import gradio as gr
from ultralytics import YOLO
import cv2
from PIL import Image
import spaces  # <--- 1. Wajib import library spaces bawaan Hugging Face

# Load model kustom YOLOv11
model = YOLO("best.pt")

# 2. Wajib tambahkan decorator ini tepat di atas fungsi yang melakukan pemrosesan gambar
@spaces.GPU  
def predict_image(img):
    # Jalankan inferensi YOLO
    results = model(img)
    # Plot hasil ke gambar (BGR)
    res_bgr = results[0].plot()
    # Konversi ke RGB untuk Gradio
    res_rgb = cv2.cvtColor(res_bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(res_rgb)

# Buat Interface Gradio yang simpel dan interaktif
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload Foto"),
    outputs=gr.Image(type="pil", label="Hasil Deteksi"),
    title="😷 Face Mask Detection",
    description="Unggah foto untuk mendeteksi penggunaan masker (mask / no-mask). Silakan upload foto proyek Anda!"
)

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