开源一个安全帽佩戴检测数据集及预训练模型

Posted 2022-09-23 supersayajin

本文开源了一个安全帽佩戴检测数据集及预训练模型，该项目已上传至github，点此链接，感觉有帮助的话请点star 。同时简要介绍下实践上如何完成一个端到端的目标检测任务。可以看下效果图：

 同时该模型也可以做人头检测，效果如下：

一、背景介绍

最近几年深度学习的发展让很多计算机视觉任务落地成为可能，这些任务渗透到了各行各业，比如工业安全，包含的任务如安全帽佩戴检测、高空坠物检测、异常事故检测（行人跌倒不起等）等等，这里以安全帽检测为例，简单介绍下如何完成一个端到端的任务，包括：

1. 业务场景分析

2. 数据收集、预处理及标注

3.算法设计及训练

4.模型转换及部署

二、业务场景分析

 怎样定义“安全帽佩戴检测”这个问题呢？正常来看是需要判断行人头部区域有没有安全帽，比如有些做法是先用目标检测模型如SSD、Yolo系列及改进系列（各种加速版及适用移动端的网络）检测行人区域，然后再设计一个小的分类网络判断区域内是否存在安全帽。这种做法的好处是比较简单：

一来，行人检测算是做的比较多的应用，数据集、算法都比较多；

二来，很多数据都是以监控视频的形式存在，标注的时候标注较粗粒度的区域即可，然后做分类。

不过这种做法的劣势也比较明显：

第一，这种做法并不是一个端到端的预测过程，显然要做到先检测出行人，再用分类网络判断是否存在安全帽；

第二，目标检测和分类的区别，如果用这种粗粒度的标注方法，也就是标注出包含安全帽的较大的区域，而不是紧紧包含安全帽的那块区域，分类效果并不会太好；

第三，严谨来说，对于用户更希望得到的是，定位出佩戴安全帽的目标位置，很多情况安全帽虽然在行人区域内，但这时候并不是出于“佩戴”状态。

基于以上分析，这里用的方法是注重安全帽“佩戴”状态，所以业务层面来看，需要的数据是有“佩戴”安全帽和未“佩戴”安全帽的图片。如下图所示：

                     

其中蓝色的bounding box是未佩戴安全帽的目标，红色的bounding box是佩戴安全帽的目标，这样可以更精准、更直接的确定是否佩戴的状态，不过需要覆盖更多场景的正负类目标数据。

三、数据收集、预处理及标注

确定了业务场景之后，需要手机大量的数据，一般来说包含两大来源，一部分是网络数据，可以通过百度、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，比如还有的检测算法不需要什么都没标注的背景图像，可以检测有没有这种“空”类别的数据；甚至是笔误敲错了类别的标签；等等这些都需要手动写自动化或半自动化的脚本来做纠错处理，这样的工具在标注时应该经常用到。也可以看出，一旦标注项目形成规模，规范的自动化流程会节省很多资源。

四.算法设计及训练

基于之前的分析，这里的任务可以看成一个常规的目标检测的情形，目标检测算法发展很快，单阶段的多阶段的，各种anchor及anchor-free的，GIoU，soft-NMS，各种金字塔多尺度预测的等等，不过很多sota的方法很难应用在实际项目中，毕竟速度、功耗、成本等因素均要考虑在内，本文仅用gluoncv自带的model-zoo中的yolo系列做检测，用了原始的darknet53、mobilenet1.0级mobilenet0.25做示例，使用的时候注意调整输入图像的尺寸参数。更多的效果图如下：

             

以及较为密集情况下的人头检测（测试的时候注意短边尺寸的大小）：

 

       

 

模型的使用清参考github上的项目：点此链接。

五.模型转换及部署

 如果说前几年深度学习主题追求的是越来越深的模型及越来越准的方法，那么最近两年关注的点是模型的轻量化及部署，这是一个比较大而且比较复杂的话题。即训练好的模型怎样部署在实际的项目中，也就是我们常说的train和inference(或者forward，test，deploy)主要关注以下几点：

(1) 场景和平台选择，是移动端还是服务器端的，CPU还是GPU的，速度和准确率要求等等。有的时候像TensorFlow、PyTorch、MXNet等框架不太适合直接拿来部署，即使用这些框架部署也要结合硬件平台手动做一些配置上的加速，比如Intel CPU可以用MKL或者SSE、AVX等SIMD指令集加速等等。

(2) 移动端有很多大厂的forward框架比如主流的NCNN，上半年阿里开源的MNN，以及比较低调的armNN等等，目前看NCNN用的最多，MNN很有潜力，移动端的GPU支持貌似不是很好，NVIDIA GPU的话主流还是他们自家的TensorRT加速。训练好的模型部署到移动端需要做模型转换，要注意有些框架只支持几个大厂的格式，比如TensorFlow，PyTorch、Caffe、onnx等，话说这里的onnx貌似发展的并不怎么样，感觉caffe本身反而是最时候做中间转换的桥梁，哈哈。。比如MXNet->Caffe->MNN；另外诸多框架的op也尚未完全统一，很多需要自己自定义op实现，还有的方法是做截断，也就是拿支持的那部分网络的op做转换，先计算得到这部分的结果，后接自己实现的剩余部分。

(3) 用TVM做编译优化，这个看起来比较有意思，很可能是深度学习部署优化的方向之一，针对不同框架、不同语言的前端（目前使用relay）生成中间层IR，然后针对不同的硬件平台做针对性的优化工作。优化的内容包括内存分配、算子融合、GEMM优化等等。之前试过mobilefacenet的TVM版本，差不多能加速30%左右，14ms->9.5ms还是非常给力的，但是目前看默认的优化选项并不是最优的，比如我也试过mxnet gluoncv中的模型还不如不用TVM来得快，更多的时候需要用tvm中auto-tuning方法搜索一个最优的配置，这个过程比较耗时，一个模型可能需要几个小时来实验。

(4) 手动写forward方法，如果做到完全脱离第三方库，需要自己手写forward方法，基本思路还是im2col+GEMM等等，做好内存分配，数据对齐重排，使用指定架构的寄存器提高访存效率、提高cache命中率等等，模型压缩和优化的坑太深，这里不多说了。。。。

六. 总结

 本文开源了一个安全帽佩戴检测数据集和预训练模型点此链接，目前来看该数据集还有几个不足之处：

(1)数据量还是少，总共接近一万个bounding box，真实场景覆盖的越多越好，比如至少十万以上；

(2)业务规则设计是否真的好呢？这里有个问题，其实安全帽佩戴的bounding box也包含人脸，因此会存在数据标签重叠的情况，标注普通人头的时候bounding box稍微大一些覆盖到整个头部边缘而非人头的一部分效果会更好；

(3)由于数据量的问题，佩戴其他种类的帽子可能会被误认为是安全帽，需要更多这些“困难”样本会达到更好的效果。

深度学习工程化是一个很复杂的任务，每一环展开都有很多的内容展开，本文只是浅尝辄止~关于目标检测算法的部分以后再补~