fix again

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import torch
 import torchvision.transforms as transforms
-from PIL import Image
+from PIL import Image, ImageOps
 import torch.nn as nn
+import torch.nn.functional as F
 
 # 如果你的模型结构与标准的torchvision模型不同，请确保在此处定义或导入你的模型结构
 # 例如，如果你有一个model.py文件：
@@ -79,19 +80,31 @@ transform = transforms.Compose([
 
 # 定义预测函数
 def classify_image(image):
-    # 如果输入是灰度图，将其转换为RGB
-    if image.mode != "RGB":
-        image = image.convert("RGB")
+    # 将 NumPy 数组转换为 PIL 图像
+    image = Image.fromarray(image).convert("L")
     
-    # 预处理图像
+    # 反转颜色
+    image = ImageOps.invert(image)
+    
+    # 调整图像大小到模型需要的输入尺寸
+    image = image.resize((224, 224))
+    
+    # 图像预处理（根据您的模型需要进行调整）
     img = transform(image).unsqueeze(0)  # 添加批次维度
     
     # 模型预测
     with torch.no_grad():
         outputs = model(img)
+        # 如果模型输出未经过 softmax，可以添加
+        probabilities = F.softmax(outputs, dim=1)
     
     # 获取预测结果
     _, predicted = torch.max(outputs, 1)
+    
+    # 如果需要返回概率
+    # return {str(predicted.item()): probabilities[0][predicted].item()}
+    
+    # 只返回预测的类别
     return str(predicted.item())
 
 # # 创建Gradio界面
@@ -103,16 +116,13 @@ def classify_image(image):
 #     description="上传一张28x28的灰度图像，模型将预测其所属的数字类别。"
 # )
 
-
 iface = gr.Interface(
     fn=classify_image,
-    inputs=gr.Sketchpad(
-        shape=(224, 224),
-        label="Draw a digit"
-    ),
+    inputs=gr.Sketchpad(crop_size=(256,256), type='numpy', image_mode='L', brush=gr.Brush()),
     outputs=gr.Label(num_top_classes=1),
     title="MNIST Digit Classification with ViT", 
     description="使用鼠标手绘一个数字，模型将预测其所属的类别。"
 )
 
+
 iface.launch()