import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image

# Load model YOLOv5 yang telah kamu latih
model = torch.hub.load('ultralytics/yolov8', 'custom', path='my_model.pt', force_reload=True)

# Fungsi deteksi wajah
def detect_face(image):
    image_cv = np.array(image)  # Konversi ke array numpy
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)

    results = model(image_cv)  # Deteksi wajah
    
    # Gambar kotak di sekitar wajah
    for *box, conf, cls in results.xyxy[0]:  
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)

    image_result = Image.fromarray(cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB))  # Konversi kembali ke PIL
    return image_result

# Buat UI dengan Gradio
iface = gr.Interface(
    fn=detect_face, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Image(type="pil"),
    title="Deteksi Wajah dengan YOLOv5",
    description="Upload gambar dan model akan mendeteksi wajah."
)

# Jalankan aplikasi
iface.launch()