Update app.py

app.py

@@ -7,6 +7,7 @@ from pathlib import Path
 import sys
 import importlib.util
 import time
+import gc
 
 # Ensure models directory exists
 MODELS_DIR = Path("models")
@@ -47,7 +48,19 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-def run_tracking_single(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, temp_dir):
+def clean_memory():
+    """Force garbage collection to free memory."""
+    gc.collect()
+    if hasattr(torch, 'cuda') and torch.cuda.is_available():
+        try:
+            import torch
+            torch.cuda.empty_cache()
+            print("🧹 Cleared CUDA memory cache")
+        except (ImportError, NameError):
+            print("⚠️ Could not clear CUDA memory (torch not available)")
+    print("🧹 Memory cleaned")
+
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, temp_dir):
     """Run object tracking with a single tracking method."""
     try:
         # Prepare input
@@ -167,66 +180,64 @@ def run_tracking_single(video_file, yolo_model, reid_model, tracking_method, cla
         traceback.print_exc()
         return None, f"Error in {tracking_method}: {str(e)}", 0
 
-def benchmark_trackers(video_path, yolo_model, reid_model, class_ids, conf_threshold):
-    """Run all tracking methods and return their results."""
+def sequential_tracker(video_path, yolo_model, reid_model, tracking_method, class_ids, conf_threshold, progress=gr.Progress()):
+    """Run a single tracker and return its results."""
     if not video_path:
-        return (None, None, None, None,  # Videos
-                "Please upload a video file", "Please upload a video file", 
-                "Please upload a video file", "Please upload a video file",  # Statuses
-                "", "", "", "",  # Times
-                None)  # DataFrame
+        return None, "Please upload a video file", "", None
     
-    print(f"Benchmarking video: {video_path}")
+    print(f"Processing video: {video_path} with {tracking_method}")
     print(f"Parameters: model={yolo_model}, reid={reid_model}, classes={class_ids}, conf={conf_threshold}")
     
-    tracking_methods = ["bytetrack", "botsort", "ocsort", "strongsort"]
-    results = [None] * 4
-    statuses = ["Processing..."] * 4
-    times = [0] * 4
+    progress(0, desc=f"Starting {tracking_method}...")
     
     try:
-        # Create a single temporary directory for all processes
+        # Create a temporary directory for processing
         with tempfile.TemporaryDirectory() as temp_dir:
-            # Process each tracking method sequentially
-            for i, method in enumerate(tracking_methods):
-                result, status, process_time = run_tracking_single(
-                    video_path,
-                    yolo_model,
-                    reid_model,
-                    method,
-                    class_ids,
-                    conf_threshold,
-                    temp_dir
-                )
-                
-                # Store results
-                if result:
-                    results[i] = os.path.abspath(result)
-                statuses[i] = status
-                times[i] = f"{process_time:.2f} seconds"
+            progress(0.1, desc=f"Running {tracking_method}...")
+            
+            result, status, process_time = run_tracking(
+                video_path,
+                yolo_model,
+                reid_model,
+                tracking_method,
+                class_ids,
+                conf_threshold,
+                temp_dir
+            )
+            
+            progress(0.9, desc="Finalizing results...")
+            
+            # Clean up memory after each tracker run
+            clean_memory()
+            
+            # Create the DataFrame data for this single tracker
+            comparison_data = [[tracking_method, f"{process_time:.2f} seconds", "Success" if result else "Failed"]]
+            
+            # Return the results
+            return result, status, f"{process_time:.2f} seconds", comparison_data
     
     except Exception as e:
         import traceback
         traceback.print_exc()
-        # On failure, fill in the error for all pending results
-        error_msg = f"Benchmark process error: {str(e)}"
-        for i in range(len(tracking_methods)):
-            if statuses[i] == "Processing...":
-                statuses[i] = error_msg
-    
-    # Create the DataFrame data
-    comparison_data = [[tracking_methods[i], times[i], "Success" if results[i] else "Failed"] 
-                      for i in range(len(tracking_methods))]
+        error_msg = f"Process error for {tracking_method}: {str(e)}"
+        
+        # Clean memory even if there was an error
+        clean_memory()
+        
+        return None, error_msg, "", [[tracking_method, "Error", "Failed"]]
+
+def update_comparison_table(current_data, new_data):
+    """Update the comparison table with results from a new tracker."""
+    if current_data is None:
+        return new_data
     
-    # Return all results as separate outputs
-    return (results[0], results[1], results[2], results[3],
-            statuses[0], statuses[1], statuses[2], statuses[3],
-            times[0], times[1], times[2], times[3],
-            comparison_data)
+    # Append the new tracker data to the existing table
+    return current_data + new_data
 
 # 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()
@@ -235,7 +246,7 @@ apply_patches()
 # Create the Gradio interface
 with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
     gr.Markdown("# 🔍 YOLO Object Tracking Benchmark")
-    gr.Markdown("Upload a video file to benchmark all four tracking methods. Processing may take several minutes depending on video length.")
+    gr.Markdown("Upload a video file and run each tracking method sequentially. Results will be displayed as they become available.")
     
     # Add class reference information
     with gr.Accordion("YOLO Class Reference", open=False):
@@ -259,6 +270,9 @@ with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
         [See full COCO class list here](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datasets/coco128.yaml)
         """)
     
+    # State variables to keep track of comparison data
+    comparison_data_state = gr.State([])
+    
     with gr.Row():
         with gr.Column(scale=1):
             input_video = gr.Video(label="Input Video", sources=["upload"])
@@ -289,8 +303,18 @@ with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
                     step=0.05, 
                     label="Confidence Threshold"
                 )
-                
-            benchmark_btn = gr.Button("Benchmark All Trackers", variant="primary")
+            
+            # Individual tracker buttons
+            with gr.Group():
+                gr.Markdown("### Run Trackers One by One")
+                with gr.Row():
+                    bytetrack_btn = gr.Button("Run ByteTrack", variant="primary")
+                with gr.Row():
+                    botsort_btn = gr.Button("Run BoTSORT", variant="primary")  
+                with gr.Row():
+                    ocsort_btn = gr.Button("Run OC-SORT", variant="primary")
+                with gr.Row():
+                    strongsort_btn = gr.Button("Run StrongSORT", variant="primary")
     
     # Output Tabs for each tracking method
     with gr.Tabs() as tabs:
@@ -323,17 +347,49 @@ with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
                 label="Performance Comparison"
             )
     
-    # Connect the benchmark button to the function and outputs
-    benchmark_btn.click(
-        fn=benchmark_trackers,
-        inputs=[input_video, yolo_model, reid_model, class_ids, conf_threshold],
-        outputs=[
-            bytetrack_video, botsort_video, ocsort_video, strongsort_video,
-            bytetrack_status, botsort_status, ocsort_status, strongsort_status,
-            bytetrack_time, botsort_time, ocsort_time, strongsort_time,
-            perf_table
-        ],
+    # Connect individual tracker buttons to their respective functions
+    bytetrack_btn.click(
+        fn=sequential_tracker,
+        inputs=[input_video, yolo_model, reid_model, gr.State("bytetrack"), class_ids, conf_threshold],
+        outputs=[bytetrack_video, bytetrack_status, bytetrack_time, comparison_data_state],
+        show_progress="full"
+    ).then(
+        fn=update_comparison_table,
+        inputs=[perf_table, comparison_data_state],
+        outputs=perf_table
+    )
+    
+    botsort_btn.click(
+        fn=sequential_tracker,
+        inputs=[input_video, yolo_model, reid_model, gr.State("botsort"), class_ids, conf_threshold],
+        outputs=[botsort_video, botsort_status, botsort_time, comparison_data_state],
         show_progress="full"
+    ).then(
+        fn=update_comparison_table,
+        inputs=[perf_table, comparison_data_state],
+        outputs=perf_table
+    )
+    
+    ocsort_btn.click(
+        fn=sequential_tracker,
+        inputs=[input_video, yolo_model, reid_model, gr.State("ocsort"), class_ids, conf_threshold],
+        outputs=[ocsort_video, ocsort_status, ocsort_time, comparison_data_state],
+        show_progress="full"
+    ).then(
+        fn=update_comparison_table,
+        inputs=[perf_table, comparison_data_state],
+        outputs=perf_table
+    )
+    
+    strongsort_btn.click(
+        fn=sequential_tracker,
+        inputs=[input_video, yolo_model, reid_model, gr.State("strongsort"), class_ids, conf_threshold],
+        outputs=[strongsort_video, strongsort_status, strongsort_time, comparison_data_state],
+        show_progress="full"
+    ).then(
+        fn=update_comparison_table,
+        inputs=[perf_table, comparison_data_state],
+        outputs=perf_table
     )
     
     # Add a debug section
@@ -362,10 +418,23 @@ with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
             for model in os.listdir(MODELS_DIR) if os.path.exists(MODELS_DIR) else []:
                 info.append(f"  - {model}")
             
+            # Check for GPU
+            try:
+                import torch
+                info.append(f"CUDA available: {torch.cuda.is_available()}")
+                if torch.cuda.is_available():
+                    info.append(f"CUDA device: {torch.cuda.get_device_name(0)}")
+                    info.append(f"CUDA memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+            except ImportError:
+                info.append("PyTorch: Not installed")
+            
             return "\n".join(info)
         
         check_btn = gr.Button("Check Environment")
         check_btn.click(fn=check_environment, outputs=debug_text)
+        
+        clean_mem_btn = gr.Button("Force Memory Cleanup")
+        clean_mem_btn.click(fn=clean_memory)
 
 # Launch the app
 if __name__ == "__main__":