This commit is contained in:
Glenn Jocher 2020-03-03 19:16:13 -08:00
parent dce753ead4
commit 308f7c8563
1 changed files with 4 additions and 4 deletions

View File

@ -524,10 +524,6 @@ def non_max_suppression(prediction, conf_thres=0.5, iou_thres=0.5, multi_cls=Tru
# Apply width-height constraint
pred = pred[(pred[:, 2:4] > min_wh).all(1) & (pred[:, 2:4] < max_wh).all(1)]
# If none remain process next image
if len(pred) == 0:
continue
# Compute conf
pred[..., 5:] *= pred[..., 4:5] # conf = obj_conf * cls_conf
@ -550,6 +546,10 @@ def non_max_suppression(prediction, conf_thres=0.5, iou_thres=0.5, multi_cls=Tru
if not torch.isfinite(pred).all():
pred = pred[torch.isfinite(pred).all(1)]
# If none remain process next image
if not pred.shape[0]:
continue
# Batched NMS
if method == 'vision_batch':
c = pred[:, 5] * 0 if agnostic else pred[:, 5] # class-agnostic NMS