tejastake commited on
Commit
3a555e5
·
verified ·
1 Parent(s): 5989673

Upload 3 files

Browse files
Files changed (3) hide show
  1. Dockerfile +19 -0
  2. app.py +44 -0
  3. best.pt +3 -0
Dockerfile ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9
2
+
3
+ WORKDIR /app
4
+
5
+ COPY . /app
6
+
7
+ RUN pip3 install fastapi uvicorn pydantic ultralytics pillow
8
+ RUN useradd -m -u 1000 user
9
+
10
+ USER user
11
+
12
+ ENV HOME=/home/user \
13
+ PATH=/home/user/.local/bin:$PATH
14
+
15
+ WORKDIR $HOME/app
16
+
17
+ COPY --chown=user . $HOME/app
18
+
19
+ CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException
2
+ from pydantic import BaseModel
3
+ from ultralytics import YOLO
4
+ from PIL import Image
5
+ import io
6
+ import base64
7
+
8
+ app = FastAPI()
9
+
10
+ # Load the YOLO model
11
+ model = YOLO(r'best.pt')
12
+
13
+ class ImageData(BaseModel):
14
+ image_base64: str
15
+
16
+ @app.post("/process_image/")
17
+ async def process_image(data: ImageData):
18
+ try:
19
+ # Decode the base64 string to an image
20
+ image_data = base64.b64decode(data.image_base64)
21
+ image = Image.open(io.BytesIO(image_data))
22
+
23
+ # Process the image with YOLO
24
+ results = model(image)
25
+ result = results[0]
26
+
27
+ # Extract bounding boxes and confidence scores
28
+ boxes = result.boxes.xyxy # Bounding box coordinates
29
+ scores = result.boxes.conf # Confidence scores
30
+
31
+ if len(boxes) > 0:
32
+ # Get the index of the bounding box with the highest score
33
+ highest_score_idx = scores.argmax()
34
+ # Extract the bounding box with the highest score
35
+ highest_score_box = boxes[highest_score_idx].tolist()
36
+ x1, y1, x2, y2 = map(int, highest_score_box) # Convert to integers
37
+ else:
38
+ # If no boxes, return the whole image dimensions
39
+ x1, y1, x2, y2 = 0, 0, image.width, image.height
40
+
41
+ return {"x1": x1, "y1": y1, "x2": x2, "y2": y2}
42
+
43
+ except Exception as e:
44
+ raise HTTPException(status_code=500, detail=str(e))
best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:10420dfe7db4c88e5bd446f604c3b97e99322962e22e1898ad4d2f325e9b2a3a
3
+ size 52177122