用opencv的dnn模块做yolov5目标检测

本文原创首发于极市平台公众号,如需转载请私信作者

最近在微信公众号里看到多篇讲解yolov5在openvino部署做目标检测文章,但是没看到过用opencv的dnn模块做yolov5目标检测的。于是,我就想着编写一套用opencv的dnn模块做yolov5目标检测的程序。在编写这套程序时,遇到的bug和解决办法,在这篇文章里讲述一下。

在yolov5之前的yolov3和yolov4的官方代码都是基于darknet框架的实现的,因此opencv的dnn模块做目标检测时,读取的是.cfg和.weight文件,那时候编写程序很顺畅,没有遇到bug。但是yolov5的官方代码(https://github.com/ultralytics/yolov5)是基于pytorch框架实现的,但是opencv的dnn模块不支持读取pytorch的训练模型文件的。如果想要把pytorch的训练模型.pth文件加载到opencv的dnn模块里,需要先把pytorch的训练模型.pth文件转换到.onnx文件,然后才能载入到opencv的dnn模块里。

因此,用opencv的dnn模块做yolov5目标检测的程序,包含两个步骤:(1).把pytorch的训练模型.pth文件转换到.onnx文件。(2).opencv的dnn模块读取.onnx文件做前向计算。

(1).把pytorch的训练模型.pth文件转换到.onnx文件

在做这一步时,我得吐槽一下官方代码:https://github.com/ultralytics/yolov5,这套程序里的代码混乱,在pytorch里,通常是在.py文件里定义网络结构的,但是官方代码是在.yaml文件定义网络结构,利用pytorch动态图特性,解析.yaml文件自动生成网络结构。在.yaml文件里有depth_multiple和width_multiple,它是控制网络的深度和宽度的参数。这么做的好处是能够灵活的配置网络结构,但是不利于理解网络结构,假如你想设断点查看某一层的参数和输出数值,那就没办法了。因此,在我编写的转换到.onnx文件的程序里,网络结构是在.py文件里定义的。其次,在官方代码里,还有一个奇葩的地方,那就是.pth文件。起初,我下载官方代码到本地运行时,torch.load读取.pth文件总是出错,后来把pytorch升级到1.7,就读取成功了。可以看到版本兼容性不好,这是它的一个不足之处。设断点查看读取的.pth文件里的内容,可以看到ultralytics的.pt文件里既存储有模型参数,也存储有网络结构,还储存了一些超参数,包括anchors,stride等等的。第一次见到有这种操作的,通常情况下,.pth文件里只存储了训练模型参数的。

查看models\yolo.py里的Detect类,在构造函数里,有这么两行代码:

我尝试过把这两行代码改成self.anchors = a 和 self.anchor_grid = a.clone().view(self.nl, 1, -1, 1, 1, 2),程序依然能正常运行,但是torch.save保存模型文件后,可以看到.pth文件里没有存储anchors和anchor_grid了,在百度搜索register_buffer,解释是:pytorch中register_buffer模型保存和加载的时候可以写入和读出。

在这两行代码的下一行:

它的作用是做特征图的输出通道对齐,通过1x1卷积把三种尺度特征图的输出通道都调整到 num_anchors*(num_classes+5)。

阅读Detect类的forward函数代码,可以看出它的作用是根据偏移公式计算出预测框的中心坐标和高宽,这里需要注意的是,计算高和宽的代码:

                                                                 pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]

没有采用exp操作,而是直接乘上anchors[i],这是yolov5与yolov3v4的一个最大区别(还有一个区别就是在训练阶段的loss函数里,yolov5采用邻域的正样本anchor匹配策略,增加了正样本。其它的是一些小区别,比如yolov5的第一个模块采用FOCUS把输入数据2倍下采样切分成4份,在channel维度进行拼接,然后进行卷积操作,yolov5的激活函数没有使用Mish)。

现在可以明白Detect类的作用是计算预测框的中心坐标和高宽,简单来说就是生成proposal,作为后续NMS的输入,进而输出最终的检测框。我觉得在Detect类里定义的1x1卷积是不恰当的,应该把它定义在Detect类的外面,紧邻着Detect类之前定义1x1卷积。

在官方代码里,有转换到onnx文件的程序: python models/export.py --weights yolov5s.pt --img 640 --batch 1

在pytorch1.7版本里,程序是能正常运行生成onnx文件的。观察export.py里的代码,在执行torch.onnx.export之前,有这么一段代码:

注意其中的for循环,我试验过注释掉它,重新运行就会出错,打印出的错误如下:

由此可见,这段for循环代码是必需的。SiLU其实就是swish激活函数,而在onnx模型里是不直接支持swish算子的,因此在转换生成onnx文件时,SiLU激活函数不能直接使用nn.Module里提供的接口,而需要自定义实现它。

(2).opencv的dnn模块读取.onnx文件做前向计算

在生成.onnx文件后,就可以用opencv的dnn模块里的cv2.dnn.readNet读取它。然而,在读取时,出现了如下错误:

我在百度搜索这个问题的解决办法,看到一篇知乎文章(Pytorch转ONNX-实战篇2(实战踩坑总结) - 知乎),文章里讲述的第一条:

于是查看yolov5的代码,在common.py文件的Focus类,torch.cat的输入里有4次切片操作,代码如下:

那么现在需要更换索引式的切片操作,观察到注释的Contract类,它就是用view和permute函数完成切片操作的,于是修改代码如下:

其次,在models\yolo.py里的Detect类里,也有切片操作,代码如下:

前面说过,Detect类的作用是计算预测框的中心坐标和高宽,生成proposal,这个是属于后处理的,因此不需要把它写入到onnx文件里。

总结一下,按照上面的截图代码,修改Focus类,把Detect类里面的1x1卷积定义在紧邻着Detect类之前的外面,然后去掉Detect类,组成新的model,作为torch.onnx.export的输入,

torch.onnx.export(model, inputs, output_onnx, verbose=False, opset_version=12, input_names=['images'], output_names=['out0', 'out1', 'out2'])

最后生成的onnx文件,opencv的dnn模块就能成功读取了,接下来对照Detect类里的forward函数,用python或者C++编写计算预测框的中心坐标和高宽的功能。

周末这两天,我在win10+cpu机器里编写了用opencv的dnn模块做yolov5目标检测的程序,包含Python和C++两个版本的。程序都调试通过了,运行结果也是正确的。我把这套代码发布在github上,地址是

https://github.com/hpc203/yolov5-dnn-cpp-python

后处理模块,python版本用numpy array实现的,C++版本的用vector和数组实现的,整套程序只依赖opencv库(opencv4版本以上的)就能正常运行,彻底摆脱对深度学习框架pytorch,tensorflow,caffe,mxnet等等的依赖。用openvino作目标检测,需要把onnx文件转换到.bin和.xml文件,相比于用dnn模块加载onnx文件做目标检测是多了一个步骤的。因此,我就想编写一套用opencv的dnn模块做yolov5目标检测的程序,用opencv的dnn模块做深度学习目标检测,在win10和ubuntu,在cpu和gpu上都能运行,可见dnn模块的通用性更好,很接地气。

生成yolov5s_param.pth 的步骤,首先下载https://github.com/ultralytics/yolov5 的源码到本地,在yolov5-master主目录(注意不是我发布的github代码目录)里新建一个.py文件,把下面的代码复制到.py文件里

import torch
from collections import OrderedDict
import pickle
import os

device = 'cuda' if torch.cuda.is_available() else 'cpu'

if __name__=='__main__':
    choices = ['yolov5s', 'yolov5l', 'yolov5m', 'yolov5x']
    modelfile = choices[0]+'.pt'
    utl_model = torch.load(modelfile, map_location=device)
    utl_param = utl_model['model'].model
    torch.save(utl_param.state_dict(), os.path.splitext(modelfile)[0]+'_param.pth')
    own_state = utl_param.state_dict()
    print(len(own_state))

    numpy_param = OrderedDict()
    for name in own_state:
        numpy_param[name] = own_state[name].data.cpu().numpy()
    print(len(numpy_param))
    with open(os.path.splitext(modelfile)[0]+'_numpy_param.pkl', 'wb') as fw:
        pickle.dump(numpy_param, fw)

运行这个.py文件,这时候就可以生成yolov5s_param.pth文件。之所以要进行这一步,我在上面讲到过:ultralytics的.pt文件里既存储有模型参数,也存储有网络结构,还储存了一些超参数,包括anchors,stride等等的。torch.load加载ultralytics的官方.pt文件,也就是utl_model = torch.load(modelfile, map_location=device)这行代码,在这行代码后设断点查看utl_model里的内容,截图如下

可以看到utl_model里含有既存储有模型参数,也存储有网络结构,还储存了一些超参数等等的,这会严重影响转onnx文件。此外,我还发现,如果pytorch的版本低于1.7,那么在torch.load加载.pt文件时就会出错的。

此外,有读者反映,在Focus类的构造函数里添加了 self.contract = Contract(gain=2) ,之后在forward函数里报错,错误信息是没有contract这个成员。之所以会出现这个错误,原因正如上面所说的在ultralytics的.pt文件里既存储有模型参数,也存储有网络结构。因而在Focus类的构造函数里添加的成员是不会生效的,做一个小实验来验证,在forward函数设断点,运行,截图如下

 可以看到self里并没有contract这个成员。想要在导出onnx文件时不出错的做法是,代码如下:

class Focus(nn.Module):  
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
        #self.contract = Contract(gain=2)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        if torch.onnx.is_in_onnx_export():
            contract = Contract(gain=2)
            return self.conv(contract(x))
        else:
            return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

因此在程序里,我把模型参数转换到cpu.numpy形式的,最后保存在.pkl文件里。这时候在win10系统cpu环境里,即使你的电脑没有安装pytorch,也能通过python程序访问到模型参数。

pytorch转onnx常见坑:

1. onnx只能输出静态图,因此不支持if-else分支。一次只能走一个分支。如果代码中有if-else语句,需要改写。
2. onnx不支持步长为2的切片。例如a[::2,::2]
3. onnx不支持对切片对象赋值。例如a[0,:,:,:]=b, 可以用torch.cat改写
4. onnx里面的resize要求output shape必须为常量。可以用以下代码解决:

if isinstance(size, torch.Size):
    size = tuple(int(x) for x in size)

此外,在torch.onnx.export(model, inputs, output_onnx)的输入参数model里,应该只包含网络结构,也就是说model里只含有nn.Conv2d, nn.MaxPool2d, nn.BatchNorm2d, F.relu等等的这些算子组件,而不应该含有后处理模块的。图像预处理和后处理模块需要自己使用C++或者Python编程实现。

在明白了这些之后,在转换生成onnx文件,你需要执行两个步骤,第一步把原始训练模型.pt文件里的参数保存到新的.pth文件里,第二步编写yolov5.py文件,把yolov5的往来结构定义在.py文件里,此时需要注意网络结构里不能包含切片对象赋值操作,F.interpolate里的size参数需要加int强制转换。在执行完这两步之后才能生成一个opencv能成功读取并且做前向推理的onnx文件。

不过,最近我发现在yolov5-pytorch程序里,其实可以直接把原始训练模型.pt文件转换生成onnx文件的,而且我在一个yolov5检测人脸+关键点的程序里实验成功了。

        5月1日,我把这套程序发布在github上,地址是 https://github.com/hpc203/yolov5-face-landmarks-opencv 和 https://github.com/hpc203/yolov5-face-landmarks-opencv-v2

这套程序只依赖opencv库就可以运行yolov5检测人脸+关键点,程序依然是包含C++和Python两个版本的,这套程序里还有一个转换生成onnx文件的python程序文件。只需运行这一个.py文件就可以生成onnx文件,而不需要之前讲的那样执行两个步骤,这样大大简化了生成onnx文件的流程,使用方法可以阅读程序里的README文档。

在这个新的转换生成onnx文件的程序里,需要重新定义yolov5网络结构,主要是修改第一个模块Focus,用Contract类替换索引式的切片操作,在最后一个模块Detect类里,只保留三个1x1卷积,剩下的make_grid和decode属于后处理,不能包含在网络结构里,代码截图如下

如果要转换生成onnx文件,需要设置export = True,这时候Detect模块的forward就只进行1x1卷积,这时的网络结构就可以作为torch.onnx.export(model, inputs, output_onnx)的输入参数model。不过由于ultralytics的yolov5代码仓库几乎每天都在更新,因此你现在看到的ultralytics的yolov5里的Detect类很有可能不是这么写的,那这是需要你手动修改程序,然后再运行。

        8月8日,看到最近旷视发布的anchor-free系列的YOLOX,而在github开源的代码里,并没有使用opencv部署的程序。因此,我就编写了一套使用OpenCV部署YOLOX的程序,支持YOLOX-S、YOLOX-M、YOLOX-L、YOLOX-X、YOLOX-Darknet53五种结构,包含C++和Python两种版本的程序实现。在今天我在github发布了这套程序,地址是

 https://github.com/hpc203/yolox-opencv-dnn

在旷视发布的YOLOX代码里,提供了在COCO数据集上训练出来的.pth模型文件,并且也提供了导出onnx模型的export_onnx.py文件,起初我运行export_onnx.py生成onnx文件之后Opencv读取onnx文件失败了,报错原因跟文章最开始的第(2)节里的一样,这说明在YOLOX的网络结构里有切片操作,经过搜索后,在 yolox\models\network_blocks.py 里有个Focus类,它跟YOLOv5里的Focus是一样的,都是把输入张量切分成4份,然后concat+conv。这时按照第(2)节里讲述的解决办法,修改Focus类,重新运行export_onnx.py生成onnx文件,Opencv读取onnx文件就不会再出错了。

        8月22日,我在github发布了一套使用OpenCV部署Yolo-FastestV2的程序,依然是包含C++和Python两种版本的程序实现。地址是

https://github.com/hpc203/yolo-fastestv2-opencv

经过运行,体验到这个Yolo-FastestV2的速度确实很快,而且onnx文件只有957kb大小,不超过1M。在官方代码https://github.com/dog-qiuqiu/Yolo-FastestV2里,学习它的网络结构。设断点调试,查看中间变量可以看到,在model/detector.py,网络输出了6个张量

 它们的形状分别是

torch.Size([1, 12, 22, 22])
torch.Size([1, 3, 22, 22])
torch.Size([1, 80, 22, 22])
torch.Size([1, 12, 11, 11])
torch.Size([1, 3, 11, 11])
torch.Size([1, 80, 11, 11])

结合配置文件data/coco.data,可以看到模型输入是352x352的图片,而输出有22x22和11x11这两种尺度的特征图,这说明Yolo-FastestV2的输出只有缩放16倍和缩放32倍这两种尺度的特征图,比yolov3,v4,v5系列的都要少一个尺度特征图。其次在配置文件data/coco.data还可以看到anchor一共有6个,分别给两个尺度特征图里的网格点分配3个。观察输出的6个张量的形状信息,很明显前3个张量是22x22尺度特征图的检测框坐标回归量bbox_reg,检测框目标置信度obj_conf,检测框类别置信度cls_conf。由于给每个网格点分配3个anchor,检测框坐标包含(center_x, center_y, width, height),因此维数是4*3=12,这也就明白了bbox_reg的第1个维度是12,obj_conf的第1个维度是3,而COCO数据集有80类,那么cls_conf的第1个维度应该是3*80=240,但是在上面调试信息里显示的是80类。继续设断点调试代码,在utils/utils.py里,第326行有这么一行代码

类别置信度复制了3份,结合这个后处理代码,可以看出类别置信度对3个anchor是共享的。

        在观察出Yolo-FastestV2的这些特性之后,可以理解为何它的速度快和模型文件小的原因了。主要是因为它的输入图片尺寸比传统yolov3,v4,v5系列的要小,它的输出特征图尺寸个数,也比传统yolo的要少,最后对网格点上的3个anchor是共享类别置信度的,这也减少了特种通道数。

        8月29日,我在github发布了一套使用OpenCV部署全景驾驶感知网络YOLOP,可同时处理交通目标检测、可驾驶区域分割、车道线检测,三项视觉感知任务,依然是包含C++和Python两种版本的程序实现。地址是:

https://github.com/hpc203/YOLOP-opencv-dnn

在这里我讲一下生成onnx文件需要注意的地方,YOLOP的官方代码地址是 https://github.com/hustvl/YOLOP  ,它是华中科技大学视觉团队发布的,它的代码是使用pytorch作为深度学习框架。仔细阅读和运行调试他的代码,可以看出,它的代码是在ultralytics的yolov5里修改的,添加了可行驶区域分割和车道线分割这两个分割头,在bdd100k数据集上的训练的,不过YOLOP的检测类别只保留了bdd100k数据集里的车辆这一个类别。生成onnx文件,第一步是把我发布的代码里的export_onnx.py拷贝到https://github.com/hustvl/YOLOP的主目录里。第二步,在https://github.com/hustvl/YOLOP的主目录里,打开lib/models/common.py,首先修改Focus类,原始的Focus类的forward函数里是由切片操作的,那么这时按照第(2)节里讲述的解决办法,修改Focus类,示例代码如下

class Contract(nn.Module):
    # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
    def __init__(self, gain=2):
        super().__init__()
        self.gain = gain

    def forward(self, x):
        N, C, H, W = x.size()  # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
        s = self.gain
        x = x.view(N, C, H // s, s, W // s, s)  # x(1,64,40,2,40,2)
        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # x(1,2,2,64,40,40)
        return x.view(N, C * s * s, H // s, W // s)  # x(1,256,40,40)
    
class Focus(nn.Module):
    # Focus wh information into c-space
    # slice concat conv
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
        self.contract = Contract(gain=2)
    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        # return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
        return self.conv(self.contract(x))

接下来修改Detect类里的forward函数,示例代码如下

def forward(self, x):
    if not torch.onnx.is_in_onnx_export():
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            # print(str(i)+str(x[i].shape))
            bs, _, ny, nx = x[i].shape  # x(bs,255,w,w) to x(bs,3,w,w,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            # print(str(i)+str(x[i].shape))

            if not self.training:  # inference
                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
                y = x[i].sigmoid()
                # print("**")
                # print(y.shape) #[1, 3, w, h, 85]
                # print(self.grid[i].shape) #[1, 3, w, h, 2]
                y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i]  # xy
                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                """print("**")
                print(y.shape)  #[1, 3, w, h, 85]
                print(y.view(bs, -1, self.no).shape) #[1, 3*w*h, 85]"""
                z.append(y.view(bs, -1, self.no))
        return x if self.training else (torch.cat(z, 1), x)
    else:
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            # print(str(i)+str(x[i].shape))
            bs, _, ny, nx = x[i].shape  # x(bs,255,w,w) to x(bs,3,w,w,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            x[i] = torch.sigmoid(x[i])
            x[i] = x[i].view(-1, self.no)
        return torch.cat(x, dim=0)

修改完之后,运行export_onnx.py就能生成onnx文件,并且opencv读取正常的。

        9月18日,我在github上发布了一套使用ONNXRuntime部署anchor-free系列的YOLOR,依然是包含C++和Python两种版本的程序。起初我是想使用OpenCV部署的,但是opencv读取onnx文件总是出错,于是我换用ONNXRuntime部署。地址是:

https://github.com/hpc203/yolor-onnxruntime

       9月19日,我在github上发布了一套使用深度学习算法实现虚拟试衣镜,结合了人体姿态估计、人体分割、几何匹配和GAN,四种模型。仅仅只依赖opencv库就能运行,除此之外不再依赖任何库。源码地址是:

https://github.com/hpc203/virtual_try_on_use_deep_learning

      10月6日,我在github上发布了一套使用ONNXRuntime部署鲁棒性视频抠图的程序,依然是包含C++和Python两种版本的程序。起初,我想使用opencv的dnn模块作为推理引擎,但是程序运行到cv2.dnn.readNet(modelpath) 这里时报错,因此使用onnxruntime 作为推理引擎,源码地址是:

https://github.com/hpc203/robustvideomatting-onnxruntime

        10月17日,我在github上发布了使用OpenCV部署SCRFD人脸检测,依然是包含C++和Python两种版本的程序实现。SCRFD是一个FCOS式的人脸检测器,2021年5月在insightface仓库里发布的,它也是检测人脸矩形框和5个关键点。我发布在github上的源码地址是:

https://github.com/hpc203/scrfd-opencv

        11月6日,我在github发布了使用OpenCV部署libface人脸检测和SFace人脸识别,包含C++和Python两种版本的程序,仅仅只依赖OpenCV库就能运行。源码地址是:

https://github.com/hpc203/libface-sface_detect-recognition-opencv

人脸检测和人脸识别模块是由人脸识别领域的两位大牛设计的, 其中人脸检测是南科大的于仕琪老师设计的,人脸识别模块是北邮的邓伟洪教授设计,其研究成果SFace发表在图像处理顶级期刊IEEE Transactions on Image Processing。 人脸检测示例程序在opencv-master/samples/dnn/face_detect.cpp里,起初我在win10系统里,在visual stdio 2019 里新建一个空项目,然后把opencv-master/samples/dnn/face_detect.cpp拷贝进来作为主程序,尝试编译,发现编译不通过。 仔细看代码可以发现face_detect.cpp里使用了类的继承和虚函数重写,这说明依赖包含了其他的.cpp和.hpp头文件的。因此我就编写一套程序, 人脸检测和人脸识别程序从opencv源码里剥离出来,只需编写一个main.cpp文件,就能运行人脸检测和人脸识别程序。于仕琪老师设计的libface人脸检测,有一个特点就是输入图像的尺寸是动态的,也就是说对输入图像不需要做resize到固定尺寸,就能输入到神经网络做推理的,此前我发布的一些人脸检测程序都没有做到这一点,而且模型文件.onnx只有336KB。因此,这套人脸检测模型是 非常有应用价值的。在下载完代码之后,在visual stdio 2019里新建一个空项目,配置opencv,然后把main.cpp和weights文件拷贝进去,接下来编译运行就可以了。

        12月4日,我在github发布了使用OpenCV部署P2PNet人群检测和计数,包含C++和Python两种版本的实现,仅仅只依赖OpenCV库就能运行。源码地址是:

https://github.com/hpc203/crowdcounting-p2pnet-opencv

        12月12日,我在github发布了使用OpenCV部署faster-rcnn检测证件照,包含C++和Python两种版本的程序,仅仅只依赖opencv库就能运行。源码地址是:

https://github.com/hpc203/faster-rcnn-card-opencv

还发布了使用OpenCV部署YOLOV3检测二维码,包含C++和Python两种版本的程序,仅仅只依赖opencv库就能运行。源码地址是:

https://github.com/hpc203/yolo-qrcode-opencv

        12月18日,我在github发布了使用ONNXRuntime部署PicoDet目标检测,包含C++和Python两个版本的程序,源码地址是:

https://github.com/hpc203/picodet-onnxruntime

起初,我是想使用opencv部署PicoDet目标检测的,但是opencv的dnn模块读取.onnx文件失败了,报错信息是这样的

可以看到在onnx文件里有一个opencv的dnn模块不支持的层HardSigmoid,在PicoDet官方代码仓库里搜索HardSigmoid,结果截图如下

可以看到它是直接调用pytorch的接口函数HardSigmoid()来实现的。如果自定义函数实现HardSigmoid,那么转换生成的onnx文件就能被opencv的dnn模块正常读取的了。可是PicoDet官方代码是基于PaddlePaddle框架编写的,我对这个框架不熟悉,因此没有去修改源码重新生成.onnx文件的。于是我就是用onnxruntime库部署PicoDet目标检测,在编写这套代码时,我发现之前编写的使用opencv部署nanodet的程序里,有百分之90的代码是可以复用的(拷贝粘贴过来),除了模型初始化的构造函数,需要重新编写。最后八卦一下在知乎上看到一个帖子,链接: 如何看待百度picodet工程中大量复制粘贴nanodet,却在各公众号和pr中只讲如何如何吊打后者? - 知乎

打开之后,可以看到帖子是关闭状态的。

12月28日,我在github发布了

使用OpenCV部署NanoDet-Plus,包含C++和Python两个版本的程序

使用ONNXRuntime部署NanoDet-Plus,包含C++和Python两个版本的程序

源码地址是:https://github.com/hpc203/nanodet-plus-opencv

2022年1月16日,我在github发布了使用ONNXRuntime部署人脸动漫化——AnimeGAN,包含C++和Python两个版本的代码实现,源码地址是

https://github.com/hpc203/AnimeGAN-onnxruntime

1月21日,我在github发布了

使用OpenCV部署ByteTrack目标跟踪,包含C++和Python两个版本的程序。

使用ONNXRuntime部署ByteTrack目标跟踪,包含C++和Python两个版本的程序。

源码地址是:

https://github.com/hpc203/bytetrack-opencv-onnxruntime
1月22日,我在github发布了使用ONNXRuntime部署U-2-Net生成人脸素描画,包含C++和Python两个版本的程序,源码地址是

https://github.com/hpc203/u2net-onnxruntime

1月28日,我在github发布了

使用OpenCV部署yolov5检测车牌和4个角点,包含C++和Python两个版本的程序

使用ONNXRuntime部署yolov5检测车牌和4个角点,包含C++和Python两个版本的程序

程序会输出车牌的水平矩形框的左上和右下顶点的坐标(x,y),车牌的4个角点的坐标(x,y)

源码地址是

https://github.com/hpc203/yolov5-detect-car_plate_corner

2月7日,我在github发布了使用ONNXRuntime部署yolov5-lite目标检测,包含C++和Python两个版本的程序,源码地址是

https://github.com/hpc203/yolov5-lite-onnxruntime

2月17日,我在github发布了

使用OpenCV部署多任务的yolov5目标检测+语义分割,包含C++和Python两个版本的程序

使用ONNXRuntime部署多任务的yolov5目标检测+语义分割,包含C++和Python两个版本的程序

源码地址是

https://github.com/hpc203/multiyolov5-opencv-onnxrun

2月19日,我在github发布了

使用OpenCV部署yolov5旋转目标检测,包含C++和Python两个版本的程序

使用ONNXRuntime部署yolov5旋转目标检测,包含C++和Python两个版本的程序

源码地址是

https://github.com/hpc203/rotate-yolov5-opencv-onnxrun

在编写这套程序的过程中,发现在python程序里,opencv的dnn模块提供了现成的计算旋转矩形框的NMS函数cv2.dnn.NMSBoxesRotated。但是在C++程序里,opencv的dnn模块提供现成的计算旋转矩形框的NMS函数NMSBoxesRotated。因此在C++程序里,需要自己编写实现计算旋转矩形框的NMS函数。经调研发现opencv库里有表示旋转矩形框的结构体RotatedRect,也有计算两个旋转矩形框的交叠区域和交叠面积的函数rotatedRectangleIntersection和contourArea,最终编写完成了计算旋转矩形框的NMS函数。附上一张C++程序的运行结果图

2月26日,对于https://github.com/ultralytics/yolov5 在最近更新的v6.1版本的,我在github发布了

使用OpenCV部署yolov5-v6.1目标检测,包含C++和Python两个版本的程序。

使用ONNXRuntime部署yolov5-v6.1目标检测,包含C++和Python两个版本的程序。

支持yolov5s,yolov5m,yolov5l,yolov5n,yolov5x, yolov5s6,yolov5m6,yolov5l6,yolov5n6,yolov5x6的十种结构的yolov5-v6.1

源码地址是:
https://github.com/hpc203/yolov5-v6.1-opencv-onnxrun

在这里特别讲解一下转换生成onnx文件的方法,首先下载yolov5-v6.1的源码和模型.pt文件之后,在程序主目录里,打开models/yolo.py,进入到Detect类的forward函数里,插入代码,示例截图如下:

 插入的代码片段是:

        if torch.onnx.is_in_onnx_export():
            for i in range(self.nl):  # 分别对三个输出层处理
                x[i] = self.m[i](x[i])  # conv
                bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
                x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
                y = x[i].sigmoid()
                z.append(y.view(bs, -1, self.no))
            return torch.cat(z, 1)

接下来,我尝试运行python export.py --weights=yolov5s.pt --include=onnx,但是生成onnx文件失败了。这时在export.py里,我自定义了一个导出onnx文件的函数,代码片段如下:

def my_export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
    print('anchors:', model.yaml['anchors'])
    wtxt = open('class.names', 'w')
    for name in model.names:
        wtxt.write(name+'\n')
    wtxt.close()
    # YOLOv5 ONNX export
    print(im.shape)
    if not dynamic:
        f = os.path.splitext(file)[0] + '.onnx'
        torch.onnx.export(model, im, f, verbose=False, opset_version=12, input_names=['images'], output_names=['output'])
    else:
        f = os.path.splitext(file)[0] + '_dynamic.onnx'
        torch.onnx.export(model, im, f, verbose=False, opset_version=12, input_names=['images'],
                          output_names=['output'], dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                                        'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
                                        })
    try:
        import cv2
        net = cv2.dnn.readNet(f)
    except:
        exit(f'export {f} failed')
    exit(f'export {f} sucess')

接下来在官方定义的export_onnx函数里插入调用这个函数,代码截图如下:

这时运行

python export.py --weights=yolov5s.pt --include=onnx --imgsz=640

python export.py --weights=yolov5s6.pt --include=onnx --imgsz=1280

就能成功生成.onnx文件,并且opencv的dnn模块能读取onnx文件做推理。

版权声明:本文为CSDN博主「nihate」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/nihate/article/details/112731327

nihate

我还没有学会写个人说明!

暂无评论

发表评论

相关推荐

Pytorch—万字入门SSD物体检测

前言 由于初入物体检测领域,我在学习SSD模型的时候遇到了很多的困难。一部分困难在于相关概念不清楚,专业词汇不知其意,相关文章不知所云;另一部分困难在于网上大部分文章要么只是简要介绍了SS