1. 📣 数据介绍

确定了业务场景之后,需要收集大量的数据(之前参加过一个安全帽识别检测的比赛,但是数据在比赛平台无法下载为己用),一般来说包含两大来源,一部分是网络数据,可以通过百度、Google图片爬虫拿到,另一部分是用户场景的视频录像,后一部分相对来说数据量更大,但出于商业因素几乎不会开放。

本项目使用开源的安全帽检测数据集(SafetyHelmetWearing-Dataset, SHWD)主要通过爬虫拿到,总共有7581张图像,包含9044个佩戴安全帽的bounding box(正类),以及111514个未佩戴安全帽的bounding box(负类),所有的图像用labelimg标注出目标区域及类别。其中每个bounding box的标签:“hat”表示佩戴安全帽,“person”表示普通未佩戴的行人头部区域的bounding box。另外本数据集中person标签的数据大多数来源于SCUT-HEAD数据集,用于判断是未佩戴安全帽的人。大致说一下数据集构造的过程:

1.数据爬取

用百度图片和Google图片的方法,百度图片用自己写的访问web页面的方式,通过不同的关键词多线程爬取数据,如果是Google图的话推荐用google-images-download,使用方法不多描述,也是爬取多个不同的关键词。关键词是个很有意思的选项,直接用“安全帽”这样的并不是一个好的选择,更多的时候可以用“建筑工人”等之类的词语;英文注意安全帽既可以是“safety Helmet”也可以是“safety hat”,“hard hat”等等。

2.数据清洗

显然用以上爬取得到的图片包含大量重复的,或者是并不包含ROI的图片,需要过滤掉大量的这些图片,这里介绍自己用到的几个方法:

(1)用已有的行人检测方法过滤掉大部分非ROI图像;

(2)可以使用深度学习模型zoo,比如ImageNet分类预训练好的模型提取特征,判断图像相似度,去除极为相似的图像;

(3)剩余的部分存在重名或者文件大小一致的图像,通常情况下这些都是不同链接下的相同图片,在数量少的情况下可以手动清洗。

3.bounding box标注

用的开源标注工具labelImg,这个没什么多说的,是个体力活,不过一个更为省力的方法是数据回灌,也就是先用标注好的一部分数据训练出一个粗糙的检测模型,精度虽然不高,不过可以拿来定位出大致的目标区域位置,然后进行手动调整bounding box位置,这样省时省力,反复这样可以减少工期。

另外标注的过程中会出不少问题比如由于手抖出现图中小圈的情形,这种情况会导致标注的xml出现bounding box的四个坐标宽或高相等,显然不符合常理,所以需要手动写脚本检查和处理有这种或者其他问题的xml的annotation,比如还有的检测算法不需要什么都没标注的背景图像,可以检测有没有这种“空”类别的数据;甚至是笔误敲错了类别的标签;等等这些都需要手动写自动化或半自动化的脚本来做纠错处理,这样的工具在标注时应该经常用到。也可以看出,一旦标注项目形成规模,规范的自动化流程会节省很多资源。

2.✨ 模型介绍

我们使用纯Tensorflow实现的YOLOv3. 包含了训练和测试自己数据集的全pipeline. 其主要的特点包括:

  • 高效的 tf.data pipeline
  • 将COCO数据集预训练的模型迁移学习
  • 支持GPU版的NMS.
  • 训练和测试推断过程全部有代码样例.
  • 使用Kmeans自己训练先验的anchor.

Python 版本: 2 or 3

Packages:

  • tensorflow >= 1.8.0 (支持tf.data的版本都可以)
  • opencv-python
  • tqdm

将预训练的darknet的权重下载,官方下载地址:https://pjreddie.com/media/files/yolov3.weights,并将该weight文件拷贝到./data/darknet_weights/

python convert_weight.py
# 注意先编译完模型再去改变最新的yolo_anchor[改变anchor请参考:使用Kmeans生成先验anchors]
3.🔰 训练数据构建

训练集的整体结构同VOC相同,可以参考VOC构建自己的数据集。

(1) annotation文件

运行

python data_pro.py

分割训练集,验证集,测试集并在 ./data/my_data/labal 下生成 train.txt/val.txt/test.txt ,对于一张图像对应一行数据,包括 image_index , image_absolute_path , img_width , img_height , box_1 , box_2 ,…, box_n,每个字段中间是用空格分隔的,其中:

  • image_index
  • image_absolute_path
  • img_width , img_height , box_1 , box_2 ,…, box_n
  • box_x的形式为:label_index,x_min,y_min,x_max,y_max
  • label_index是label对应的index(取值为[0~class_num-1]),这里要注意YOLO系列的模型训练与SSD不同,label不包含background

例子:

0 xxx/xxx/a.jpg 1920,1080,0 453 369 473 391 1 588 245 608 268
1 xxx/xxx/b.jpg 1920,1080,1 466 403 485 422 2 793 300 809 320
...

(2) class_names文件:

coco.names 文件在 ./data/

hat
person

(3) 先验anchor文件:

使用Kmeans生成先验anchors:

python get_kmeans.py




PyTorch 安全帽监测 tensorflow安全帽检测_检测安全帽佩戴


可以得到9个anchors和平均的IOU,把anchors保存在文本文件:./data/yolo_anchors.txt,

注意: Kmeans计算出的YOLO Anchors是在调整大小的图像比例的,默认的调整大小方法是保持图像的纵横比。

4.📝 训练

修改arg.py中的一些参数,如下:

  • 修改arg.py

运行:

CUDA_VISIBLE_DEVICES=GPU_ID python train.py

我们训练的环境为:

  • ubuntu 16.04
  • Tesla V100 32G
5.🔖 推断

我们使用 test_single_image.py 和 video_test.py 推断单张图片和视频,测试Demo在 6.⛏Demo

python3 test_single_image.py /home/myuser/xujing/YOLO_V3_hat/data/my_data/JPEGImages/000002.jpg
6.⛏Demo


PyTorch 安全帽监测 tensorflow安全帽检测_YOLO V3_02


PyTorch 安全帽监测 tensorflow安全帽检测_PyTorch 安全帽监测_03


PyTorch 安全帽监测 tensorflow安全帽检测_检测安全帽佩戴_04


7.⛏训练的一些Trick

(1) 使用two-stage训练或one-stage训练:

  • Two-stage training:
  • 第一阶段:在COCO数据集训练的ckeckpoints上加载darknet53_body部分的weights,训练YOLO V3的head部分,使用较大的学习率比如0.001,直到损失降下来;
  • 第二阶段:加载第一阶段训练的模型,训练整个模型的参数,使用较小的学习率比如0.0001。
  • One-stage training:

直接加载除Conv_6,Conv_14和Conv_22(这三层是输出层需要根据自己训练数据调整)的预训练模型,这种情况需要注意Loss的nan问题,对于该项目为了简单,我们采用One-stage training。

(2) args.py中有很多有用的训练参数调整策略:

  • 学习率的decay(Cosine decay of lr (SGDR))
  • 多尺度训练(Multi-scale training)
  • 标签平滑(Label smoothing)
  • 数据增强(Mix up data augmentation)
  • Focal loss(来源于RetinaNet主要修正目标检测中的unblance问题)

这么多策略,不一定都能提升你的模型性能,根据自己的数据集自行调整选择.