Update app.py

app.py

@@ -1,71 +1,57 @@
 import cv2
 import torch
-import gradio as gr
+import gradio
 import numpy as np
 from ultralytics import YOLO
 
-# Load YOLOv8 model and set device (GPU if available)
+# Load YOLOv8 model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = YOLO('./data/best.pt')  # Path to your model
-model.to(device)
+print(f"Using device: {device}")  # Debug
+model = YOLO('./data/best.pt').to(device)
 
-# Define the function that processes the uploaded video
 def process_video(video):
-    # video is now the file path string, not a file object
-    input_video = cv2.VideoCapture(video)  # Directly pass the path to cv2.VideoCapture
-
-    # Get frame width, height, and fps from input video
-    frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = input_video.get(cv2.CAP_PROP_FPS)
-
-    # Resize to reduce computation (optional)
-    new_width, new_height = 640, 480  # Resize to 640x480 resolution
-    frame_width, frame_height = new_width, new_height
-
-    # Create a VideoWriter object to write processed frames to an output file
-    output_video_path = "processed_output.mp4"
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4 format
-    output_video = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
-
-    frame_skip = 1  # Skip 10 frames between each processed frame
+    cap = cv2.VideoCapture(video)
+    frame_width, frame_height = 320, 240  # Smaller resolution
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    output_path = "processed_output.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fourcc, fps, (frame_width, frame_height))
+    
     frame_count = 0
-
+    frame_skip = 5  # Process every 5th frame
+    max_frames = 100  # Limit to 100 frames for testing
+    
     while True:
-        # Read a frame from the video
-        ret, frame = input_video.read()
-        if not ret:
-            break  # End of video
-
+        ret, frame = cap.read()
+        if not ret or frame_count > max_frames:
+            break
+            
         frame_count += 1
         if frame_count % frame_skip != 0:
-            continue  # Skip frames
-
-        # Resize the frame to reduce computational load
-        frame = cv2.resize(frame, (new_width, new_height))
-
-        # Perform inference on the frame
-        results = model(frame)  # Automatically uses GPU if available
-
-        # 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)
-
-    # Release resources
-    input_video.release()
-    output_video.release()
-
-    # Return the processed video file path
-    return output_video_path
-
-# Create a Gradio interface for video upload
-iface = gr.Interface(fn=process_video, 
-                     inputs=gr.Video(label="Upload Video"),  # Updated line
-                     outputs=gr.Video(label="Processed Video"),  # This will display the output video directly
-                     title="YOLOv8 Object Detection on Video",
-                     description="Upload a video for object detection using YOLOv8")
-
-# Launch the interface
-iface.launch()
+            continue
+            
+        frame = cv2.resize(frame, (frame_width, frame_height))
+        print(f"Processing frame {frame_count}")  # Debug
+        
+        # Inference
+        results = model(frame)
+        annotated_frame = results[0].plot()
+        
+        # Write frame
+        out.write(annotated_frame)
+    
+    cap.release()
+    out.release()
+    return output_path
+
+# Gradio interface
+iface = gr.Interface(
+    fn=process_video,
+    inputs=gr.Video(label="Upload Video"),
+    outputs=gr.Video(label="Processed Video"),
+    title="YOLOv8 Object Detection",
+    description="Upload a short video for testing"
+)
+
+iface.launch()