xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactored `mask_non_max_suppression`

This commit is contained in:
SkalskiP 2024-02-07 09:18:41 +01:00
parent 35c9ebdd60
commit 239f7333f1
1 changed files with 5 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
supervision/detection/utils.py
View File

@ -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(