本博客将从标图到最终采用tensorrt部署yolov5_obb(用于斜框目标检测的yolov5框架)，一步一步手把手教你如何成为一个合格的算法搬运工。yolov5_obb是一款用于斜框目标检测的神经网络，具体原理我就不说了，关于这个网络的中文博客百度一大堆，作者也是在知乎发了一系列文章关于如何改写ultralytics大佬的yolov5用于斜框检测。
这个代码实际上是ultralytics/yolov5的6.0版本的基础上修改的，所以预训练模型采用ultralytics/yolov5的6.0的预训练模型即可
这个博客原则上需要具有一定深度学习能力的搬运工食用，指能够熟练torch，opencv，cuda，cudnn，tensorrt的安装与部分使用。
本文环境：GTX1080TI，cuda10.2 cudnn8.2.4 Tensorrt8.0.1.6 Opencv4.5.4，文中代码地址为：这里
一、模型训练与模型转换
1.数据标注
标注软件为rolabelimg，标注方法见这里，直接百度爬取到番茄来搞，我们检测番茄及其方向。标注如图所示。
2.标签转换
将标注好的xml文件，转换为DOTA_devkit的dota格式，如下：
txt的格式为：Format: x1, y1, x2, y2, x3, y3, x4, y4, category, difficulty
275.0 463.0 411.0 587.0 312.0 600.0 222.0 532.0 tomato 0
341.0 376.0 487.0 487.0 434.0 556.0 287.0 444.0 tomato 0
428.0 6.0 519.0 66.0 492.0 108.0 405.0 50.0 tomato 0
代码如下：
# *_* coding : UTF-8 *_*
# 功能描述 ：把旋转框 cx,cy,w,h,angle，转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4,class,difficulty
import os
import xml.etree.ElementTree as ET
import math
label=['tomato']
def edit_xml(xml_file):
"""
修改xml文件
:param xml_file:xml文件的路径
:return:
"""
print(xml_file)
tree = ET.parse(xml_file)
f=open(xml_file.replace('xml','txt').replace('anns','labelTxt'),'w')
objs = tree.findall('object')
for ix, obj in enumerate(objs):
obj_type = obj.find('type')
type = obj_type.text
if type == 'bndbox':
obj_bnd = obj.find('bndbox')
obj_xmin = obj_bnd.find('xmin')
obj_ymin = obj_bnd.find('ymin')
obj_xmax = obj_bnd.find('xmax')
obj_ymax = obj_bnd.find('ymax')
xmin = float(obj_xmin.text)
ymin = float(obj_ymin.text)
xmax = float(obj_xmax.text)
ymax = float(obj_ymax.text)
obj_bnd.remove(obj_xmin) # 删除节点
obj_bnd.remove(obj_ymin)
obj_bnd.remove(obj_xmax)
obj_bnd.remove(obj_ymax)
x0 = xmin
y0 = ymin
x1 = xmax
y1 = ymin
x2 = xmin
y2 = ymax
x3 = xmax
y3 = ymax
elif type == 'robndbox':
obj_bnd = obj.find('robndbox')
obj_bnd.tag = 'bndbox' # 修改节点名
obj_cx = obj_bnd.find('cx')
obj_cy = obj_bnd.find('cy')
obj_w = obj_bnd.find('w')
obj_h = obj_bnd.find('h')
obj_angle = obj_bnd.find('angle')
cx = float(obj_cx.text)
cy = float(obj_cy.text)
w = float(obj_w.text)
h = float(obj_h.text)
angle = float(obj_angle.text)
x0, y0 = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
x1, y1 = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
x2, y2 = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
x3, y3 = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)
classes=int(obj.find('name').text)
axis=list([str(x0),str(y0),str(x1), str(y1),str(x2), str(y2),str(x3), str(y3),label[classes],'0'])
bb = " ".join(axis)
f.writelines(bb)
f.writelines("\n")
f.close()
# 转换成四点坐标
def rotatePoint(xc, yc, xp, yp, theta):
xoff = xp - xc;
yoff = yp - yc;
cosTheta = math.cos(theta)
sinTheta = math.sin(theta)
pResx = cosTheta * xoff + sinTheta * yoff
pResy = - sinTheta * xoff + cosTheta * yoff
return int(xc + pResx), int(yc + pResy)
if __name__ == '__main__':
for path in os.listdir('anns/'):
edit_xml('anns/'+path)
转换好之后，新建images与labelTxt文件夹，把图片与标签复制过去，最终目录为：
data_dir/images/*.jpg
data_dir/labelTxt/*.txt
3.模型训练
3.1 编译安装
这里默认你已经下好了yolov5_obb代码，以及预训练模型（为了方便，我把训练的代码也放repo下了，可以直接用我的，戳这里），现在有两个地方需要编译
1.安装nms_rotated
pip install -r requirements.txt
cd utils/nms_rotated
python setup.py develop #or "pip install -v -e ."2.安装DOTA_devkit
sudo apt-get install swig
swig -c++ -python polyiou.i
python setup.py build_ext --inplace
3.2 数据分割
yolov5_obb要求输入的图片是矩形的，为此需要采用DOTA_devkit的ImgSplit_multi_process.py将转换为512×512的图像输入（不限于512，只要是能被32整除即可，比如）。最终得到的目录如下：
data_dir/split/images/*.jpg
data_dir/split/labelTxt/*.txt
3.3 划分训练集与验证集
主要是生成train.txt与val.txt,代码如下
# -*- coding: utf-8 -*-
import os
import random
# obb data split
annfilepath=r'/split/labelTxt/'
saveBasePath=r'split/'
train_percent=0.95
total_file = os.listdir(annfilepath)
num=len(total_file)
list=range(num)
tr=int(num*train_percent)
train=random.sample(list,tr)
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
for i in list:
name=total_file[i].split('.')[0]+'\n'
if i in train:
ftrain.write(name)
else:
fval.write(name)
ftrain.close()
fval.close()
print("train size",tr)
print("valid size",num-tr)
最终得到的目录如下，我的test.txt是直接复制的val.txt：
data_dir/split/images/*.jpg
data_dir/split/labelTxt/*.txt
data_dir/split/train.txt
data_dir/split/val.txt
data_dir/split/test.txt
3.4 模型训练—以yolov5s模型为例
如果你要更换数据集，需要改以下几个地方
1.data/dotav15_poly.yaml 改成自己的数据集路径。以及类别数和类别标签 ，我是一分类，类别数改成1即可，标签为tomato
2.models/yolov5s.yaml 中改成自己的类别数，我是一分类，改成1即可
3.开始训练
python train.py --weights yolov5s.pt --cfg models/yolov5s.yaml --data data/dotav15_poly.yaml --hyp data/hyps/obb/hyp.finetune_dota.yaml --imgsz 5124.生成wts模型
前面说过这个代码是基于yolov5-6.0版本修改的，wangxinyu大佬他的 tensorrtx repo中已经实现了tensort的c++版本推理。在此，我也在他的基础上修改最后的后处理，添加旋转角度计算与斜框NMS代码，以此来实现最终的斜框推理框架。
wangxinyu大佬和一般的过程不一样，我们大部分时候是pt-onnx-engine，而wangxinyu大佬是pt-wts-engine，区别是后者需要在c++里面重写网络结构，然后用wts的权重参数来初始化这个网络，相当于有难度，之前我自己写个简单的deeplabv3+结构写了500多行，像yolov5这种复杂的结构我肯定是不会的，所以直接采用他的代码做后处理的修改。
这里默认你已经下载好了我的代码【如果没有下载的话，戳这里】
第一步，复制gen_wts.py到yolov5_obb的根目录下
第二步，生成wts文件
python gen_wts.py -w {runs下你训练好的pt模型路径} -o yolov5s.wts二、Tensorrt模型转换
首先将yolov5s.wts复制到我那个c++的工程下面{Yolov5_obb_Tensorrt_Infer目录下}
1.编译安装
这里默认你已经装好了cuda10.2、cudnn8.2、opencv4.5、tensorrt8.0，cmake3.15，等一些其他组件，采用下面方法进行编译
cd 到下来解压后的目录
mkdir build
cd build
cmake ..
make 2.生成engine文件
在1中make之后，会在build生成yolov5_gen 这个可以执行的文件（windows中的exe），在终端运行
sudo ./yolov5_gen -s ../yolov5s.wts ../yolov5s.engine s不出意外的，经过一些时间的运行，生成yolov5s.engine这个文件（最好是是按照我的软件版本安装）。
3.tensorrt推理
主要就是用到opencv与tensorrt中nvinfer这个库了，部分代码如下：
在1中make之后，会在build生成yolov5_use这个可以执行的文件（windows中的exe），在终端运行
sudo ./yolov5_use ../yolovs.engine ../images/test.jpg不出意外的话，会显示下面这个结果，可以看到经过tensorrt推理之后的速度是4ms(我的1080ti熵)左右
三、参考
1.hukaixuan大佬的yolov5_obb
2.wangxinyu大佬的tensorrtx
3.再次上的链接：代码地址
文章出处登录后可见！
