Pytoorch轻松学–RetinaNet自定义对象检测

原创 OpenCV学堂 2023-11-15 08:43

【直播】创新芯片重塑机器人未来 入门级示波器的属性缺陷

点击上方蓝字关注我们

微信公众号：OpenCV学堂
关注获取更多计算机视觉与深度学习知识

RetinaNet网络介绍

2017发布，实现了一阶段网络首次在精度方面超过二阶段网络的经典网络，作者最大的一个创新就是在训练损失函数方面，论文比较了CE、BCE、以及论文提出感知损失函数（FL），最后说明感知损失可以有效解决一阶段网络训练中的样本不平衡现象，从而取得更佳的训练效果。论文中提出的感知损失函数如下：

最终RetinaNet网络结构如下：

数据集准备与制作

自己百度收集了一个无人机与飞鸟的数据集，其中训练集270张图像，测试集26张图像。

使用labelImg工具完成标注，工具下载地址：

https://gitee.com/opencv_ai/opencv_tutorial_data/tree/master/tools

三：模型训练

制作好数据集之后，模型训练就成为一件很简单事情，基于OpenMV工具软件，零代码即可实现模型训练。运行下面界面如下：

总计训练了25个轮次以后，发现效果已经是相当的不错了，直接导出ONNX格式RetinaNet模型文件。

把模型转换为ONNX格式，Pytorch是原生支持的，只需要把通过torch.onnx.export接口，填上相关的参数，然后直接运行就可以生成ONNX模型文件。相关的转换代码如下：

model = tv.models.detection.retinanet_resnet50_fpn(pretrained=True)
dummy_input = torch.randn(1, 3, 1333, 800)
model.eval()
model(dummy_input)
im = torch.zeros(1, 3, 1333, 800).to("cpu")
torch.onnx.export(model, im,
                    "retinanet_resnet50_fpn.onnx",
                    verbose=False,
                    opset_version=11,
                    training=torch.onnx.TrainingMode.EVAL,
                    do_constant_folding=True,
                    input_names=['input'],
                    output_names=['output'],
                    dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}}
                  )

运行时候控制台会有一系列的警告输出，但是绝对不影响模型转换，影响不影响精度我还没做个仔细的对比。

模型转换之后，可以直接查看模型的输入与输出结构，图示如下：

推理运行

推理部分的代码很简单，只有三十几行，Python就是方便使用，这里最需要注意的是输入图像的预处理必须是RGB格式，需要归一化到0~1之间。对得到的三个输出层分别解析，就可以获取到坐标（boxes里面包含的实际坐标，无需转换）。基于OpenMV工具软件，可以实现一键零代码推理演示，效果如下：