트랜스포머"

app.py

@@ -1,8 +1,35 @@
+# 모델로딩
 
+# 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)
 
-def image_classifier(img):
-    return {'cat': 0.3, 'dog': 0.7}
+# 이미지 예측 분류함수
 
+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
-demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")  # greet이라는 함수
-demo.launch(debug=True) # share=True : 다른 사람과 공유
+interface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.components.Image(type="pil", label="Upload an Image"),
+    outputs="text",
+    live=True
+)
+interface.launch()