Update app.py

app.py

@@ -1,51 +1,45 @@
-import gradio as gr
+from ultralytics import YOLO  # Use Ultralytics' YOLO module
 import torch
 import cv2
 import numpy as np
+import gradio as gr
 from PIL import Image
-from torchvision.transforms import functional as F
-from ultralytics.yolo.utils.ops import non_max_suppression
-from ultralytics.yolo.engine.model import Model
-
 
 # Load YOLOv5 model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
+model = YOLO("yolov5s.pt")  # Load pre-trained YOLOv5s model
+model.to(device)
 model.eval()
 
 def preprocess_image(image):
     image = Image.fromarray(image)
-    image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
-    return image_tensor
+    image = image.convert("RGB")
+    return image
 
-def draw_boxes(image, outputs, threshold=0.3):
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    h, w, _ = image.shape
+def detect_objects(image):
+    image = preprocess_image(image)
+    results = model.predict(image)  # Run YOLOv5 inference
 
-    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}"
-            cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+    # Convert results to bounding box format
+    detections = []
+    for result in results:
+        for box in result.boxes.xyxy:
+            x1, y1, x2, y2 = map(int, box[:4])
+            detections.append([x1, y1, x2, y2])
 
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    # Draw bounding boxes
+    image = np.array(image)
+    for x1, y1, x2, y2 in detections:
+        cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
 
-def detect_objects(image):
-    image_tensor = preprocess_image(image)
-    outputs = model(image_tensor)
-    outputs = non_max_suppression(outputs)[0]
-    result_image = draw_boxes(image, outputs.cpu().numpy())
-    return result_image
+    return image
 
+# Gradio UI
 iface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Image(type="numpy"),
     outputs=gr.Image(type="numpy"),
-    title="YOLOv5 Object Detection",
-    description="Upload an image to detect objects using the YOLOv5 model."
+    live=True,
 )
 
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch()