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