app.py

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

# The model we will use
MODEL_ID = "umm-maybe/AI-image-detector"

# Load the model pipeline (image-classification)
pipe = pipeline("image-classification", model=MODEL_ID)

def predict_image(image: Image.Image):
    """
    image: PIL image from Gradio
    We focus on the fine-tuned ai_detector output if available.
    """
    try:
        # run the pipeline
        res = pipe(image)

        # typical pipeline returns list of dicts: [{"label": "...", "score": ...}, ...]
        if isinstance(res, list) and isinstance(res[0], dict) and "label" in res[0]:
            top = res[0]
            label = top.get("label", "")
            score = float(top.get("score", 0.0))

            if label.lower().startswith("artificial") or "ai" in label.lower():
                verdict = "AI-generated"
            elif label.lower().startswith("human") or "real" in label.lower():
                verdict = "Human-made"
            else:
                verdict = label

            return f"{verdict} — {score:.2f}"

        # Handle nested response structures
        if isinstance(res, dict):
            imr = res.get("individual_model_results") or res.get("models") or {}
            ai_det = imr.get("ai_detector") or None
            if ai_det:
                ai_prob = ai_det.get("ai_probability") or ai_det.get("score") or ai_det.get("confidence")
                ai_prob = float(ai_prob)
                if ai_prob >= 0.6:
                    return f"AI-generated — {ai_prob:.2f}"
                elif ai_prob <= 0.4:
                    return f"Human-made — {(1 - ai_prob):.2f}"
                else:
                    return f"Uncertain — {ai_prob:.2f}"

        return str(res)

    except Exception as e:
        traceback.print_exc()
        return f"Error: {e}"

# Build Gradio app
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="AI Image Detector",
    description="Upload an image to detect whether it is AI-generated or human-made."
)

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