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() pred = pred.float()
# Apply NMS # 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 # Apply Classifier
if classify: 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('--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('--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('--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('--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('--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') 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, batch_size=16,
img_size=416, img_size=416,
conf_thres=0.001, conf_thres=0.001,
nms_thres=0.5, iou_thres=0.5,
save_json=False, save_json=False,
model=None, model=None,
dataloader=None): dataloader=None):
@ -88,7 +88,7 @@ def test(cfg,
loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls loss += compute_loss(train_out, targets, model)[1][:3].cpu() # GIoU, obj, cls
# Run NMS # 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 # Statistics per image
for si, pred in enumerate(output): 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('--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('--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('--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('--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('--task', default='test', help="'test', 'study', 'benchmark'")
parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu') 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.batch_size,
opt.img_size, opt.img_size,
opt.conf_thres, opt.conf_thres,
opt.nms_thres, opt.iou_thres,
opt.save_json) opt.save_json)
elif opt.task == 'benchmark': elif opt.task == 'benchmark':
@ -262,6 +262,6 @@ if __name__ == '__main__':
ax[2].set_ylabel('time (s)') ax[2].set_ylabel('time (s)')
for i in range(3): for i in range(3):
ax[i].legend() ax[i].legend()
ax[i].set_xlabel('nms_thr') ax[i].set_xlabel('iou_thr')
fig.tight_layout() fig.tight_layout()
plt.savefig('study.jpg', dpi=200) plt.savefig('study.jpg', dpi=200)

View File

@ -497,7 +497,7 @@ def build_targets(model, targets):
return tcls, tbox, indices, av 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' Removes detections with lower object confidence score than 'conf_thres'
Non-Maximum Suppression to further filter detections. 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 # Batched NMS
if method == 'vision_batch': 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 continue
# Sort by confidence # 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 dc = dc[:500] # limit to first 500 boxes: https://github.com/ultralytics/yolov3/issues/117
if method == 'vision': 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 elif method == 'or': # default
# METHOD1 # METHOD1
@ -570,7 +570,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru
# while len(ind): # while len(ind):
# j = ind[0] # j = ind[0]
# det_max.append(dc[j:j + 1]) # save highest conf detection # 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)] # [ind.pop(i) for i in reversed(reject)]
# METHOD2 # 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 if len(dc) == 1: # Stop if we're at the last detection
break break
iou = bbox_iou(dc[0], dc[1:]) # iou with other boxes 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 elif method == 'and': # requires overlap, single boxes erased
while len(dc) > 1: while len(dc) > 1:
iou = bbox_iou(dc[0], dc[1:]) # iou with other boxes iou = bbox_iou(dc[0], dc[1:]) # iou with other boxes
if iou.max() > 0.5: if iou.max() > 0.5:
det_max.append(dc[:1]) 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 elif method == 'merge': # weighted mixture box
while len(dc): while len(dc):
if len(dc) == 1: if len(dc) == 1:
det_max.append(dc) det_max.append(dc)
break 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] weights = dc[i, 4:5]
dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum() dc[0, :4] = (weights * dc[i, :4]).sum(0) / weights.sum()
det_max.append(dc[:1]) det_max.append(dc[:1])