import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image
import requests

# Load pretrained general image classification model
model_name = "microsoft/beit-large-patch16-224"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Load ImageNet labels
labels = model.config.id2label

def predict(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_id = logits.argmax(-1).item()
    label = labels[predicted_class_id]
    confidence = torch.softmax(logits, dim=1)[0][predicted_class_id].item()
    return f"### Top Prediction: `{label}`\n**Confidence:** `{confidence:.2%}`"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Markdown(),
    title="Image Classifier (Realism Check)",
    description="This uses Microsoft's BEiT model to classify the uploaded image. Useful for assessing whether the image has real-world consistency.",
)

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