From 13e857cf0010eb2d1e7274101bc429e37647befd Mon Sep 17 00:00:00 2001 From: AI-Hao <2865467769@qq.com> Date: Mon, 7 Jul 2025 15:21:28 +0800 Subject: [PATCH] =?UTF-8?q?refactor(detection):=20=E9=87=8D=E6=9E=84?= =?UTF-8?q?=E5=90=88=E5=B9=B6=E6=A3=80=E6=B5=8B=E5=AF=B9=E8=B1=A1=E7=9A=84?= =?UTF-8?q?=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 移除 merge_inner_detections_objects 函数中的 threshold 和 match_metric 参数 - 新增 merge_inner_detections_objects_without_iou 函数，用于合并检测对象 - 优化非最大值抑制（NMM）的逻辑，提高效率 - 调整 mask_iou_batch 函数的实现，简化代码 --- supervision/detection/core.py | 17 ++++++-- supervision/detection/overlap_filter.py | 30 +++++++++----- supervision/detection/utils.py | 53 +++++++++++++------------ 3 files changed, 63 insertions(+), 37 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 5dcb46ea..ab9bddbb 100644 --- a/supervision/detection/core.py +++ b/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 diff --git a/supervision/detection/overlap_filter.py b/supervision/detection/overlap_filter.py index dbe1159a..b5aecbe2 100644 --- a/supervision/detection/overlap_filter.py +++ b/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 diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index a915c96d..8bfdce84 100644 --- a/supervision/detection/utils.py +++ b/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 )