app.py

import os
import gradio as gr
import subprocess
import tempfile
import shutil
from pathlib import Path
import sys
import importlib.util

# 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 = [
        "ultralytics",
        "boxmot"
    ]
    
    for package in required_packages:
        try:
            importlib.import_module(package)
            print(f"✅ {package} is installed")
        except ImportError:
            print(f"⚠️ {package} is not installed, attempting to install...")
            subprocess.run([sys.executable, "-m", "pip", "install", package], check=True)

# Apply tracker patches if tracker_patch.py exists
def apply_patches():
    patch_path = Path("tracker_patch.py")
    if patch_path.exists():
        spec = importlib.util.spec_from_file_location("tracker_patch", patch_path)
        if spec:
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            if hasattr(module, "patch_trackers"):
                module.patch_trackers()
                print("✅ Applied tracker patches")
            else:
                print("⚠️ tracker_patch.py exists but has no patch_trackers function")
    else:
        print("⚠️ tracker_patch.py not found, skipping patches")

def run_tracking(video_file, yolo_model, reid_model, tracking_method, class_ids, conf_threshold):
    """Run object tracking on the uploaded video."""
    try:
        # 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, "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"
            ]
            
            # 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, "Invalid class IDs. Please enter comma-separated numbers (e.g., '0,1,2')."
            
            # 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:
        import traceback
        traceback.print_exc()
        return None, f"Error: {str(e)}"

# 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}")
    
    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 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="🚀 Object Tracking") as app:
    gr.Markdown("# 🚀  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):
        gr.Markdown("""
        # YOLO Class IDs Reference
        
        Enter the class IDs as comma-separated numbers in the "Target Classes" field.
        Leave empty to track all classes.
        
        ## Common Class IDs:
        - 0: person
        - 1: bicycle
        - 2: car
        - 3: motorcycle
        - 5: bus
        - 7: truck
        - 16: dog
        - 17: horse
        - 67: cell phone
        
        [See full COCO class list here](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datasets/coco128.yaml)
        """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_video = gr.Video(label="Input Video", sources=["upload"])
            
            with gr.Group():
                yolo_model = gr.Dropdown(
                    choices=yolo_models, 
                    value="yolov8n.pt", 
                    label="YOLO Model"
                )
                reid_model = gr.Dropdown(
                    choices=reid_models, 
                    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(
                    value="",
                    label="Target Classes (comma-separated IDs, e.g. '0,2,5', leave empty for all classes)",
                    placeholder="e.g. 0,2,5"
                )
                
                conf_threshold = gr.Slider(
                    minimum=0.1, 
                    maximum=0.9, 
                    value=0.3, 
                    step=0.05, 
                    label="Confidence Threshold"
                )
                
            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")
    
    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
    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__":
    app.launch(debug=True, share=True)