Update app.py

app.py

@@ -6,8 +6,6 @@ import shutil
 from pathlib import Path
 import sys
 import importlib.util
-import time
-import gc
 
 # Ensure models directory exists
 MODELS_DIR = Path("models")
@@ -48,193 +46,150 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-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."""
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold):
+    """Run object tracking on the uploaded video."""
     try:
-        # Prepare input
-        input_path = os.path.join(temp_dir, "input_video.mp4")
-        if not os.path.exists(input_path):
+        # Create temporary workspace
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Prepare input
+            input_path = os.path.join(temp_dir, "input_video.mp4")
             shutil.copy(video_file, input_path)
-        
-        # Prepare output directory
-        output_dir = os.path.join(temp_dir, tracking_method)
-        os.makedirs(output_dir, exist_ok=True)
-        
-        start_time = time.time()
-        
-        # 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"
-        ]
-        
-        # Add class filtering if specific classes are provided
-        if class_ids and class_ids.strip():
-            # Parse the comma-separated class IDs
-            try:
-                # Split by comma and convert to integers to validate
-                class_list = [int(c.strip()) for c in class_ids.split(",") if c.strip()]
-                # Add each class ID as a separate argument
-                if class_list:
-                    cmd.append("--classes")
-                    cmd.extend(str(c) for c in class_list)
-            except ValueError:
-                return None, f"Invalid class IDs for {tracking_method}. Please enter comma-separated numbers.", 0
-        
-        # Special handling for OcSort
-        if tracking_method == "ocsort":
-            cmd.append("--per-class")
-        
-        # Execute tracking with error handling
-        print(f"Executing command for {tracking_method}: {' '.join(cmd)}")
-        process = subprocess.run(
-            cmd,
-            capture_output=True,
-            text=True
-        )
-        
-        end_time = time.time()
-        processing_time = end_time - start_time
-        
-        # Check for errors in output
-        if process.returncode != 0:
-            error_message = process.stderr or process.stdout
-            print(f"Process for {tracking_method} failed with return code {process.returncode}")
-            print(f"Error: {error_message}")
-            return None, f"Error in {tracking_method}: {error_message}", processing_time
-        
-        print(f"Process for {tracking_method} completed with return code {process.returncode}")
-        
-        # 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))
-        
-        print(f"Found output files for {tracking_method}: {output_files}")
-        
-        if not output_files:
-            print(f"No output video files found for {tracking_method}")
-            return None, f"No output video was generated for {tracking_method}.", processing_time
-        
-        output_file = output_files[0]
-        print(f"Selected output file for {tracking_method}: {output_file}")
-        
-        # Verify file exists and has size
-        if os.path.exists(output_file):
-            file_size = os.path.getsize(output_file)
-            print(f"Output file for {tracking_method} exists with size: {file_size} bytes")
             
-            if file_size == 0:
-                return None, f"Output video for {tracking_method} was generated but has zero size.", processing_time
+            # Prepare output directory
+            output_dir = os.path.join(temp_dir, "output")
+            os.makedirs(output_dir, exist_ok=True)
             
-            # Copy to permanent location with unique name
-            permanent_path = os.path.join(OUTPUT_DIR, f"{tracking_method}_{os.path.basename(video_file)}")
-            shutil.copy(output_file, permanent_path)
-            print(f"Copied {tracking_method} output to permanent location: {permanent_path}")
+            # 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"
+            ]
             
-            # Ensure the file is in MP4 format for better compatibility with Gradio
-            if not permanent_path.lower().endswith('.mp4'):
-                mp4_path = os.path.splitext(permanent_path)[0] + '.mp4'
+            # Add class filtering if specific classes are provided
+            if class_ids and class_ids.strip():
+                # Parse the comma-separated class IDs
                 try:
-                    print(f"Converting to MP4 format: {mp4_path}")
-                    subprocess.run([
-                        'ffmpeg', '-i', permanent_path, 
-                        '-c:v', 'libx264', '-preset', 'fast',
-                        '-c:a', 'aac', mp4_path
-                    ], check=True, capture_output=True)
-                    os.remove(permanent_path)  # Remove the original file
-                    permanent_path = mp4_path
-                except Exception as e:
-                    print(f"Failed to convert to MP4: {str(e)}")
-                    # Continue with original file if conversion fails
+                    # Split by comma and convert to integers to validate
+                    class_list = [int(c.strip()) for c in class_ids.split(",") if c.strip()]
+                    # Add each class ID as a separate argument
+                    if class_list:
+                        cmd.append("--classes")
+                        cmd.extend(str(c) for c in class_list)
+                except ValueError:
+                    return None, "Invalid class IDs. Please enter comma-separated numbers (e.g., '0,1,2')."
             
-            return permanent_path, f"{tracking_method} completed successfully in {processing_time:.2f} seconds!", processing_time
-        else:
-            print(f"Output file for {tracking_method} not found at {output_file}")
-            return None, f"Output file for {tracking_method} was referenced but doesn't exist on disk.", processing_time
-        
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        return None, f"Error in {tracking_method}: {str(e)}", 0
-
-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, "Please upload a video file", "", None
-    
-    print(f"Processing video: {video_path} with {tracking_method}")
-    print(f"Parameters: model={yolo_model}, reid={reid_model}, classes={class_ids}, conf={conf_threshold}")
-    
-    progress(0, desc=f"Starting {tracking_method}...")
-    
-    try:
-        # Create a temporary directory for processing
-        with tempfile.TemporaryDirectory() as temp_dir:
-            progress(0.1, desc=f"Running {tracking_method}...")
+            # Special handling for OcSort
+            if tracking_method == "ocsort":
+                cmd.append("--per-class")
             
-            result, status, process_time = run_tracking(
-                video_path,
-                yolo_model,
-                reid_model,
-                tracking_method,
-                class_ids,
-                conf_threshold,
-                temp_dir
+            # Execute tracking with error handling
+            print(f"Executing command: {' '.join(cmd)}")
+            process = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True
             )
             
-            progress(0.9, desc="Finalizing results...")
+            # Check for errors in output
+            if process.returncode != 0:
+                error_message = process.stderr or process.stdout
+                print(f"Process failed with return code {process.returncode}")
+                print(f"Error: {error_message}")
+                return None, f"Error in tracking process: {error_message}"
             
-            # Clean up memory after each tracker run
-            clean_memory()
+            print(f"Process completed with return code {process.returncode}")
             
-            # Create the DataFrame data for this single tracker
-            comparison_data = [[tracking_method, f"{process_time:.2f} seconds", "Success" if result else "Failed"]]
+            # 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))
+            
+            print(f"Found output files: {output_files}")
+            
+            if not output_files:
+                print("No output video files found")
+                return None, "No output video was generated. Check if tracking was successful."
+            
+            output_file = output_files[0]
+            print(f"Selected output file: {output_file}")
+            
+            # Verify file exists and has size
+            if os.path.exists(output_file):
+                file_size = os.path.getsize(output_file)
+                print(f"Output file exists with size: {file_size} bytes")
+                
+                if file_size == 0:
+                    return None, "Output video was generated but has zero size."
+                
+                # Copy to permanent location with unique name
+                permanent_path = os.path.join(OUTPUT_DIR, f"output_{os.path.basename(video_file)}")
+                shutil.copy(output_file, permanent_path)
+                print(f"Copied output to permanent location: {permanent_path}")
+                
+                # Ensure the file is in MP4 format for better compatibility with Gradio
+                if not permanent_path.lower().endswith('.mp4'):
+                    mp4_path = os.path.splitext(permanent_path)[0] + '.mp4'
+                    try:
+                        print(f"Converting to MP4 format: {mp4_path}")
+                        subprocess.run([
+                            'ffmpeg', '-i', permanent_path, 
+                            '-c:v', 'libx264', '-preset', 'fast',
+                            '-c:a', 'aac', mp4_path
+                        ], check=True, capture_output=True)
+                        os.remove(permanent_path)  # Remove the original file
+                        permanent_path = mp4_path
+                    except Exception as e:
+                        print(f"Failed to convert to MP4: {str(e)}")
+                        # Continue with original file if conversion fails
+                
+                return permanent_path, "Processing completed successfully!"
+            else:
+                print(f"Output file not found at {output_file}")
+                return None, "Output file was referenced but doesn't exist on disk."
             
-            # Return the results
-            return result, status, f"{process_time:.2f} seconds", comparison_data
-    
     except Exception as e:
         import traceback
         traceback.print_exc()
-        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"]]
+        return None, f"Error: {str(e)}"
 
-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
+# Define the Gradio interface
+def process_video(video_path, yolo_model, reid_model, tracking_method, class_ids, conf_threshold):
+    # Validate inputs
+    if not video_path:
+        return None, "Please upload a video file"
+    
+    print(f"Processing video: {video_path}")
+    print(f"Parameters: model={yolo_model}, reid={reid_model}, tracker={tracking_method}, classes={class_ids}, conf={conf_threshold}")
     
-    # Append the new tracker data to the existing table
-    return current_data + new_data
+    output_path, status = run_tracking(
+        video_path,
+        yolo_model,
+        reid_model,
+        tracking_method,
+        class_ids,
+        conf_threshold
+    )
+    
+    if output_path:
+        print(f"Returning output path: {output_path}")
+        # Make sure the path is absolute for Gradio
+        abs_path = os.path.abspath(output_path)
+        return abs_path, status
+    else:
+        print(f"No output path available. Status: {status}")
+        return None, status
 
-# Available models
+# Available models and tracking methods
 yolo_models = ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"]
 reid_models = ["osnet_x0_25_msmt17.pt"]
 tracking_methods = ["bytetrack", "botsort", "ocsort", "strongsort"]
@@ -244,9 +199,9 @@ ensure_dependencies()
 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 and run each tracking method sequentially. Results will be displayed as they become available.")
+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.")
     
     # Add class reference information
     with gr.Accordion("YOLO Class Reference", open=False):
@@ -270,9 +225,6 @@ 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"])
@@ -288,6 +240,11 @@ with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
                     value="osnet_x0_25_msmt17.pt", 
                     label="ReID Model"
                 )
+                tracking_method = gr.Dropdown(
+                    choices=tracking_methods, 
+                    value="bytetrack", 
+                    label="Tracking Method"
+                )
                 
                 # Class ID input field
                 class_ids = gr.Textbox(
@@ -303,93 +260,17 @@ with gr.Blocks(title="YOLO Object Tracking Benchmark") as app:
                     step=0.05, 
                     label="Confidence Threshold"
                 )
-            
-            # 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:
-        with gr.TabItem("ByteTrack"):
-            bytetrack_video = gr.Video(label="ByteTrack Result")
-            bytetrack_status = gr.Textbox(label="Status", value="Ready to process")
-            bytetrack_time = gr.Textbox(label="Processing Time", value="")
-            
-        with gr.TabItem("BoTSORT"):
-            botsort_video = gr.Video(label="BoTSORT Result")
-            botsort_status = gr.Textbox(label="Status", value="Ready to process")
-            botsort_time = gr.Textbox(label="Processing Time", value="")
-            
-        with gr.TabItem("OC-SORT"):
-            ocsort_video = gr.Video(label="OC-SORT Result")
-            ocsort_status = gr.Textbox(label="Status", value="Ready to process")
-            ocsort_time = gr.Textbox(label="Processing Time", value="")
-            
-        with gr.TabItem("StrongSORT"):
-            strongsort_video = gr.Video(label="StrongSORT Result")
-            strongsort_status = gr.Textbox(label="Status", value="Ready to process")
-            strongsort_time = gr.Textbox(label="Processing Time", value="")
-    
-    # Comparison Tab
-    with gr.Tabs() as comparison_tab:
-        with gr.TabItem("Performance Comparison"):
-            perf_table = gr.DataFrame(
-                headers=["Tracker", "Processing Time", "Status"],
-                datatype=["str", "str", "str"],
-                label="Performance Comparison"
-            )
-    
-    # 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
-    )
+                
+            process_btn = gr.Button("Process Video", variant="primary")
+        
+        with gr.Column(scale=1):
+            output_video = gr.Video(label="Output Video with Tracking")
+            status_text = gr.Textbox(label="Status", value="Ready to process video")
     
-    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
+    process_btn.click(
+        fn=process_video,
+        inputs=[input_video, yolo_model, reid_model, tracking_method, class_ids, conf_threshold],
+        outputs=[output_video, status_text]
     )
     
     # Add a debug section
@@ -418,23 +299,10 @@ 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__":