Update app.py

app.py

@@ -9,14 +9,8 @@ model_name = "microsoft/beit-base-patch16-224"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
-# ImageNet 1000 类别名称（手动加载）
-imagenet_labels = [
-    "tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", 
-    "electric ray", "stingray", "cock", "hen", "ostrich",
-    "brambling", "goldfinch", "house finch", "junco", "indigo bunting",
-    # ...... (省略 900 多个类别)
-    "sports car", "convertible", "minivan", "pickup", "SUV"
-]
+# 获取模型内置的类别标签
+labels = model.config.id2label
 
 # 定义分类函数
 def classify_image(image):
@@ -34,13 +28,11 @@ def classify_image(image):
     predicted_class_idx = logits.argmax(-1).item()
 
     # 获取类别名称
-    if predicted_class_idx < len(imagenet_labels):
-        class_name = imagenet_labels[predicted_class_idx]
-    else:
-        class_name = f"Unknown Class (ID: {predicted_class_idx})"
-    
+    class_name = labels.get(predicted_class_idx, f"Unknown Class (ID: {predicted_class_idx})")
+
     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")
 demo.launch()
+