# 모델로딩

# ImageNet-1k에 훈련된 모델과 특징 추출기 로드
from transformers import ViTImageProcessor, ViTForImageClassification
model_name = "google/vit-base-patch16-224"
model = ViTForImageClassification.from_pretrained(model_name)
image_processor = ViTImageProcessor.from_pretrained(model_name)

# 이미지 예측 분류함수

import torch
def classify_image(inp):
    # 이미지를 특징 벡터로 변환
    inputs = image_processor(images=inp, return_tensors="pt")
    pixel_values = inputs["pixel_values"]

    # 예측 수행
    outputs = model(pixel_values)
    logits = outputs.logits
    predicted_index = torch.argmax(logits, 1)[0].item()

    # 가장 확률이 높은 라벨 반환``
    label = model.config.id2label[predicted_index]
    return label

# Gradio 인터페이스 설정
from PIL import Image
import gradio as gr
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.components.Image(type="pil", label="Upload an Image"),
    outputs="text",
    live=True
)
interface.launch()