向AI转型的程序员都关注了这个号????????????
机器学习AI算法工程 公众号:datayx
该项目是使用 YOLOv5 来训练在智能工地安全领域中头盔目标检测的应用
代码及运行教程,训练模型,数据集 获取:
关注微信公众号 datayx 然后回复 安全帽 即可获取。
AI项目体验地址 https://loveai.tech
创建自己的数据集配置文件
因为我这里只是判断 【人没有带安全帽】、【人有带安全帽】、【人体】 3个类别 ,基于 data/coco128.yaml
文件,创建自己的数据集配置文件 custom_data.yaml
创建每个图片对应的标签文件
使用标注工具类似于 Labelbox 、CVAT 、精灵标注助手 标注之后,需要生成每个图片对应的 .txt
文件,其规范如下:
- 每一行都是一个目标
- 类别序号是零索引开始的(从0开始)
- 每一行的坐标
class x_center y_center width height
格式 - 框坐标必须采用归一化的 xywh格式(从0到1)。如果您的框以像素为单位,则将
x_center
和width
除以图像宽度,将y_center
和height
除以图像高度。代码如下:
生成的 .txt
文件放置的名字是图片的名字,放置在 label 文件夹中,例如:
聚类得出先验框(可选)
选择一个你需要的模型
在文件夹 ./models
下选择一个你需要的模型然后复制一份出来,将文件开头的 nc =
修改为数据集的分类数,下面是借鉴 ./models/yolov5s.yaml
来修改的
开始训练
这里选择了 yolov5s
模型进行训练,权重也是基于 yolov5s.pt
来训练
其中,yolov5s.pt
需要自行下载放在本工程的根目录即可,下载地址 官方权重
看训练之后的结果
训练之后,权重会保存在 ./runs
文件夹里面的每个 exp
文件里面的 weights/best.py
,里面还可以看到训练的效果
侦测图片会保存在 ./inferenct/output/
文件夹下
运行命令:
检测危险区域内是否有人危险区域标注方式,我这里使用的是 精灵标注助手 标注,生成了对应图片的 json 文件
侦测图片会保存在 ./inferenct/output/
文件夹下
运行命令:
效果:在危险区域里面的人体会被 红色框 选出来
生成 ONNX安装 onnx 库
pip install onnx
执行生成
增加数据集的分类关于增加数据集分类的方法:
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
类型的合并
不断更新资源
深度学习、机器学习、数据分析、python
搜索公众号添加: datayx
机大数据技术与机器学习工程
搜索公众号添加: datanlp
长按图片,识别二维码