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
import os

# Verify paths and Hugging Face repository details
REPO_ID = "StephanST/WALDO30"  # Replace with the correct repo ID if different
MODEL_FILENAME = "WALDO30_yolov8m_640x640.pt"  # Replace if the filename is different

# Download the model from Hugging Face
try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
    print(f"Model downloaded successfully to: {model_path}")
except Exception as e:
    raise RuntimeError(f"Failed to download model from Hugging Face. Verify `repo_id` and `filename`. Error: {e}")

# Load the YOLOv8 model
try:
    model = YOLO(model_path)  # Load the YOLOv8 model
    print("Model loaded successfully!")
except Exception as e:
    raise RuntimeError(f"Failed to load the YOLO model. Verify the model file at `{model_path}`. Error: {e}")

# Detection function for images
def detect_on_image(image):
    try:
        results = model(image)  # Perform detection
        annotated_frame = results[0].plot()  # Get annotated image
        return Image.fromarray(annotated_frame)
    except Exception as e:
        raise RuntimeError(f"Error during image processing: {e}")

# Detection function for videos
def detect_on_video(video):
    try:
        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
    except Exception as e:
        raise RuntimeError(f"Error during video processing: {e}")

# Gradio Interface using Blocks
with gr.Blocks() as app:
    gr.Markdown("# Sat ESPR View")
    gr.Markdown("Upload an image or video to see object detection results.")

    # Image processing block
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            image_button = gr.Button("Detect on Image")
        with gr.Column():
            image_output = gr.Image(type="pil", label="Detected Image")

    # Video processing block
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Upload Video")
            video_button = gr.Button("Detect on Video")
        with gr.Column():
            video_output = gr.Video(label="Detected Video")

    # Set up events
    image_button.click(detect_on_image, inputs=image_input, outputs=image_output)
    video_button.click(detect_on_video, inputs=video_input, outputs=video_output)

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