This commit is contained in:
Glenn Jocher 2020-03-05 09:54:41 -08:00
parent 8b6c8a5318
commit b8b89a3132
1 changed files with 4 additions and 0 deletions

View File

@ -525,6 +525,10 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_cls=Tru
# Apply width-height constraint # Apply width-height constraint
pred = pred[((pred[:, 2:4] > min_wh) & (pred[:, 2:4] < max_wh)).all(1)] 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 # Compute conf
pred[..., 5:] *= pred[..., 4:5] # conf = obj_conf * cls_conf pred[..., 5:] *= pred[..., 4:5] # conf = obj_conf * cls_conf