N_M / app.py
lohanajk-56's picture
Create app.py
2caad62 verified
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()