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

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

# ---------------- IMAGE DETECTION ----------------
def detect(image):
    results = model.predict(image, conf=0.4, iou=0.5)
    return results[0].plot()


# ---------------- VIDEO DETECTION ----------------
def detect_video(video):

    input_path = video
    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name

    cap = cv2.VideoCapture(input_path)

    fps = cap.get(cv2.CAP_PROP_FPS)
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

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

        results = model.predict(frame, conf=0.4, iou=0.5)
        r = results[0]

        frame = r.plot()
        out.write(frame)

    cap.release()
    out.release()

    return output_path


# ---------------- UI ----------------
image_ui = gr.Interface(
    fn=detect,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(label="Detection Result"),
    title="Spectacles and Teapots Object Detection - 4124139E",
    description="Upload an image or choose any example at the bottom for object detection.",
    examples=[
        ["image_test1.jpg"],
        ["image_test2.jpg"],
   ], 
    cache_examples=False,
)

video_ui = gr.Interface(
    fn=detect_video,
    inputs=gr.Video(label="Upload Video (.mp4)"),
    outputs=gr.Video(label="Processed Video"),
    title="Spectacles and Teapots Object Detection - 4124139E",
    description="Upload a short video or choose any example at the bottom for object detection.",
    examples=[
        ["harry_potter_short.mp4"]
    ]
)

demo = gr.TabbedInterface(
    [image_ui, video_ui],
    ["Images", "Video"]
)

demo.launch(server_name="0.0.0.0", server_port=7860)