Update app.py

app.py

@@ -3,15 +3,15 @@ import torch
 from torchvision import transforms
 from PIL import Image
 from timm import create_model
+import json
+
+with open('class_mapping.json', 'r') as json_file:
+    class_mapping = json.load(json_file)
 
-# 定义类别名称
-CLASSES = [
-    "Class 1", "Class 2", "Class 3",  # 替换为你的类别名称
-]
 
 # 加载模型
 def load_model(model_path):
-    model = create_model('resnet18', pretrained=False, num_classes=len(CLASSES))
+    model = create_model('resnet18', pretrained=False, num_classes=len(class_mapping))
     model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
     model.eval()
     return model
@@ -32,9 +32,12 @@ def classify_image(image):
     image = preprocess_image(image)
     with torch.no_grad():
         outputs = model(image)
-        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-        confidences = {CLASSES[i]: float(probabilities[i]) for i in range(len(CLASSES))}
-    return confidences
+        
+        _, predicted_class = torch.max(outputs, 1)
+        predicted_class_idx = predicted_class.item()
+        predicted_class_name = class_mapping[str(predicted_class_idx)]
+
+    return predicted_class_name
 
 # 创建 Gradio 接口
 title = "Bird Species Classifier"
@@ -43,7 +46,7 @@ description = "Upload an image of a bird, and the model will predict its species
 interface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=3),
+    outputs="text",
     title=title,
     description=description,
 )