updates
This commit is contained in:
parent
8ae06ad7c3
commit
fea54c4a85
|
@ -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
10
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)
|
||||
|
|
|
@ -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])
|
||||
|
|
Loading…
Reference in New Issue