點擊上方「小白學視覺」，選擇加"星標"或「置頂」javascript
重磅乾貨，第一時間送達
css
轉自 | AI深度前沿視線java
1、YOLO v5訓練本身數據集教程 python
1.1 建立本身的數據集配置文件 nginx
1.2 建立每一個圖片對應的標籤文件 git
1.3 文件放置規範 github
1.4 聚類得出先驗框（可選） web
1.5 選擇一個你須要的模型 json
1.6 開始訓練 安全
1.7 看訓練以後的結果
2、偵測
3、檢測危險區域內是否有人
3.1 危險區域標註方式
3.2 執行偵測
3.3 效果：在危險區域裏面的人體會被 紅色框 選出來
4、生成 ONNX
5、增長數據集的分類
該項目是使用 YOLOv5 v2.x 來訓練在智能工地安全領域中頭盔目標檢測的應用,先來一波演示！
指標
yolov5s 爲基礎訓練，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 |
對應的權重文件：https://pan.baidu.com/share/init?surl=ELPhtW-Q4G8UqEr4YrV_5A，提取碼: b981
yolov5m 爲基礎訓練，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 |
對應的權重文件：https://pan.baidu.com/share/init?surl=0hlKrgpxVsw4d_vHnPHwEA，提取碼: psst
yolov5l 爲基礎訓練，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 |
對應的權重文件：https://pan.baidu.com/share/init?surl=iMZkRNXY1fowpQCcapFDqw，提取碼: a66e
使用的數據集：Safety-Helmet-Wearing-Dataset ，感謝這位大神的開源數據集！
https://github.com/njvisionpower/Safety-Helmet-Wearing-Dataset
本文結合 YOLOv5官方教程 來寫
環境準備
首先確保本身的環境：
Python >= 3.7Pytorch == 1.5.x
訓練本身的數據
提示：
關於增長數據集分類的方法，請看【5. 增長數據集的分類】
1.1 建立本身的數據集配置文件
由於我這裏只是判斷 【人沒有帶安全帽】、【人有帶安全帽】、【人體】 3個類別 ，基於data/coco128.yaml 文件，建立本身的數據集配置文件 custom_data.yaml
# 訓練集和驗證集的 labels 和 image 文件的位置train: ./score/images/trainval: ./score/images/val# number of classesnc: 3
# class namesnames: ['person', 'head', 'helmet']
使用標註工具相似於 Labelbox 、CVAT 、精靈標註助手 標註以後，須要生成每一個圖片對應的.txt 文件，其規範以下：
每一行都是一個目標
類別序號是零索引開始的（從0開始）
每一行的座標
class x_center y_center width height格式框座標必須採用歸一化的 xywh格式（從0到1）。若是您的框以像素爲單位，則將
x_center和width除以圖像寬度，將y_center和height除以圖像高度。代碼以下：
import numpy as npdef 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.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 文件放置規範
文件樹以下
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來修改的
# parametersnc: 3 # number of classes <============ 修改這裏爲數據集的分類數depth_multiple: 0.33 # model depth multiplewidth_multiple: 0.50 # layer channel multiple# anchorsanchors: # <============ 根據 ./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 開始訓練
這裏選擇了 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 須要自行下載放在本工程的根目錄便可，下載地址 官方權重
1.7 看訓練以後的結果
訓練以後，權重會保存在 ./runs 文件夾裏面的每一個 exp 文件裏面的 weights/best.py ，裏面還能夠看到訓練的效果
偵測圖片會保存在 ./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.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
會推理出全部的分類，並在 中生成對應圖片的 標籤文件；yolov5inference/output.txt
修改 ./data/gen_data/merge_data.py 中的本身數據集標籤所在的路徑，執行這個python腳本，會進行 person 類型的合併
本文分享自微信公衆號 - 小白學視覺（NoobCV）。
若有侵權，請聯繫 support@oschina.cn 刪除。
本文參與「OSC源創計劃」，歡迎正在閱讀的你也加入，一塊兒分享。