Update app.py

app.py

@@ -1,46 +1,59 @@
-import os
+!pip install -U torch torchvision cython
+!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
+!pip install gradio
+
 import torch
-import gradio as gr
-import cv2
-import numpy as np
-from PIL import Image
+import torchvision
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
 from torchvision.transforms import functional as F
+from PIL import Image
+import cv2
+from google.colab.patches import cv2_imshow
+import gradio as gr
 
-# Clone yolov5 if not present
-if not os.path.exists("yolov5"):
-    os.system("git clone https://github.com/ultralytics/yolov5.git")
-    os.system("pip install -r yolov5/requirements.txt")
+!git clone https://github.com/ultralytics/yolov5
+%cd yolov5
+!pip install -r requirements.txt
 
-from yolov5.utils.general import non_max_suppression
+import torch
+from yolov5.models.yolo import Model
 
-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()
 
+from yolov5.utils.general import non_max_suppression
+
 def preprocess_image(image):
-    image = image.convert("RGB")
-    image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
-    return image_tensor
+    image_tensor = F.to_tensor(image)
+    return image_tensor.unsqueeze(0).to(device)
 
 def draw_boxes(image, outputs, threshold=0.3):
-    image = np.array(image)
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    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:
             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}"
+            text = f"{model.names[label]:s}: {score:.2f}"
             cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
-    return Image.fromarray(image)
+
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 def detect_objects(image):
     image_tensor = preprocess_image(image)
     outputs = model(image_tensor)
     outputs = non_max_suppression(outputs)[0]
-    return draw_boxes(image, outputs)
+    result_image = draw_boxes(image, outputs)
+    return result_image
 
-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.")
+iface = gr.Interface(
+    fn=detect_objects,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Image(type="pil"),
+    live=True
+)
 
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()