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

chore: make enum comparisons

This commit is contained in:
soumik12345 2025-07-15 15:34:24 +05:30
parent b3024e780d
commit 007d8f5dad
4 changed files with 75 additions and 252 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
supervision/detection/core.py
View File

@ -12,7 +12,6 @@ from supervision.config import (
ORIENTED_BOX_COORDINATES,
)
from supervision.detection.overlap_filter import (
OverlapMetric,
box_non_max_merge,
box_non_max_suppression,
mask_non_max_merge,
@ -36,6 +35,7 @@ from supervision.detection.utils import (
merge_metadata,
process_roboflow_result,
xywh_to_xyxy,
OverlapMetric,
)
from supervision.detection.vlm import (
LMM,

83
supervision/detection/overlap_filter.py
View File

@ -7,6 +7,7 @@ import numpy as np
import numpy.typing as npt
from supervision.detection.utils import box_iou_batch, mask_iou_batch
from supervision.detection.utils import OverlapMetric
def resize_masks(masks: np.ndarray, max_dimension: int = 640) -> np.ndarray:
@ -42,7 +43,7 @@ def mask_non_max_suppression(
predictions: np.ndarray,
masks: np.ndarray,
iou_threshold: float = 0.5,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
mask_dimension: int = 640,
) -> np.ndarray:
"""
@ -58,8 +59,7 @@ def mask_non_max_suppression(
dimensions of each mask.
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
mask_dimension (int): The dimension to which the masks should be
resized before computing IOU values. Defaults to 640.
@ -84,7 +84,7 @@ def mask_non_max_suppression(
predictions = predictions[sort_index]
masks = masks[sort_index]
masks_resized = resize_masks(masks, mask_dimension)
ious = mask_iou_batch(masks_resized, masks_resized, match_metric)
ious = mask_iou_batch(masks_resized, masks_resized, overlap_metric)
categories = predictions[:, 5]
keep = np.ones(rows, dtype=bool)
@ -97,7 +97,7 @@ def mask_non_max_suppression(
def box_non_max_suppression(
predictions: np.ndarray, iou_threshold: float = 0.5, match_metric: str = "IOU"
predictions: np.ndarray, iou_threshold: float = 0.5, overlap_metric: OverlapMetric = OverlapMetric.IOU
) -> np.ndarray:
"""
Perform Non-Maximum Suppression (NMS) on object detection predictions.
@ -108,8 +108,7 @@ def box_non_max_suppression(
or `(x_min, y_min, x_max, y_max, score, class)`.
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
np.ndarray: A boolean array indicating which predictions to keep after n
@ -135,7 +134,7 @@ def box_non_max_suppression(
boxes = predictions[:, :4]
categories = predictions[:, 5]
ious = box_iou_batch(boxes, boxes, match_metric)
ious = box_iou_batch(boxes, boxes, overlap_metric)
ious = ious - np.eye(rows)
keep = np.ones(rows, dtype=bool)
@ -155,7 +154,7 @@ def box_non_max_suppression(
def group_overlapping_boxes(
predictions: npt.NDArray[np.float64],
iou_threshold: float = 0.5,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> List[List[int]]:
"""
Apply greedy version of non-maximum merging to avoid detecting too many
@ -167,8 +166,7 @@ def group_overlapping_boxes(
and the confidence scores.
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression. Defaults to 0.5.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
List[List[int]]: Groups of prediction indices be merged.
@ -189,7 +187,7 @@ def group_overlapping_boxes(
merge_candidate = np.expand_dims(predictions[idx], axis=0)
ious = box_iou_batch(
predictions[order][:, :4], merge_candidate[:, :4], match_metric
predictions[order][:, :4], merge_candidate[:, :4], overlap_metric
)
ious = ious.flatten()
@ -205,7 +203,7 @@ def mask_non_max_merge(
masks: np.ndarray,
iou_threshold: float = 0.5,
mask_dimension: int = 640,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> List[List[int]]:
"""
Perform Non-Maximum Merging (NMM) on segmentation predictions.
@ -222,8 +220,7 @@ def mask_non_max_merge(
to use for non-maximum suppression.
mask_dimension (int): The dimension to which the masks should be
resized before computing IOU values. Defaults to 640.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
np.ndarray: A boolean array indicating which predictions to keep after
@ -236,7 +233,7 @@ def mask_non_max_merge(
masks_resized = resize_masks(masks, mask_dimension)
if predictions.shape[1] == 5:
return group_overlapping_masks(
predictions, masks_resized, iou_threshold, match_metric
predictions, masks_resized, iou_threshold, overlap_metric
)
category_ids = predictions[:, 5]
@ -247,7 +244,7 @@ def mask_non_max_merge(
predictions[curr_indices],
masks_resized[curr_indices],
iou_threshold,
match_metric,
overlap_metric,
)
for merge_class_group in merge_class_groups:
@ -264,7 +261,7 @@ def mask_non_max_merge(
def box_non_max_merge(
predictions: npt.NDArray[np.float64],
iou_threshold: float = 0.5,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> List[List[int]]:
"""
Apply greedy version of non-maximum merging per category to avoid detecting
@ -277,22 +274,21 @@ def box_non_max_merge(
detections of different classes to be merged.
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression. Defaults to 0.5.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
List[List[int]]: Groups of prediction indices be merged.
Each group may have 1 or more elements.
"""
if predictions.shape[1] == 5:
return group_overlapping_boxes(predictions, iou_threshold, match_metric)
return group_overlapping_boxes(predictions, iou_threshold, overlap_metric)
category_ids = predictions[:, 5]
merge_groups = []
for category_id in np.unique(category_ids):
curr_indices = np.where(category_ids == category_id)[0]
merge_class_groups = group_overlapping_boxes(
predictions[curr_indices], iou_threshold, match_metric
predictions[curr_indices], iou_threshold, overlap_metric
)
for merge_class_group in merge_class_groups:
@ -310,7 +306,7 @@ def group_overlapping_masks(
predictions: npt.NDArray[np.float64],
masks: npt.NDArray[np.float64],
iou_threshold: float = 0.5,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> List[List[int]]:
"""
Apply greedy version of non-maximum merging to avoid detecting too many
@ -323,8 +319,7 @@ def group_overlapping_masks(
the predictions.
iou_threshold (float): The intersection-over-union threshold
to use for non-maximum suppression. Defaults to 0.5.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
List[List[int]]: Groups of prediction indices be merged.
@ -348,7 +343,7 @@ def group_overlapping_masks(
candidate_groups = [idx]
while len(order) > 0:
# 'IOU or IOS' of the calculate mask and the remaining mask
ious = mask_iou_batch(masks[order], merge_candidate, match_metric)
ious = mask_iou_batch(masks[order], merge_candidate, overlap_metric)
above_threshold: np.ndarray = ious.flatten() >= iou_threshold
# if no mask is above threshold, break
if not above_threshold.any():
@ -406,39 +401,3 @@ class OverlapFilter(Enum):
f"Invalid value type: {type(value)}. Must be an instance of "
f"{cls.__name__} or str."
)
class OverlapMetric(Enum):
"""
Enum specifying the metric for measuring overlap between detections.
Attributes:
IOU: Intersection over Union. A region-overlap metric that compares
two shapes (usually bounding boxes or masks) by normalising the
shared area with the area of their union.
IOS: Intersection over Smaller, a region-overlap metric that compares
two shapes (usually bounding boxes or masks) by normalising the
shared area with the smaller of the two shapes.
"""
IOU = "IOU"
IOS = "IOS"
@classmethod
def list(cls):
return list(map(lambda c: c.value, cls))
@classmethod
def from_value(cls, value: Union[OverlapFilter, str]) -> OverlapFilter:
if isinstance(value, cls):
return value
if isinstance(value, str):
value = value.lower()
try:
return cls(value)
except ValueError:
raise ValueError(f"Invalid value: {value}. Must be one of {cls.list()}")
raise ValueError(
f"Invalid value type: {type(value)}. Must be an instance of "
f"{cls.__name__} or str."
)

70
supervision/detection/utils.py
View File

@ -1,4 +1,7 @@
from __future__ import annotations
from itertools import chain
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union
import cv2
@ -11,6 +14,42 @@ from supervision.geometry.core import Vector
MIN_POLYGON_POINT_COUNT = 3
class OverlapMetric(Enum):
"""
Enum specifying the metric for measuring overlap between detections.
Attributes:
IOU: Intersection over Union. A region-overlap metric that compares
two shapes (usually bounding boxes or masks) by normalising the
shared area with the area of their union.
IOS: Intersection over Smaller, a region-overlap metric that compares
two shapes (usually bounding boxes or masks) by normalising the
shared area with the smaller of the two shapes.
"""
IOU = "IOU"
IOS = "IOS"
@classmethod
def list(cls):
return list(map(lambda c: c.value, cls))
@classmethod
def from_value(cls, value: Union[OverlapMetric, str]) -> OverlapMetric:
if isinstance(value, cls):
return value
if isinstance(value, str):
value = value.lower()
try:
return cls(value)
except ValueError:
raise ValueError(f"Invalid value: {value}. Must be one of {cls.list()}")
raise ValueError(
f"Invalid value type: {type(value)}. Must be an instance of "
f"{cls.__name__} or str."
)
def xyxy_to_polygons(box: np.ndarray) -> np.ndarray:
"""
Convert an array of boxes to an array of polygons.
@ -109,7 +148,7 @@ def box_iou(
def box_iou_batch(
boxes_true: np.ndarray, boxes_detection: np.ndarray, match_metric: str = "IOU"
boxes_true: np.ndarray, boxes_detection: np.ndarray, overlap_metric: OverlapMetric = OverlapMetric.IOU
) -> np.ndarray:
"""
Compute Intersection over Union (IoU) of two sets of bounding boxes -
@ -126,8 +165,7 @@ def box_iou_batch(
`shape = (N, 4)` where `N` is number of true objects.
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes.
`shape = (M, 4)` where `M` is number of detected objects.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
@ -167,7 +205,7 @@ def box_iou_batch(
area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
if match_metric.upper() == "IOU":
if overlap_metric == OverlapMetric.IOU:
union_area = area_true[:, None] + area_detection - area_inter
ious = np.divide(
area_inter,
@ -175,7 +213,7 @@ def box_iou_batch(
out=np.zeros_like(area_inter, dtype=float),
where=union_area != 0,
)
elif match_metric.upper() == "IOS":
elif overlap_metric == OverlapMetric.IOS:
small_area = np.minimum(area_true[:, None], area_detection)
ious = np.divide(
area_inter,
@ -185,7 +223,7 @@ def box_iou_batch(
)
else:
raise ValueError(
f"match_metric {match_metric} is not supported, "
f"overlap_metric {overlap_metric} is not supported, "
"only 'IOU' and 'IOS' are supported"
)
@ -196,7 +234,7 @@ def box_iou_batch(
def _mask_iou_batch_split(
masks_true: np.ndarray,
masks_detection: np.ndarray,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> np.ndarray:
"""
Internal function.
@ -206,8 +244,7 @@ def _mask_iou_batch_split(
Args:
masks_true (np.ndarray): 3D `np.ndarray` representing ground-truth masks.
masks_detection (np.ndarray): 3D `np.ndarray` representing detection masks.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
Returns:
np.ndarray: Pairwise IoU of masks from `masks_true` and `masks_detection`.
@ -219,7 +256,7 @@ def _mask_iou_batch_split(
masks_true_area = masks_true.sum(axis=(1, 2)) # (area1, area2, ...)
masks_detection_area = masks_detection.sum(axis=(1, 2)) # (area1)
if match_metric.upper() == "IOU":
if overlap_metric == OverlapMetric.IOU:
union_area = masks_true_area[:, None] + masks_detection_area - intersection_area
ious = np.divide(
intersection_area,
@ -227,7 +264,7 @@ def _mask_iou_batch_split(
out=np.zeros_like(intersection_area, dtype=float),
where=union_area != 0,
)
elif match_metric.upper() == "IOS":
elif overlap_metric == OverlapMetric.IOS:
# ios = intersection_area / min(area1, area2)
small_area = np.minimum(masks_true_area[:, None], masks_detection_area)
ious = np.divide(
@ -238,7 +275,7 @@ def _mask_iou_batch_split(
)
else:
raise ValueError(
f"match_metric {match_metric} is not supported, "
f"overlap_metric {overlap_metric} is not supported, "
"only 'IOU' and 'IOS' are supported"
)
@ -249,7 +286,7 @@ def _mask_iou_batch_split(
def mask_iou_batch(
masks_true: np.ndarray,
masks_detection: np.ndarray,
match_metric: str = "IOU",
overlap_metric: OverlapMetric = OverlapMetric.IOU,
memory_limit: int = 1024 * 5,
) -> np.ndarray:
"""
@ -259,8 +296,7 @@ def mask_iou_batch(
Args:
masks_true (np.ndarray): 3D `np.ndarray` representing ground-truth masks.
masks_detection (np.ndarray): 3D `np.ndarray` representing detection masks.
match_metric (str): Metric used for matching detections in slices.
"IOU" or "IOS". Defaults "IOU".
overlap_metric (OverlapMetric): Metric used for matching detections in slices.
memory_limit (int): memory limit in MB, default is 1024 * 5 MB (5GB).
Returns:
@ -275,7 +311,7 @@ def mask_iou_batch(
/ 1024
)
if memory <= memory_limit:
return _mask_iou_batch_split(masks_true, masks_detection, match_metric)
return _mask_iou_batch_split(masks_true, masks_detection, overlap_metric)
ious = []
step = max(
@ -292,7 +328,7 @@ def mask_iou_batch(
for i in range(0, masks_true.shape[0], step):
ious.append(
_mask_iou_batch_split(
masks_true[i : i + step], masks_detection, match_metric
masks_true[i : i + step], masks_detection, overlap_metric
)
)

172
test/detection/test_overlap_filter.py
View File

@ -7,7 +7,6 @@ import pytest
from supervision.detection.overlap_filter import (
box_non_max_suppression,
group_overlapping_boxes,
mask_non_max_merge,
mask_non_max_suppression,
)
@ -448,174 +447,3 @@ def test_mask_non_max_suppression(
predictions=predictions, masks=masks, iou_threshold=iou_threshold
)
assert np.array_equal(result, expected_result)
@pytest.mark.parametrize(
"predictions, masks, iou_threshold, expected_result, exception",
[
(
np.empty((0, 6)),
np.empty((0, 5, 5)),
0.5,
[],
DoesNotRaise(),
), # empty predictions and masks
(
np.array([[0, 0, 0, 0, 0.8]]),
np.array(
[
[
[False, False, False, False, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, False, False, False, False],
]
]
),
0.5,
[[0]],
DoesNotRaise(),
), # single mask with no category
(
np.array([[0, 0, 0, 0, 0.8, 0]]),
np.array(
[
[
[False, False, False, False, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, False, False, False, False],
]
]
),
0.5,
[[0]],
DoesNotRaise(),
), # single mask with category
(
np.array([[0, 0, 0, 0, 0.9], [0, 0, 0, 0, 0.8]]),
np.array(
[
[
[False, False, False, False, False],
[False, True, True, False, False],
[False, True, True, False, False],
[False, False, False, False, False],
[False, False, False, False, False],
],
[
[False, False, False, False, False],
[False, False, False, False, False],
[False, False, False, True, True],
[False, False, False, True, True],
[False, False, False, False, False],
],
]
),
0.5,
[[0], [1]],
DoesNotRaise(),
), # two masks non-overlapping with no category
(
np.array([[0, 0, 0, 0, 0.9], [0, 0, 0, 0, 0.8]]),
np.array(
[
[
[False, False, False, False, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, False, False, False, False],
],
[
[False, False, False, False, False],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, False, False, False],
],
]
),
0.2,
[[0, 1]],
DoesNotRaise(),
), # two masks partially overlapping with no category
(
np.array([[0, 0, 0, 0, 0.9, 0], [0, 0, 0, 0, 0.8, 1]]),
np.array(
[
[
[False, False, False, False, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, False, False, False, False],
],
[
[False, False, False, False, False],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, False, False, False],
],
]
),
0.2,
[[0], [1]],
DoesNotRaise(),
), # two masks partially overlapping with different category
(
np.array(
[
[0, 0, 0, 0, 0.9, 0],
[0, 0, 0, 0, 0.8, 0],
[0, 0, 0, 0, 0.85, 1],
]
),
np.array(
[
[ # mask 0, class 0
[False, False, False, False, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, True, True, True, False],
[False, False, False, False, False],
],
[ # mask 1, class 0
[False, False, False, False, False],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, False, False, False],
],
[ # mask 2, class 1
[False, False, False, False, False],
[False, False, False, True, True],
[False, False, False, True, True],
[False, False, False, False, False],
[False, False, False, False, False],
],
]
),
0.2,
[[0, 1], [2]],
DoesNotRaise(),
), # three masks, two overlapping with same class
],
)
def test_mask_non_max_merge(
predictions: np.ndarray,
masks: np.ndarray,
iou_threshold: float,
expected_result: List[List[int]],
exception: Exception,
) -> None:
with exception:
result = mask_non_max_merge(
predictions=predictions, masks=masks, iou_threshold=iou_threshold
)
result = sorted([sorted(group) for group in result])
expected_result = sorted([sorted(group) for group in expected_result])
assert result == expected_result