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

# Load the YOLOv8 model ONCE (faster processing)
MODEL_PATH = "yolov8l.pt"
model = YOLO(MODEL_PATH)

# Get sample images from Space directory
valid_extensions = (".jpg", ".jpeg", ".png")
preloaded_images = [img for img in os.listdir() if img.lower().endswith(valid_extensions)]

# Object detection function
def predict(image):
    if isinstance(image, str):  # Sample image selected
        image = Image.open(image)
    else:  # Uploaded image
        image = Image.fromarray(image)

    results = model(image)  # YOLO detection

    # Convert image for OpenCV processing
    image_cv = np.array(image)
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
    overlay = image_cv.copy()

    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            label = model.names[int(box.cls)]
            confidence = float(box.conf)

            # Translucent blue overlay
            cv2.rectangle(overlay, (x1, y1), (x2, y2), (255, 0, 0), -1)
            cv2.rectangle(image_cv, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(image_cv, f"{label}: {confidence:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    # Blend overlay with transparency
    alpha = 0.4
    image_masked = cv2.addWeighted(overlay, alpha, image_cv, 1 - alpha, 0)

    return Image.fromarray(cv2.cvtColor(image_masked, cv2.COLOR_BGR2RGB))

# Processing function with correct priority
def process_image(image, sample_name):
    if image is not None:  # Prioritize uploaded image
        return predict(image)
    elif sample_name:  # Otherwise, use the selected sample
        return predict(sample_name)
    return None  # No input provided

# Gradio Interface
with gr.Blocks() as interface:
    gr.Markdown("# 🪨 Moon Rock Detection")
    gr.Markdown("Upload a moon surface image or select a sample.")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="numpy", label="Upload Image")
            sample_dropdown = gr.Dropdown(
                choices=preloaded_images, label="Or Select a Sample Image", interactive=True
            )

        with gr.Column(scale=2):
            output_image = gr.Image(type="pil", label="Detection Result")

    # Automatically detect when an image is uploaded or selected
    image_input.change(process_image, inputs=[image_input, sample_dropdown], outputs=output_image)
    sample_dropdown.change(process_image, inputs=[image_input, sample_dropdown], outputs=output_image)

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