From 239f7333f10b54bec59ec9d4686e8dc0d544948d Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 7 Feb 2024 09:18:41 +0100 Subject: [PATCH] refactored `mask_non_max_suppression` --- supervision/detection/utils.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index e74facb0..2aaa23f0 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -211,18 +211,12 @@ def mask_non_max_suppression_2( ious = mask_iou_batch(masks_resized, masks_resized) categories = predictions[:, 5] - keep = np.ones(num_predictions, dtype=bool) + for i in range(num_predictions): + if keep[i]: + condition = (ious[i] > iou_threshold) & (categories == category) + keep[i + 1:] = np.where(condition[i + 1:], False, keep[i + 1:]) - for index, (iou, category) in enumerate(zip(ious, categories)): - if not keep[index]: - continue - - # drop detections with iou > iou_threshold and - # same category as current detections - condition = (iou > iou_threshold) & (categories == category) - keep = keep & ~condition - - return keep[sort_index.argsort()] + return keep def box_non_max_suppression(