Update app.py

app.py

@@ -4,8 +4,10 @@ import gradio as gr
 import numpy as np
 from ultralytics import YOLO
 
-# Load YOLOv8 model
+# Load YOLOv8 model and set device (GPU if available)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = YOLO('./data/best.pt')  # Path to your model
+model.to(device)
 
 # Define the function that processes the uploaded video
 def process_video(video):
@@ -17,19 +19,33 @@ def process_video(video):
     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 = 10  # Skip 10 frames between each processed frame
+    frame_count = 0
+
     while True:
         # Read a frame from the video
         ret, frame = input_video.read()
         if not ret:
             break  # End of video
 
+        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)
+        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