Upload app.py

app.py

@@ -5,15 +5,15 @@ import torchvision.transforms as transforms
 
 
 def img2label(left, right):
-    left_img = Image.open(left).convert('RGB')
-    right_img = Image.open(right).convert('RGB')
+    left = Image.fromarray(left.astype('uint8'), 'RGB')
+    right = Image.fromarray(right.astype('uint8'), 'RGB')
     # 将右眼底镜像反转
     r2l = transforms.RandomHorizontalFlip(p=1)
-    right_img = r2l(right_img)
+    right = r2l(right)
 
     # 调整图片
-    left_img = my_transforms(left_img).to(device)
-    right_img = my_transforms(right_img).to(device)
+    left_img = my_transforms(left).to(device)
+    right_img = my_transforms(right).to(device)
 
     # 读取模型
     model = torch.load('densenet_FD_e4_l5e-4_b32.pkl', map_location='cpu').to(device)
@@ -22,9 +22,19 @@ def img2label(left, right):
         output = model(left=left_img.unsqueeze(0), right=right_img.unsqueeze(0))
 
     output = torch.sigmoid(output.squeeze(0))
-    pred = output.cpu().numpy().tolist()
+    # pred = output.cpu().numpy().tolist()
+    # return {LABELS[i]: pred[i] for i in range(len(pred))}
+    pred = torch.nonzero(output > 0.4).view(-1)
+    pred = pred.cpu().numpy().tolist()
 
-    return {LABELS[i]: pred[i] for i in range(len(pred))}
+    if len(pred) == 0 or (len(pred) == 1 and pred[0] == 0):
+        return LABELS[0]
+    res = ''
+    for i in pred:
+        if i == 0:
+            continue
+        res += ', ' + LABELS[i]
+    return '目前的身体状态：' + res[2:]
 
 
 if __name__ == '__main__':
@@ -54,9 +64,11 @@ if __name__ == '__main__':
               6: '病理性近视',
               7: '其他疾病'}
 
-    # left_img_dir = 'left.jpg'
-    # right_img_dir = 'right.jpg'
+    left_img_dir = 'left.jpg'
+    right_img_dir = 'right.jpg'
+    examples = [[left_img_dir, right_img_dir]]
     # r = img2label(left_img_dir, right_img_dir)
-    demo = gr.Interface(fn=img2label, inputs=[gr.inputs.Image(), gr.inputs.Image()], outputs='label',
+    demo = gr.Interface(fn=img2label, inputs=[gr.inputs.Image(), gr.inputs.Image()],
+                        outputs="text", examples=examples,
                         title=title, description=description)
-    demo.launch()
+    demo.launch(share=True)