點(diǎn)擊左上方藍(lán)字關(guān)注我們

全網(wǎng)搜集目標(biāo)檢測(cè)相關(guān)，人工篩選最優(yōu)價(jià)值內(nèi)容

編者薦語

文章以智能工地安全領(lǐng)域中的頭盔為檢測(cè)目標(biāo)，系統(tǒng)性的介紹了yolov5從環(huán)境搭建到模型訓(xùn)練的整個(gè)過程。

轉(zhuǎn)載自 | 機(jī)器學(xué)習(xí)算法那些事

一、YOLO v5訓(xùn)練自己數(shù)據(jù)集教程

1.1 創(chuàng)建自己的數(shù)據(jù)集配置文件
1.2 創(chuàng)建每個(gè)圖片對(duì)應(yīng)的標(biāo)簽文件
1.3 文件放置規(guī)范
1.4 聚類得出先驗(yàn)框（可選）
1.5 選擇一個(gè)你需要的模型
1.6 開始訓(xùn)練
1.7 看訓(xùn)練之后的結(jié)果

二、偵測(cè)

三、檢測(cè)危險(xiǎn)區(qū)域內(nèi)是否有人

3.1 危險(xiǎn)區(qū)域標(biāo)注方式
3.2 執(zhí)行偵測(cè)
3.3 效果：在危險(xiǎn)區(qū)域里面的人體會(huì)被紅色框選出來

四、生成 ONNX

五、增加數(shù)據(jù)集的分類

該項(xiàng)目是使用YOLOv5 v2.x來訓(xùn)練在智能工地安全領(lǐng)域中頭盔目標(biāo)檢測(cè)的應(yīng)用,先來一波演示！

指標(biāo)

yolov5s 為基礎(chǔ)訓(xùn)練，`epoch = 50`

分類	P	R	mAP0.5
總體	0.884	0.899	0.888
人體	0.846	0.893	0.877
頭	0.889	0.883	0.871
安全帽	0.917	0.921	0.917

對(duì)應(yīng)的權(quán)重文件：https://pan.baidu.com/share/init?surl=ELPhtW-Q4G8UqEr4YrV_5A，提取碼:b981

yolov5m 為基礎(chǔ)訓(xùn)練，`epoch = 100`

分類	P	R	mAP0.5
總體	0.886	0.915	0.901
人體	0.844	0.906	0.887
頭	0.9	0.911	0.9
安全帽	0.913	0.929	0.916

對(duì)應(yīng)的權(quán)重文件：https://pan.baidu.com/share/init?surl=0hlKrgpxVsw4d_vHnPHwEA，提取碼:psst

yolov5l 為基礎(chǔ)訓(xùn)練，`epoch = 100`

分類	P	R	mAP0.5
總體	0.892	0.919	0.906
人體	0.856	0.914	0.897
頭	0.893	0.913	0.901
安全帽	0.927	0.929	0.919

對(duì)應(yīng)的權(quán)重文件：https://pan.baidu.com/share/init?surl=iMZkRNXY1fowpQCcapFDqw，提取碼:a66e

YOLOv5訓(xùn)練自己數(shù)據(jù)集教程

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

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

本文結(jié)合 YOLOv5官方教程來寫

環(huán)境準(zhǔn)備

首先確保自己的環(huán)境：

Python >= 3.7Pytorch == 1.5.x

訓(xùn)練自己的數(shù)據(jù)

提示：

關(guān)于增加數(shù)據(jù)集分類的方法，請(qǐng)看【5. 增加數(shù)據(jù)集的分類】

1.1 創(chuàng)建自己的數(shù)據(jù)集配置文件

因?yàn)槲疫@里只是判斷【人沒有帶安全帽】、【人有帶安全帽】、【人體】 3個(gè)類別，基于data/coco128.yaml文件，創(chuàng)建自己的數(shù)據(jù)集配置文件custom_data.yaml

# 訓(xùn)練集和驗(yàn)證集的 labels 和 image 文件的位置train: ./score/images/trainval: ./score/images/val
# number of classesnc: 3
# class namesnames: ['person', 'head', 'helmet']

1.2 創(chuàng)建每個(gè)圖片對(duì)應(yīng)的標(biāo)簽文件

使用標(biāo)注工具類似于 Labelbox 、CVAT 、精靈標(biāo)注助手標(biāo)注之后，需要生成每個(gè)圖片對(duì)應(yīng)的.txt文件，其規(guī)范如下：

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

import numpy as npdef convert(size, box):    """    將標(biāo)注的 xml 文件生成的【左上角x,左上角y,右下角x，右下角y】標(biāo)注轉(zhuǎn)換為yolov5訓(xùn)練的坐標(biāo)    :param size: 圖片的尺寸：[w,h]    :param box: anchor box 的坐標(biāo) [左上角x,左上角y,右下角x,右下角y,]    :return: 轉(zhuǎn)換后的 [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.159159163013100621 0.5240000248886645 0.29129129834473133 0.0800000037997961 0.168168172240257261 0.6060000287834555 0.29579580295830965 0.08400000398978591 0.17717718146741391 0.6760000321082771 0.25375375989824533 0.10000000474974513 0.213213218376040460 0.39300001866649836 0.2552552614361048 0.17800000845454633 0.28228228911757470 0.7200000341981649 0.5570570705458522 0.25200001196935773 0.42942943982779980 0.7720000366680324 0.2567567629739642 0.1520000072196126 0.23123123683035374

1.3 文件放置規(guī)范

文件樹如下

1.4 聚類得出先驗(yàn)框（Yolov5 內(nèi)部已做適配，可選）

使用代碼./data/gen_anchors/clauculate_anchors.py，修改數(shù)據(jù)集的路徑

FILE_ROOT = r"xxx" # 根路徑ANNOTATION_ROOT = r"xxx"  # 數(shù)據(jù)集標(biāo)簽文件夾路徑ANNOTATION_PATH = FILE_ROOT + ANNOTATION_ROOT

跑完會(huì)生成一個(gè)文件anchors.txt，里面有得出的建議先驗(yàn)框：

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 選擇一個(gè)您需要的模型

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

# parametersnc: 3  # number of classes     <============ 修改這里為數(shù)據(jù)集的分類數(shù)depth_multiple: 0.33  # model depth multiplewidth_multiple: 0.50  # layer channel multiple
# anchorsanchors: # <============ 根據(jù) ./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 backbonebackbone:  # [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 headhead:  [[-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 開始訓(xùn)練

這里選擇了yolov5s模型進(jìn)行訓(xùn)練，權(quán)重也是基于yolov5s.pt來訓(xùn)練

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需要自行下載放在本工程的根目錄即可，下載地址官方權(quán)重

1.7 看訓(xùn)練之后的結(jié)果

訓(xùn)練之后，權(quán)重會(huì)保存在./runs文件夾里面的每個(gè)exp文件里面的weights/best.py，里面還可以看到訓(xùn)練的效果

推斷

偵測(cè)圖片會(huì)保存在./inferenct/output/文件夾下

運(yùn)行命令：

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權(quán)重檢測(cè)圖片，可以運(yùn)行以下命令，偵測(cè)圖片會(huì)保存在./inferenct/output/文件夾下

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

檢測(cè)危險(xiǎn)區(qū)域內(nèi)是否有人

3.1 危險(xiǎn)區(qū)域標(biāo)注方式

我這里使用的是精靈標(biāo)注助手標(biāo)注，生成了對(duì)應(yīng)圖片的 json 文件

3.2 執(zhí)行偵測(cè)

偵測(cè)圖片會(huì)保存在./inferenct/output/文件夾下

運(yùn)行命令：

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

3.3 效果：在危險(xiǎn)區(qū)域里面的人體會(huì)被紅色框選出來

生成 ONNX

4.1 安裝`onnx`庫

pip install onnx

4.2 執(zhí)行生成

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

onnx和torchscript文件會(huì)生成在./weights文件夾中

增加數(shù)據(jù)集的分類

關(guān)于增加數(shù)據(jù)集分類的方法：

SHWD數(shù)據(jù)集里面沒有person的類別，先將現(xiàn)有的自己的數(shù)據(jù)集執(zhí)行腳本生成yolov5需要的標(biāo)簽文件.txt，之后再用yolov5x.pt加上yolov5x.yaml，使用指令檢測(cè)出人體

python detect.py --save-txt --source ./自己數(shù)據(jù)集的文件目錄 --weights ./weights/yolov5x.pt

yolov5 會(huì)推理出所有的分類，并在 inference/output 中生成對(duì)應(yīng)圖片的 .txt 標(biāo)簽文件；

修改./data/gen_data/merge_data.py中的自己數(shù)據(jù)集標(biāo)簽所在的路徑，執(zhí)行這個(gè)python腳本，會(huì)進(jìn)行person類型的合并。

END

雙一流大學(xué)研究生團(tuán)隊(duì)創(chuàng)建，專注于目標(biāo)檢測(cè)與深度學(xué)習(xí)，希望可以將分享變成一種習(xí)慣！

整理不易，點(diǎn)贊三連↓

（附代碼）實(shí)戰(zhàn) | YOLOv5在建筑工地中安全帽佩戴檢測(cè)的應(yīng)用