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

# Resolve base directory (folder where this script is located)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Load model
model = YOLO(os.path.join(BASE_DIR, "best.pt"))

# --- Helper function for safe image loading ---
def load_image_case_insensitive(folder, filename):
    """Find and open an image file ignoring case."""
    target = filename.lower()
    for f in os.listdir(folder):
        if f.lower() == target:
            return Image.open(os.path.join(folder, f))
    raise FileNotFoundError(f"{filename} not found in {folder}")

# --- YOLO Prediction Functions ---
def predict_image(image):
    """Predict on an image (PIL)."""
    results = model.predict(image)
    return Image.fromarray(results[0].plot()[:, :, ::-1])  # Convert BGR to RGB

def predict_video(video):
    """Predict on a video file and return annotated video."""
    cap = cv2.VideoCapture(video)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out_path = os.path.join(BASE_DIR, "output.mp4")
    out = None

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        results = model.predict(frame)
        annotated_frame = results[0].plot()
        
        if out is None:
            h, w, _ = annotated_frame.shape
            out = cv2.VideoWriter(out_path, fourcc, 20.0, (w, h))
        
        out.write(annotated_frame)

    cap.release()
    if out:
        out.release()

    return out_path  # Return path to processed video

# --- Load example images safely ---
examples_folder = os.path.join(BASE_DIR, "examples")
image_examples = [
    [load_image_case_insensitive(examples_folder, "test1.jpg")],
    [load_image_case_insensitive(examples_folder, "test2.jpg")],
    [load_image_case_insensitive(examples_folder, "test3.jpg")]
]

# --- Build interface ---
with gr.Blocks(title="My YOLO Demo") as demo:
    gr.Markdown("## YOLO Object Detection Demo")
    gr.Markdown("Upload an image, a video, or use your webcam.")
    
    with gr.Tab("Image / Webcam"):
        gr.Interface(
            fn=predict_image,
            inputs=gr.Image(type="pil", sources=["upload", "webcam"], label="Upload or Capture Image"),
            outputs=gr.Image(type="pil", label="Output"),
            examples=image_examples
        )
    
    with gr.Tab("Video"):
        gr.Interface(
            fn=predict_video,
            inputs=gr.Video(sources=["upload"], label="Upload Video"),
            outputs=gr.Video(label="Processed Video")
        )

demo.launch()