refactor(detection): 重构合并检测对象的逻辑

- 移除 merge_inner_detections_objects 函数中的 threshold 和 match_metric 参数
- 新增 merge_inner_detections_objects_without_iou 函数，用于合并检测对象
- 优化非最大值抑制（NMM）的逻辑，提高效率
- 调整 mask_iou_batch 函数的实现，简化代码

supervision/detection/core.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass, field
 from enum import Enum
+from functools import reduce
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import numpy as np
@@ -1455,9 +1456,7 @@ class Detections:
         result = []
         for merge_group in merge_groups:
             unmerged_detections = [self[i] for i in merge_group]
-            merged_detections = merge_inner_detections_objects(
-                unmerged_detections, threshold, match_metric
-            )
+            merged_detections = merge_inner_detections_objects_without_iou(unmerged_detections)
             result.append(merged_detections)
 
         return Detections.merge(result)
@@ -1577,6 +1576,18 @@ def merge_inner_detections_objects(
         detections_1 = merge_inner_detection_object_pair(detections_1, detections_2)
     return detections_1
 
+def merge_inner_detections_objects_without_iou(
+    detections: List[Detections]
+) -> Detections:
+    """
+    Given N detections each of length 1 (exactly one object inside), combine them into a
+    single detection object of length 1. The contained inner object will be the merged
+    result of all the input detections.
+
+    For example, this lets you merge N boxes into one big box, N masks into one mask,
+    etc.
+    """
+    return reduce(merge_inner_detection_object_pair, detections)
 
 def validate_fields_both_defined_or_none(
     detections_1: Detections, detections_2: Detections

supervision/detection/overlap_filter.py

@@ -91,7 +91,7 @@ def mask_non_max_suppression(
     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 :])
+            keep[i + 1:] = np.where(condition[i + 1:], False, keep[i + 1:])
 
     return keep[sort_index.argsort()]
 
@@ -206,7 +206,7 @@ def mask_non_max_merge(
     iou_threshold: float = 0.5,
     mask_dimension: int = 640,
     match_metric: str = "IOU",
-) -> np.ndarray:
+) -> List[List[int]]:
     """
     Perform Non-Maximum Merging (NMM) on segmentation predictions.
 
@@ -342,14 +342,26 @@ def group_overlapping_masks(
             merge_groups.append([idx])
             break
 
-        merge_candidate = np.expand_dims(masks[idx], axis=0)
-        ious = mask_iou_batch(masks[order], merge_candidate, match_metric)
-        ious = ious.flatten()
+        # calculate mask
+        merge_candidate = masks[idx][None, ...]
+        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)
+            above_threshold: np.ndarray = ious.flatten() >= iou_threshold
+            # if no mask is above threshold, break
+            if not above_threshold.any():
+                break
+            # get indexes that meet the threshold
+            above_idx = order[above_threshold]
+            # update merge_candidate
+            merge_candidate = np.logical_or.reduce(np.concatenate([masks[above_idx], merge_candidate]), axis=0, keepdims=True)
+            # add indexes that meet the criteria to the candidate_groups
+            candidate_groups.extend(np.flip(above_idx).tolist())
+            # update order, masks
+            order = order[~above_threshold]
 
-        above_threshold = ious >= iou_threshold
-        merge_group = [idx, *np.flip(order[above_threshold]).tolist()]
-        merge_groups.append(merge_group)
-        order = order[~above_threshold]
+        merge_groups.append(candidate_groups)
     return merge_groups
 
 

supervision/detection/utils.py

@@ -49,7 +49,9 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n
 
 
 def box_iou_batch(
-    boxes_true: np.ndarray, boxes_detection: np.ndarray, match_metric: str = "IOU"
+    boxes_true: np.ndarray,
+    boxes_detection: np.ndarray,
+    match_metric: str = "IOU",
 ) -> np.ndarray:
     """
     Compute Intersection over Union (IoU) of two sets of bounding boxes -
@@ -107,7 +109,9 @@ def box_iou_batch(
 
 
 def _mask_iou_batch_split(
-    masks_true: np.ndarray, masks_detection: np.ndarray, match_metric: str = "IOU"
+    masks_true: np.ndarray,
+    masks_detection: np.ndarray,
+    match_metric: str = "IOU",
 ) -> np.ndarray:
     """
     Internal function.
@@ -123,12 +127,10 @@ def _mask_iou_batch_split(
     Returns:
         np.ndarray: Pairwise IoU of masks from `masks_true` and `masks_detection`.
     """
-    intersection_area = np.logical_and(masks_true[:, None], masks_detection).sum(
-        axis=(2, 3)
-    )
+    intersection_area = np.logical_and(masks_true[:, None], masks_detection).sum(axis=(2, 3))
 
-    masks_true_area = masks_true.sum(axis=(1, 2))
-    masks_detection_area = masks_detection.sum(axis=(1, 2))
+    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":
         union_area = masks_true_area[:, None] + masks_detection_area - intersection_area
@@ -139,7 +141,8 @@ def _mask_iou_batch_split(
             where=union_area != 0,
         )
     elif match_metric.upper() == "IOS":
-        small_area = np.minimum(intersection_area, masks_detection_area)
+        # ios = intersection_area / min(area1, area2)
+        small_area = np.minimum(masks_true_area[:, None], masks_detection_area)
         ious = np.divide(
             intersection_area,
             small_area,
@@ -176,12 +179,12 @@ def mask_iou_batch(
         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
+            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, match_metric)
@@ -192,16 +195,16 @@ def mask_iou_batch(
         * 1024
         * 1024
         // (
-            masks_detection.shape[0]
-            * masks_detection.shape[1]
-            * masks_detection.shape[2]
+                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, match_metric
+                masks_true[i: i + step], masks_detection, match_metric
             )
         )
 
@@ -581,7 +584,7 @@ def filter_polygons_by_area(
         polygon
         for polygon, area in zip(polygons, ares)
         if (min_area is None or area >= min_area)
-        and (max_area is None or area <= max_area)
+           and (max_area is None or area <= max_area)
     ]
 
 
@@ -918,9 +921,9 @@ def move_masks(
 
     if source_x_end > source_x_start and source_y_end > source_y_start:
         mask_array[
-            :,
-            destination_y_start:destination_y_end,
-            destination_x_start:destination_x_end,
+        :,
+        destination_y_start:destination_y_end,
+        destination_x_start:destination_x_end,
         ] = masks[:, source_y_start:source_y_end, source_x_start:source_x_end]
 
     return mask_array
@@ -1022,11 +1025,11 @@ def is_metadata_equal(metadata_a: Dict[str, Any], metadata_b: Dict[str, Any]) ->
     return set(metadata_a.keys()) == set(metadata_b.keys()) and all(
         np.array_equal(metadata_a[key], metadata_b[key])
         if (
-            isinstance(metadata_a[key], np.ndarray)
-            and isinstance(metadata_b[key], np.ndarray)
+                isinstance(metadata_a[key], np.ndarray)
+                and isinstance(metadata_b[key], np.ndarray)
         )
         else metadata_a[key] == metadata_b[key]
-        for key in metadata_a
+            for key in metadata_a
     )