done

app.py

@@ -6,12 +6,15 @@ import shutil
 from pathlib import Path
 import sys
 import importlib.util
-from concurrent.futures import ThreadPoolExecutor
 
 # Ensure models directory exists
 MODELS_DIR = Path("models")
 os.makedirs(MODELS_DIR, exist_ok=True)
 
+# Create permanent output directory
+OUTPUT_DIR = Path("outputs")
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
 def ensure_dependencies():
     """Ensure all required dependencies are installed."""
     required_packages = [
@@ -44,125 +47,151 @@ def apply_patches():
     else:
         print("⚠️ tracker_patch.py not found, skipping patches")
 
-def run_single_tracker(video_file, yolo_model, reid_model, tracking_method, conf_threshold, temp_dir):
-    """Run object tracking with a single tracking method."""
+def run_tracking(video_file, yolo_model, reid_model, tracking_method, 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 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}"
             
-        return output_files[0], f"{tracking_method} completed successfully!"
-        
+            # 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
+            print(f"Executing command: {' '.join(cmd)}")
+            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
+                print(f"Process failed with return code {process.returncode}")
+                print(f"Error: {error_message}")
+                return None, f"Error in tracking process: {error_message}"
+            
+            print(f"Process 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: {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."
+            
     except Exception as e:
-        return None, f"Error in {tracking_method}: {str(e)}"
+        import traceback
+        traceback.print_exc()
+        return None, f"Error: {str(e)}"
 
-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 = []
-    
+# Define the Gradio interface
+def process_video(video_path, yolo_model, reid_model, tracking_method, conf_threshold):
     # Validate inputs
     if not video_path:
-        return None, None, None, None, "Please upload a video file"
+        return None, "Please upload a video file"
     
-    # Create temporary workspace
-    temp_dir = tempfile.mkdtemp()
+    print(f"Processing video: {video_path}")
+    print(f"Parameters: model={yolo_model}, reid={reid_model}, tracker={tracking_method}, conf={conf_threshold}")
     
-    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(
+    output_path, status = run_tracking(
         video_path,
         yolo_model,
         reid_model,
+        tracking_method,
         conf_threshold
     )
     
-    return bytetrack_output, botsort_output, ocsort_output, strongsort_output, status
+    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"]
 
 # Ensure dependencies and apply patches at startup
 ensure_dependencies()
 apply_patches()
 
 # Create the Gradio interface
-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.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.Row():
-        with gr.Column(scale=1):
+        with gr.Column():
             input_video = gr.Video(label="Input Video", sources=["upload"])
             
             with gr.Group():
@@ -176,6 +205,11 @@ with gr.Blocks(title="YOLO Multi-Tracker Comparison") 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, 
@@ -184,33 +218,48 @@ with gr.Blocks(title="YOLO Multi-Tracker Comparison") as app:
                     label="Confidence Threshold"
                 )
                 
-            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)
-            
+            process_btn = gr.Button("Process Video", variant="primary")
+        
         with gr.Column():
-            gr.Markdown("### StrongSORT")
-            strongsort_video = gr.Video(label="StrongSORT Result", autoplay=False)
+            output_video = gr.Video(label="Output Video with Tracking")
+            status_text = gr.Textbox(label="Status", value="Ready to process video")
     
     process_btn.click(
         fn=process_video,
-        inputs=[input_video, yolo_model, reid_model, conf_threshold],
-        outputs=[bytetrack_video, botsort_video, ocsort_video, strongsort_video, status_text]
+        inputs=[input_video, yolo_model, reid_model, tracking_method, conf_threshold],
+        outputs=[output_video, status_text]
     )
+    
+    # Add a debug section
+    with gr.Accordion("Debug Information", open=False):
+        debug_text = gr.Textbox(label="Debug Log", lines=10, interactive=False)
+        
+        def check_environment():
+            info = []
+            # Check Python version
+            info.append(f"Python version: {sys.version}")
+            # Check Gradio version
+            info.append(f"Gradio version: {gr.__version__}")
+            # Check for ffmpeg
+            try:
+                ffmpeg_version = subprocess.run(["ffmpeg", "-version"], capture_output=True, text=True)
+                info.append("ffmpeg: Installed")
+            except:
+                info.append("ffmpeg: Not found")
+            # Check tracking directory
+            if os.path.exists("tracking"):
+                info.append("tracking directory: Found")
+            else:
+                info.append("tracking directory: Not found")
+            # Check models
+            info.append("Models:")
+            for model in os.listdir(MODELS_DIR) if os.path.exists(MODELS_DIR) else []:
+                info.append(f"  - {model}")
+            
+            return "\n".join(info)
+        
+        check_btn = gr.Button("Check Environment")
+        check_btn.click(fn=check_environment, outputs=debug_text)
 
 # Launch the app
 if __name__ == "__main__":