diff --git a/supervision/detection/core.py b/supervision/detection/core.py index c428356a..3b343b46 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -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, diff --git a/supervision/detection/overlap_filter.py b/supervision/detection/overlap_filter.py index 9cff643f..34bfb024 100644 --- a/supervision/detection/overlap_filter.py +++ b/supervision/detection/overlap_filter.py @@ -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." - ) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index b44cd2b6..63670144 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -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 ) ) diff --git a/test/detection/test_overlap_filter.py b/test/detection/test_overlap_filter.py index 08cd870f..f628c30f 100644 --- a/test/detection/test_overlap_filter.py +++ b/test/detection/test_overlap_filter.py @@ -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