• 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格式，如下：

    1. txt的格式为：Format: x1, y1, x2, y2, x3, y3, x4, y4, category, difficulty
    2.  
    3. 275.0 463.0 411.0 587.0 312.0 600.0 222.0 532.0 tomato 0
    4. 341.0 376.0 487.0 487.0 434.0 556.0 287.0 444.0 tomato 0
    5. 428.0 6.0 519.0 66.0 492.0 108.0 405.0 50.0 tomato 0

    代码如下：

    1. # *_* coding : UTF-8 *_*
    2. # 功能描述   ：把旋转框 cx,cy,w,h,angle，转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4,class,difficulty
    3.  
    4. import os
    5. import xml.etree.ElementTree as ET
    6. import math
    7.  
    8. label=['tomato']
    9. def edit_xml(xml_file):
    10.     """
    11.     修改xml文件
    12.     :param xml_file:xml文件的路径
    13.     :return:
    14.     """
    15.     print(xml_file)
    16.     tree = ET.parse(xml_file)
    17.     f=open(xml_file.replace('xml','txt').replace('anns','labelTxt'),'w')
    18.     objs = tree.findall('object')
    19.     for ix, obj in enumerate(objs):
    20.         obj_type = obj.find('type')
    21.         type = obj_type.text
    22.         
    23.         if type == 'bndbox':
    24.             obj_bnd = obj.find('bndbox')
    25.             obj_xmin = obj_bnd.find('xmin')
    26.             obj_ymin = obj_bnd.find('ymin')
    27.             obj_xmax = obj_bnd.find('xmax')
    28.             obj_ymax = obj_bnd.find('ymax')
    29.             xmin = float(obj_xmin.text)
    30.             ymin = float(obj_ymin.text)
    31.             xmax = float(obj_xmax.text)
    32.             ymax = float(obj_ymax.text)
    33.             obj_bnd.remove(obj_xmin)  # 删除节点
    34.             obj_bnd.remove(obj_ymin)
    35.             obj_bnd.remove(obj_xmax)
    36.             obj_bnd.remove(obj_ymax)
    37.             x0 = xmin
    38.             y0 = ymin
    39.             x1 = xmax
    40.             y1 = ymin
    41.             x2 = xmin
    42.             y2 = ymax
    43.             x3 = xmax
    44.             y3 = ymax
    45.         elif type == 'robndbox':
    46.             obj_bnd = obj.find('robndbox')
    47.             obj_bnd.tag = 'bndbox'   # 修改节点名
    48.             obj_cx = obj_bnd.find('cx')
    49.             obj_cy = obj_bnd.find('cy')
    50.             obj_w = obj_bnd.find('w')
    51.             obj_h = obj_bnd.find('h')
    52.             obj_angle = obj_bnd.find('angle')
    53.             cx = float(obj_cx.text)
    54.             cy = float(obj_cy.text)
    55.             w = float(obj_w.text)
    56.             h = float(obj_h.text)
    57.             angle = float(obj_angle.text)
    58.  
    59.             x0, y0 = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
    60.             x1, y1 = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
    61.             x2, y2 = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
    62.             x3, y3 = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)
    63.         classes=int(obj.find('name').text)
    64.         axis=list([str(x0),str(y0),str(x1), str(y1),str(x2), str(y2),str(x3), str(y3),label[classes],'0'])
    65.         bb = " ".join(axis)
    66.         f.writelines(bb)
    67.         f.writelines("\n")
    68.     f.close()
    69. # 转换成四点坐标
    70. def rotatePoint(xc, yc, xp, yp, theta):
    71.     xoff = xp - xc;
    72.     yoff = yp - yc;
    73.     cosTheta = math.cos(theta)
    74.     sinTheta = math.sin(theta)
    75.     pResx = cosTheta * xoff + sinTheta * yoff
    76.     pResy = - sinTheta * xoff + cosTheta * yoff
    77.     return int(xc + pResx), int(yc + pResy)
    78.  
    79. if __name__ == '__main__':
    80.     for path in os.listdir('anns/'):
    81.         edit_xml('anns/'+path)

    转换好之后，新建images与labelTxt文件夹，把图片与标签复制过去，最终目录为：

    1. data_dir/images/*.jpg
    2. data_dir/labelTxt/*.txt

    3.模型训练

    3.1 编译安装

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

    1.安装nms_rotated

    1. pip install -r requirements.txt
    2. cd utils/nms_rotated
    3. python setup.py develop  #or "pip install -v -e ."

    2.安装DOTA_devkit

    1. sudo apt-get install swig
    2. swig -c++ -python polyiou.i
    3. python setup.py build_ext --inplace

    3.2 数据分割

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

    1. data_dir/split/images/*.jpg
    2. data_dir/split/labelTxt/*.txt

    3.3 划分训练集与验证集

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

    1. # -*- coding: utf-8 -*-
    2.  
    3. import os
    4. import random
    5.  
    6. # obb data split
    7. annfilepath=r'/split/labelTxt/'
    8. saveBasePath=r'split/'
    9. train_percent=0.95
    10. total_file = os.listdir(annfilepath)
    11. num=len(total_file)
    12. list=range(num)
    13. tr=int(num*train_percent)
    14. train=random.sample(list,tr)
    15. ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
    16. fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
    17. for i  in list:
    18.     name=total_file[i].split('.')[0]+'\n'
    19.     if i in train:
    20.         ftrain.write(name)
    21.     else:
    22.         fval.write(name)
    23. ftrain.close()
    24. fval.close()
    25. print("train size",tr)
    26. print("valid size",num-tr)

    最终得到的目录如下，我的test.txt是直接复制的val.txt：

    1. data_dir/split/images/*.jpg
    2. data_dir/split/labelTxt/*.txt
    3. data_dir/split/train.txt
    4. data_dir/split/val.txt
    5. 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 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目录下}

    1.编译安装

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

    1. cd 到下来解压后的目录
    2. mkdir build
    3. cd build
    4. cmake ..
    5. 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.再次上的链接：代码地址

  • 相关阅读:
    Midjourney如何实现人物角色的一致性？
    【postgresql 数据库运维文档】
    STM32CubeMX教程31 USB_DEVICE - HID外设_模拟键盘或鼠标
    linux安装UnZip
    【PostgreSql本地备份为dump文件与恢复】使用脚本一键备份为dump文件
    vcruntime140d.dl丢失怎么办？
    化工管理杂志化工管理杂志社化工管理编辑部2023年第30期目录
    牛X，试用了下 GitHub 上 2 万 Star 的第一抢票神器，3 秒钟抢到！
    数据库设计规范
    报错TypeError: Got an unexpected keyword argument ‘keep_dims‘
  • 原文地址：https://blog.csdn.net/qq_41043389/article/details/127777272