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

# === Load 3 model publik khusus AI vs Real ===
model1 = pipeline("image-classification", model="prithivMLmods/AI-vs-Deepfake-vs-Real")
model2 = pipeline("image-classification", model="dima806/ai_vs_real_image_detection")
model3 = pipeline("image-classification", model="Hemg/AI-VS-REAL-IMAGE-DETECTION")

def detect_image(img: Image.Image):
    results = {}

    # Model 1
    res1 = model1(img)
    results["Model 1 (prithivMLmods)"] = res1

    # Model 2
    res2 = model2(img)
    results["Model 2 (dima806)"] = res2

    # Model 3
    res3 = model3(img)
    results["Model 3 (Hemg)"] = res3

    # --- Ensemble dengan threshold agar lebih stabil ---
    ai_score, real_score = 0, 0
    for res in [res1, res2, res3]:
        for item in res:
            label = item["label"].lower()
            score = item["score"]
            if "real" in label:
                real_score += score
            elif "fake" in label or "deepfake" in label or "artificial" in label:
                ai_score += score

    # Normalisasi
    total = ai_score + real_score + 1e-9
    ai_percent = (ai_score / total) * 100
    real_percent = (real_score / total) * 100

    # Keputusan akhir dengan threshold
    if real_percent > ai_percent and real_percent > 55:  
        verdict = f"✅ Foto Asli ({real_percent:.2f}%)"
    elif ai_percent > real_percent and ai_percent > 55:
        verdict = f"⚠️ AI Generated ({ai_percent:.2f}%)"
    else:
        verdict = f"⚠️ Tidak Pasti (Real {real_percent:.2f}%, AI {ai_percent:.2f}%)"

    # Format output
    output = f"## 📊 Ringkasan Deteksi\n"
    for name, res in results.items():
        output += f"\n🔹 **{name}**: {res}\n"

    output += f"\n=== 🧠 ENSEMBLE HASIL AKHIR ===\n{verdict}\n"

    return output

iface = gr.Interface(
    fn=detect_image,
    inputs=gr.Image(type="pil"),
    outputs="markdown",
    title="Deteksi AI vs Foto Asli (Ensemble 3 Model Publik)",
    description="Menggunakan 3 model publik Hugging Face (prithivMLmods, dima806, Hemg) dengan threshold untuk mendeteksi apakah gambar AI-generated atau asli."
)

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