# -*- coding: utf-8 -*-
"""image classification general.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1IKVz7857OWmXGom1Z2zq884jI84gkSUx
"""

#!pip install gradio
#!pip install transformers torch

import gradio as gr
from transformers import ViTImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load pre-trained model and feature extractor
model_name = "google/vit-base-patch16-224"
feature_extractor = ViTImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.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()