Update app.py

app.py

@@ -9,14 +9,12 @@ model_name = "microsoft/beit-base-patch16-224"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
-# ImageNet 1000 类别名称（从 Hugging Face 官方下载）
-imagenet_classes = [
-    "tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark",  # 0-4
-    "electric ray", "stingray", "cock", "hen", "ostrich",  # 5-9
-    "brambling", "goldfinch", "house finch", "junco", "indigo bunting",  # 10-14
-    # 省略中间 900 多个类别...
-    "sports car", "convertible", "minivan", "pickup", "SUV"  # 817-821（汽车类）
-]
+# 加载 ImageNet 1000 类别名称
+imagenet_labels = {
+    idx: entry.strip() for idx, entry in enumerate(
+        open("https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt").readlines()
+    )
+}
 
 # 定义分类函数
 def classify_image(image):
@@ -34,8 +32,8 @@ def classify_image(image):
     predicted_class_idx = logits.argmax(-1).item()
 
     # 获取类别名称
-    class_name = imagenet_classes[predicted_class_idx] if predicted_class_idx < len(imagenet_classes) else f"Unknown Class (ID: {predicted_class_idx})"
-    
+    class_name = imagenet_labels.get(predicted_class_idx, f"Unknown Class (ID: {predicted_class_idx})")
+
     return f"Predicted class: {class_name} (ID: {predicted_class_idx})"
 
 # 创建 Gradio 界面