本文展示了两种非极大值抑制（NMS）的Python实现，一种避免了使用for循环，另一种则在计算IoU时使用了for循环。这两种方法都是为了处理图像检测中的边界框重叠问题，通过计算IoU来决定保留哪些边界框。代码包括了NMS算法的详细步骤，并在最后给出了运行时间的比较。
代码结合了 NMS 和 IOU 两个部分。
下面展示 代码。
思路一
// JasonL
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time
'''
不同于“result_for_loop”脚本
这个脚本在遍历bbox2的时候 不使用 for loop
'''
def covert_wh(x, y, w, h) -> float:
'''
Args:
w: wide
h: height
Return:
The float values indicate the two coordinate axes of the rectangle (Diagonal coordinates)
'''
xmin, ymin = int(x - w / 2.0), int(y - h / 2.0)
xmax, ymax = int(x + w / 2.0), int(y + h / 2.0)
return xmin, ymin, xmax, ymax
def iou(bbox1, bbox2) -> float:
'''
Args:
bbox1: a list of some attributes of a boundingbox (bbox), bbox.format = [conf, class, x_center, y_center, w, h]
bbox2: another list of some attributes of a boundingbox (bbox), bbox.format = [conf, class, x_center, y_center, w, h]
Return:
a float value indicate the Intersection over Union (iou) between two inputs.
'''
_, _, x1, y1, w1, h1 = bbox1
xmin1, ymin1, xmax1, ymax1 = covert_wh(x1, y1, w1, h1)
areas1 = (xmax1 - xmin1 + 1) * (ymax1 - ymin1 + 1) # 计算面积
x2, y2 = bbox2[:, 2], bbox2[:, 3]
w2, h2 = bbox2[:, 4], bbox2[:, 5]
xmin2, ymin2 = x2 - w2 / 2.0, y2 - h2 / 2.0
xmax2, ymax2 = x2 + w2 / 2.0, y2 + h2 / 2.0
areas2 = (xmax2 - xmin2 + 1) * (ymax2 - ymin2 + 1) # 计算面积
# 将得分最高的边框与剩余边框进行比较
xx1 = np.maximum(xmin1, xmin2)
yy1 = np.maximum(ymin1, ymin2)
xx2 = np.minimum(xmax1, xmax2)
yy2 = np.minimum(ymax1, ymax2)
# 计算交集
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
intersection = w * h
# 计算并集
union = areas1 + areas2 - intersection
# 计算交并比
iou_ans = intersection / union
return iou_ans
def nms(bbox_ls, iou_threshold = 0.5) -> list:
'''
Args:
bbox_ls: a list of boundingboxes (bbox), shuffled order
bbox: a list of some attributes, bbox = [conf, class, x_center, y_center, w, h]
iou_threshold: a float indicate the iou iou_threshold to judge whether two boundingboxes are overlapped
Return:
a new list of bboxed after non-maximum-suppression
'''
conf = bbox_ls[:,0]
order = conf.argsort()[::-1] # 按score降序排序，argsort返回降序后的索引
record_ans = [] # 用于记录符合要求的值
while order.size > 0:
if order.size == 1:
top1_idx = order[0] # 选取得分最高的边框
record_ans.append(top1_idx) # 添加到候选列表
break
else:
top1_idx = order[0]
record_ans.append(top1_idx)
# bbox1 bbox2 操作
bbox1 = bbox_ls[top1_idx] # 放入iou进行对比
bbox2 = bbox_ls[order[1:]] # 与bbox1进行对比
iou_ans = iou(bbox1, bbox2) # 返回iou计算数值
inds = np.where(iou_ans <= iou_threshold)[0] # np.where筛选符合条件的选项
order = order[inds + 1] # 因为第一项是top1_idx，不在temp_order中，导入order需要修补索引之间的差值
return record_ans
if __name__ == '__main__':
time_begin = time.time()
#读取图片（任意放一张图片到当前目录命名为“img.jpg”即可）
img = cv2.imread("img.jpg")
img_cp = np.copy(img)
thickness = 2
# 模拟标注框的参数
# format = [conf, class, x_center, y_center, w, h]
info = np.array([
[0.95, 1, 150, 150, 100, 100],
[0.98, 1, 155, 155, 100, 100],
[0.96, 1, 145, 145, 110, 101],
[0.96, 2, 230, 230, 88, 99],
])
colors = [[0, 0, 255], [0, 255, 0], [255, 0, 0], [255, 255, 0]]
# 绘制使用NMS和IOS前的图片
plt.subplot(121)
plt.axis('off')
plt.title("Input image")
for i in range(len(colors)):
_, _, x_center, y_center, w, h = info[i]
x1, y1, x2, y2 = covert_wh(x_center, y_center, w, h)
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), colors[i], thickness=thickness)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img)
plt.subplot(122)
plt.axis('off')
plt.title("After NMS")
indx = nms(info)
for i in indx:
_, _, x_center, y_center, w, h = info[i]
x1, y1, x2, y2 = covert_wh(x_center, y_center, w, h)
cv2.rectangle(img_cp, (int(x1), int(y1)), (int(x2), int(y2)), colors[i], thickness=thickness)
img_cp = cv2.cvtColor(img_cp, cv2.COLOR_BGR2RGB)
plt.imshow(img_cp)
# 保存并显示图片
plt.savefig('results.png')
plt.show()
# 记录运行时间
time_end = time.time()
time = time_end - time_begin
print('time:', time)
另一种思路
调用iou的时候使用了一个for循环
// JasonL
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time
'''
不同于“result”脚本
这个脚本在遍历bbox2的时候 尝试使用 for loop
'''
def covert_wh(x, y, w, h) -> float:
'''
Args:
w: wide
h: height
Return:
The float values indicate the two coordinate axes of the rectangle (Diagonal coordinates)
'''
xmin, ymin = int(x - w / 2.0), int(y - h / 2.0)
xmax, ymax = int(x + w / 2.0), int(y + h / 2.0)
return xmin, ymin, xmax, ymax
def iou(bbox1, bbox2) -> float:
'''
Args:
bbox1: a list of some attributes of a boundingbox (bbox), bbox.format = [conf, class, x_center, y_center, w, h]
bbox2: another list of some attributes of a boundingbox (bbox), bbox.format = [conf, class, x_center, y_center, w, h]
Return:
a float value indicate the Intersection over Union (iou) between two inputs.
'''
_, _, x1, y1, w1, h1 = bbox1
xmin1, ymin1, xmax1, ymax1 = covert_wh(x1, y1, w1, h1)
_, _, x2, y2, w2, h2 = bbox2
xmin2, ymin2, xmax2, ymax2 = covert_wh(x2, y2, w2, h2)
# 计算交集的对角坐标点
xx1 = np.max([xmin1, xmin2])
yy1 = np.max([ymin1, ymin2])
xx2 = np.min([xmax1, xmax2])
yy2 = np.min([ymax1, ymax2])
# 计算交集面积
w = np.max([0.0, xx2 - xx1 + 1])
h = np.max([0.0, yy2 - yy1 + 1])
area_intersection = w * h
# 计算并集面积
area1 = (xmax1 - xmin1 + 1) * (ymax1 - ymin1 + 1)
area2 = (xmax2 - xmin2 + 1) * (ymax2 - ymin2 + 1)
area_union = area1 + area2 - area_intersection
# 计算两个边框的交并比
iou_ans = area_intersection / (area_union + 1e-8) # 避免分母为0
return iou_ans
def nms(bbox_ls, iou_threshold = 0.5) -> list:
'''
Args:
bbox_ls: a list of boundingboxes (bbox), shuffled order
bbox: a list of some attributes, bbox = [conf, class, x_center, y_center, w, h]
iou_threshold: a float indicate the iou iou_threshold to judge whether two boundingboxes are overlapped
Return:
a new list of bboxed after non-maximum-suppression
'''
# score_thresh=0.5 如果允许，score_thresh可以作为nms函数的参数，预先筛选掉一些比较差的数据。
conf = bbox_ls[:,0]
order = conf.argsort()[::-1] # 按conf降序排序，用argsort返回降序后的索引
record_ans = [] # 用于记录符合要求的值
while order.size > 0:
if order.size == 1:
top1_idx = order[0] # 选取得分最高的边框
record_ans.append(top1_idx) # 添加到候选列表
break
else:
top1_idx = order[0]
record_ans.append(top1_idx)
#
temp_record = [] # temp_record 用于存储iou结果
bbox1 = bbox_ls[top1_idx] # 放入iou进行对比
temp_order = order[1:] # temp_order 除了top1_idx之外的index
for index in temp_order: # 依次将temp_order中的数值与top1_idx进行iou计算
bbox2 = bbox_ls[index]
iou_ans = iou(bbox1, bbox2)
temp_record.append(iou_ans)
# 将计算的iou导出到temp_record，使用np.where筛选符合条件的选项
temp_record = np.array(temp_record)
inds = np.where(temp_record <= iou_threshold)[0]
order = order[inds + 1] # 因为第一项是top1_idx，不在temp_order中，导入order需要修补索引之间的差值
return record_ans
# 测试
if __name__ == '__main__':
time_begin = time.time()
#读取图片（任意放一张图片到当前目录命名为“img.jpg”即可）
img = cv2.imread("img.jpg")
img_cp = np.copy(img)
thickness = 2
# 模拟标注框的参数
# format = [conf, class, x_center, y_center, w, h]
info = np.array([
[0.95, 1, 150, 150, 100, 100],
[0.98, 1, 155, 155, 100, 100],
[0.96, 1, 145, 145, 110, 101],
[0.96, 2, 230, 230, 88, 99],
# [0.3, 1, 60, 60, 90, 90],
])
colors = [[0, 0, 255], [0, 255, 0], [255, 0, 0], [255, 255, 0]]
# 绘制使用NMS和IOS前的图片
plt.subplot(121)
plt.axis('off')
plt.title("Input image")
for i in range(len(colors)):
_, _, x_center, y_center, w, h = info[i]
x1, y1, x2, y2 = covert_wh(x_center, y_center, w, h)
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), colors[i], thickness=thickness)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img)
# 绘制使用NMS和IOS后的图片
plt.subplot(122)
plt.axis('off')
plt.title("After NMS")
indx = nms(info)
for i in indx:
_, _, x_center, y_center, w, h = info[i]
x1, y1, x2, y2 = covert_wh(x_center, y_center, w, h)
cv2.rectangle(img_cp, (int(x1), int(y1)), (int(x2), int(y2)), colors[i], thickness=thickness)
img_cp = cv2.cvtColor(img_cp, cv2.COLOR_BGR2RGB)
plt.imshow(img_cp)
# 保存并显示图片
plt.savefig('results.png')
plt.show()
# 记录运行时间
time_end = time.time()
time = time_end - time_begin
print('time:', time)
代码参考：
CVHub
TeddyZhang
扫一扫
1252
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
