说明：
本例程使用YOLOv3进行昆虫检测。例程分为数据处理、模型设计、损失函数、训练模型、模型预测和测试模型六个部分。本篇为第五部分，使用非极大值抑制来消除预测出的重叠面积过大的边框，然后显示预测结果图像。
实验代码：
模型预测：
import paddle.fluid as fluid
from paddle.fluid.dygraph.base import to_variable
from source.data import single_test_reader, display_infer
from source.model import YOLOv3
from source.infer import get_nms_infer
num_classes = 7 # 类别数量
anchor_size = [10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326] # 锚框大小
anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]] # 锚框掩码
downsample_ratio = 32 # 下采样率
image_path = './dataset/test/images/1872.jpeg' # 预测图像路径
model_path = './output/darknet53-yolov3' # 网络权重路径
sco_threshold = 0.70 # 预测得分阈值：根据测试的平均精度在准确率和召回率之间取一个平衡值
nms_threshold = 0.45 # 非极大值阈值
with fluid.dygraph.guard():
# 读取图像
image, image_size = single_test_reader(image_path) # 读取图像
image = to_variable(image) # 转换格式
image_size = to_variable(image_size) # 转换格式
# 加载模型
model = YOLOv3(num_classes=num_classes, anchor_mask=anchor_mask) # 加载模型
model_dict, _ = fluid.load_dygraph(model_path) # 加载权重
model.load_dict(model_dict) # 设置权重
model.eval（) # 设置验证
# 前向传播
infer = model(image)
# 获取结果
infer = get_nms_infer(infer, image_size, num_classes, anchor_size, anchor_mask, downsample_ratio,
sco_threshold, nms_threshold)
# 显示结果
print('image infer:', infer[0].shape[0]) # 显示图像预测结果数量
display_infer(infer[0], image_path) # 显示一张图像预测结果结果：
image infer: 6
infer.py文件
import numpy as np
def sigmoid(x):
"""
功能:
计算sigmoid函数
输入:
x - 输入数值
输出:
y - 输出数值
"""
return 0.5 * (1.0 + np.tanh(0.5 * x))
# def sigmoid(x):
# return 1.0 / (1.0 + np.exp(-x))
def get_infer(infer, image_size, num_classes, anchor_size, anchor_mask, downsample_ratio):
"""
功能:
计算每个特征图像的预测边框和得分
输入:
infer - 特征图像
image_size - 图像高宽
num_classes - 类别数量
anchor_size - 锚框大小
anchor_mask - 锚框掩码
downsample_ratio - 下采样率
输出:
pdbox - 预测边框
pdsco - 预测得分
"""
# 调整特征形状
batch_size = infer.shape[0] # 特征批数
num_rows = infer.shape[2] # 特征行数
num_cols = infer.shape[3] # 特征列数
num_anchor = len(anchor_mask) # 锚框数量
infer = infer.numpy()
infer = infer.reshape([-1, num_anchor, 5 + num_classes, num_rows, num_cols]) # 转换特征形状
# 计算预测边框
pdloc = infer[:, :, 0:4, :, :] # 获取预测位置：[b,c,4,n,m]
pdbox = np.zeros(pdloc.shape) # 预测边框数组：[b,c,4,n,m]
image_h = num_rows * downsample_ratio # 预测图像高度
image_w = num_cols * downsample_ratio # 预测图像宽度
for m in range(batch_size): # 遍历图像
for i in range(num_rows): # 遍历行数
for j in range(num_cols): # 遍历列数
for k in range(num_anchor): # 遍历锚框
# 获取边框大小
anchor_w = anchor_size[2 * anchor_mask[k]] # 锚框宽度
anchor_h = anchor_size[2 * anchor_mask[k] + 1] # 锚框高度
# 设置预测边框
pdbox[m, k, 0, i, j] = j # 预测边框cx
pdbox[m, k, 1, i, j] = i # 预测边框cy
pdbox[m, k, 2, i, j] = anchor_w # 预测边框pw
pdbox[m, k, 3, i, j] = anchor_h # 预测边框ph
pdbox[:, :, 0, :, :] = (pdbox[:, :, 0, :, :] + sigmoid(pdloc[:, :, 0, :, :])) / num_cols # 预测边框x=cx + dx
pdbox[:, :, 1, :, :] = (pdbox[:, :, 1, :, :] + sigmoid(pdloc[:, :, 1, :, :])) / num_rows # 预测边框y=cy + dy
pdbox[:, :, 2, :, :] = (pdbox[:, :, 2, :, :] * np.exp(pdloc[:, :, 2, :, :])) / image_w # 预测边框w=pw * exp(tw)
pdbox[:, :, 3, :, :] = (pdbox[:, :, 3, :, :] * np.exp(pdloc[:, :, 3, :, :])) / image_h # 预测边框h=ph * exp(th)
pdbox = np.clip(pdbox, 0.0, 1.0) # 限制预测边框范围为[0,1]
pdbox = pdbox.transpose((0, 1, 3, 4, 2)) # 调整数据维度：[b,c,n,m,4]
pdbox = pdbox.reshape((pdbox.shape[0], -1, pdbox.shape[-1])) # 调整数据形状：[b,c*n*m,4]
# 调整坐标格式
pdbox[:, :, 0] = pdbox[:, :, 0] - pdbox[:, :, 2] / 2.0 # 预测边框x1
pdbox[:, :, 1] = pdbox[:, :, 1] - pdbox[:, :, 3] / 2.0 # 预测边框y1
pdbox[:, :, 2] = pdbox[:, :, 0] + pdbox[:, :, 2] # 预测边框x2
pdbox[:, :, 3] = pdbox[:, :, 1] + pdbox[:, :, 3] # 预测边框y2
# 计算原图坐标
scale = image_size.numpy() # 原图高宽
for m in range(batch_size):
pdbox[m, :, 0] = pdbox[m, :, 0] * scale[m, 1] # 预测边框x1
pdbox[m, :, 1] = pdbox[m, :, 1] * scale[m, 0] # 预测边框y1
pdbox[m, :, 2] = pdbox[m, :, 2] * scale[m, 1] # 预测边框x2
pdbox[m, :, 3] = pdbox[m, :, 3] * scale[m, 0] # 预测边框y2
# 计算预测得分
pdobj = sigmoid(infer[:, :, 4, :, :]) # 预测物体概率：[b,c,n,m]，对损失函数计算结果求sigmoid
pdcls = sigmoid(infer[:, :, 5:5+num_classes, :, :]) # 预测类别概率：[b,c,7,n,m]，对损失函数计算结果求sigmoid
pdobj = np.expand_dims(pdobj, axis=2) # 添加数据维度：[b,c,1,n,m]
pdsco = pdobj * pdcls # 计算预测得分：[b,c,7,n,m]
pdsco = pdsco.transpose((0, 1, 3, 4, 2)) # 调整数据维度：[b,c,n,m,7]
pdsco = pdsco.reshape((pdsco.shape[0], -1, pdsco.shape[-1])) # 调整数据形状：[b,c*n*m,7]
pdsco = pdsco.transpose((0, 2, 1)) # 调整数据维度：[b,7,c*n*m]
return pdbox, pdsco
# def get_infer(infer, image_size, num_classes, anchor_size, anchor_mask, downsample_ratio):
# # 获取锚框大小
# anchor_list = [] # 锚框列表
# for i in anchor_mask: # 遍历锚框
# anchor_list.append(anchor_size[2 * i]) # 锚框宽度
# anchor_list.append(anchor_size[2 * i + 1]) # 锚框高度
# # 计算预测结果
# pdbox, pdsco = fluid.layers.yolo_box(
# x=infer,
# img_size=image_size,
# class_num=num_classes,
# anchors=anchor_list,
# conf_thresh=0.01,
# downsample_ratio=downsample_ratio)
# pdsco = fluid.layers.transpose(pdsco, perm=[0, 2, 1])
# return pdbox.numpy(), pdsco.numpy()
def get_sum_infer(infer, image_size, num_classes, anchor_size, anchor_mask, downsample_ratio):
"""
功能:
计算三个输出的预测结果的边框和得分
输入:
infer - 特征列表
image_size - 图像高宽
num_classes - 类别数量
anchor_size - 锚框大小
anchor_mask - 锚框掩码
downsample_ratio - 下采样率
输出:
pdbox - 预测边框
pdsco - 预测得分
"""
# 计算预测结果
pdbox_list = [] # 预测边框列表
pdsco_list = [] # 预测得分列表
for i in range(len(infer)): # 遍历特征列表
pdbox, pdsco = get_infer(infer[i], image_size, num_classes, anchor_size, anchor_mask[i], downsample_ratio)
pdbox_list.append(pdbox) # 添加边框列表
pdsco_list.append(pdsco) # 添加得分列表
# 减小下采样率
downsample_ratio //= 2 # 减小下采样率
# 合并预测结果
pdbox = np.concatenate(pdbox_list, axis=1) # 连接预测边框列表第一维
pdsco = np.concatenate(pdsco_list, axis=2) # 连接预测得分列表第二维
return pdbox, pdsco
##############################################################################################################
def get_box_iou_xyxy(box1, box2):
"""
功能:
计算边框交并比值
输入:
box1 - 边界框1
box2 - 边界框2
输出:
iou - 交并比值
"""
# 计算交集面积
x1_min, y1_min, x1_max, y1_max = box1[0], box1[1], box1[2], box1[3]
x2_min, y2_min, x2_max, y2_max = box2[0], box2[1], box2[2], box2[3]
x_min = np.maximum(x1_min, x2_min)
y_min = np.maximum(y1_min, y2_min)
x_max = np.minimum(x1_max, x2_max)
y_max = np.minimum(y1_max, y2_max)
w = np.maximum(x_max - x_min + 1.0, 0)
h = np.maximum(y_max - y_min + 1.0, 0)
intersection = w * h # 交集面积
# 计算并集面积
s1 = (y1_max - y1_min + 1.0) * (x1_max - x1_min + 1.0)
s2 = (y2_max - y2_min + 1.0) * (x2_max - x2_min + 1.0)
union = s1 + s2 - intersection # 并集面积
# 计算交并比
iou = intersection / union
return iou
def get_nms_index(pdbox, pdsco, sco_threshold, nms_threshold):
"""
功能:
获取非极大值抑制预测索引
输入:
pdbox - 预测边框
pdsco - 预测得分
sco_threshold - 预测得分阈值
nms_threshold - 非极大值阈值
输出:
nms_index - 预测索引
"""
# 获取得分索引
sco_index = np.argsort(pdsco)[::-1] # 对得分逆向排序，获取预测得分索引
# 非极大值抑制
nms_index = [] # 预测索引列表
while(len(sco_index) > 0): # 如果剩余得分索引数量大于0，则进行非极大值抑制
# 获取最大得分
max_index = sco_index[0] # 获取最大得分索引
max_score = pdsco[max_index] # 获取最大得分
if max_score < sco_threshold: # 如果最大得分小于预测得分阈值，则不处理剩余得分索引
break
# 设置保留标识
keep_flag = True # 保留标识为真
for i in nms_index: # 遍历保留索引
# 计算交并比值
box1 = pdbox[max_index] # 第一个边框坐标
box2 = pdbox[i] # 保留的边框坐标
iou = get_box_iou_xyxy(box1, box2) # 计算交并比值
if iou > nms_threshold: # 如果交并比值大于非极大值阈值，则不处理剩余保留索引
keep_flag = False # 保留标识为假
break
# 添加保留索引
if keep_flag: # 如果保留标识为真，则添加预测索引
nms_index.append(max_index) # 添加预测索引列表
# 获取剩余索引
sco_index = sco_index[1:]
# 转换数据格式
nms_index = np.array(nms_index)
return nms_index
def get_nms_class(pdbox, pdsco, sco_threshold, nms_threshold):
"""
功能:
获取非极大值抑制的预测结果
输入:
pdbox - 预测边框
pdsco - 预测得分
sco_threshold - 预测得分阈值
nms_threshold - 非极大值阈值
输出:
infer_list - 预测结果列表
"""
# 获取批次结果
batch_size = pdbox.shape[0] # 预测批数数量
class_numb = pdsco.shape[1] # 总的类别数量
infer_list = [] # 预测结果列表
for i in range(batch_size): # 遍历批次
# 获取预测结果
infer = [] # 每批预测列表
for c in range(class_numb): # 遍历类别
# 获取预测索引
nms_index = get_nms_index(pdbox[i], pdsco[i][c], sco_threshold, nms_threshold)
if len(nms_index) < 1: # 如果预测索引为0，则计算下一个类别索引
continue
# 设置预测结果
nms_pdsco = pdsco[i][c][nms_index] # 预测得分
nms_pdbox = pdbox[i][nms_index] # 预测边框
nms_infer = np.zeros([nms_pdsco.shape[0], 6]) # 预测结果
nms_infer[:, 0] = c # 设置预测类别
nms_infer[:, 1] = nms_pdsco[:] # 设置预测得分
nms_infer[:, 2:6] = nms_pdbox[:, :] # 设置预测边框
infer.append(nms_infer) # 添加每类结果
# 添加预测列表
if len(infer) > 0:
infer = np.concatenate(infer, axis=0) # 合并各批预测结果
infer_list.append(infer) # 添加预测结果列表
else:
infer_list.append(infer) # 添加空的预测结果
return infer_list
##############################################################################################################
def get_nms_infer(infer, image_size, num_classes, anchor_size, anchor_mask, downsample_ratio,
sco_threshold, nms_threshold):
"""
功能:
获取三个输出的非极大值抑制的预测结果
输入:
infer - 特征列表
image_size - 原图高宽
num_classes - 类别数量
anchor_size - 锚框大小
anchor_mask - 锚框掩码
downsample_ratio - 下采样率
sco_threshold - 预测得分阈值
nms_threshold - 非极大值阈值
输出:
infer - 预测结果
"""
# 计算预测结果
pdbox, pdsco = get_sum_infer(infer, image_size, num_classes, anchor_size, anchor_mask, downsample_ratio)
# 非极大值抑制
infer = get_nms_class(pdbox, pdsco, sco_threshold, nms_threshold)
return infer
参考资料：
https://aistudio.baidu.com/aistudio/projectdetail/742781
https://aistudio.baidu.com/aistudio/projectdetail/672017
https://aistudio.baidu.com/aistudio/projectdetail/868589
https://aistudio.baidu.com/aistudio/projectdetail/122277
本文章为转载内容，我们尊重原作者对文章享有的著作权。如有内容错误或侵权问题，欢迎原作者联系我们进行内容更正或删除文章。
