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:24:45 +01:00
parent c46a2e9c08
commit f2e65cf6fc
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
supervision/detection/utils.py
View File

@ -199,10 +199,10 @@ def mask_non_max_suppression_2(
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
f"{iou_threshold} given."
)
num_predictions, columns = predictions.shape
rows, columns = predictions.shape
if columns == 5:
predictions = np.c_[predictions, np.zeros(num_predictions)]
predictions = np.c_[predictions, np.zeros(rows)]
sort_index = predictions[:, 4].argsort()[::-1]
predictions = predictions[sort_index]
@ -212,7 +212,7 @@ def mask_non_max_suppression_2(
categories = predictions[:, 5]
keep = np.ones(rows, dtype=bool)
for i in range(num_predictions):
for i in range(rows):
if keep[i]:
condition = (ious[i] > iou_threshold) & (categories == category)
keep[i + 1:] = np.where(condition[i + 1:], False, keep[i + 1:])