IoU 和 NMS 工具

OverlapFilter

supervision.detection.utils.iou_and_nms.OverlapFilter

Bases: Enum

Enum specifying the strategy for filtering overlapping detections.

Attributes:

Name	Type	Description
NONE		Do not filter detections based on overlap.
NON_MAX_SUPPRESSION		Filter detections using non-max suppression. This means, detections that overlap by more than a set threshold will be discarded, except for the one with the highest confidence.
NON_MAX_MERGE		Merge detections with non-max merging. This means, detections that overlap by more than a set threshold will be merged into a single detection.

Source code in supervision/detection/utils/iou_and_nms.py

class OverlapFilter(Enum):
    """
    Enum specifying the strategy for filtering overlapping detections.

    Attributes:
        NONE: Do not filter detections based on overlap.
        NON_MAX_SUPPRESSION: Filter detections using non-max suppression. This means,
            detections that overlap by more than a set threshold will be discarded,
            except for the one with the highest confidence.
        NON_MAX_MERGE: Merge detections with non-max merging. This means,
            detections that overlap by more than a set threshold will be merged
            into a single detection.
    """

    NONE = "none"
    NON_MAX_SUPPRESSION = "non_max_suppression"
    NON_MAX_MERGE = "non_max_merge"

    @classmethod
    def list(cls):
        return list(map(lambda c: c.value, cls))

    @classmethod
    def from_value(cls, value: 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."
        )

OverlapMetric

supervision.detection.utils.iou_and_nms.OverlapMetric

Bases: Enum

Enum specifying the metric for measuring overlap between detections.

Attributes:

Name	Type	Description
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.

Source code in supervision/detection/utils/iou_and_nms.py

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: 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."
        )

box_iou

supervision.detection.utils.iou_and_nms.box_iou(box_true, box_detection)

Compute the Intersection over Union (IoU) between two bounding boxes.

\[ \text{IoU} = \frac{|\text{box}_{\text{true}} \cap \text{box}_{\text{detection}}|}{|\text{box}_{\text{true}} \cup \text{box}_{\text{detection}}|} \]

Note

Use box_iou when computing IoU between two individual boxes. For comparing multiple boxes (arrays of boxes), use box_iou_batch for better performance.

Parameters:

Name	Type	Description	Default
box_true	Union[List[float], ndarray]	A single bounding box represented as [x_min, y_min, x_max, y_max].	required
box_detection	Union[List[float], ndarray]	A single bounding box represented as [x_min, y_min, x_max, y_max].	required

Returns:

Name	Type	Description
IoU	float	IoU score between the two boxes. Ranges from 0.0 (no overlap) to 1.0 (perfect overlap).

Examples:

import numpy as np
import supervision as sv

box_true = np.array([100, 100, 200, 200])
box_detection = np.array([150, 150, 250, 250])

sv.box_iou(box_true=box_true, box_detection=box_detection)
# 0.14285814285714285

Source code in supervision/detection/utils/iou_and_nms.py

def box_iou(
    box_true: list[float] | np.ndarray,
    box_detection: list[float] | np.ndarray,
) -> float:
    r"""
    Compute the Intersection over Union (IoU) between two bounding boxes.

    \[
    \text{IoU} = \frac{|\text{box}_{\text{true}} \cap \text{box}_{\text{detection}}|}{|\text{box}_{\text{true}} \cup \text{box}_{\text{detection}}|}
    \]

    Note:
        Use `box_iou` when computing IoU between two individual boxes.
        For comparing multiple boxes (arrays of boxes), use `box_iou_batch` for better
        performance.

    Args:
        box_true (Union[List[float], np.ndarray]): A single bounding box represented as
            [x_min, y_min, x_max, y_max].
        box_detection (Union[List[float], np.ndarray]):
            A single bounding box represented as [x_min, y_min, x_max, y_max].

    Returns:
        IoU (float): IoU score between the two boxes. Ranges from 0.0 (no overlap)
            to 1.0 (perfect overlap).

    Examples:
        ```python
        import numpy as np
        import supervision as sv

        box_true = np.array([100, 100, 200, 200])
        box_detection = np.array([150, 150, 250, 250])

        sv.box_iou(box_true=box_true, box_detection=box_detection)
        # 0.14285814285714285
        ```
    """  # noqa: E501
    box_true = np.array(box_true)
    box_detection = np.array(box_detection)

    inter_x1 = max(box_true[0], box_detection[0])
    inter_y1 = max(box_true[1], box_detection[1])
    inter_x2 = min(box_true[2], box_detection[2])
    inter_y2 = min(box_true[3], box_detection[3])

    inter_w = max(0, inter_x2 - inter_x1)
    inter_h = max(0, inter_y2 - inter_y1)

    inter_area = inter_w * inter_h

    area_true = (box_true[2] - box_true[0]) * (box_true[3] - box_true[1])
    area_detection = (box_detection[2] - box_detection[0]) * (
        box_detection[3] - box_detection[1]
    )

    union_area = area_true + area_detection - inter_area

    return inter_area / union_area + 1e-6

box_iou_batch

supervision.detection.utils.iou_and_nms.box_iou_batch(boxes_true, boxes_detection, overlap_metric=OverlapMetric.IOU)

Compute Intersection over Union (IoU) of two sets of bounding boxes - boxes_true and boxes_detection. Both sets of boxes are expected to be in (x_min, y_min, x_max, y_max) format.

Note

Use box_iou when computing IoU between two individual boxes. For comparing multiple boxes (arrays of boxes), use box_iou_batch for better performance.

Parameters:

Name	Type	Description	Default
boxes_true	ndarray	2D np.ndarray representing ground-truth boxes. shape = (N, 4) where N is number of true objects.	required
boxes_detection	ndarray	2D np.ndarray representing detection boxes. shape = (M, 4) where M is number of detected objects.	required
overlap_metric	OverlapMetric	Metric used for matching detections in slices.	IOU

Returns:

Type	Description
ndarray	np.ndarray: Pairwise IoU of boxes from boxes_true and boxes_detection. shape = (N, M) where N is number of true objects and M is number of detected objects.

Examples:

import numpy as np
import supervision as sv

boxes_true = np.array([
    [100, 100, 200, 200],
    [300, 300, 400, 400]
])
boxes_detection = np.array([
    [150, 150, 250, 250],
    [320, 320, 420, 420]
])

sv.box_iou_batch(boxes_true=boxes_true, boxes_detection=boxes_detection)
# array([
#     [0.14285714, 0.        ],
#     [0.        , 0.47058824]
# ])

Source code in supervision/detection/utils/iou_and_nms.py

def box_iou_batch(
    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 -
        `boxes_true` and `boxes_detection`. Both sets
        of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.

    Note:
        Use `box_iou` when computing IoU between two individual boxes.
        For comparing multiple boxes (arrays of boxes), use `box_iou_batch` for better
        performance.

    Args:
        boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes.
            `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.
        overlap_metric (OverlapMetric): Metric used for matching detections in slices.

    Returns:
        np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
            `shape = (N, M)` where `N` is number of true objects and
            `M` is number of detected objects.

    Examples:
        ```python
        import numpy as np
        import supervision as sv

        boxes_true = np.array([
            [100, 100, 200, 200],
            [300, 300, 400, 400]
        ])
        boxes_detection = np.array([
            [150, 150, 250, 250],
            [320, 320, 420, 420]
        ])

        sv.box_iou_batch(boxes_true=boxes_true, boxes_detection=boxes_detection)
        # array([
        #     [0.14285714, 0.        ],
        #     [0.        , 0.47058824]
        # ])
        ```
    """

    def box_area(box):
        return (box[2] - box[0]) * (box[3] - box[1])

    area_true = box_area(boxes_true.T)
    area_detection = box_area(boxes_detection.T)

    top_left = np.maximum(boxes_true[:, None, :2], boxes_detection[:, :2])
    bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])

    area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)

    if overlap_metric == OverlapMetric.IOU:
        union_area = area_true[:, None] + area_detection - area_inter
        ious = np.divide(
            area_inter,
            union_area,
            out=np.zeros_like(area_inter, dtype=float),
            where=union_area != 0,
        )
    elif overlap_metric == OverlapMetric.IOS:
        small_area = np.minimum(area_true[:, None], area_detection)
        ious = np.divide(
            area_inter,
            small_area,
            out=np.zeros_like(area_inter, dtype=float),
            where=small_area != 0,
        )
    else:
        raise ValueError(
            f"overlap_metric {overlap_metric} is not supported, "
            "only 'IOU' and 'IOS' are supported"
        )

    ious = np.nan_to_num(ious)
    return ious

box_iou_batch_with_jaccard

supervision.detection.utils.iou_and_nms.box_iou_batch_with_jaccard(boxes_true, boxes_detection, is_crowd)

Calculate the intersection over union (IoU) between detection bounding boxes (dt) and ground-truth bounding boxes (gt). Reference: https://github.com/rafaelpadilla/review_object_detection_metrics

Parameters:

Name	Type	Description	Default
boxes_true	List[List[float]]	List of ground-truth bounding boxes in the format [x, y, width, height].	required
boxes_detection	List[List[float]]	List of detection bounding boxes in the format [x, y, width, height].	required
is_crowd	List[bool]	List indicating if each ground-truth bounding box is a crowd region or not.	required

Returns:

Type	Description
ndarray	np.ndarray: Array of IoU values of shape (len(dt), len(gt)).

Examples:

import numpy as np
import supervision as sv

boxes_true = [
    [10, 20, 30, 40],  # x, y, w, h
    [15, 25, 35, 45]
]
boxes_detection = [
    [12, 22, 28, 38],
    [16, 26, 36, 46]
]
is_crowd = [False, False]

ious = sv.box_iou_batch_with_jaccard(
    boxes_true=boxes_true,
    boxes_detection=boxes_detection,
    is_crowd=is_crowd
)
# array([
#     [0.8866..., 0.4960...],
#     [0.4000..., 0.8622...]
# ])

Source code in supervision/detection/utils/iou_and_nms.py

def box_iou_batch_with_jaccard(
    boxes_true: list[list[float]],
    boxes_detection: list[list[float]],
    is_crowd: list[bool],
) -> np.ndarray:
    """
    Calculate the intersection over union (IoU) between detection bounding boxes (dt)
    and ground-truth bounding boxes (gt).
    Reference: https://github.com/rafaelpadilla/review_object_detection_metrics

    Args:
        boxes_true (List[List[float]]): List of ground-truth bounding boxes in the \
            format [x, y, width, height].
        boxes_detection (List[List[float]]): List of detection bounding boxes in the \
            format [x, y, width, height].
        is_crowd (List[bool]): List indicating if each ground-truth bounding box \
            is a crowd region or not.

    Returns:
        np.ndarray: Array of IoU values of shape (len(dt), len(gt)).

    Examples:
        ```python
        import numpy as np
        import supervision as sv

        boxes_true = [
            [10, 20, 30, 40],  # x, y, w, h
            [15, 25, 35, 45]
        ]
        boxes_detection = [
            [12, 22, 28, 38],
            [16, 26, 36, 46]
        ]
        is_crowd = [False, False]

        ious = sv.box_iou_batch_with_jaccard(
            boxes_true=boxes_true,
            boxes_detection=boxes_detection,
            is_crowd=is_crowd
        )
        # array([
        #     [0.8866..., 0.4960...],
        #     [0.4000..., 0.8622...]
        # ])
        ```
    """
    assert len(is_crowd) == len(boxes_true), (
        "`is_crowd` must have the same length as `boxes_true`"
    )
    if len(boxes_detection) == 0 or len(boxes_true) == 0:
        return np.array([])
    ious = np.zeros((len(boxes_detection), len(boxes_true)), dtype=np.float64)
    for g_idx, g in enumerate(boxes_true):
        for d_idx, d in enumerate(boxes_detection):
            ious[d_idx, g_idx] = _jaccard(d, g, is_crowd[g_idx])
    return ious

mask_iou_batch

supervision.detection.utils.iou_and_nms.mask_iou_batch(masks_true, masks_detection, overlap_metric=OverlapMetric.IOU, memory_limit=1024 * 5)

Compute Intersection over Union (IoU) of two sets of masks - masks_true and masks_detection.

Parameters:

Name	Type	Description	Default
masks_true	ndarray	3D np.ndarray representing ground-truth masks.	required
masks_detection	ndarray	3D np.ndarray representing detection masks.	required
overlap_metric	OverlapMetric	Metric used for matching detections in slices.	IOU
memory_limit	int	memory limit in MB, default is 1024 * 5 MB (5GB).	1024 * 5

Returns:

Type	Description
ndarray	np.ndarray: Pairwise IoU of masks from masks_true and masks_detection.

Source code in supervision/detection/utils/iou_and_nms.py

def mask_iou_batch(
    masks_true: np.ndarray,
    masks_detection: np.ndarray,
    overlap_metric: OverlapMetric = OverlapMetric.IOU,
    memory_limit: int = 1024 * 5,
) -> np.ndarray:
    """
    Compute Intersection over Union (IoU) of two sets of masks -
        `masks_true` and `masks_detection`.

    Args:
        masks_true (np.ndarray): 3D `np.ndarray` representing ground-truth masks.
        masks_detection (np.ndarray): 3D `np.ndarray` representing detection masks.
        overlap_metric (OverlapMetric): Metric used for matching detections in slices.
        memory_limit (int): memory limit in MB, default is 1024 * 5 MB (5GB).

    Returns:
        np.ndarray: Pairwise IoU of masks from `masks_true` and `masks_detection`.
    """
    memory = (
        masks_true.shape[0]
        * masks_true.shape[1]
        * masks_true.shape[2]
        * masks_detection.shape[0]
        / 1024
        / 1024
    )
    if memory <= memory_limit:
        return _mask_iou_batch_split(masks_true, masks_detection, overlap_metric)

    ious = []
    step = max(
        memory_limit
        * 1024
        * 1024
        // (
            masks_detection.shape[0]
            * masks_detection.shape[1]
            * masks_detection.shape[2]
        ),
        1,
    )
    for i in range(0, masks_true.shape[0], step):
        ious.append(
            _mask_iou_batch_split(
                masks_true[i : i + step], masks_detection, overlap_metric
            )
        )

    return np.vstack(ious)

oriented_box_iou_batch

supervision.detection.utils.iou_and_nms.oriented_box_iou_batch(boxes_true, boxes_detection)

Compute Intersection over Union (IoU) of two sets of oriented bounding boxes - boxes_true and boxes_detection. Both sets of boxes are expected to be in ((x1, y1), (x2, y2), (x3, y3), (x4, y4)) format.

Parameters:

Name	Type	Description	Default
boxes_true	ndarray	a np.ndarray representing ground-truth boxes. shape = (N, 4, 2) where N is number of true objects.	required
boxes_detection	ndarray	a np.ndarray representing detection boxes. shape = (M, 4, 2) where M is number of detected objects.	required

Returns:

Type	Description
ndarray	np.ndarray: Pairwise IoU of boxes from boxes_true and boxes_detection. shape = (N, M) where N is number of true objects and M is number of detected objects.

Source code in supervision/detection/utils/iou_and_nms.py

def oriented_box_iou_batch(
    boxes_true: np.ndarray, boxes_detection: np.ndarray
) -> np.ndarray:
    """
    Compute Intersection over Union (IoU) of two sets of oriented bounding boxes -
    `boxes_true` and `boxes_detection`. Both sets of boxes are expected to be in
    `((x1, y1), (x2, y2), (x3, y3), (x4, y4))` format.

    Args:
        boxes_true (np.ndarray): a `np.ndarray` representing ground-truth boxes.
            `shape = (N, 4, 2)` where `N` is number of true objects.
        boxes_detection (np.ndarray): a `np.ndarray` representing detection boxes.
            `shape = (M, 4, 2)` where `M` is number of detected objects.

    Returns:
        np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
            `shape = (N, M)` where `N` is number of true objects and
            `M` is number of detected objects.
    """

    boxes_true = boxes_true.reshape(-1, 4, 2)
    boxes_detection = boxes_detection.reshape(-1, 4, 2)

    max_height = int(max(boxes_true[:, :, 0].max(), boxes_detection[:, :, 0].max()) + 1)
    # adding 1 because we are 0-indexed
    max_width = int(max(boxes_true[:, :, 1].max(), boxes_detection[:, :, 1].max()) + 1)

    mask_true = np.zeros((boxes_true.shape[0], max_height, max_width))
    for i, box_true in enumerate(boxes_true):
        mask_true[i] = polygon_to_mask(box_true, (max_width, max_height))

    mask_detection = np.zeros((boxes_detection.shape[0], max_height, max_width))
    for i, box_detection in enumerate(boxes_detection):
        mask_detection[i] = polygon_to_mask(box_detection, (max_width, max_height))

    ious = mask_iou_batch(mask_true, mask_detection)
    return ious

box_non_max_suppression

supervision.detection.utils.iou_and_nms.box_non_max_suppression(predictions, iou_threshold=0.5, overlap_metric=OverlapMetric.IOU)

Perform Non-Maximum Suppression (NMS) on object detection predictions.

Parameters:

Name	Type	Description	Default
predictions	ndarray	An array of object detection predictions in the format of (x_min, y_min, x_max, y_max, score) or (x_min, y_min, x_max, y_max, score, class).	required
iou_threshold	float	The intersection-over-union threshold to use for non-maximum suppression.	0.5
overlap_metric	OverlapMetric	Metric used for matching detections in slices.	IOU

Returns:

Type	Description
ndarray	np.ndarray: A boolean array indicating which predictions to keep after n on-maximum suppression.

Raises:

Type	Description
AssertionError	If iou_threshold is not within the closed range from 0 to 1.

Source code in supervision/detection/utils/iou_and_nms.py

def box_non_max_suppression(
    predictions: np.ndarray,
    iou_threshold: float = 0.5,
    overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> np.ndarray:
    """
    Perform Non-Maximum Suppression (NMS) on object detection predictions.

    Args:
        predictions (np.ndarray): An array of object detection predictions in
            the format of `(x_min, y_min, x_max, y_max, score)`
            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.
        overlap_metric (OverlapMetric): Metric used for matching detections in slices.

    Returns:
        np.ndarray: A boolean array indicating which predictions to keep after n
            on-maximum suppression.

    Raises:
        AssertionError: If `iou_threshold` is not within the
            closed range from `0` to `1`.
    """
    assert 0 <= iou_threshold <= 1, (
        "Value of `iou_threshold` must be in the closed range from 0 to 1, "
        f"{iou_threshold} given."
    )
    rows, columns = predictions.shape

    # add column #5 - category filled with zeros for agnostic nms
    if columns == 5:
        predictions = np.c_[predictions, np.zeros(rows)]

    # sort predictions column #4 - score
    sort_index = np.flip(predictions[:, 4].argsort())
    predictions = predictions[sort_index]

    boxes = predictions[:, :4]
    categories = predictions[:, 5]
    ious = box_iou_batch(boxes, boxes, overlap_metric)
    ious = ious - np.eye(rows)

    keep = np.ones(rows, dtype=bool)

    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()]

mask_non_max_suppression

supervision.detection.utils.iou_and_nms.mask_non_max_suppression(predictions, masks, iou_threshold=0.5, overlap_metric=OverlapMetric.IOU, mask_dimension=640)

Perform Non-Maximum Suppression (NMS) on segmentation predictions.

Parameters:

Name	Type	Description	Default
predictions	ndarray	A 2D array of object detection predictions in the format of (x_min, y_min, x_max, y_max, score) or (x_min, y_min, x_max, y_max, score, class). Shape: (N, 5) or (N, 6), where N is the number of predictions.	required
masks	ndarray	A 3D array of binary masks corresponding to the predictions. Shape: (N, H, W), where N is the number of predictions, and H, W are the dimensions of each mask.	required
iou_threshold	float	The intersection-over-union threshold to use for non-maximum suppression.	0.5
overlap_metric	OverlapMetric	Metric used for matching detections in slices.	IOU
mask_dimension	int	The dimension to which the masks should be resized before computing IOU values. Defaults to 640.	640

Returns:

Type	Description
ndarray	np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.

Raises:

Type	Description
AssertionError	If iou_threshold is not within the closed range from 0 to 1.

Source code in supervision/detection/utils/iou_and_nms.py

def mask_non_max_suppression(
    predictions: np.ndarray,
    masks: np.ndarray,
    iou_threshold: float = 0.5,
    overlap_metric: OverlapMetric = OverlapMetric.IOU,
    mask_dimension: int = 640,
) -> np.ndarray:
    """
    Perform Non-Maximum Suppression (NMS) on segmentation predictions.

    Args:
        predictions (np.ndarray): A 2D array of object detection predictions in
            the format of `(x_min, y_min, x_max, y_max, score)`
            or `(x_min, y_min, x_max, y_max, score, class)`. Shape: `(N, 5)` or
            `(N, 6)`, where N is the number of predictions.
        masks (np.ndarray): A 3D array of binary masks corresponding to the predictions.
            Shape: `(N, H, W)`, where N is the number of predictions, and H, W are the
            dimensions of each mask.
        iou_threshold (float): The intersection-over-union threshold
            to use for non-maximum suppression.
        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.

    Returns:
        np.ndarray: A boolean array indicating which predictions to keep after
            non-maximum suppression.

    Raises:
        AssertionError: If `iou_threshold` is not within the closed
            range from `0` to `1`.
    """
    assert 0 <= iou_threshold <= 1, (
        "Value of `iou_threshold` must be in the closed range from 0 to 1, "
        f"{iou_threshold} given."
    )
    rows, columns = predictions.shape

    if columns == 5:
        predictions = np.c_[predictions, np.zeros(rows)]

    sort_index = predictions[:, 4].argsort()[::-1]
    predictions = predictions[sort_index]
    masks = masks[sort_index]
    masks_resized = resize_masks(masks, mask_dimension)
    ious = mask_iou_batch(masks_resized, masks_resized, overlap_metric)
    categories = predictions[:, 5]

    keep = np.ones(rows, dtype=bool)
    for i in range(rows):
        if keep[i]:
            condition = (ious[i] > iou_threshold) & (categories[i] == categories)
            keep[i + 1 :] = np.where(condition[i + 1 :], False, keep[i + 1 :])

    return keep[sort_index.argsort()]

box_non_max_merge

supervision.detection.utils.iou_and_nms.box_non_max_merge(predictions, iou_threshold=0.5, overlap_metric=OverlapMetric.IOU)

Apply greedy version of non-maximum merging per category to avoid detecting too many overlapping bounding boxes for a given object.

Parameters:

Name	Type	Description	Default
predictions	NDArray[float64]	An array of shape (n, 5) or (n, 6) containing the bounding boxes coordinates in format [x1, y1, x2, y2], the confidence scores and class_ids. Omit class_id column to allow detections of different classes to be merged.	required
iou_threshold	float	The intersection-over-union threshold to use for non-maximum suppression. Defaults to 0.5.	0.5
overlap_metric	OverlapMetric	Metric used for matching detections in slices.	IOU

Returns:

Type	Description
list[list[int]]	list[list[int]]: Groups of prediction indices be merged. Each group may have 1 or more elements.

Source code in supervision/detection/utils/iou_and_nms.py

def box_non_max_merge(
    predictions: npt.NDArray[np.float64],
    iou_threshold: float = 0.5,
    overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> list[list[int]]:
    """
    Apply greedy version of non-maximum merging per category to avoid detecting
    too many overlapping bounding boxes for a given object.

    Args:
        predictions (npt.NDArray[np.float64]): An array of shape `(n, 5)` or `(n, 6)`
            containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`,
            the confidence scores and class_ids. Omit class_id column to allow
            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.
        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, 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, overlap_metric
        )

        for merge_class_group in merge_class_groups:
            merge_groups.append(curr_indices[merge_class_group].tolist())

    for merge_group in merge_groups:
        if len(merge_group) == 0:
            raise ValueError(
                f"Empty group detected when non-max-merging detections: {merge_groups}"
            )
    return merge_groups

mask_non_max_merge

supervision.detection.utils.iou_and_nms.mask_non_max_merge(predictions, masks, iou_threshold=0.5, mask_dimension=640, overlap_metric=OverlapMetric.IOU)

Perform Non-Maximum Merging (NMM) on segmentation predictions.

Parameters:

Name	Type	Description	Default
predictions	ndarray	A 2D array of object detection predictions in the format of (x_min, y_min, x_max, y_max, score) or (x_min, y_min, x_max, y_max, score, class). Shape: (N, 5) or (N, 6), where N is the number of predictions.	required
masks	ndarray	A 3D array of binary masks corresponding to the predictions. Shape: (N, H, W), where N is the number of predictions, and H, W are the dimensions of each mask.	required
iou_threshold	float	The intersection-over-union threshold to use for non-maximum suppression.	0.5
mask_dimension	int	The dimension to which the masks should be resized before computing IOU values. Defaults to 640.	640
overlap_metric	OverlapMetric	Metric used for matching detections in slices.	IOU

Returns:

Type	Description
list[list[int]]	np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.

Raises:

Type	Description
AssertionError	If iou_threshold is not within the closed range from 0 to 1.

Source code in supervision/detection/utils/iou_and_nms.py

def mask_non_max_merge(
    predictions: np.ndarray,
    masks: np.ndarray,
    iou_threshold: float = 0.5,
    mask_dimension: int = 640,
    overlap_metric: OverlapMetric = OverlapMetric.IOU,
) -> list[list[int]]:
    """
    Perform Non-Maximum Merging (NMM) on segmentation predictions.

    Args:
        predictions (np.ndarray): A 2D array of object detection predictions in
            the format of `(x_min, y_min, x_max, y_max, score)`
            or `(x_min, y_min, x_max, y_max, score, class)`. Shape: `(N, 5)` or
            `(N, 6)`, where N is the number of predictions.
        masks (np.ndarray): A 3D array of binary masks corresponding to the predictions.
            Shape: `(N, H, W)`, where N is the number of predictions, and H, W are the
            dimensions of each mask.
        iou_threshold (float): The intersection-over-union threshold
            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.
        overlap_metric (OverlapMetric): Metric used for matching detections in slices.

    Returns:
        np.ndarray: A boolean array indicating which predictions to keep after
            non-maximum suppression.

    Raises:
        AssertionError: If `iou_threshold` is not within the closed
            range from `0` to `1`.
    """
    masks_resized = resize_masks(masks, mask_dimension)
    if predictions.shape[1] == 5:
        return _group_overlapping_masks(
            predictions, masks_resized, 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_masks(
            predictions[curr_indices],
            masks_resized[curr_indices],
            iou_threshold,
            overlap_metric,
        )

        for merge_class_group in merge_class_groups:
            merge_groups.append(curr_indices[merge_class_group].tolist())

    for merge_group in merge_groups:
        if len(merge_group) == 0:
            raise ValueError(
                f"Empty group detected when non-max-merging detections: {merge_groups}"
            )
    return merge_groups

Comments