Update app.py

app.py

@@ -9,12 +9,14 @@ model_name = "microsoft/beit-base-patch16-224"
 model = AutoModelForImageClassification.from_pretrained(model_name)
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
-# 加载 ImageNet 1000 类别名称
-imagenet_labels = {
-    idx: entry.strip() for idx, entry in enumerate(
-        open("https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt").readlines()
-    )
-}
+# 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"
+]
 
 # 定义分类函数
 def classify_image(image):
@@ -32,8 +34,11 @@ def classify_image(image):
     predicted_class_idx = logits.argmax(-1).item()
 
     # 获取类别名称
-    class_name = imagenet_labels.get(predicted_class_idx, f"Unknown Class (ID: {predicted_class_idx})")
-
+    if predicted_class_idx < len(imagenet_labels):
+        class_name = imagenet_labels[predicted_class_idx]
+    else:
+        class_name = f"Unknown Class (ID: {predicted_class_idx})"
+    
     return f"Predicted class: {class_name} (ID: {predicted_class_idx})"
 
 # 创建 Gradio 界面