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

# --- CONFIGURATION ---
MODEL_PATH = "best.pt" # Pastikan file ini ada di Files

# Load Model
try:
    model = YOLO(MODEL_PATH)
    print("✅ Model berhasil dimuat!")
except Exception as e:
    model = None
    print(f"❌ Error Load Model: {e}")

def process_image(image):
    """
    Fungsi ini menerima gambar RGB dari Gradio,
    mendeteksi objek, dan mengembalikan gambar RGB dengan kotak & label.
    """
    if image is None: return None
    if model is None: return image

    # 1. Konversi dari RGB (Gradio) ke BGR (OpenCV)
    # Ini penting agar warna tidak aneh saat diproses YOLO
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # 2. Deteksi
    # conf=0.25 artinya hanya tampilkan jika yakin > 25%
    results = model.predict(image_bgr, conf=0.25)

    # 3. Gambar Kotak & Label Otomatis
    # PERBAIKAN DI SINI:
    # Gunakan 'labels=True' dan 'conf=True' (bukan show_labels)
    annotated_bgr = results[0].plot(labels=True, conf=True)

    # 4. Konversi Balik dari BGR ke RGB (Agar kulit tidak biru di Web)
    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)

    return annotated_rgb

# --- TAMPILAN WEB ---
with gr.Blocks() as demo:
    gr.Markdown("# 👁️ Tes Model YOLO (Final Fix)")
    gr.Markdown("Pastikan wajah tidak biru dan label muncul.")
    
    with gr.Row():
        # Input kamera (Mode Foto agar lebih ringan & akurat untuk tes)
        inp = gr.Image(sources=["webcam"], label="Ambil Foto", streaming=False)
        out = gr.Image(label="Hasil Deteksi")
    
    btn = gr.Button("🔍 Deteksi Sekarang", variant="primary")
    
    # Jalankan fungsi saat tombol ditekan
    btn.click(fn=process_image, inputs=inp, outputs=out)

# Matikan SSR agar lebih stabil
if __name__ == "__main__":
    demo.launch(ssr_mode=False)