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

# Model ID from Hugging Face
MODEL_ID = "Ateeqq/ai-vs-human-image-detector"

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

def detect_ai_image(image: Image.Image):
    try:
        results = pipe(image)
        # Sort and take top results
        output = {r['label']: float(r['score']) for r in results}
        return output
    except Exception as e:
        return {"error": str(e)}

# Create Gradio interface
demo = gr.Interface(
    fn=detect_ai_image,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="AI vs Human Image Detector",
    description="Detect whether an image is AI-generated or human-made using the Ateeqq model.",
    theme="soft"
)

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