jjzjj

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

机器鱼 2023-05-27 原文

         本博客将从标图到最终采用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将转换为512x512的图像输入(不限于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熵)左右

三、参考

1.hukaixuan大佬的yolov5_obb

2.wangxinyu大佬的tensorrtx

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

有关Yolov5旋转框(斜框)检测tensorrt部署(C++)从入门到入坟的更多相关文章

  1. ruby - RuntimeError(自动加载常量 Apps 多线程时检测到循环依赖 - 2

    我收到这个错误:RuntimeError(自动加载常量Apps时检测到循环依赖当我使用多线程时。下面是我的代码。为什么会这样?我尝试多线程的原因是因为我正在编写一个HTML抓取应用程序。对Nokogiri::HTML(open())的调用是一个同步阻塞调用,需要1秒才能返回,我有100,000多个页面要访问,所以我试图运行多个线程来解决这个问题。有更好的方法吗?classToolsController0)app.website=array.join(',')putsapp.websiteelseapp.website="NONE"endapp.saveapps=Apps.order("

  2. ruby-on-rails - 每次我尝试部署时,我都会得到 - (gcloud.preview.app.deploy) 错误响应 : [4] DEADLINE_EXCEEDED - 2

    我是Google云的新手,我正在尝试对其进行首次部署。我的第一个部署是RubyonRails项目。我基本上是在关注thisguideinthegoogleclouddocumentation.唯一的区别是我使用的是我自己的项目,而不是他们提供的“helloworld”项目。这是我的app.yaml文件runtime:customvm:trueentrypoint:bundleexecrackup-p8080-Eproductionconfig.ruresources:cpu:0.5memory_gb:1.3disk_size_gb:10当我转到我的项目目录并运行gcloudprevie

  3. ruby-on-rails - Ruby on Rails 可以部署在 Azure 网站上吗? - 2

    我可以在Azure网站上部署RubyonRails吗? 最佳答案 还没有。目前仅支持.NET和PHP。 关于ruby-on-rails-RubyonRails可以部署在Azure网站上吗?,我们在StackOverflow上找到一个类似的问题: https://stackoverflow.com/questions/12964010/

  4. 旋转矩阵的几何意义 - 2

    点向量坐标矩阵的几何意义介绍旋转矩阵的几何含义之前,先介绍一下点向量坐标矩阵的几何含义点:在一维空间下就是一个标量,如同一条直线上,以任意某一个位置为0点,以一定的尺度间隔为1,2,3...,相反方向为-1,-2,-3...;如此就形成了一维坐标系,这时候任何一个点都可以用一个数值表示,如点p1=5,即即从原点出发沿着x轴正方向移动5个尺度;点p2=-3,负方向移动3个尺度;     在一维坐标系上过原点做垂直于一维坐标系的直线,则形成了二维坐标系,此时描述一个点需要两个数值来表示点p3=(3,2),即从原点出发沿着x轴正方向移动3个尺度,在此基础上沿着y轴正方向移动两个尺度的位置就是点p3。

  5. Unity 3D 制作开关门动画,旋转门制作,推拉门制作,门把手动画制作 - 2

    Unity自动旋转动画1.开门需要门把手先动,门再动2.关门需要门先动,门把手再动3.中途播放过程中不可以再次进行操作觉得太复杂?查看我的文章开关门简易进阶版效果:如果这个门可以直接打开的话,就不需要放置"门把手"如果门把手还有钥匙需要旋转,那就可以把钥匙放在门把手的"门把手",理论上是可以无限套娃的可调整参数有:角度,反向,轴向,速度运行时点击Test进行测试自己写的代码比较垃圾,命名与结构比较拉,高手轻点喷,新手有类似的需求可以拿去做参考上代码usingSystem.Collections;usingSystem.Collections.Generic;usingUnityEngine;u

  6. jenkins部署1--jenkins+gitee持续集成 - 2

    前置步骤我们都操作完了,这篇开始介绍jenkins的集成。话不多说,看操作1、登录进入jenkins后会让你选择安装插件,选择第一个默认的就行。安装完成后设置账号密码,重新登录。2、配置JDK和Git都需要执行路径,所以需要先把执行路径找到,先进入服务器的docker容器,2.1JDK的路径root@69eef9ee86cf:/usr/bin#echo$JAVA_HOME/usr/local/openjdk-82.2Git的路径root@69eef9ee86cf:/#whichgit/usr/bin/git3、先配置JDK和Git。点击:ManageJenkins>>GlobalToolCon

  7. LC滤波器设计学习笔记(一)滤波电路入门 - 2

    目录前言滤波电路科普主要分类实际情况单位的概念常用评价参数函数型滤波器简单分析滤波电路构成低通滤波器RC低通滤波器RL低通滤波器高通滤波器RC高通滤波器RL高通滤波器部分摘自《LC滤波器设计与制作》,侵权删。前言最近需要学习放大电路和滤波电路,但是由于只在之前做音乐频谱分析仪的时候简单了解过一点点运放,所以也是相当从零开始学习了。滤波电路科普主要分类滤波器:主要是从不同频率的成分中提取出特定频率的信号。有源滤波器:由RC元件与运算放大器组成的滤波器。可滤除某一次或多次谐波,最普通易于采用的无源滤波器结构是将电感与电容串联,可对主要次谐波(3、5、7)构成低阻抗旁路。无源滤波器:无源滤波器,又称

  8. 深度学习部署:Windows安装pycocotools报错解决方法 - 2

    深度学习部署:Windows安装pycocotools报错解决方法1.pycocotools库的简介2.pycocotools安装的坑3.解决办法更多Ai资讯:公主号AiCharm本系列是作者在跑一些深度学习实例时,遇到的各种各样的问题及解决办法,希望能够帮助到大家。ERROR:Commanderroredoutwithexitstatus1:'D:\Anaconda3\python.exe'-u-c'importsys,setuptools,tokenize;sys.argv[0]='"'"'C:\\Users\\46653\\AppData\\Local\\Temp\\pip-instal

  9. 微信小程序开发入门与实战(Behaviors使用) - 2

    @作者:SYFStrive @博客首页:HomePage📜:微信小程序📌:个人社区(欢迎大佬们加入)👉:社区链接🔗📌:觉得文章不错可以点点关注👉:专栏连接🔗💃:感谢支持,学累了可以先看小段由小胖给大家带来的街舞👉微信小程序(🔥)目录自定义组件-behaviors    1、什么是behaviors    2、behaviors的工作方式    3、创建behavior    4、导入并使用behavior    5、behavior中所有可用的节点    6、同名字段的覆盖和组合规则总结最后自定义组件-behaviors    1、什么是behaviorsbehaviors是小程序中,用于实现

  10. 【Java入门】使用Java实现文件夹的遍历 - 2

    遍历文件夹我们通常是使用递归进行操作,这种方式比较简单,也比较容易理解。本文为大家介绍另一种不使用递归的方式,由于没有使用递归,只用到了循环和集合,所以效率更高一些!一、使用递归遍历文件夹整体思路1、使用File封装初始目录,2、打印这个目录3、获取这个目录下所有的子文件和子目录的数组。4、遍历这个数组,取出每个File对象4-1、如果File是否是一个文件,打印4-2、否则就是一个目录,递归调用代码实现publicclassSearchFile{publicstaticvoidmain(String[]args){//初始目录Filedir=newFile("d:/Dev");Datebeg

随机推荐