轉載自：3D視覺工坊
編輯：AI深度學習視線

教程來自：https://github.com/PeterH0323/Smart/_Construction

一、YOLO v5訓練自己數據集教程

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

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

1.3 文件放置規范

1.4 聚類得出先驗框（可選）

1.5 選擇一個你需要的模型

1.6 開始訓練

1.7 看訓練之后的結果

二、偵測
三、檢測危險區域內是否有人

3.1 危險區域標注方式

3.2 執行偵測

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

四、生成 ONNX
五、增加數據集的分類
該項目是使用 YOLOv5 v2.x 來訓練在智能工地安全領域中頭盔目標檢測的應用,先來一波演示！

指標

yolov5s 為基礎訓練，epoch = 50

對應的權重文件：https://pan.baidu.com/share/i.../_5A，提取碼: b981

yolov5m 為基礎訓練，epoch = 100

對應的權重文件：https://pan.baidu.com/share/i.../_vHnPHwEA，提取碼: psst

yolov5l 為基礎訓練，epoch = 100

對應的權重文件：https://pan.baidu.com/share/i...，提取碼: a66e

1

YOLOv5訓練自己數據集教程

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

https://github.com/njvisionpo...

本文結合 YOLOv5官方教程 來寫

環境準備

首先確保自己的環境：

Python >= 3.7

訓練自己的數據

提示：

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

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

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

# 訓練集和驗證集的 labels 和 image 文件的位置

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

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

./score/images/train/00001.jpg  # image

生成的.txt例子

1 0.1830000086920336 0.1396396430209279 0.13400000636465847 0.15915916301310062

1.3文件放置規范

文件樹如下

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

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

FILE_ROOT = r"xxx" # 根路徑

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

Best Accuracy = 79.72%

1.5 選擇一個您需要的模型

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

# parameters

1.6 開始訓練

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

pythontrain.py--img640--batch16--epochs10--data./data/custom_data.yaml--cfg./models/custom_yolov5.yaml--weights./weights/yolov5s.pt

其中，yolov5s.pt 需要自行下載放在本工程的根目錄即可，下載地址 官方權重

1.7 看訓練之后的結果

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

2

推斷
偵測圖片會保存在 ./inferenct/output/ 文件夾下
運行命令：

python detect.py --source   0  # webcam

例如使用我的 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 效果：在危險區域里面的人體會被 紅色框 選出來

生成 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 文件夾中

增加數據集的分類

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

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

yolov5

修改 ./data/gen_data/merge_data.py 中的自己數據集標簽所在的路徑，執行這個python腳本，會進行 person 類型的合并
本文僅做學術分享，如有侵權，請聯系刪文。

審核編輯 黃昊宇