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

model = YOLO("best.pt")

def detect_potholes(video):
    if video is None:
        return "Please upload a video.", None
    
    video_path = video  # gr.Video gives file path

    # Remove old results
    if os.path.exists("runs"):
        shutil.rmtree("runs")

    model.predict(
        source=video_path,
        save=True,
        conf=0.35,
        iou=0.5
    )

    output_files = glob.glob("runs/detect/*/*.mp4")

    if len(output_files) == 0:
        return "❌ No output generated.", None

    output_video = output_files[0]

    return output_video, output_video   # ⬅️ return video for preview AND download


app = gr.Interface(
    fn=detect_potholes,
    inputs=gr.Video(label="Upload Dashcam Video"),
    outputs=[
        gr.Video(label="Processed Video ✅"),
        gr.File(label="Download Video ⬇️")
    ],
    title="🚧 Pothole Detection AI",
    description="Upload a dashcam video and detect potholes using YOLO."
)
app.launch()