diff --git a/utils/utils.py b/utils/utils.py index 5416455a..482f8974 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -525,6 +525,10 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_cls=Tru # Apply width-height constraint pred = pred[((pred[:, 2:4] > min_wh) & (pred[:, 2:4] < max_wh)).all(1)] + # If none remain process next image + if not pred.shape[0]: + continue + # Compute conf pred[..., 5:] *= pred[..., 4:5] # conf = obj_conf * cls_conf