Create app.py

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+import torch
+from torchvision import transforms
+from PIL import Image
+from timm import create_model
+
+# 定义类别名称
+CLASSES = [
+    "Class 1", "Class 2", "Class 3",  # 替换为你的类别名称
+]
+
+# 加载模型
+def load_model(model_path):
+    model = create_model('resnet18', pretrained=False, num_classes=len(CLASSES))
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+    model.eval()
+    return model
+
+model = load_model("model.pth")
+
+# 定义图像预处理
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    return transform(image).unsqueeze(0)
+
+# 定义推理函数
+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
+
+# 创建 Gradio 接口
+title = "Bird Species Classifier"
+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),
+    title=title,
+    description=description,
+)
+
+if __name__ == "__main__":
+    interface.launch()