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,标注方法见这里,直接百度爬取到番茄来搞,我们检测番茄及其方向。标注如图所示。Yolov5旋转框(斜框)检测tensorrt部署(C++)从入门到入坟

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 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,等一些其他组件,采用下面方法进行编译

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熵)左右

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

三、参考

1.hukaixuan大佬的yolov5_obb

2.wangxinyu大佬的tensorrtx

3.再次上的链接:代码地址

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
社会演员多的头像社会演员多普通用户
上一篇 2023年3月5日 下午12:08
下一篇 2023年3月5日 下午12:10

相关推荐