點擊上方“計算機視覺cv”即可“進入公眾號”

重磅干貨第一時間送達

前言

隨著人工智能的發展，現在越來越多的場景需要人工智能。在工廠的廠區中以安全為首，但工人普遍缺乏佩戴安全帽意識；工廠環境復雜，有各種各樣的禁止進入的區域，普通的圖像識別算法很難實現；加上使用傳統的人工監管存在諸多缺點。

基于計算機視覺的安全帽自動識別技術設計通過在施工現場布設視頻監控設備或利用現有的施工監控設備，采用機器視覺的相關方法進行安全帽的自動識別，可以實現對作業人員安全帽佩戴情況信息的全程快速識別，在降低監管費用的同時提升了監管信息化水平。

項目鏈接：

https://github.com/PeterH0323/Smart_Construction

數據集鏈接：

https://github.com/njvisionpower/Safety-Helmet-Wearing-Dataset

Smart_Construction

該項目是使用 YOLOv5 來訓練在智能工地安全領域中頭盔目標檢測的應用, 先來一波演示！

指標

yolov5s 為基礎訓練，epoch = 50

對應的 yolov5s 權重文件：百度云：https://pan.baidu.com/s/1ELPhtW-Q4G8UqEr4YrV_5A，提取碼: b981

yolov5m 為基礎訓練，epoch = 100

對應的 yolov5m 權重文件：百度云：https://pan.baidu.com/s/10hlKrgpxVsw4d_vHnPHwEA，提取碼: psst

yolov5l 為基礎訓練，epoch = 100

對應的 yolov5l 權重文件：百度云：https://pan.baidu.com/s/1iMZkRNXY1fowpQCcapFDqw，提取碼: a66e

1.YOLO v5訓練自己數據集教程

使用的數據集：Safety-Helmet-Wearing-Dataset，感謝這位大神的開源數據集！

本文結合 YOLOv5官方教程 來寫

環境準備

首先確保自己的環境：

    Python >= 3.7
    Pytorch == 1.5.x

訓練自己的數據

提示：

關于增加數據集分類的方法，請看【5. 增加數據集的分類】

1.1 創建自己的數據集配置文件

因為我這里只是判斷 【人沒有帶安全帽】、【人有帶安全帽】、【人體】 3個類別 ，基于 data/coco128.yaml 文件，創建自己的數據集配置文件 custom_data.yaml

#?訓練集和驗證集的?labels?和?image?文件的位置
train:?./score/images/train
val:?./score/images/val

#?number?of?classes
nc:?3

#?class?names
names:?['person',?'head',?'helmet']

1.2 創建每個圖片對應的標簽文件

使用標注工具類似于 Labelbox 、CVAT、精靈標注助手 標注之后，需要生成每個圖片對應的 .txt 文件，其規范如下：

• 每一行都是一個目標
• 類別序號是零索引開始的(從0開始)
• 每一行的坐標 class x_center y_center width height 格式
• 框坐標必須采用歸一化的 xywh格式(從0到1)。如果您的框以像素為單位，則將x_center和width除以圖像寬度，將y_center和height除以圖像高度。代碼如下：

import?numpy?as?np
def?convert(size,?box):
????"""
????將標注的?xml?文件生成的【左上角x,左上角y,右下角x，右下角y】標注轉換為yolov5訓練的坐標
????:param size:?圖片的尺寸：?[w,h]
????:param?box:?anchor?box?的坐標?[左上角x,左上角y,右下角x,右下角y,]
????:return:?轉換后的?[x,y,w,h]
????"""

????x1?=?int(box[0])
????y1?=?int(box[1])
????x2?=?int(box[2])
????y2?=?int(box[3])

????dw?=?np.float32(1.?/?int(size[0]))
????dh?=?np.float32(1.?/?int(size[1]))

????w?=?x2?-?x1
????h?=?y2?-?y1
????x?=?x1?+?(w?/?2)
????y?=?y1?+?(h?/?2)

????x?=?x?*?dw
????w?=?w?*?dw
????y?=?y?*?dh
????h?=?h?*?dh
????return?[x,?y,?w,?h]

生成的 .txt 文件放置的名字是圖片的名字，放置在 label 文件夾中，例如：

./score/images/train/00001.jpg  # image
./score/labels/train/00001.txt  # label

生成的 .txt 例子：

1 0.1830000086920336 0.1396396430209279 0.13400000636465847 0.15915916301310062
1 0.5240000248886645 0.29129129834473133 0.0800000037997961 0.16816817224025726
1 0.6060000287834555 0.29579580295830965 0.08400000398978591 0.1771771814674139
1 0.6760000321082771 0.25375375989824533 0.10000000474974513 0.21321321837604046
0 0.39300001866649836 0.2552552614361048 0.17800000845454633 0.2822822891175747
0 0.7200000341981649 0.5570570705458522 0.25200001196935773 0.4294294398277998
0 0.7720000366680324 0.2567567629739642 0.1520000072196126 0.23123123683035374

1.3 文件放置規范

文件樹如下

1.4 聚類得出先驗框(Yolov5 內部已做適配，可選)

使用代碼 ./data/gen_anchors/clauculate_anchors.py ，修改數據集的路徑

FILE_ROOT?=?r"xxx"?#?根路徑
ANNOTATION_ROOT?=?r"xxx"??#?數據集標簽文件夾路徑
ANNOTATION_PATH?=?FILE_ROOT?+?ANNOTATION_ROOT

跑完會生成一個文件 anchors.txt，里面有得出的建議先驗框：

Best Accuracy = 79.72%

Best Anchors = [[14.74, 27.64], [23.48, 46.04], [28.88, 130.0], [39.33, 148.07], [52.62, 186.18], [62.33, 279.11], [85.19, 237.87], [88.0, 360.89], [145.33, 514.67]]

1.5 選擇一個您需要的模型

在文件夾 ./models 下選擇一個你需要的模型然后復制一份出來，將文件開頭的 nc = 修改為數據集的分類數，下面是借鑒 ./models/yolov5s.yaml來修改的

#?parameters
nc:?3??#?number?of?classes?????<============?修改這里為數據集的分類數
depth_multiple:?0.33??#?model?depth?multiple
width_multiple:?0.50??#?layer?channel?multiple

#?anchors
anchors:?#?<============?根據?./data/gen_anchors/anchors.txt?中的?Best?Anchors?修改，需要取整(可選)
??-?[14,27,?23,46,?28,130]?
??-?[39,148,?52,186,?62.,279]?
??-?[85,237,?88,360,?145,514]

#?YOLOv5?backbone
backbone:
??#?[from,?number,?module,?args]
??[[-1,?1,?Focus,?[64,?3]],??#?0-P1/2
???[-1,?1,?Conv,?[128,?3,?2]],??#?1-P2/4
???[-1,?3,?BottleneckCSP,?[128]],
???[-1,?1,?Conv,?[256,?3,?2]],??#?3-P3/8
???[-1,?9,?BottleneckCSP,?[256]],
???[-1,?1,?Conv,?[512,?3,?2]],??#?5-P4/16
???[-1,?9,?BottleneckCSP,?[512]],
???[-1,?1,?Conv,?[1024,?3,?2]],??#?7-P5/32
???[-1,?1,?SPP,?[1024,?[5,?9,?13]]],
???[-1,?3,?BottleneckCSP,?[1024,?False]],??#?9
??]

#?YOLOv5?head
head:
??[[-1,?1,?Conv,?[512,?1,?1]],
???[-1,?1,?nn.Upsample,?[None,?2,?'nearest']],
???[[-1,?6],?1,?Concat,?[1]],??#?cat?backbone?P4
???[-1,?3,?BottleneckCSP,?[512,?False]],??#?13

???[-1,?1,?Conv,?[256,?1,?1]],
???[-1,?1,?nn.Upsample,?[None,?2,?'nearest']],
???[[-1,?4],?1,?Concat,?[1]],??#?cat?backbone?P3
???[-1,?3,?BottleneckCSP,?[256,?False]],??#?17

???[-1,?1,?Conv,?[256,?3,?2]],
???[[-1,?14],?1,?Concat,?[1]],??#?cat?head?P4
???[-1,?3,?BottleneckCSP,?[512,?False]],??#?20

???[-1,?1,?Conv,?[512,?3,?2]],
???[[-1,?10],?1,?Concat,?[1]],??#?cat?head?P5
???[-1,?3,?BottleneckCSP,?[1024,?False]],??#?23

???[[17,?20,?23],?1,?Detect,?[nc,?anchors]],??#?Detect(P3,?P4,?P5)
??]

1.6 開始訓練

這里選擇了 yolov5s 模型進行訓練，權重也是基于 yolov5s.pt 來訓練

python?train.py?--img?640?--batch?16?--epochs?10?--data?./data/custom_data.yaml?--cfg?./models/custom_yolov5.yaml?--weights?./weights/yolov5s.pt

其中，yolov5s.pt 需要自行下載放在本工程的根目錄即可，下載地址：https://drive.google.com/open?id=1Drs_Aiu7xx6S-ix95f9kNsA6ueKRpN2J

1.7 看訓練之后的結果

訓練之后，權重會保存在 ./runs 文件夾里面的每個 exp 文件里面的 weights/best.py ，里面還可以看到訓練的效果

2. 偵測

偵測圖片會保存在 ./inferenct/output/ 文件夾下

運行命令：

python?detect.py?--source???0??#?webcam
????????????????????????????file.jpg??#?image?
????????????????????????????file.mp4??#?video
????????????????????????????path/??#?directory
????????????????????????????path/*.jpg??#?glob
????????????????????????????rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa??#?rtsp?stream
????????????????????????????http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8??#?http?stream

例如使用我的 s 權重檢測圖片，可以運行以下命令，偵測圖片會保存在 ./inferenct/output/ 文件夾下

python?detect.py?--source?圖片路徑?--weights?./weights/helmet_head_person_s.pt

3. 檢測危險區域內是否有人

3.1 危險區域標注方式

我這里使用的是 精靈標注助手標注，生成了對應圖片的 json 文件

3.2 執行偵測

偵測圖片會保存在 ./inferenct/output/ 文件夾下

運行命令：

python?area_detect.py?--source?./area_dangerous?--weights?./weights/helmet_head_person_s.pt

3.3 效果：在危險區域里面的人體會被 紅色框 選出來

4. 生成 ONNX

4.1 安裝 onnx 庫

pip?install?onnx

4.2 執行生成

python?./models/export.py?--weights?./weights/helmet_head_person_s.pt?--img?640?--batch?1

onnx 和 torchscript 文件會生成在 ./weights 文件夾中

5. 增加數據集的分類

關于增加數據集分類的方法：

SHWD 數據集里面沒有 person 的類別，先將現有的自己的數據集執行腳本生成 yolov5 需要的標簽文件 .txt，之后再用 yolov5x.pt 加上 yolov5x.yaml ，使用指令檢測出人體

python?detect.py?--save-txt?--source?./自己數據集的文件目錄?--weights?./weights/yolov5x.pt

yolov5 會推理出所有的分類，并在 inference/output 中生成對應圖片的 .txt 標簽文件；

修改 ./data/gen_data/merge_data.py 中的自己數據集標簽所在的路徑，執行這個python腳本，會進行 person 類型的合并

總結

以上就是 YOLOv5 實現 智慧工地安全帽和危險區域檢測系統 的全部內容。歡迎 star ???

本文轉載自csdn博客，文章鏈接：https://blog.csdn.net/hxj0323/article/details/109257255