Upload app.py

app.py

@@ -22,19 +22,20 @@ 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()
-    # return {LABELS[i]: pred[i] for i in range(len(pred))}
+    output_ = output.cpu().numpy().tolist()
+    res_dict = {LABELS[i]: output_[i] for i in range(len(output_))}
+
     pred = torch.nonzero(output > 0.4).view(-1)
     pred = pred.cpu().numpy().tolist()
 
     if len(pred) == 0 or (len(pred) == 1 and pred[0] == 0):
-        return LABELS[0]
+        return LABELS[0], res_dict
     res = ''
     for i in pred:
         if i == 0:
             continue
         res += ', ' + LABELS[i]
-    return '目前的身体状态：' + res[2:]
+    return '目前的身体状态：' + res[2:], res_dict
 
 
 if __name__ == '__main__':
@@ -68,6 +69,8 @@ if __name__ == '__main__':
     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="text", examples=examples, title=title, description=description)
+    demo = gr.Interface(fn=img2label,
+                        inputs=[gr.inputs.Image(), gr.inputs.Image()],
+                        outputs=["text", "label"],
+                        examples=examples, title=title, description=description)
     demo.launch()