Update app.py

app.py

@@ -1,34 +1,62 @@
 import gradio as gr
 import cv2
-import numpy as np
 from ultralytics import YOLO
+import tempfile
+import os
 
-# Load the YOLOv8 model
-model = YOLO("best.pt")  # Ensure this file is in the same directory as app.py on Hugging Face
+# Load the YOLO model
+model = YOLO("best.pt")  # Replace with the path to your model
 
 # Define the inference function
-def predict(image):
-    # Convert the input image from RGB to BGR (OpenCV format)
-    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    
-    # Run the model on the input image
-    results = model(image_bgr)
-    
-    # Extract the result image with detections
-    annotated_image = results[0].plot()  # Returns a BGR image with annotations
-    
-    # Convert the image back to RGB for displaying in Gradio
-    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
-    
-    return annotated_image_rgb
+def yolo_inference(input_file):
+    # Check if the input is an image or a video
+    if input_file.endswith((".jpg", ".jpeg", ".png")):
+        # Process as an image
+        img = cv2.imread(input_file)
+        results = model(img)
+        annotated_img = results[0].plot()
+
+        # Save the annotated image
+        output_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+        cv2.imwrite(output_path, annotated_img)
+        return output_path
+
+    elif input_file.endswith((".mp4", ".avi", ".mov")):
+        # Process as a video
+        cap = cv2.VideoCapture(input_file)
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        
+        # Create a temporary output video path
+        output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+        
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            
+            # Run YOLO on each frame
+            results = model(frame)
+            annotated_frame = results[0].plot()
+            out.write(annotated_frame)
+        
+        cap.release()
+        out.release()
+        return output_video_path
+
+    else:
+        raise ValueError("Unsupported file format. Please upload an image or video.")
 
 # Define the Gradio interface
 interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="numpy", label="Upload an Image"),
-    outputs=gr.Image(type="numpy", label="Detected Objects"),
-    title="YOLOv8 Object Detection",
-    description="Upload an image to detect objects with YOLOv8 model."
+    fn=yolo_inference,
+    inputs=gr.File(label="Upload an Image or Video"),
+    outputs=gr.File(label="Processed Image/Video"),
+    title="YOLO Object Detection",
+    description="Upload an image or video for object detection using YOLO."
 )
 
 # Launch the app