YOLOv5预处理高性能实现
目录
-
-
- 前言
- 1. YOLOv5推理
- 2. 高性能预处理
- 结语
- 下载链接
- 参考
-
前言
梳理下YOLOv5预处理和后处理流程，并实现GPU版预处理(仅供自己参考)
视频讲解:抖音/手写AI
代码参考:preprocess_kernel.cu
1. YOLOv5推理
YOLOv5推理包括预处理和后处理两部分，其中预处理主要包括warpAffine和双线性插值，不了解的可以查看here，后处理主要包括decode解码和NMS两部分。
废话少说直接上代码
import numpy as np
import cv2
def preprocess(img, dst_width=640, dst_height=640):
'''
:param img: 输入的图片
:param dst_width: 预处理后的图像宽
:param dst_height: 预处理后的图像高
:return: 预处理后的图片,仿射变换矩阵的逆变换矩阵IM
'''
scale = min((dst_width / img.shape[1]), (dst_height / img.shape[0]))
ox = (-scale * img.shape[1] + dst_width) / 2
oy = (-scale * img.shape[0] + dst_height) / 2
M = np.array([
[scale, 0, ox],
[0, scale, oy]
], dtype=np.float32)
# img_pre为仿射变换后的图即原始图像缩放到[dst_width,dst_height]
img_pre = cv2.warpAffine(img, M, dsize=[dst_width, dst_height], flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT, borderValue=(114,114,114))
IM = cv2.invertAffineTransform(M)
# -----------------------------------------------------------------------#
# 需要进行的预处理
# 1. BGR -> RGB
# 2. /255.0
# 3. 通道数变换 H,W,C -> C,H,W
# 4. 添加batch维度 C,H,W -> B,C,H,W
# -----------------------------------------------------------------------#
img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32)
img_pre = img_pre.transpose(2,0,1)[None]
return img_pre, IM
def iou(box1, box2):
def area_box(box):
return (box[2] - box[0]) * (box[3] - box[1])
# box -> [x1,y1,x2,y2,...]
left, top = max(box1[:2], box2[:2])
right, bottom = min(box1[2:4], box2[2:4])
union = max((right-left), 0) * max((bottom-top), 0)
cross = area_box(box1) + area_box(box2) - union
if cross == 0 or union == 0:
return 0
return union / cross
def NMS(boxes, iou_thresh=0.45):
'''
:param boxes: decode解码排序后的boxes [n,7] 7 = x1,y1,x2,y2,conf,img_id,label
:param iou_thresh: iou阈值
:return: 经过NMS的boxes
'''
# 利用remove_flags标记需要去除的box
remove_flags = [False] * len(boxes)
# 保留下的box
keep_boxes = []
for i in range(len(boxes)):
if remove_flags[i]:
continue
ibox = boxes[i]
keep_boxes.append(ibox)
for j in range(len(boxes)):
if remove_flags[j]:
continue
jbox = boxes[j]
# 只有同一张图片中的同一个类别的box才计算iou
if ibox[5] != jbox[5] or ibox[6] != jbox[6]:
continue
# 计算iou,若大于阈值则标记去除
if iou(ibox, jbox) > iou_thresh:
remove_flags[j] = True
return keep_boxes
def postprocess(pred, IM, iou_thresh=0.45, conf_thresh=0.25):
'''
:param pred: 模型推理的结果 [1,25200,85] 85 = cx,cy,w,h,conf + 80
:param IM: 仿射变换矩阵的逆变换，主要用来将box映射回原图
:param iou_thresh: iou阈值
:param cof_thresh: 置信度阈值
:return: 经过NMS的boxes
'''
# 保存decode解码后的boxes
boxes = []
for img_id, box_id in zip(*np.where(pred[...,4] >= conf_thresh)):
item = pred[img_id][box_id]
cx, cy, w, h, obj_conf = item[:5]
label = item[5:].argmax()
confidence = obj_conf * item[5+label]
if confidence < conf_thresh:
continue
left = cx - w * 0.5
top = cy - h * 0.5
right = cx + w * 0.5
bottom = cy + h * 0.5
boxes.append([left, top, right, bottom, confidence, img_id, label])
# 利用IM将box映射回原图
boxes = np.array(boxes)
lr = boxes[..., [0, 2]]
tb = boxes[..., [1, 3]]
boxes[..., [0, 2]] = lr * IM[0][0] + IM[0][2]
boxes[..., [1, 3]] = tb * IM[1][1] + IM[1][2]
# 将boxes按照置信度高低排序
boxes = sorted(boxes.tolist(), key= lambda x : x[4], reverse=True)
# 将排序后的boxes作NMS
return NMS(boxes, iou_thresh=iou_thresh)
if __name__ == '__main__':
img = cv2.imread("bus.jpg")
# 预处理
img_pre, IM = preprocess(img, dst_width=640, dst_height=640)
# 加载模型推理的结果
pred = np.load("pred.npz")['arr_0']
# 后处理
boxes = postprocess(pred, IM)
for obj in boxes:
x1, y1, x2, y2 = map(int, obj[:4])
label = int(obj[6])
confidence = obj[4]
cv2.rectangle(img, (x1, y1), (x2, y2), (0,255,0), 1, 8)
cv2.putText(img, f"{label}:{confidence:.3f}", (x1, y1-6), 0, 1, (0,0,255), 2, 8)
cv2.imshow("img_pre", img)
cv2.waitKey(0)
效果如下图
2. 高性能预处理
主要将YOLOv5的预处理(即warpAffine和双线性插值)以核函数的形式实现，代码如下
#include 结语
代码仅供自己参考，大家也可查看对应的小视频观看讲解，后续有时间再和大家详细分享。
下载链接
- 高性能预处理[password:yolo]
参考
- 抖音/手写AI
- preprocess_kernel.cu
- 图像预处理之warpaffine与双线性插值