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

# Load the pre-trained AI image detector model
detector = pipeline(
    "image-classification",
    model="umm-maybe/AI-image-detector",
    device=0 if torch.cuda.is_available() else -1  # Use GPU if available
)

def detect_image(image):
    if image is None:
        return "Please upload an image."
    
    # Run inference
    results = detector(image)
    
    # Get top prediction
    top_pred = results[0]
    label = top_pred['label']
    prob = top_pred['score'] * 100
    
    # Format output
    if label == 'real':
        return f"Real Image (Confidence: {prob:.2f}%)"
    else:
        return f"AI-Generated Image (Confidence: {prob:.2f}%)"

# Create Gradio interface
with gr.Blocks(title="AI Image Detector – CBSE Board Project 2025-26") as demo:
    gr.Markdown("# AI Image Detector by [Dipika Singh, Rashvi Singh, Rupanjali Rai, Khusboo Yadav] – CBSE Board Project 2025-26")
    gr.Markdown("Upload any image → Instantly know if it's Real or AI-Generated!")
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Drop Image Here - or - Click to Upload")
        output = gr.Textbox(label="Detection Result")
    
    with gr.Row():
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear")
    
    # Event handlers
    submit_btn.click(fn=detect_image, inputs=image_input, outputs=output)
    clear_btn.click(fn=lambda: (None, "Upload an image to start!"), inputs=None, outputs=[image_input, output])

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