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

model = YOLO("best.pt")  # make sure model file name matches

def detect_potholes(video):
    if video is None:
        return "Please upload a video."
    
    temp_dir = tempfile.mkdtemp()
    video_path = os.path.join(temp_dir, "input.mp4")

    with open(video_path, "wb") as f:
        f.write(video.read())

    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 "❌ Error: No output generated."
    
    return output_files[0]

app = gr.Interface(
    fn=detect_potholes,
    inputs=gr.Video(label="Upload Dashcam Video"),
    outputs=gr.Video(label="Detected Video Output"),
    title="🚧 Pothole Detection AI",
    description="Upload a dashcam video and detect potholes using YOLO."
)

app.launch()