Update app.py

app.py

@@ -1,22 +1,38 @@
-from fastapi import FastAPI, UploadFile, File
-from transformers import YolosForObjectDetection, YolosImageProcessor
-from PIL import Image
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
 import torch
-import io
+from PIL import Image
+from torchvision.transforms import functional as F
+from yolov5.models.yolo import Model
+from yolov5.utils.general import non_max_suppression
 
 app = FastAPI()
 
-# Load the YOLOv5 model and processor from Hugging Face
-model = YolosForObjectDetection.from_pretrained('ultralytics/yolov5')
-processor = YolosImageProcessor.from_pretrained('ultralytics/yolov5')
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
 model.eval()
 
+def preprocess_image(image):
+    image_tensor = F.to_tensor(image)
+    return image_tensor.unsqueeze(0).to(device)
+
+def draw_boxes(outputs, threshold=0.3):
+    boxes = []
+    for box in outputs:
+        score, label, x1, y1, x2, y2 = box[4].item(), int(box[5].item()), box[0].item(), box[1].item(), box[2].item(), box[3].item()
+        if score > threshold:
+            boxes.append({
+                "label": model.names[label],
+                "score": score,
+                "box": [x1, y1, x2, y2]
+            })
+    return boxes
+
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents))
-    inputs = processor(images=image, return_tensors="pt").to(device)
-    outputs = model(**inputs)
-    return outputs
+    image = Image.open(file.file)
+    image_tensor = preprocess_image(image)
+    outputs = model(image_tensor)
+    outputs = non_max_suppression(outputs)[0]
+    boxes = draw_boxes(outputs)
+    return JSONResponse(content={"boxes": boxes})