用CenterNet对旋转目标进行检测 | 附源码
共 2214字，需浏览 5分钟
·
2020-12-04 11:19
点击上方“AI算法与图像处理”，选择加"星标"或“置顶”
重磅干货，第一时间送达
导读
该项目主要内容包括对原版的CenterNet目标检测代码进行了精简，使用更加方便，同时作者也在持续维护中，有问题的开源去GitHub上提issue，后续会开放工程化工作 C++落地
前言
前段时间纯粹为了论文凑字数做的一个工作，本文不对CenterNet原理进行详细解读，如果你对CenterNet原理不了解，建议简单读一下原论文然后对照本文代码理解（对原版CenterNet目标检测代码进行了极大程度精简）。
基本想法就是直接修改CenterNet的head部分，但是是直接在长宽图上加一层通道表示角度，还是多引出一 路feature-map呢？实测是多引出一张feature map比较好，直接在长宽图上加一层通道很难收敛，具体原因我也是猜测，角度和尺度特征基本无共享特征，且会相互干扰（角度发生些许变化，目标的长宽可能就变了，如果角度是错的，长宽本来是对的呢？反之亦然）引出的feature-map只经历了一层卷积层就开始计算loss，对于这种复杂的关系表征能力不够，可能造成弄巧成拙。网络结构如下：
代码
代码主要分为五个部分：
train.py：模型的训练
predict.py：模型的前向推理
backbone：模型的主干网，给了DLA和ResNet的DCN与普通版本，层数可以自定义
loss.py：模型的损失函数
dataset.py：模型的数据处理
不是很重要：
data/airplane：示例训练数据与图片
dcn：编译好的dcn，说明一下，这里与原版CenterNet编译dcn一样，直接文件夹复制过来即可，如果你不会编译dcn，就用backbone内的dlanet.py与resnet.py
原版CenterNet代码较多，我只需要做目标检测，所以把各种3D检测等都删了，模型架构也拆了只保留了有用部分，方便自己阅读以及魔改。
其次，因为只是加了一个角度检测，所以主要是修改了一下数据处理部分，用的还是VOC格式，只是在.josn文件里面加了一个角度信息，打标签的时候用[0,π]表示，后续在loss内添加了角度的feature-map损失，用的Smooth-L1 loss，打标签方法已在下方更新。
2020.1021代码已更新(不是很重要)
labelGenerator：生成自己的训练数据
evaluation.py：性能指标计算
imgs：性能指标计算示例图片
鉴于一些同学想知道怎么对自己的数据打标签以及生成R-CenterNet可以训练的数据，所以更新一个labelGenerator文件夹，内包含转换函数以及用例。注意，这个文件夹以及其内部函数不是网络必须的，只是想训练自己打标签的数据时用的。
鉴于一些同学想知道模型训练完毕，怎么对比性能，比如旋转框怎么计算mIOU等，所以更新一个evaluation.py以及对应的案例图片和文件夹imgs。注意，这个.py和imgs文件夹不是必须的，只是模型训练结束计算性能用的。
注：
每个label里面的目标五个数值：
分别为目标中心点（x,y），以及宽度，长度，角度，角度是以12点钟为0°，顺时针旋转的角度，最大为179.99999°（旋转180°，相当于没转）。
Demo
训练自己的多分类网络
理论上torch >1.0即可，如果报了显存不足的问题就是torch版本低了
打标签用labelGenerator文件夹里面的代码
修改代码中所有num_classes为你的类别数目
增加predict.py中方框颜色，我这里只检测单目标，所以只有红蓝框。
修改back_bone中hm的数目为你的类别数，如：
def DlaNet(num_layers=34, heads = {'hm': your classes num, 'wh': 2, 'ang':1, 'reg': 2}, head_conv=256, plot=False):最后，有问题可以去github上看一下issue，或者私信留言。
后续有时间会将上面的工作工程化，用C++落地。
感谢每个开源的作者给我们提供了大量可以参考的案例，让我们能够有条件去使用，不用花费大量的时间去做重复的工作，未来会我会继续分享优质的工作内容和项目！希望小伙伴看到时，可以顺手点个“在看”或许能帮到其他人，当人这也是我坚持的动力之一！！！感谢
下载1：leetcode 开源书
在「AI算法与图像处理」公众号后台回复：leetcode，即可下载。每题都 runtime beats 100% 的开源好书，你值得拥有！
下载2 CVPR2020 在「AI算法与图像处理」公众号后台回复：CVPR2020，即可下载1467篇CVPR 2020论文 个人微信（如果没有备注不拉群！） 请注明：地区+学校/企业+研究方向+昵称
觉得不错就点亮在看吧
