Python中如何快速去除高度重叠的目标检测边界框

原代码性能差的核心原因 #

• 时间复杂度高：双层嵌套循环的时间复杂度为O(n²)，且每个框对会被重复计算2次，无意义开销大
• 无分组逻辑：不同图片的边界框不需要计算IoU，原代码未按Image name分组，产生大量无效计算
• 遍历方式低效：iterrows是pandas中最慢的遍历方式之一，性能远低于numpy向量化操作
• 频繁修改DataFrame：多次调用df.drop(inplace=True)会反复修改DataFrame结构，额外开销极高

优化方案 #

方案1：使用底层优化的现成NMS接口（推荐，性能提升上千倍） #

非极大值抑制（NMS）是目标检测的通用算子，现有工具库的实现都是C++底层优化，2400条标注只需几毫秒即可处理完成，以下是两种常用实现示例：

基于torchvision的实现

import pandas as pd
import torch
from torchvision.ops import nms

result = []
# 按图片名称分组，仅同一张图的框计算IoU
for img_name, group_df in df.groupby('Image name'):
    # 转换为张量，CPU运行即可获得极高性能
    boxes = torch.tensor(group_df[['xmin', 'ymin', 'xmax', 'ymax']].values, dtype=torch.float32)
    scores = torch.tensor(group_df['confidence'].values, dtype=torch.float32)
    # 传入IoU阈值，得到保留的框索引
    keep_idx = nms(boxes, scores, iou_threshold=0.5)
    result.append(group_df.iloc[keep_idx.numpy()])

# 合并所有图片的结果得到最终去重数据
df_clean = pd.concat(result, ignore_index=True)

基于OpenCV的实现（无需安装PyTorch）

import pandas as pd
import cv2
import numpy as np

result = []
for img_name, group_df in df.groupby('Image name'):
    boxes = group_df[['xmin', 'ymin', 'xmax', 'ymax']].values
    # OpenCV的NMS要求输入格式为[xmin, ymin, width, height]，先做格式转换
    boxes_cv = np.hstack([boxes[:, :2], (boxes[:, 2:] - boxes[:, :2])])
    scores = group_df['confidence'].values
    # 阈值设为0表示不过滤低置信度框，仅按IoU去重
    keep_idx = cv2.dnn.NMSBoxes(boxes_cv.tolist(), scores.tolist(), score_threshold=0, nms_threshold=0.5)
    result.append(group_df.iloc[keep_idx.flatten()])

df_clean = pd.concat(result, ignore_index=True)

方案2：手动实现优化版NMS（无额外依赖） #

如果不想引入第三方库的NMS实现，可按以下逻辑优化手动实现：

1. 先按置信度降序排序所有框
2. 从最高置信度的框开始遍历，每次保留当前框，再删除所有和当前框IoU≥0.5的低置信度框
3. 所有计算使用numpy向量化操作，避免循环遍历DataFrame

import pandas as pd
import numpy as np

def calculate_iou(boxes, box):
    # boxes是所有框的数组，shape为(n,4)，box是当前对比的框，shape为(4,)，向量化计算所有IoU
    x1 = np.maximum(boxes[:, 0], box[0])
    y1 = np.maximum(boxes[:, 1], box[1])
    x2 = np.minimum(boxes[:, 2], box[2])
    y2 = np.minimum(boxes[:, 3], box[3])
    inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    area_box = (box[2] - box[0]) * (box[3] - box[1])
    area_boxes = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    iou = inter / (area_box + area_boxes - inter + 1e-6)
    return iou

result = []
for img_name, group_df in df.groupby('Image name'):
    # 按置信度降序排序
    group_df = group_df.sort_values('confidence', ascending=False).reset_index(drop=True)
    boxes = group_df[['xmin', 'ymin', 'xmax', 'ymax']].values
    keep = []
    while len(boxes) > 0:
        # 保留第一个（置信度最高）的框
        keep.append(0)
        # 计算剩下所有框和当前框的IoU
        ious = calculate_iou(boxes[1:], boxes[0])
        # 只保留IoU<0.5的框进入下一轮
        boxes = boxes[1:][ious < 0.5]
    # 取出保留的行
    result.append(group_df.iloc[keep])

df_clean = pd.concat(result, ignore_index=True)

内容的提问来源于stack exchange，提问作者Mohammad M. Sa'adoun