Update app.py

app.py

@@ -1,7 +1,51 @@
+
 import gradio as gr
+import torch
+from torchvision import transforms
+from PIL import Image
+import torch.nn.functional as F
+
+# 加载模型
+class SimpleModel(torch.nn.Module):
+    def __init__(self, num_classes=3):
+        super(SimpleModel, self).__init__()
+        self.fc = torch.nn.Linear(3 * 224 * 224, num_classes)
+
+    def forward(self, x):
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+# 初始化模型并加载权重
+model = SimpleModel(num_classes=3)
+model_path = "model.pth"
+model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+model.eval()
+
+
+def preprocess(image: 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.Image):
+    input_tensor = preprocess(image)
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = F.softmax(output[0], dim=0)
+    top_class = probabilities.argmax().item()
+    return f"Predicted: {class_names[top_class]} (Confidence: {probabilities[top_class]:.2f})"
+
 
-def greet(name):
-    return "Hello " + name + "!!"
+app = gr.Interface(
+    fn=classify_image,  # 推理函数
+    inputs=gr.Image(type="pil"),  # 接收输入图片，返回 PIL 格式
+    outputs="text",  # 输出分类结果
+    title="Image Classification with PyTorch",
+    description="Upload an image to classify it using a pre-trained PyTorch model."
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+app.launch()