Update app.py

app.py

@@ -3,17 +3,20 @@ 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()
+# 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（汽车类）
+]
 
 # 定义分类函数
 def classify_image(image):
@@ -31,7 +34,7 @@ def classify_image(image):
     predicted_class_idx = logits.argmax(-1).item()
 
     # 获取类别名称
-    class_name = imagenet_classes.get(str(predicted_class_idx), "Unknown")
+    class_name = imagenet_classes[predicted_class_idx] if predicted_class_idx < len(imagenet_classes) else f"Unknown Class (ID: {predicted_class_idx})"
     
     return f"Predicted class: {class_name} (ID: {predicted_class_idx})"