import gradio as gr
import cv2
from PIL import Image
import numpy as np
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# Download the model from Hugging Face
model_path = hf_hub_download(repo_id="StephanST/WALDO30", filename="WALDO30_yolov8m_640x640.pt")
model = YOLO(model_path)  # Load YOLOv8 model

# Detection function for images
def detect_on_image(image):
    results = model(image)  # Perform detection
    annotated_frame = results[0].plot()  # Get annotated image
    return Image.fromarray(annotated_frame)

# Detection function for videos
def detect_on_video(video):
    temp_video_path = "processed_video.mp4"
    cap = cv2.VideoCapture(video)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(temp_video_path, fourcc, cap.get(cv2.CAP_PROP_FPS),
                          (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model(frame)  # Perform detection
        annotated_frame = results[0].plot()  # Get annotated frame
        out.write(annotated_frame)

    cap.release()
    out.release()
    return temp_video_path

# Gradio Interface
app = gr.Interface(
    fn=[detect_on_image, detect_on_video],
    inputs=[gr.inputs.Image(type="pil", label="Upload Image"), gr.inputs.Video(type="file", label="Upload Video")],
    outputs=[gr.outputs.Image(type="pil", label="Detected Image"), gr.outputs.Video(label="Detected Video")],
    title="WALDO30 YOLOv8 Object Detection",
    description="Upload an image or video to see object detection results using the WALDO30 YOLOv8 model."
)

if __name__ == "__main__":
    app.launch()