done

app.py

@@ -6,6 +6,7 @@ import shutil
 from pathlib import Path
 import sys
 import importlib.util
+from concurrent.futures import ThreadPoolExecutor
 
 # Ensure models directory exists
 MODELS_DIR = Path("models")
@@ -43,96 +44,125 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-def run_tracking(video_file, yolo_model, reid_model, tracking_method, conf_threshold):
-    """Run object tracking on the uploaded video."""
+def run_single_tracker(video_file, yolo_model, reid_model, tracking_method, conf_threshold, temp_dir):
+    """Run object tracking with a single tracking method."""
     try:
-        # Create temporary workspace
-        with tempfile.TemporaryDirectory() as temp_dir:
-            # Prepare input
-            input_path = os.path.join(temp_dir, "input_video.mp4")
+        # Prepare input
+        input_path = os.path.join(temp_dir, "input_video.mp4")
+        if not os.path.exists(input_path):
             shutil.copy(video_file, input_path)
+        
+        # Prepare output directory for this tracker
+        output_dir = os.path.join(temp_dir, f"output_{tracking_method}")
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Build command
+        cmd = [
+            "python", "tracking/track.py",
+            "--yolo-model", str(MODELS_DIR / yolo_model),
+            "--reid-model", str(MODELS_DIR / reid_model),
+            "--tracking-method", tracking_method,
+            "--source", input_path,
+            "--conf", str(conf_threshold),
+            "--save",
+            "--project", output_dir,
+            "--name", tracking_method,
+            "--exist-ok"
+        ]
+        
+        # Special handling for OcSort
+        if tracking_method == "ocsort":
+            cmd.append("--per-class")
+        
+        # Execute tracking
+        process = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True
+        )
+        
+        # Check for errors
+        if process.returncode != 0:
+            error_message = process.stderr or process.stdout
+            return None, f"Error in {tracking_method}: {error_message}"
+        
+        # Find output video
+        output_files = []
+        for root, _, files in os.walk(output_dir):
+            for file in files:
+                if file.lower().endswith((".mp4", ".avi", ".mov")):
+                    output_files.append(os.path.join(root, file))
+        
+        if not output_files:
+            return None, f"No output video generated for {tracking_method}"
             
-            # Prepare output directory
-            output_dir = os.path.join(temp_dir, "output")
-            os.makedirs(output_dir, exist_ok=True)
-            
-            # Build command
-            cmd = [
-                "python", "tracking/track.py",
-                "--yolo-model", str(MODELS_DIR / yolo_model),
-                "--reid-model", str(MODELS_DIR / reid_model),
-                "--tracking-method", tracking_method,
-                "--source", input_path,
-                "--conf", str(conf_threshold),
-                "--save",
-                "--project", output_dir,
-                "--name", "track",
-                "--exist-ok"
-            ]
-            
-            # Special handling for OcSort
-            if tracking_method == "ocsort":
-                cmd.append("--per-class")
-            
-            # Execute tracking with error handling
-            process = subprocess.run(
-                cmd,
-                capture_output=True,
-                text=True
-            )
-            
-            # Check for errors in output
-            if process.returncode != 0:
-                error_message = process.stderr or process.stdout
-                return None, f"Error in tracking process: {error_message}"
-            
-            # Find output video
-            output_files = []
-            for root, _, files in os.walk(output_dir):
-                for file in files:
-                    if file.lower().endswith((".mp4", ".avi", ".mov")):
-                        output_files.append(os.path.join(root, file))
-            
-            if not output_files:
-                return None, "No output video was generated. Check if tracking was successful."
-                
-            return output_files[0], "Processing completed successfully!"
-            
+        return output_files[0], f"{tracking_method} completed successfully!"
+        
     except Exception as e:
-        return None, f"Error: {str(e)}"
+        return None, f"Error in {tracking_method}: {str(e)}"
 
-# Define the Gradio interface
-def process_video(video_path, yolo_model, reid_model, tracking_method, conf_threshold):
+def run_all_trackers(video_path, yolo_model, reid_model, conf_threshold):
+    """Run all tracking methods and return results."""
+    tracking_methods = ["bytetrack", "botsort", "ocsort", "strongsort"]
+    results = {}
+    status_messages = []
+    
     # Validate inputs
     if not video_path:
-        return None, "Please upload a video file"
+        return None, None, None, None, "Please upload a video file"
+    
+    # Create temporary workspace
+    temp_dir = tempfile.mkdtemp()
     
-    output_path, status = run_tracking(
+    try:
+        for method in tracking_methods:
+            # Run tracking for this method
+            output_path, status = run_single_tracker(
+                video_path,
+                yolo_model,
+                reid_model,
+                method,
+                conf_threshold,
+                temp_dir
+            )
+            
+            results[method] = output_path
+            status_messages.append(status)
+        
+        combined_status = "\n".join(status_messages)
+        return results.get("bytetrack"), results.get("botsort"), results.get("ocsort"), results.get("strongsort"), combined_status
+        
+    except Exception as e:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        return None, None, None, None, f"Error: {str(e)}"
+
+# Define the Gradio interface
+def process_video(video_path, yolo_model, reid_model, conf_threshold):
+    # Run all tracking methods
+    bytetrack_output, botsort_output, ocsort_output, strongsort_output, status = run_all_trackers(
         video_path,
         yolo_model,
         reid_model,
-        tracking_method,
         conf_threshold
     )
     
-    return output_path, status
+    return bytetrack_output, botsort_output, ocsort_output, strongsort_output, status
 
-# Available models and tracking methods
+# Available models
 yolo_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"]
 reid_models = ["osnet_x0_25_msmt17.pt"]
-tracking_methods = ["bytetrack", "botsort", "ocsort", "strongsort"]
 
 # Ensure dependencies and apply patches at startup
 ensure_dependencies()
 apply_patches()
 
 # Create the Gradio interface
-with gr.Blocks(title="YOLO Object Tracking") as app:
-    gr.Markdown("# 🚀 YOLO Object Tracking")
-    gr.Markdown("Upload a video file to detect and track objects. Processing may take a few minutes depending on video length.")
+with gr.Blocks(title="YOLO Multi-Tracker Comparison") as app:
+    gr.Markdown("# 🚀 YOLO Multi-Tracker Comparison")
+    gr.Markdown("Upload a video file to see tracking results from all four trackers side by side.")
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             input_video = gr.Video(label="Input Video", sources=["upload"])
             
             with gr.Group():
@@ -146,11 +176,6 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
                     value="osnet_x0_25_msmt17.pt", 
                     label="ReID Model"
                 )
-                tracking_method = gr.Dropdown(
-                    choices=tracking_methods, 
-                    value="bytetrack", 
-                    label="Tracking Method"
-                )
                 conf_threshold = gr.Slider(
                     minimum=0.1, 
                     maximum=0.9, 
@@ -159,16 +184,32 @@ with gr.Blocks(title="YOLO Object Tracking") as app:
                     label="Confidence Threshold"
                 )
                 
-            process_btn = gr.Button("Process Video", variant="primary")
-        
+            process_btn = gr.Button("Process Video with All Trackers", variant="primary")
+            status_text = gr.Textbox(label="Status", value="Ready to process video", lines=5)
+    
+    gr.Markdown("## Tracking Results")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### ByteTrack")
+            bytetrack_video = gr.Video(label="ByteTrack Result", autoplay=False)
+            
+        with gr.Column():
+            gr.Markdown("### BOT-SORT")
+            botsort_video = gr.Video(label="BOT-SORT Result", autoplay=False)
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### OC-SORT")
+            ocsort_video = gr.Video(label="OC-SORT Result", autoplay=False)
+            
         with gr.Column():
-            output_video = gr.Video(label="Output Video with Tracking", autoplay=True)
-            status_text = gr.Textbox(label="Status", value="Ready to process video")
+            gr.Markdown("### StrongSORT")
+            strongsort_video = gr.Video(label="StrongSORT Result", autoplay=False)
     
     process_btn.click(
         fn=process_video,
-        inputs=[input_video, yolo_model, reid_model, tracking_method, conf_threshold],
-        outputs=[output_video, status_text]
+        inputs=[input_video, yolo_model, reid_model, conf_threshold],
+        outputs=[bytetrack_video, botsort_video, ocsort_video, strongsort_video, status_text]
     )
 
 # Launch the app