Update app.py

app.py

@@ -1,38 +1,45 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
+import gradio as gr
 import torch
+import cv2
+import numpy as np
 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()
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 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)
+    image = image.convert("RGB")
+    image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
+    return image_tensor
 
-def draw_boxes(outputs, threshold=0.3):
-    boxes = []
+def draw_boxes(image, outputs, threshold=0.3):
+    image = np.array(image)
+    h, w, _ = image.shape
     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
+            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+            cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+            text = f"{model.names[label]}: {score:.2f}"
+            cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+    return Image.fromarray(image)
 
-@app.post("/predict/")
-async def predict(file: UploadFile = File(...)):
-    image = Image.open(file.file)
+def detect_objects(image):
     image_tensor = preprocess_image(image)
     outputs = model(image_tensor)
     outputs = non_max_suppression(outputs)[0]
-    boxes = draw_boxes(outputs)
-    return JSONResponse(content={"boxes": boxes})
+    return draw_boxes(image, outputs)
+
+iface = gr.Interface(
+    fn=detect_objects,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="YOLO Object Detector",
+    description="Upload an image to detect objects using YOLOv5."
+)
+
+if __name__ == "__main__":
+    iface.launch()