import gradio as gr
from transformers import ViTForImageClassification, ViTProcessor

# Load the model and processor
model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
processor = ViTProcessor.from_pretrained("google/vit-base-patch16-224")

def predict_image(img):
    inputs = processor(img, return_tensors="pt")
    outputs = model(**inputs)
    predictions = outputs.logits.argmax(-1)
    return model.config.id2label[predictions.item()]

# Create the interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(shape=(224, 224)),
    outputs="text",
    live=True,
    capture_session=True,
    title="Image recognition",
    description="Upload an image you want to categorize.",
    theme="Monochrome"
)

iface.launch()