非最大值抑制（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.4
def 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] # x
boxes_xywh[:, 1] = boxes[:, 1] # y
boxes_xywh[:, 2] = boxes[:, 2] - boxes[:, 0] # w
boxes_xywh[:, 3] = boxes[:, 3] - boxes[:, 1] # h
return boxes_xywh
def 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 np
def 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 np
scores = 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 = box
width = x2 - x1
height = 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 patches
def plot_iou(boxA, boxB):
fig, ax = plt.subplots(1, figsize=(8, 8))
# 绘制 Box A
x1A, y1A, x2A, y2A = boxA
widthA = x2A - x1A
heightA = y2A - y1A
rectA = patches.Rectangle((x1A, y1A), widthA, heightA, linewidth=2, edgecolor='blue', facecolor='blue', label='Box A')
ax.add_patch(rectA)
# 绘制 Box B
x1B, y1B, x2B, y2B = boxB
widthB = x2B - x1B
heightB = y2B - y1B
rectB = 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
# 计算 IoU
iou = 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)
![[leetcode]24-game](https://img-blog.csdnimg.cn/direct/849046fed9904a6c8ae7301fe55eef2a.png)
