Update app.py

app.py

@@ -1,6 +1,7 @@
 import cv2
 import torch
 import gradio as gr
+import numpy as np
 from ultralytics import YOLO
 
 # Load YOLOv8 model
@@ -16,10 +17,8 @@ def process_video(video):
     frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = input_video.get(cv2.CAP_PROP_FPS)
 
-    # Define output video writer
-    output_video_path = "/mnt/data/output_video.mp4"  # Path to save the output video
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    output_video = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
+    # Create an empty list to store processed frames
+    processed_frames = []
 
     while True:
         # Read a frame from the video
@@ -33,19 +32,22 @@ def process_video(video):
         # The results object contains annotations for the frame
         annotated_frame = results[0].plot()  # Plot the frame with bounding boxes
 
-        # Write the annotated frame to the output video
-        output_video.write(annotated_frame)
+        # Convert the annotated frame to RGB format
+        annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
+
+        # Append the frame to the list
+        processed_frames.append(annotated_frame_rgb)
 
     # Release resources
     input_video.release()
-    output_video.release()
 
-    return output_video_path
+    # Return the processed frames as an output video in Gradio
+    return processed_frames
 
 # Create a Gradio interface for video upload
 iface = gr.Interface(fn=process_video, 
                      inputs=gr.Video(label="Upload Video"),  # Updated line
-                     outputs="file", 
+                     outputs=gr.Video(),  # This will display the output video directly
                      title="YOLOv8 Object Detection on Video",
                      description="Upload a video for object detection using YOLOv8")