Update app.py

app.py

@@ -1,24 +1,30 @@
-
+```python
 import os
-os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'  # Set Ultralytics config path
 import gradio as gr
 import cv2
 import numpy as np
 from ultralytics import YOLO
 import tempfile
-from moviepy.editor import ImageSequenceClip
-from PIL import Image
+
+# Set Ultralytics config path
+os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
+
+# Custom function to load model with weights_only=False
+def load_model_with_trusted_weights(model_path):
+    import torch
+    with torch.serialization.safe_globals():
+        return YOLO(model_path)
 
 # Load both YOLO models
-model_yolo11 = YOLO('./data/yolo11n.pt')
-model_best = YOLO('./data/best.pt')
+model_yolo11 = load_model_with_trusted_weights('./data/yolo11n.pt')
+model_best = load_model_with_trusted_weights('./data/best.pt')
 
 def process_video(video_path, model_name, conf_threshold=0.4):
     """
     Process the input video frame by frame using the selected YOLO model,
     draw bounding boxes, and return the processed video path.
     """
-    # Select model to use
+    # Select model
     model = model_yolo11 if model_name == "YOLO11n" else model_best
 
     # Open video capture
@@ -31,8 +37,10 @@ def process_video(video_path, model_name, conf_threshold=0.4):
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-    # Store processed frames
-    processed_frames = []
+    # Define output video path
+    temp_video_path = os.path.join(tempfile.gettempdir(), "output.mp4")
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
+    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
 
     while cap.isOpened():
         ret, frame = cap.read()
@@ -50,15 +58,12 @@ def process_video(video_path, model_name, conf_threshold=0.4):
 
         # Draw bounding boxes
         for result in results:
-            im_array = result.plot() # Plot boxes
-            processed_frames.append(im_array[..., ::-1]) # Convert BGR to RGB
+            im_array = result.plot()  # Plot boxes
+            out.write(im_array)  # Write frame to output video
 
     cap.release()
-
-    # Save processed frames to temp video
-    temp_video_path = os.path.join(tempfile.gettempdir(), "output.mp4")
-    clip = ImageSequenceClip(processed_frames, fps=fps)
-    clip.write_videofile(temp_video_path, codec='libx264')
+    out.release()
+    cv2.destroyAllWindows()
 
     return temp_video_path
 
@@ -102,3 +107,4 @@ with gr.Blocks() as app:
 
 if __name__ == "__main__":
     app.launch()
+```