非最大值抑制（NMS）函数
flyfish
非最大值抑制（Non-Maximum Suppression, NMS）是计算机视觉中常用的一种后处理技术，主要用于目标检测任务。其作用是从一组可能存在大量重叠的候选边界框中，筛选出最具代表性的边界框，即通过置信度分数和重叠区域的过滤，保留最具代表性的边界框。
边界框（Bounding Boxes）：一组表示候选目标区域的矩形框，每个框由左上角和右下角的坐标（x1, y1, x2, y2）表示。
置信度分数（Confidence Scores）：每个边界框对应的一个置信度分数，表示该框内包含目标的可能性。
执行步骤
初始化：
boxes：输入的边界框列表。
scores：每个边界框对应的置信度得分列表。
confidence_threshold：过滤边界框的最低置信度阈值。
iou_threshold：用于确定边界框是否重叠的 IOU 阈值。
过滤低置信度边界框：
根据 confidence_threshold 过滤掉置信度低于该阈值的边界框。
按置信度排序：
对剩余的边界框按照置信度从高到低排序。
非极大值抑制：
从排序后的列表中选择置信度最高的边界框，并计算其与其他边界框的 Intersection-over-Union (IoU)。
如果 IoU大于 iou_threshold，则移除该边界框（表示重叠太多）。
重复该过程直到处理完所有边界框。
返回结果：
返回保留的边界框的索引。
可视化 Intersection-over-Union (IoU)
蓝色矩形表示 Box A，红色矩形表示 Box B，绿色矩形表示它们的交集区域,剩余的红色和蓝色是并集区域。
torchvision.ops.nms 和 cv2.dnn.NMSBoxes 的调用
import numpy as np
import torch
import torchvision.ops as ops
import cv2# 输入数据
boxes = np.array([[100, 100, 210, 210], [220, 220, 320, 330], [300, 300, 400, 400],[50, 50, 150, 200], [200, 150, 280, 320], [280, 280, 380, 380],[80, 90, 190, 210], [250, 250, 350, 370], [290, 290, 390, 390]
])# (x1, y1, x2, y2)格式
scores = np.array([0.9, 0.8, 0.75, 0.85, 0.7, 0.65, 0.82, 0.78, 0.6])
score_threshold = 0.5
nms_threshold = 0.4def convert_to_xywh(boxes): #opencv用 (x, y, w, h)格式"""将边界框从 (x1, y1, x2, y2) 格式转换为 (x, y, w, h) 格式。参数:- boxes: 形状为 (N, 4) 的数组，其中 N 是边界框的数量返回:- boxes_xywh: 形状为 (N, 4) 的数组，包含转换后的边界框"""boxes_xywh = np.zeros_like(boxes)boxes_xywh[:, 0] = boxes[:, 0] # xboxes_xywh[:, 1] = boxes[:, 1] # yboxes_xywh[:, 2] = boxes[:, 2] - boxes[:, 0] # wboxes_xywh[:, 3] = boxes[:, 3] - boxes[:, 1] # hreturn boxes_xywhdef nms_torchvision(boxes, scores, nms_threshold):boxes_tensor = torch.tensor(boxes, dtype=torch.float32)scores_tensor = torch.tensor(scores, dtype=torch.float32)keep = ops.nms(boxes_tensor, scores_tensor, nms_threshold)return keep.numpy()def nms_opencv(boxes, scores, score_threshold, nms_threshold):boxes = convert_to_xywh(boxes)indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), score_threshold, nms_threshold)return np.array(indices).flatten()# 调用 NMS
keep_torchvision = nms_torchvision(boxes, scores, nms_threshold)
keep_opencv = nms_opencv(boxes, scores, score_threshold, nms_threshold)print("使用 torchvision.ops.nms 保留的边界框索引: ", keep_torchvision)
print("使用 cv2.dnn.NMSBoxes 保留的边界框索引: ", keep_opencv)
输出
使用 torchvision.ops.nms 保留的边界框索引: [0 3 1 7 2 4]
使用 cv2.dnn.NMSBoxes 保留的边界框索引: [0 3 1 7 2 4]
用纯 NumPy 实现的非最大值抑制（NMS）函数
import numpy as npdef nms(boxes, scores, score_threshold, nms_threshold):"""单类 NMS 使用 NumPy 实现。"""# 过滤掉低置信度的框indices = np.where(scores > score_threshold)[0]boxes = boxes[indices]scores = scores[indices]# 如果没有剩余的框，返回空列表if len(boxes) == 0:return []# 提取每个边界框的坐标x1 = boxes[:, 0]y1 = boxes[:, 1]x2 = boxes[:, 2]y2 = boxes[:, 3]# 计算每个边界框的面积areas = (x2 - x1 + 1) * (y2 - y1 + 1)# 根据分数进行排序（从高到低）order = scores.argsort()[::-1]keep = []while order.size > 0:i = order[0]keep.append(indices[i])# 计算当前边界框与其余边界框的交集坐标xx1 = np.maximum(x1[i], x1[order[1:]])yy1 = np.maximum(y1[i], y1[order[1:]])xx2 = np.minimum(x2[i], x2[order[1:]])yy2 = np.minimum(y2[i], y2[order[1:]])# 计算交集的宽度和高度w = np.maximum(0.0, xx2 - xx1 + 1)h = np.maximum(0.0, yy2 - yy1 + 1)# 计算交集面积inter = w * h# 计算交并比（IOU）ovr = inter / (areas[i] + areas[order[1:]] - inter)# 只保留 IOU 小于阈值的边界框inds = np.where(ovr <= nms_threshold)[0]order = order[inds + 1]return keep# 示例数据
boxes = np.array([[100, 100, 210, 210], [220, 220, 320, 330], [300, 300, 400, 400],[50, 50, 150, 200], [200, 150, 280, 320], [280, 280, 380, 380],[80, 90, 190, 210], [250, 250, 350, 370], [290, 290, 390, 390]
])
scores = np.array([0.9, 0.8, 0.75, 0.85, 0.7, 0.65, 0.82, 0.78, 0.6])
score_threshold = 0.5
nms_threshold = 0.4# 调用NMS
keep_indices = nms(boxes, scores, score_threshold, nms_threshold)
print("使用 NumPy 实现的 NMS 保留的边界框索引: ", keep_indices)
使用 NumPy 实现的 NMS 保留的边界框索引: [0, 3, 1, 7, 2, 4]
关于语法的解释
在 NumPy 中，冒号 : 用于数组切片。它们可以用来提取数组的子集、重排数组或选取特定的元素。
示例1
scores.argsort()[::-1]
scores.argsort()：返回 scores 中元素的索引数组，这些索引会将 scores 排序。
[::-1]：表示反转数组。
在这个例子中，[::-1] 表示从开始到结束，步长为 -1，因此数组会被反转。这里的两个冒号是为了清楚地表示切片的完整语法 [start:stop:step]，其中省略了 start 和 stop，只指定了 step 为 -1。
import numpy as npscores = np.array([0.9, 0.8, 0.75, 0.85, 0.7, 0.65, 0.82, 0.78, 0.6])
sorted_indices = scores.argsort() # 升序排序的索引
print("sorted_indices:", sorted_indices)# 反转排序索引（降序排序）
reversed_indices = sorted_indices[::-1]
print("reversed_indices:", reversed_indices)
sorted_indices: [8 5 4 2 7 1 6 3 0]
reversed_indices: [0 3 6 1 7 2 4 5 8]
示例2
boxes[:, 0]
boxes[:, 0]：选取 boxes 数组中第 0 列的所有元素。
: 表示选择所有行，0 表示选择第 0 列。
这段代码的作用是提取 boxes 数组中每个边界框的 x1 坐标（左上角的 x 坐标）。
import numpy as np
boxes = np.array([[100, 100, 210, 210],[220, 220, 320, 330],[300, 300, 400, 400],[50, 50, 150, 200]
])x1 = boxes[:, 0]
print("x1:", x1)
x1: [100 220 300 50]
可视化数据的代码
def plot_boxes(boxes, keep_indices):fig, ax = plt.subplots(1, figsize=(12, 12))for i, box in enumerate(boxes):x1, y1, x2, y2 = boxwidth = x2 - x1height = y2 - y1# 所有输入框用蓝色绘制edgecolor = 'blue'if i in keep_indices:# NMS 保留的框用绿色绘制edgecolor = 'green'else:# 被抑制的框用红色绘制edgecolor = 'red'rect = patches.Rectangle((x1, y1), width, height, linewidth=2, edgecolor=edgecolor, facecolor='none')ax.add_patch(rect)# 设置坐标范围ax.set_xlim(0, np.max(boxes[:, [0, 2]]) + 50)ax.set_ylim(0, np.max(boxes[:, [1, 3]]) + 50)ax.invert_yaxis() # 图像坐标系和实际坐标系相反时需要plt.show()# 示例数据
boxes = np.array([[100, 100, 210, 210], [220, 220, 320, 330], [300, 300, 400, 400],[50, 50, 150, 200], [200, 150, 280, 320], [280, 280, 380, 380],[80, 90, 190, 210], [250, 250, 350, 370], [290, 290, 390, 390]
])
scores = np.array([0.9, 0.8, 0.75, 0.85, 0.7, 0.65, 0.82, 0.78, 0.6])
score_threshold = 0.5
nms_threshold = 0.4# 调用NMS
keep_indices = nms(boxes, scores, score_threshold, nms_threshold)
print("使用 NumPy 实现的 NMS 保留的边界框索引: ", keep_indices)# 绘图
plot_boxes(boxes, keep_indices)
可视化 Intersection-over-Union (IoU)的代码
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patchesdef plot_iou(boxA, boxB):fig, ax = plt.subplots(1, figsize=(8, 8))# 绘制 Box Ax1A, y1A, x2A, y2A = boxAwidthA = x2A - x1AheightA = y2A - y1ArectA = patches.Rectangle((x1A, y1A), widthA, heightA, linewidth=2, edgecolor='blue', facecolor='blue', label='Box A')ax.add_patch(rectA)# 绘制 Box Bx1B, y1B, x2B, y2B = boxBwidthB = x2B - x1BheightB = y2B - y1BrectB = patches.Rectangle((x1B, y1B), widthB, heightB, linewidth=2, edgecolor='red', facecolor='red', label='Box B')ax.add_patch(rectB)# 计算交集xx1 = np.maximum(x1A, x1B)yy1 = np.maximum(y1A, y1B)xx2 = np.minimum(x2A, x2B)yy2 = np.minimum(y2A, y2B)w = np.maximum(0, xx2 - xx1)h = np.maximum(0, yy2 - yy1)intersection_area = w * h# 计算并集areaA = (x2A - x1A) * (y2A - y1A)areaB = (x2B - x1B) * (y2B - y1B)union_area = areaA + areaB - intersection_area# 计算 IoUiou = intersection_area / union_area# 绘制交集if w > 0 and h > 0:rect_intersection = patches.Rectangle((xx1, yy1), w, h, linewidth=2, edgecolor='green', facecolor='green', linestyle='--', label='Intersection')ax.add_patch(rect_intersection)# 显示图例handles, labels = ax.get_legend_handles_labels()plt.legend(handles=handles)plt.xlim(0, 500)plt.ylim(0, 500)plt.gca().set_aspect('equal', adjustable='box')plt.title(f'IoU = {iou:.2f}')plt.show()# 示例框
boxA = [100, 100, 300, 300]
boxB = [200, 200, 400, 400]plot_iou(boxA, boxB)
