Yolov5旋转框(斜框)检测tensorrt部署(C++)从入门到入坟

         本博客将从标图到最终采用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, difficulty275.0 463.0 411.0 587.0 312.0 600.0 222.0 532.0 tomato 0341.0 376.0 487.0 487.0 434.0 556.0 287.0 444.0 tomato 0428.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,difficultyimport osimport xml.etree.ElementTree as ETimport mathlabel=['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/*.jpgdata_dir/labelTxt/*.txt

3.模型训练

3.1 编译安装

        这里默认你已经下好了yolov5_obb代码，以及预训练模型（为了方便，我把训练的代码也放repo下了，可以直接用我的，戳这里），现在有两个地方需要编译

安装nms_rotated

pip install -r requirements.txtcd utils/nms_rotatedpython setup.py develop  #or "pip install -v -e ."

安装DOTA_devkit

sudo apt-get install swigswig -c++ -python polyiou.ipython setup.py build_ext --inplace

3.2 数据分割

        yolov5_obb要求输入的图片是矩形的，为此需要采用DOTA_devkit的ImgSplit_multi_process.py将转换为512x512的图像输入（不限于512，只要是能被32整除即可，比如）。最终得到的目录如下：

data_dir/split/images/*.jpgdata_dir/split/labelTxt/*.txt

3 划分训练集与验证集

        主要是生成train.txt与val.txt,代码如下

# -*- coding: utf-8 -*-import osimport random# obb data splitannfilepath=r'/split/labelTxt/'saveBasePath=r'split/'train_percent=0.95total_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/*.jpgdata_dir/split/labelTxt/*.txtdata_dir/split/train.txtdata_dir/split/val.txtdata_dir/split/test.txt

3.4 模型训练---以yolov5s模型为例

如果你要更换数据集，需要改以下几个地方

data/dotav15_poly.yaml 改成自己的数据集路径。以及类别数和类别标签 ，我是一分类，类别数改成1即可，标签为tomato

models/yolov5s.yaml 中改成自己的类别数，我是一分类，改成1即可

开始训练

python train.py --weights yolov5s.pt --cfg models/yolov5s.yaml --data data/dotav15_poly.yaml --hyp data/hyps/obb/hyp.finetune_dota.yaml --imgsz 512

4.生成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目录下}

编译安装

        这里默认你已经装好了cuda10.2、cudnn8.2、opencv4.5、tensorrt8.0，cmake3.15，等一些其他组件，采用下面方法进行编译

cd 到下来解压后的目录mkdir buildcd buildcmake ..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熵)左右

​

三、参考

hukaixuan大佬的yolov5_obb

wangxinyu大佬的tensorrtx

再次上的链接：代码地址

来源地址：https://blog.csdn.net/qq_41043389/article/details/127777272