fix infer bug and adjust thres

deploy/ONNX/export_onnx.py

@@ -34,7 +34,7 @@
     parser.add_argument('--max-wh', type=int, default=None, help='None for tensorrt nms, int value for onnx-runtime nms')
     parser.add_argument('--topk-all', type=int, default=100, help='topk objects for every images')
     parser.add_argument('--iou-thres', type=float, default=0.45, help='iou threshold for NMS')
-    parser.add_argument('--conf-thres', type=float, default=0.25, help='conf threshold for NMS')
+    parser.add_argument('--conf-thres', type=float, default=0.5, help='conf threshold for NMS')
     parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     args = parser.parse_args()
     args.img_size *= 2 if len(args.img_size) == 1 else 1  # expand

tools/infer.py

@@ -21,7 +21,7 @@ def get_args_parser(add_help=True):
     parser.add_argument('--source', type=str, default='data/images', help='the source path, e.g. image-file/dir.')
     parser.add_argument('--yaml', type=str, default='data/coco.yaml', help='data yaml file.')
     parser.add_argument('--img-size', type=int, default=640, help='the image-size(h,w) in inference size.')
-    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold for inference.')
+    parser.add_argument('--conf-thres', type=float, default=0.5, help='confidence threshold for inference.')
     parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold for inference.')
     parser.add_argument('--max-det', type=int, default=1000, help='maximal inferences per image.')
     parser.add_argument('--device', default='0', help='device to run our model i.e. 0 or 0,1,2,3 or cpu.')
@@ -47,7 +47,7 @@ def run(weights=osp.join(ROOT, 'yolov6s.pt'),
         source=osp.join(ROOT, 'data/images'),
         yaml=None,
         img_size=640,
-        conf_thres=0.25,
+        conf_thres=0.5,
         iou_thres=0.45,
         max_det=1000,
         device='',

yolov6/core/evaler.py

@@ -251,8 +251,8 @@ def check_thres(conf_thres, iou_thres, task):
                     LOGGER.warning(f'The best conf_thresh when evaluate the model is less than 0.03, while you set it to: {conf_thres}')
                 if iou_thres != 0.65:
                     LOGGER.warning(f'The best iou_thresh when evaluate the model is 0.65, while you set it to: {iou_thres}')
-            if task == 'speed' and conf_thres < 0.25:
-                LOGGER.warning(f'The best conf_thresh when test the speed of the model is larger than 0.25, while you set it to: {conf_thres}')
+            if task == 'speed' and conf_thres < 0.5:
+                LOGGER.warning(f'The best conf_thresh when test the speed of the model is larger than 0.5, while you set it to: {conf_thres}')
 
     @staticmethod
     def reload_device(device, model, task):

yolov6/layers/common.py

@@ -357,7 +357,7 @@ def __init__(self, weights='yolov6s.pt', device=None, dnn=True):
         self.__dict__.update(locals())  # assign all variables to self
 
     def forward(self, im, val=False):
-        y = self.model(im)
+        y, _ = self.model(im)
         if isinstance(y, np.ndarray):
             y = torch.tensor(y, device=self.device)
         return y