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

# โหลดโมเดลและตัวแปลงภาพ
model_name = "google/vit-base-patch16-224"
model = ViTForImageClassification.from_pretrained(model_name)
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)

# ฟังก์ชันประมวลผลภาพ
def classify_image(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class = logits.argmax(-1).item()
    label = model.config.id2label[predicted_class]
    return label

# สร้าง Gradio UI
gr.Interface(fn=classify_image,
             inputs=gr.Image(type="pil"),
             outputs="label",
             title="ViT Image Classifier (google/vit-base-patch16-224)").launch()