import gradio as gr
from transformers import pipeline
from PIL import Image
import numpy as np

# ----------------------------
# Inisialisasi model
# ----------------------------

# Model A: AI detector wajah
model_a = pipeline("image-classification", model="rodrigoberriel/real-or-fake-face")

# Model B: AI detector general (misal: resnet-50)
model_b = pipeline("image-classification", model="microsoft/resnet-50")

# Model C: General image classifier
model_c = pipeline("image-classification", model="google/vit-base-patch16-224")

# ----------------------------
# Fungsi prediksi dengan ensemble
# ----------------------------
def detect_image(image):
    results = []
    
    # Model A
    res_a = model_a(image)[0]
    label_a, score_a = res_a['label'].lower(), res_a['score']
    results.append((label_a, score_a))
    
    # Model B
    res_b = model_b(image)[0]
    label_b, score_b = res_b['label'].lower(), res_b['score']
    # Simplifikasi: jika label mirip 'human' atau 'person' = human, lainnya = AI
    if 'human' in label_b or 'person' in label_b:
        label_b_final = 'human'
    else:
        label_b_final = 'ai'
    results.append((label_b_final, score_b))
    
    # Model C
    res_c = model_c(image)[0]
    label_c, score_c = res_c['label'].lower(), res_c['score']
    if 'human' in label_c or 'person' in label_c:
        label_c_final = 'human'
    else:
        label_c_final = 'ai'
    results.append((label_c_final, score_c))
    
    # ----------------------------
    # Voting + Threshold
    # ----------------------------
    votes = [r[0] for r in results]
    scores = [r[1] for r in results]
    
    # Mayoritas voting
    final_label = max(set(votes), key=votes.count)
    
    # Rata-rata confidence untuk final label
    relevant_scores = [s for (l, s) in results if l == final_label]
    avg_confidence = np.mean(relevant_scores) * 100  # ke persen
    
    # Tentukan hasil akhir
    if avg_confidence < 80:
        return "⚠️ Tidak Pasti (cek manual)", round(avg_confidence, 2)
    elif final_label == 'human':
        return f"Foto Asli ✅", round(avg_confidence, 2)
    else:
        return f"AI Detected 🤖", round(avg_confidence, 2)

# ----------------------------
# Gradio Interface
# ----------------------------
iface = gr.Interface(
    fn=detect_image,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Textbox(label="Hasil Deteksi"), gr.Number(label="Confidence (%)")],
    title="AI vs Foto Asli Detector",
    description="Unggah gambar, sistem akan mendeteksi apakah gambar asli atau dihasilkan AI menggunakan ensemble 3 model + voting."
)

iface.launch()