深度学习目标检测系列论文阅读和Pytorch实现(三)——SSD代码实现

本章介绍SSD的Pytorch代码实现。

具体原理参考上一章节:https://blog.csdn.net/IEEE_FELLOW/article/details/104241995

代码我已经上传我的github,欢迎star:https://github.com/midasklr/SSD.Pytorch

主要参考了https://github.com/amdegroot/ssd.pytorch代码的实现,但是原作者代码年久失修,并且仅仅实现了SSD300的结构。

在此基础上,我修改支持最近几个版本的Pytorch,并且参考caffe原实现实现了相应的SSD512。

我的环境如下:

  1. Ubuntu1604;
  2. python3.7 ;
  3. pytorch1.3;
  4. opencv等

数据集

我在PascalVOC数据集上训练了SSD300和SSD512,并和原论文结果进行比较:

model paper this implements
SSD300 77.2 77.43
SSD512 79.8 79.80

模型文件:链接: https://pan.baidu.com/s/1DxlkOQzkFkkdYdNYsDx_MQ 提取码: dd7m 复制这段内容后打开百度网盘手机App,操作更方便哦

当然你也可以训练自己的数据集,这里我使用SSD512训练了一个车牌检测,效果如下:

效果还行,就是延时还不理想,毕竟SSD结构还是比较老。但是这里未做特征融合还能识别出场景中小目标的车牌还是很惊人的。

训练

具体见README。

测试

python eval.py --input 512 --trained_model weights/ssd512_VOC_73000_mAP79.80.pth

结果如下:

AP for aeroplane = 0.8861

AP for bicycle = 0.8694

AP for bird = 0.8078

AP for boat = 0.7698

AP for bottle = 0.6407

AP for bus = 0.8625

AP for car = 0.8825

AP for cat = 0.8671

AP for chair = 0.6424

AP for cow = 0.8712

AP for diningtable = 0.6781

AP for dog = 0.8572

AP for horse = 0.8781

AP for motorbike = 0.8531

AP for person = 0.8091

AP for pottedplant = 0.5479

AP for sheep = 0.8327

AP for sofa = 0.7562

AP for train = 0.8654

AP for tvmonitor = 0.7824

Mean AP = 0.7980

Demo

提供了demo.py,运行文件可以测试单张图片:

结果

在VOC2007 测试集上:

model paper this implements
SSD300 77.2 77.43
SSD512 79.8 79.80

训练过程:

训练中MAP和epoch关系如下:

可以看到大约在250个循环左右达到最优秀,对比原作者Caffe中的参数发现基本吻合。

代码已经上传github:https://github.com/midasklr/SSD.Pytorch

欢迎大家star和训练自己的数据集。有问题直接在issue上留言即可。

Guess you like

Origin blog.csdn.net/IEEE_FELLOW/article/details/104242232