This commit is contained in:
Glenn Jocher 2019-12-26 12:30:51 -08:00
parent 8ae06ad7c3
commit fea54c4a85
3 changed files with 14 additions and 14 deletions

View File

@ -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')

10
test.py
View File

@ -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)

View File

@ -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])