Update app.py

app.py

@@ -1,113 +1,60 @@
-
-import os
 import gradio as gr
+import torch
 import cv2
-import numpy as np
 from ultralytics import YOLO
-import tempfile
-import torch
-import ultralytics.nn.tasks
-
-# Set Ultralytics config path
-os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
-
-# Custom function to load model with trusted weights
-def load_model_with_trusted_weights(model_path):
-    with torch.serialization.safe_globals([
-        ultralytics.nn.tasks.DetectionModel,
-        torch.nn.modules.container.Sequential
-    ]):
-        return YOLO(model_path)
-
-# Load both YOLO models
-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
-    model = model_yolo11 if model_name == "YOLO11n" else model_best
 
-    # Open video capture
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError("Could not open video file")
+# Load YOLO models
+model_yolo11 = YOLO('./data/yolo11n.pt')
+model_best = YOLO('./data/best.pt')
 
-    # Get video properties
-    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))
+def process_video(video):
+    # Read video input
+    cap = cv2.VideoCapture(video.name)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-    # 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))
+    # Create a VideoWriter object to save the output video
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4
+    out = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height))
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-
-        # Perform detection
-        results = model.predict(
-            source=frame,
-            conf=conf_threshold,
-            imgsz=640,
-            show_labels=True,
-            show_conf=True
-        )
-
-        # Draw bounding boxes
-        for result in results:
-            im_array = result.plot()  # Plot boxes
-            out.write(im_array)  # Write frame to output video
+        
+        # Use both YOLO models for detection
+        results_yolo11 = model_yolo11(frame)
+        results_best = model_best(frame)
+        
+        # Combine the results from both models
+        # For simplicity, we will overlay bounding boxes and labels from both models
+        for result in results_yolo11:
+            boxes = result.boxes
+            for box in boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                label = f"YOLOv11: {box.cls[0]} - {box.conf[0]:.2f}"
+                cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+
+        for result in results_best:
+            boxes = result.boxes
+            for box in boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
+                label = f"Best: {box.cls[0]} - {box.conf[0]:.2f}"
+                cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+
+        # Write the processed frame to the output video
+        out.write(frame)
 
     cap.release()
     out.release()
-    cv2.destroyAllWindows()
 
-    return temp_video_path
+    return 'output_video.mp4'
 
 # Gradio interface
-with gr.Blocks() as app:
-    gr.HTML("""
-        <h1 style='text-align: center'>
-            Video Object Detection with YOLO Models
-        </h1>
-    """)
-    
-    with gr.Row():
-        with gr.Column():
-            video_input = gr.Video(label="Upload Video")
-            model_choice = gr.Dropdown(
-                choices=["YOLO11n", "Best Model"],
-                label="Select Model",
-                value="YOLO11n"
-            )
-            conf_threshold = gr.Slider(
-                label="Confidence Threshold",
-                minimum=0.0,
-                maximum=1.0,
-                step=0.05,
-                value=0.4
-            )
-            process_button = gr.Button("Process Video")
-        
-        with gr.Column():
-            video_output = gr.Video(
-                label="Processed Video",
-                streaming=True,
-                autoplay=True
-            )
-
-    process_button.click(
-        fn=process_video,
-        inputs=[video_input, model_choice, conf_threshold],
-        outputs=[video_output]
-    )
+iface = gr.Interface(fn=process_video, inputs=gr.Video(), outputs=gr.Video(), live=True)
 
-if __name__ == "__main__":
-    app.launch()
+# Launch the app
+iface.launch()