瀏覽代碼

修改返回类型

leon 3 周之前
父節點
當前提交
9d314d4b3c
共有 1 個文件被更改,包括 1 次插入1 次删除
  1. 1 1
      infer.py

+ 1 - 1
infer.py

@@ -186,7 +186,7 @@ class DoorInference(object):
         # cv2.waitKey(30000)
         res = self.door_model(door_image)
         class_idx = res[0].probs.top1
-        class_conf = res[0].probs.top1conf.cpu().numpy()
+        class_conf = res[0].probs.top1conf.cpu().numpy().item()
         return class_idx == 1, class_conf
 
     """