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

# Load 3 model Hugging Face
model1 = pipeline("image-classification", model="prithivMLmods/deepfake-vs-real-image-detection")
model2 = pipeline("image-classification", model="dima806/ai-image-detector")
model3 = pipeline("image-classification", model="Hemg/A-real-and-fake-image-detection")

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

    # Prediksi masing-masing model
    res1 = model1(img)
    res2 = model2(img)
    res3 = model3(img)

    # Ambil skor Real/Artificial sesuai output
    score1_real = next((r["score"] for r in res1 if "real" in r["label"].lower()), 0)
    score1_fake = next((r["score"] for r in res1 if "fake" in r["label"].lower() or "artificial" in r["label"].lower()), 0)

    score2_real = next((r["score"] for r in res2 if "real" in r["label"].lower()), 0)
    score2_fake = next((r["score"] for r in res2 if "fake" in r["label"].lower()), 0)

    score3_real = next((r["score"] for r in res3 if "real" in r["label"].lower()), 0)
    score3_fake = next((r["score"] for r in res3 if "fake" in r["label"].lower()), 0)

    # Voting: jika 2 model yakin REAL (>0.5), maka final REAL
    votes_real = sum([score1_real > 0.5, score2_real > 0.5, score3_real > 0.5])

    if votes_real >= 2:
        final_label = "✅ Foto Asli"
    else:
        # Weighted ensemble (Hemg lebih berat)
        weighted_real = (0.25 * score1_real) + (0.25 * score2_real) + (0.5 * score3_real)
        weighted_fake = (0.25 * score1_fake) + (0.25 * score2_fake) + (0.5 * score3_fake)

        if weighted_real >= weighted_fake:
            final_label = "✅ Foto Asli"
        else:
            final_label = "❌ Foto AI / Hasil Buatan"

    # Tampilkan hasil detail
    results["Model1"] = res1
    results["Model2"] = res2
    results["Model3"] = res3
    results["Final"] = final_label

    return results

# Interface Gradio
demo = gr.Interface(
    fn=detect_image,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="Deteksi Foto Asli vs AI",
    description="Menggabungkan 3 model Hugging Face dengan hybrid voting + bobot."
)

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