app.py

import os
import glob
import gradio as gr
import shutil
from ultralytics import YOLO

# Initialize YOLO model
model = YOLO('./best.pt')

def check_ffmpeg():
    """Check if ffmpeg is installed."""
    if shutil.which("ffmpeg") is None:
        raise EnvironmentError("ffmpeg is not installed or not found in the system PATH. Please install ffmpeg to proceed.")

def process_video(video_path):
    # Check if ffmpeg is installed
    check_ffmpeg()

    # Run the YOLO model on the video, specifying the tracker configuration
    results = model.track(video_path, save=True, tracker="bytetrack.yaml", persist=True)

    # Get the video file name without extension and directory
    video_name = os.path.splitext(os.path.basename(video_path))[0]
    
    # Find the latest directory created in 'runs/detect/'
    output_dir = max(glob.glob('./runs/detect/*'), key=os.path.getmtime)
    
    # Find the saved video file in the latest directory with the specific video name
    video_files = glob.glob(os.path.join(output_dir, f'{video_name}*.avi'))
    
    if not video_files:
        raise Exception(f"No .avi video files found in directory {output_dir} for {video_name}")

    # Convert the video to mp4
    mp4_files = []
    for file in video_files:
        mp4_file = f"{file[:-4]}.mp4"
        # Quote the file paths to handle spaces and overwrite without prompt
        os.system(f"ffmpeg -y -i \"{file}\" -vcodec libx264 \"{mp4_file}\"")
        os.remove(file)  # Remove the original .avi file after conversion
        mp4_files.append(mp4_file)

    # Check again for the converted .mp4 video files specifically for the processed video
    matched_mp4_files = [file for file in mp4_files if video_name in file]
    
    if not matched_mp4_files:
        raise Exception(f"No .mp4 video files found in directory {output_dir} after conversion for {video_name}.")

    # Return the first video file found (should correspond to the latest processed video)
    return matched_mp4_files[0]

# Define Gradio inputs
video_input = gr.Video()

video_interface = gr.Interface(
    fn=process_video,
    inputs=video_input,
    outputs="video",  # Display the video output
    description="YOLO Video Tracking"
)

# Deploy the interface with share enabled
gr.TabbedInterface([video_interface], ["Track Video"]).launch(share=True)