??一直想做一點(diǎn)3D目標(biāo)檢測，先來一篇單目3D目標(biāo)檢測Monodle（基于centernet的），訓(xùn)練代碼參考官方【代碼】，這里只講講如何部署。

??模型和完整仿真測試代碼，放在github上參考鏈接【模型和完整代碼】。

1 模型訓(xùn)練

??訓(xùn)練參考官方代碼 https://github.com/xinzhuma/monodle

2 導(dǎo)出onnx

??如果按照官方代碼導(dǎo)出的onnx，后處理寫起來比較復(fù)雜且后處理時耗比較長，這里將后處理的部分代碼放到模型中。原始官方導(dǎo)出的onnx模型如下圖：

本示例導(dǎo)出的onnx模型如下圖，這樣導(dǎo)出便于寫后處理代碼，模型+后處理整個時耗也比較優(yōu)：

把centernet3d.py 文件拷貝一份，命名為export_onnx.py，并進(jìn)行如下修改：


export_onnx.py 修改后的完整代碼：

import os
import cv2
import torch
import torch.nn as nn
import numpy as npfrom lib.backbones import dla
from lib.backbones.dlaup import DLAUp
from lib.backbones.hourglass import get_large_hourglass_net
from lib.backbones.hourglass import load_pretrian_modelclass CenterNet3D(nn.Module):def __init__(self, backbone='dla34', neck='DLAUp', num_class=3, downsample=4):"""CenterNet for monocular 3D object detection.:param backbone: the backbone of pipeline, such as dla34.:param neck: the necks of detection, such as dla_up.:param downsample: the ratio of down sample. [4, 8, 16, 32]:param head_conv: the channels of convolution in head. default: 256"""assert downsample in [4, 8, 16, 32]super().__init__()self.heads = {'heatmap': num_class, 'offset_2d': 2, 'size_2d': 2, 'depth': 2, 'offset_3d': 2, 'size_3d': 3,'heading': 24}self.backbone = getattr(dla, backbone)(pretrained=True, return_levels=True)channels = self.backbone.channels  # channels list for feature maps generated by backboneself.first_level = int(np.log2(downsample))scales = [2 ** i for i in range(len(channels[self.first_level:]))]self.neck = DLAUp(channels[self.first_level:], scales_list=scales)  # feature fusion [such as DLAup, FPN]self.heatmapmaxpool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)# initialize the head of pipeline, according to heads setting.for head in self.heads.keys():if head != 'heatmap':output_channels = self.heads[head]fc = nn.Sequential(nn.Conv2d(channels[self.first_level], 256, kernel_size=3, padding=1, bias=True),nn.ReLU(inplace=True),nn.Conv2d(256, output_channels, kernel_size=1, stride=1, padding=0, bias=True))else:output_channels = self.heads[head]fc = nn.Sequential(nn.Conv2d(channels[self.first_level], 256, kernel_size=3, padding=1, bias=True),nn.ReLU(inplace=True),nn.Conv2d(256, output_channels, kernel_size=1, stride=1, padding=0, bias=True),nn.Sigmoid())# initializationif 'heatmap' in head:fc[-2].bias.data.fill_(-2.19)else:self.fill_fc_weights(fc)self.__setattr__(head, fc)def forward(self, input):feat = self.backbone(input)feat = self.neck(feat[self.first_level:])ret = {}for head in self.heads:ret[head] = self.__getattr__(head)(feat)if head == 'heatmap':heatmapmax = self.heatmapmaxpool(ret[head])ret.update({'heatmapmax': heatmapmax})return retdef fill_fc_weights(self, layers):for m in layers.modules():if isinstance(m, nn.Conv2d):nn.init.normal_(m.weight, std=0.001)if m.bias is not None:nn.init.constant_(m.bias, 0)def export_onnx(model):print('===========  onnx =========== ')dummy_input = torch.randn(1, 3, 384, 1280)input_names = ['data']output_names = ['heatmap', 'offset_2d', 'size_2d', 'depth', 'offset_3d', 'size_3d', 'heading', 'heatmapmax']torch.onnx.export(model, dummy_input, './Monodle_epoch_140.onnx', verbose=False, input_names=input_names,output_names=output_names, opset_version=11)print('======================== convert onnx Finished! .... ')if __name__ == '__main__':print('This is main ...')CLASSES = ['Pedestrian', 'Car', 'Cyclist']net = CenterNet3D(backbone='dla34')checkpoint = torch.load('./weights/checkpoint_epoch_140.pth',map_location='cpu')net.load_state_dict(checkpoint['model_state'], strict=True)net.eval()export_onnx(net)input = torch.randn((1, 3, 1280, 384))print('input1:', input.shape, input.dtype)output = net(input)print(output.keys())print(output['heatmap'].shape)print(output['offset_2d'].shape)print(output['size_2d'].shape)print(output['depth'].shape)print(output['offset_3d'].shape)print(output['size_3d'].shape)print(output['heading'].shape)

運(yùn)行 python export_onnx.py 生成.onn文件。

3 測試效果

官方pytorch 測試效果

onnx 測試效果
特別說明： 由于官方代碼的2d框是用3d框計算得到的，而本博客是直接解碼的模型預(yù)測出的2d框，所以2d框有所出入。

4 onnx、rknn、horizon、tensorRT測試轉(zhuǎn)完整代碼

??模型和完整仿真測試代碼、測試圖片參考【模型和完整代碼】。