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

# Model utama untuk deteksi real vs AI
detector = pipeline("image-classification", model="umm-maybe/AI-image-detector")

# Model tambahan general classifier (backup)
general = pipeline("image-classification", model="google/vit-base-patch16-224")

def detect_image(img):
    try:
        # Prediksi dengan AI detector
        result1 = detector(img)
        label1 = result1[0]['label']
        conf1 = round(result1[0]['score'] * 100, 2)

        # Prediksi dengan model general (untuk cek ganda)
        result2 = general(img)
        label2 = result2[0]['label']
        conf2 = round(result2[0]['score'] * 100, 2)

        # Logika sederhana untuk memutuskan hasil
        if "fake" in label1.lower() or "artificial" in label1.lower():
            final = f"⚠️ Kemungkinan Besar AI Generated ({conf1}%)"
        elif "real" in label1.lower():
            final = f"✅ Kemungkinan Besar Foto Asli ({conf1}%)"
        else:
            final = f"⚠️ Tidak Pasti (cek manual)"

        output = f"""
### Hasil Deteksi:
{final}

**Model AI-detector:** {label1} ({conf1}%)
**Model General (ViT):** {label2} ({conf2}%)
"""
        return output
    except Exception as e:
        return f"Terjadi error: {str(e)}"

# UI Gradio
iface = gr.Interface(
    fn=detect_image,
    inputs=gr.Image(type="pil"),
    outputs="markdown",
    title="AI vs Real Image Detector",
    description="Upload foto untuk mendeteksi apakah gambar kemungkinan besar asli atau hasil AI."
)

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