from ultralytics import YOLO
import gradio as gr
import cv2
import numpy as np

# =========================
# Load trained YOLO model
# =========================
# Make sure best.pt is in the same folder
model = YOLO("best.pt")

# =========================
# Prediction Function
# =========================
def detect_mask(image):
    """
    Takes an image from Gradio,
    runs YOLO prediction,
    returns annotated image
    """
    # Convert RGB to BGR (OpenCV format)
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # YOLO prediction
    results = model(img)

    # Draw bounding boxes
    annotated_img = results[0].plot()

    # Convert back to RGB for Gradio
    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)

    return annotated_img

# =========================
# Gradio Interface
# =========================
interface = gr.Interface(
    fn=detect_mask,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detection Result"),
    title="😷 Face Mask Detection System",
    description="YOLO-based Face Mask Detection (Mask / No Mask)"
)

# =========================
# Launch App
# =========================
interface.launch()