import gradio as gr
from ultralytics import YOLO
import tempfile
import cv2

# Load YOLOv8 model once
model = "best.pt"

# Inference on image
def predict_image(image):
    results = model.predict(image)
    return results[0].plot()

# Inference on video
def predict_video(video_path):
    # OpenCV video capture
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2  # Resize for performance
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2

    # Output video
    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height))

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame = cv2.resize(frame, (width, height))  # Resize to speed up inference
        results = model.predict(frame, imgsz=480, conf=0.5, verbose=False)  # Lower imgsz = faster
        annotated = results[0].plot()
        out.write(annotated)

    cap.release()
    out.release()
    return temp_output.name

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Optimized YOLOv8 Detection\nFast & Accurate on Images and Videos")

    with gr.Tab("Image"):
        img_input = gr.Image(type="pil")
        img_output = gr.Image(label="Detected")
        img_btn = gr.Button("Run Detection")
        img_btn.click(predict_image, inputs=img_input, outputs=img_output)

    with gr.Tab("Video"):
        vid_input = gr.Video()
        vid_output = gr.Video()
        vid_btn = gr.Button("Run Detection on Video")
        vid_btn.click(predict_video, inputs=vid_input, outputs=vid_output)

demo.launch()