基于SSD算法的电动车头盔检测【附代码】

本项目是基于SSD算法实现的电动车头盔检测。完整的项目是基于SSD的改进-->知识蒸馏-->自蒸馏-->剪枝-->trt推理。本想用来发论文的,但可能没那时间和精力了,这里仅选择项目中的一部分内容进行开源。

目录

开源说明

项目功能说明

项目代码的参数说明

预测

两次结果在同窗口显示

分窗口显示

单目标检测显示

剪枝

训练

训练自己的数据集

剪枝后进行模型的微调训练

FPS测试

mAP测试

torch2onnx

onnx2engine[tensorRT推理]

代码链接:

权重链接:

报错


开源说明

本项目是开源项目,近期发现有些人拿我的其他项目去发布收费资源,这种情况是不允许的,我希望我发布的开源项目是帮助大家学习交流的。后期再发现有拿我开源项目进行收费的我将举报,同时未来本人将考虑关闭所有开源项目。

另:如想用本项目发表论文或者毕设等,均需经过本人同意并且有偿或挂名,否则将会举报!

项目功能说明

1.支持训练自己的数据集

2.支持剪枝

3.支持剪枝后的微调训练

4.支持mAP测试

5.支持FPS测试

6.支持tensorRT推理

7.支持检测结果在同一窗口显示

8.支持检测结果分窗口显示

9.支持单目标检测

注:我这里实现的电动车头盔检测,是分两步进行的(所以有多窗口或者单窗口显示):第一步先对骑乘人员和电动车整体进行检测;第二步是在前一次检测结果的基础上检测是否佩戴头盔。如果想一步检测,可以用单目标检测。

这种检测方法的好处有:

        1.可以与普通行人进行区分,不会将普通行人也检测进去;

        2.如果存在一辆车上有多个人员也可以实现头盔检测;

这种检测方法的缺点:

        由于分两次检测,因此占据内存比较高,在实际应用中实时检测有影响,因为会过两次GPU。


项目代码的参数说明

(helmet_dectection.py)

--model:目标检测模型选择,现只支持SSD

--cuda:是否使用GPU,默认为True

--output:一些功能需要指定的输出路径

--conf_thres:置信度阈值,默认0.6

--iou_thres:iou阈值,默认0.5

--target_weights:第一次j检测(人与电动车整体)权重路径

--helmet_weights:第二次(只对头部)权重路径,默认为EC/model_data/pruning_2th_model.pth

--fps:测试FPS,默认False

# 支持300*300,512*512分辨率的输入

--input_shape:第一次检测输入分辨率,默认512

--input_shape2:第二次检测输入分辨率,默认512

# 剪枝,在剪枝之前需要确定是否完整保存了模型而不是仅保存权重,如果没有保存完整的网络结果需要去tools/savemodel.py保存完整的网络和权重才能进行剪枝

--pruning_model:开启剪枝(默认对vgg主干剪枝)

--pruning_weights:需要剪枝的权重

# 训练配置

--train:是否开启训练

--is_fine:是否开启剪枝后的微调训练

--pruned_model_path:剪枝后的权重路径【在微调训练的时候传入】

--batch_size:batch大小

--Init_Epoch:初始化Epoch

--Freeze_Epoch:冻结训练Epoch,默认前50Epoch冻结训练

--UnFreeze_epoch:解冻训练,默认50-100解冻训练

--UnFreeze_lr:解冻后学习率

# 预测配置

--predict:是否开启预测模式,开启后默认两次检测在一个窗口显示

--predict_2windows:检测结果在两个窗口显示(分窗显示)

--predict_single:只检测一个目标

--video:是否检测视频

--video_path:视频路径

--image:是否检测图像

--mAP:是否测试mAP

--trt:是否采用trt进行推理


预测

两次结果在同窗口显示

python helmet_dectection.py --predict --video --video_path 【your video path】  --target_weights [一次检测权重路径] --helmet_weights [二次检测权重路径]

效果如下(对头部的检测会在同意一画面的右下角显示):

电动车头盔检测

图像预测可结合上面的命令输入:

python helmet_dectection.py --predict --image

分窗口显示

--predict_2windows --video --video_path 【视频路径】 --target_weights 【第一次检测权重】 --helmet_weights 【第二次检测权重】 --input_shape 512 --input_shape2 512

此时会显示两个窗口显示两次检测的结果(该功能暂时不支持trt)

我这里附上我论文中的效果图

基于SSD算法的电动车头盔检测【附代码】_第1张图片 第一次检测
基于SSD算法的电动车头盔检测【附代码】_第2张图片 第二次头部的检测

 

单目标检测显示

如果只想对第一次或者第二次单独检测显示,可以使用该功能【注意修改config.py和classes.txt中的内容】。例如:

python helmet_dectection.py --predict_single --video --video_path [视频路径] --trt

剪枝

输入以下命令:

python helmet_dectection.py --pruning_model --pruning_weights=你需要剪枝的权重  --output=剪枝以后的权重保存在哪里

比如:

python helmet_dectection.py --pruning_model --pruning_weights=model_data/whole_model.pth  --output=model_data/

出现以下信息说明剪枝成功,剪枝后的模型将保存在model_data文件下,名字为pruning_model.pth,如果需要修改保存后权重名字和剪枝细节,去tools/Model_Pruning.py中修改

[ prune_conv on vgg.24 (Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))>, Index=[1, 2, 3, 5, 8, 10, 11, 13, 17, 18, 20, 25, 28, 29, 34, 36, 38, 40
, 44, 46, 50, 53, 58, 59, 60, 63, 65, 66, 68, 69, 74, 76, 81, 82, 83, 85, 87, 91, 95, 96, 97, 99, 101, 103, 109, 113, 115, 117, 119, 123, 124, 131, 132, 133, 134, 138, 139, 141, 143, 145,
 146, 147, 148, 150, 155, 158, 162, 163, 166, 173, 175, 176, 177, 179, 182, 187, 194, 195, 198, 199, 200, 201, 203, 205, 208, 212, 217, 221, 225, 228, 229, 235, 237, 240, 245, 248, 250, 2
52, 253, 255, 256, 260, 265, 267, 269, 270, 271, 273, 275, 277, 278, 280, 282, 284, 286, 287, 288, 289, 293, 294, 296, 302, 307, 308, 309, 312, 313, 317, 320, 323, 325, 330, 337, 339, 341
, 347, 349, 350, 353, 356, 359, 361, 362, 364, 365, 370, 371, 372, 379, 381, 382, 383, 384, 385, 387, 388, 396, 400, 406, 415, 418, 423, 426, 427, 431, 432, 436, 437, 438, 439, 440, 441,
442, 444, 447, 450, 452, 453, 457, 458, 459, 460, 462, 464, 465, 466, 468, 472, 475, 482, 484, 485, 486, 487, 490, 491, 494, 495, 496, 499, 502, 503, 507, 511], NumPruned=940236]

训练

训练自己的数据集

1.如果想训练自己的数据集。将数据集图片放在VOCdevkit/VOC2007/JPEGImages/下,xml标签放在VOCdevkit/VOC2007/Annotations/下。然后运行voc2ssd.py。

2.然后修改voc_annotation.py中的classes并运行。

3.运行下面代码:

python helmet_dectection.py --train --target_weights 权重路径

剪枝后进行模型的微调训练

python helmet_dectection.py --train --is_fine --pruned_model_path=你剪枝后的权重路径

比如:

python helmet_dectection.py --train --is_fine --pruned_model_path=model_data/pruning_model.pth 

【注意:上述命令batchsize,学习率,冻结网络训练的epoch都采用默认值,需要更改直接传参就可以】

训练的时候,需要根据自己的类,去utils/config.py中修改参数,主要修改类别数,类别数量为自己的类的数+1(含有背景类),比如自己的数据集是20个类,那么num_classes=21,也可以修改min_dim【SSD支持分辨率为300和512,如果修改了分辨率,feature_map也需要修改】

如果只是正常的训练,就不需要传入is_fine即可开启正常的训练。

FPS测试

python helmet_dectection.py --fps

因为FPS测试代码是继承了tools/ssd.py中的SSD类,如果需要预测自己的类,需要注意以下ssd.py中classes_path路径中要写自己的类,同时注意传入的权重,比如输入以下命令【前提已经将classes_path写为自己的类】

mAP测试

mAP测试也是继承了tools/ssd.py,classes_path修改为自己的类,然后输入:

python helmet_dectection.py --mAP --target_weights=权重路径 --conf_thres 0.5 --iou_thres 0.5

会在当前目录中生成results文件,里面有mAP结果和F1值,P值和Recall

torch2onnx

修改torch2onnx.py中的ckpt_path,样例x的输入大小即可生成onnx权重

注意:由于本部分代码是使用蒸馏后的网络,因此有个del操作删除不需要的层,如果想使用最原始的网络,用net_student/ssd_student.py即可。

onnx2engine[tensorRT推理]

修改onnx2trt.py的权重路径再运行即可生成对应的engine。有关详细的tensorrt推理教程看我另一篇文章:

SSD目标检测网络tensorRT推理【附代码】_ssd tensorrt_爱吃肉的鹏的博客-CSDN博客

代码链接:

https://github.com/YINYIPENG-EN/E-bike_helmet_dectection.git

权重链接:

链接:https://pan.baidu.com/s/1mBEi5PJl7ibfwzJ8PBrAqQ 
提取码:yypn

报错

报错1:

conf_preds = conf_data.view(num, num_priors, self.num_classes).transpose(2, 1)
RuntimeError: shape '[1, 24656, 2]' is invalid for input of size 17464

在SSD中,共有6个预测特征层。每个特征层分配的先验框数量分别为:4,6,6,6,4,4

以输入大小为:300×300为例,共有先验框38×38×4+19×19×6+10×10×6+5×5×6+3×3×4+1×1×4=8732个先验框

如果是输入大小为512×512,共有先验框64×64×4+32×32×6+16×16×6+8×8×6+4×4×4+2×2×4=24656个先验框

因此上面的错误是输入大小和真正的模型输入大小不匹配,需要检查输入大小

 

你可能感兴趣的:(算法)