Pytorch-SSD目标检测算法(Single Shot MultiBox Detector)(简单,明了,易用,全中文注释,单机多卡训练,训练过程可视化,视频检测)
SSD-Pytorch
作者为pytorch中文网维护者之一,网站会发许多深度学习相关文章,欢迎pytorch使用者多多支持.
项目Github地址:Github 欢迎star, fork
一. 项目简介
-
项目是在实际项目上整理精简而来,主要针对单机多卡环境实现.
-
项目整体结构清晰明了,对训练过程在不失灵活性的前提下进行了包装.训练,测试,检测过程均只需简单几行即可实现.
-
如需深入了解,可对照项目文件结构进行理解,项目整体均有详细中文备注.
-
完整项目地址:Github,欢迎star, fork.
-
github连接较慢的,可以去Gitee(国内的代码托管网站),也有完整项目.
1.1 项目特点
1.结构化实现.模型各部件均分离实现,模型结构更为清晰明了,便于理解以及后续修改使用.
2.针对单机多卡环境实现.简单直观的多GPU调用实现. 只需通过列表指明需调用的GPU序号即可实现多GPU环境下的训练与测试.如: train_devices=[0,1,3]即可调用第1,2,4三块GPU用于训练.
3.全中文注释.每部分均有详细的中文注释,是你学习SSD,加深理解的不二之选.
1.2 Requirements
1.pytorch
2.opencv-python
3.torchvision >= 0.3.0
4.Vizer (在代码实现上也借鉴了lufficc大神的一些思想)
5.visdom
1.3 项目文件结构
| 文件夹 | 文件 | 说明 |
|---|---|---|
| Data | 数据相关 | |
| Dataloader | 数据加载器类’Our_Dataloader’, ‘Our_Dataloader_test’ | |
| Dataset_VOC | VOC格式数据集类 | |
| Transfroms | 数据Transfroms | |
| Transfroms_tuils | Transfroms子方法 | |
| Model | 模型相关 | |
| base_models/VGG | 现只有vgg模型,后期会更新其他模型 | |
| structs/Anchors | 默认检测框生成器 | |
| structs/MutiBoxLoss | SSD损失函数 | |
| structs/PostProcess | 后处理 | |
| structs/Predictor | 分类及回归网络 | |
| evaler | 验证器,用于在数据集上对模型进行验证(测试),计算ap,map | |
| ssd_model | SSD模型类 | |
| trainer | 训练器,用于在数据集上训练模型 | |
| Utils | 各种工具 | |
| boxs_op | 各种框体操作,编码解码,IOU计算,框体格式转换等 | |
| Weights | 模型权重存放处 | |
| pretrained | 预训练模型权重存放处 | |
| trained | 训练过程中默认模型存放处 | |
| ---- | Configs.py | 配置文件,包含了模型定义,数据以及训练过程,测试过程等的全部参数,建议备份一份再进行修改 |
| ---- | Demo_train.py | 模型训练的例子,训练过程中的模型会保存在Weights/Our/ |
| ---- | Demo_eval.py | 模型测试的例子,计算模型ap,map |
| ---- | Demo_detect_one_image.py | 检测单张图片例子 |
| ---- | Demo_detect_video.py | 视频检测例子,传入一个视频,进行检测 |
二. Demo例子
本项目Demo均是完整的训练,测试,检测过程,可直接运行.
2.1 训练train
本项目采用visdom进行训练过程可视化
visdom安装 及运行:
visdom安装
pip install visdom 安装
python -m visdom.server 运行
打开网页`127.0.0.1:8097` 进行训练过程可视化
(如未运行visdom,会报警告,但是不影响训练过程.建议通过visdom进行可视化)
训练过程通过训练器进行配置,训练器支持单机单卡,单机多卡训练;不支持多机多卡以及cpu.
# 实例化模型. 模型的具体各种参数在Config文件中进行配置
net = SSD(cfg)
# 将模型移动到gpu上,cfg.DEVICE.MAINDEVICE定义了模型所使用的主GPU
# 模型的参数更新会在主GPU上进行.
net.to(cfg.DEVICE.MAINDEVICE)
# 初始化训练器,训练器参数已通过cfg进行配置;也可传入参数进行配置,但不建议
trainer = Trainer(cfg, max_iter=None, batch_size=None,
train_devices=None, model_save_step=None,
model_save_root=None, vis = None, vis_step=None)
# 训练器开始在 数据集上训练模型
trainer(net, train_dataset)
2.2 测试Eval
对模型进行测试,测试器支持单机单卡,单机多卡测试,但不支持多机多卡,cpu.
net = SSD(cfg)
# 将模型移动到gpu上,cfg.DEVICE.MAINDEVICE定义了模型所使用的主GPU
net.to(cfg.DEVICE.MAINDEVICE)
# 模型从权重文件中加载权重
net.load_pretrained_weight('Weights/pretrained/vgg_ssd300_voc0712.pkl')
# 初始化验证器,验证器参数通过cfg进行配置;也可传入参数进行配置,但不建议
evaler = Evaler(cfg, eval_devices=None)
# 验证器开始在数据集上验证模型
ap, map = evaler(model=net, test_dataset=test_dataset)
print(ap)
print(map)
2.3 检测Detect
检测过程本项目给出俩个例子,分别为单张图片检测与视频检测
2.3.1 图片检测
单张图片检测,通过调用模型Detect_single_img方法实现.支持gpu与cpu,依模型运行环境而定.
# 实例化模型
net = SSD(cfg)
# 使用cpu或gpu
net.to('cuda')
# 模型从权重文件中加载权重
net.load_pretrained_weight('Weights/pretrained/vgg_ssd300_voc0712.pkl')
# 打开图片
image = Image.open("Images/000133.jpg")
# 进行检测, 分别返回 绘制了检测框的图片数据/回归框/标签/分数.
drawn_image, boxes, labels, scores = net.Detect_single_img(image=image,score_threshold=0.5)
plt.imshow(drawn_image)
plt.show()
2.3.2 视频检测
# 实例化模型
net = SSD(cfg)
# 使用cpu或gpu
net.to('cuda')
# 模型从权重文件中加载权重
net.load_pretrained_weight('Weights/pretrained/vgg_ssd300_voc0712.pkl')
video_path = 'aaa.mp4' # 现在忙于其他事,项目后期会上传视频检测样例
# 进行检测,
# save_video_path=None,则不保存视频; save_video_path=bbb.mp4 将检测后视频保存为bbb.mp4
# show=True,实时显示检测结果
net.Detect_video(video_path=video_path, score_threshold=0.02, save_video_path=None, show=True)
说明
本项目是在实际项目上整理精简而来,现配置文件为在voc2007上进行训练,如有其他需求,请自行修改.
完整项目地址:Github,欢迎star, fork.
仅限用于交流学习,如需引用,请联系作者.