import gradio as gr
from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import torch

# Load pre-trained model and feature extractor
model_name = "google/vit-base-patch16-224"
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# Define the prediction function
def classify_image(img):
    inputs = feature_extractor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_class_idx]
    return predicted_label

# Build the Gradio interface
interface = gr.Interface(fn=classify_image,
                         inputs=gr.Image(type="pil"),
                         outputs="text",
                         title="Image Classification with ViT",
                         description="Upload an image and classify it using Vision Transformer (ViT)")

# Launch the app
interface.launch()