diff --git a/detect.py b/detect.py index 9df3a1c6..74838a49 100644 --- a/detect.py +++ b/detect.py @@ -86,7 +86,7 @@ def detect(save_txt=False, save_img=False): pred = pred.float() # Apply NMS - pred = non_max_suppression(pred, opt.conf_thres, opt.nms_thres) + pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres) # Apply Classifier if classify: @@ -162,7 +162,7 @@ if __name__ == '__main__': parser.add_argument('--output', type=str, default='output', help='output folder') # output folder parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold') - parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression') + parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)') parser.add_argument('--half', action='store_true', help='half precision FP16 inference') parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu') diff --git a/test.py b/test.py index 4315d826..19aca3c7 100644 --- a/test.py +++ b/test.py @@ -14,7 +14,7 @@ def test(cfg, batch_size=16, img_size=416, conf_thres=0.001, - nms_thres=0.5, + iou_thres=0.5, save_json=False, model=None, dataloader=None): @@ -88,7 +88,7 @@ def test(cfg, loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls # Run NMS - output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres) + output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres) # Statistics per image for si, pred in enumerate(output): @@ -212,7 +212,7 @@ if __name__ == '__main__': parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') - parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression') + parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') parser.add_argument('--task', default='test', help="'test', 'study', 'benchmark'") parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu') @@ -228,7 +228,7 @@ if __name__ == '__main__': opt.batch_size, opt.img_size, opt.conf_thres, - opt.nms_thres, + opt.iou_thres, opt.save_json) elif opt.task == 'benchmark': @@ -262,6 +262,6 @@ if __name__ == '__main__': ax[2].set_ylabel('time (s)') for i in range(3): ax[i].legend() - ax[i].set_xlabel('nms_thr') + ax[i].set_xlabel('iou_thr') fig.tight_layout() plt.savefig('study.jpg', dpi=200) diff --git a/utils/utils.py b/utils/utils.py index 38ff9fd3..6a9d0861 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -497,7 +497,7 @@ def build_targets(model, targets): return tcls, tbox, indices, av -def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=True, method='vision_batch'): +def non_max_suppression(prediction, conf_thres=0.5, iou_thres=0.5, multi_cls=True, method='vision_batch'): """ Removes detections with lower object confidence score than 'conf_thres' Non-Maximum Suppression to further filter detections. @@ -542,7 +542,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru # Batched NMS if method == 'vision_batch': - output[image_i] = pred[torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], nms_thres)] + output[image_i] = pred[torchvision.ops.boxes.batched_nms(pred[:, :4], pred[:, 4], pred[:, 5], iou_thres)] continue # Sort by confidence @@ -562,7 +562,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru dc = dc[:500] # limit to first 500 boxes: https://github.com/ultralytics/yolov3/issues/117 if method == 'vision': - det_max.append(dc[torchvision.ops.boxes.nms(dc[:, :4], dc[:, 4], nms_thres)]) + det_max.append(dc[torchvision.ops.boxes.nms(dc[:, :4], dc[:, 4], iou_thres)]) elif method == 'or': # default # METHOD1 @@ -570,7 +570,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru # while len(ind): # j = ind[0] # det_max.append(dc[j:j + 1]) # save highest conf detection - # reject = (bbox_iou(dc[j], dc[ind]) > nms_thres).nonzero() + # reject = (bbox_iou(dc[j], dc[ind]) > iou_thres).nonzero() # [ind.pop(i) for i in reversed(reject)] # METHOD2 @@ -579,21 +579,21 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru if len(dc) == 1: # Stop if we're at the last detection break iou = bbox_iou(dc[0], dc[1:]) # iou with other boxes - dc = dc[1:][iou < nms_thres] # remove ious > threshold + dc = dc[1:][iou < iou_thres] # remove ious > threshold elif method == 'and': # requires overlap, single boxes erased while len(dc) > 1: iou = bbox_iou(dc[0], dc[1:]) # iou with other boxes if iou.max() > 0.5: det_max.append(dc[:1]) - dc = dc[1:][iou < nms_thres] # remove ious > threshold + dc = dc[1:][iou < iou_thres] # remove ious > threshold elif method == 'merge': # weighted mixture box while len(dc): if len(dc) == 1: det_max.append(dc) break - i = bbox_iou(dc[0], dc) > nms_thres # iou with other boxes + i = bbox_iou(dc[0], dc) > iou_thres # iou with other boxes weights = dc[i, 4:5] dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum() det_max.append(dc[:1])