開源一個安全帽佩戴檢測數據集及預訓練模型

  • 2019 年 10 月 3 日
  • 筆記

本文開源了一個安全帽佩戴檢測數據集及預訓練模型,該項目已上傳至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)由於數據量的問題,佩戴其他種類的帽子可能會被誤認為是安全帽,需要更多這些“困難”樣本會達到更好的效果。

深度學習工程化是一個很複雜的任務,每一環展開都有很多的內容展開,本文只是淺嘗輒止~關於目標檢測算法的部分以後再補~