app.py

import gradio as gr
from ultralytics import YOLO
from PIL import Image
import tempfile
import os

# Load model once
model = YOLO("best.pt")

def detect_image(image, conf):
    results = model.predict(image, conf=conf)
    return results[0].plot()

def detect_video(video, conf):
    # Save uploaded video to temp file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
        tmp.write(video.read())
        tmp_path = tmp.name

    # Run YOLO inference on video
    results = model.predict(source=tmp_path, conf=conf, save=True)

    # Find output video
    save_dir = results[0].save_dir
    output_path = None
    for f in os.listdir(save_dir):
        if f.endswith(".mp4"):
            output_path = os.path.join(save_dir, f)
            break

    return output_path if output_path else None

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🚧 Pothole Detector (YOLOv8)")

    conf = gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Confidence Threshold")

    with gr.Tab("📷 Image Detection"):
        img_in = gr.Image(type="pil", label="Upload Image")
        img_out = gr.Image(type="numpy", label="Detections")
        img_btn = gr.Button("Detect Potholes")
        img_btn.click(fn=detect_image, inputs=[img_in, conf], outputs=img_out)

    with gr.Tab("🎥 Video Detection"):
        vid_in = gr.File(file_types=[".mp4", ".avi", ".mov"], label="Upload Video")
        vid_out = gr.Video(label="Processed Video")
        vid_btn = gr.Button("Process Video")
        vid_btn.click(fn=detect_video, inputs=[vid_in, conf], outputs=vid_out)

demo.launch()