Update app.py

app.py

@@ -2,20 +2,38 @@ import gradio as gr
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 import torch
 from PIL import Image
+import numpy as np
+import json
+import requests
 
 # 加载模型和特征提取器
 model_name = "microsoft/beit-base-patch16-224"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
+# 获取 ImageNet 类别映射
+LABELS_URL = "https://storage.googleapis.com/bit_models/imagenet21k_wordnet_id_map.json"
+imagenet_classes = requests.get(LABELS_URL).json()
+
 # 定义分类函数
 def classify_image(image):
-    image = feature_extractor(images=image, return_tensors="pt")
+    # 转换 PIL Image 为 numpy 数组
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+
+    # 进行特征提取
+    inputs = feature_extractor(images=image, return_tensors="pt")
+
+    # 预测类别
     with torch.no_grad():
-        outputs = model(**image)
+        outputs = model(**inputs)
     logits = outputs.logits
-    predicted_class = logits.argmax(-1).item()
-    return f"Predicted class: {predicted_class}"
+    predicted_class_idx = logits.argmax(-1).item()
+
+    # 获取类别名称
+    class_name = imagenet_classes.get(str(predicted_class_idx), "Unknown")
+    
+    return f"Predicted class: {class_name} (ID: {predicted_class_idx})"
 
 # 创建 Gradio 界面
 demo = gr.Interface(fn=classify_image, inputs="image", outputs="text", title="Image Classification Demo")