Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 import cv2
-import torch
 import numpy as np
+from ultralytics import YOLO  # Import YOLO from ultralytics
 
 # Load the YOLOv5 model
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+model = YOLO('yolov5s')  # Use the YOLOv5s pre-trained model
 
 # Function to run inference on an image
 def run_inference(image):
@@ -12,11 +12,10 @@ def run_inference(image):
     image = np.array(image)
 
     # Run YOLOv5 inference
-    results = model(image)
+    results = model.predict(source=image, save=False, conf=0.25, stream=False)
 
-    # Convert the annotated image from BGR to RGB for display
-    annotated_image = results.render()[0]
-    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+    # Annotate the image with detected objects
+    annotated_image = results[0].plot()  # Use YOLO's built-in plotting function
 
     return annotated_image
 
@@ -24,7 +23,7 @@ def run_inference(image):
 interface = gr.Interface(
     fn=run_inference,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="numpy"),
     title="YOLOv5 Object Detection",
     description="Upload an image to run YOLOv5 object detection and see the results."
 )