本科阶段最后一次竞赛Vlog——2024年智能车大赛智慧医疗组准备全过程——2Yolo使用之ONNX模型准备
大家好,因为板端BPU环境,可以加速目标检测的速度,所以今天在此先给大家带来如何准备一个模型,下一期会给大家带来如何在板端部署这一期的目标识别模型
本文给大家带来的是yolov5-2.0 版本,在我最先接触的就是这个板子,当时也是参考了地平线社区的大佬带来一些方法,当然这个参考是在下一期哈哈、
这一期其方法上是通用的,出这一期也是能够方便大家以后学会yolo任何版本的使用。
因为我没有研究很深入,我写是针对比赛,至于其他版本会不会在使用过程中有小坑,在地平线官网也是有教程的,大家需要其他版本的可以去看看其他大佬的文章。
1.获取目标检测源码
首先,是要获得yolov5-2.0得源码,正常情况是要教大家去github进行下载源码的,但是咳咳好像不太容易,因为“科学上网”
当然下载完成2.0版本还有有很多bug的,这里给大家以附件形式给出我的源码吧
下载的源码里主要由一下文件夹
2.数据准备与标记
要训练一个模型,首先需要一定的图片数据,然后标记,接下来开始告诉大家如何进行这些准备
当然你首先要有Pyhton环境,这里就不教大家配置Python了,我这里使用的是Anaconda+Pycharm
之前发现哔哩哔哩有好多讲解使用,大家可以看下
2.1 数据集准备
首先建立一个文件夹dataset,里面放置image和label,每个子文件夹里面放置train和test。然后把你的jpg或者png图片放在image的train里面
2.2 数据标记
打开你的命令行,输入如下
pip install labelimg
稍等安装完毕,然后在命令行输入,如下内容,你会发现弹出如下窗口
labelimg
接下来点击View-选择第一行设置自动保存模式
接下来点击 Open Dir——>选择你刚才image所在的目录
接下来选择 Change Save Dir 找到你的label文件夹里的train
最后保证选择Yolo模式
此刻你就可以看到你的图片已经加载进来了
这里把鼠标移动到页面里,按下键盘里w,会出现一个十字
这时候你就可以,进行绘制矩形了,绘制完成后,就会出现一个提示框,这里写上你的标签
当你打完这个标签之后,使用键盘键A和D,可以切换上一张,下一张,最后你标记完成你所有图片,应该是下面这个样子
一个文件夹有图片,另外一个文件夹里面是txt格式文件
3.配置文件
3.1目录结构
在你下载的源码中创建下面这样的文件夹,接下来,我将挨个进行刚才没有提及的文件进行解释
3.2 test
对于test里面的文件夹,这里下面我给大家提供一个分割脚本,其实本质是就是将刚才打标签的所有内容,进行划分,分成两部分
当然下面代码是要大家改一些东西的,一些路径字符串我已经给大家写好了,大家改成自己对应的就可以了
import os
import shutil
from sklearn.model_selection import train_test_split
def move_files(src_dir, dst_dir, files, extension):
"""将指定的文件从源目录移动到目标目录。"""
os.makedirs(dst_dir, exist_ok=True) # 如果目标目录不存在,则创建
for base_name in files:
src_path = os.path.join(src_dir, f"{base_name}{extension}")
dst_path = os.path.join(dst_dir, f"{base_name}{extension}")
shutil.move(src_path, dst_path)
def split_dataset(src_images_dir, src_labels_dir, dst_dirs, train_ratio=0.8, val_ratio=0.1):
"""
将数据集分为训练集、验证集和测试集,并移动到指定目录。
参数:
- src_images_dir: 图片的源目录。
- src_labels_dir: 标签的源目录。
- dst_dirs: 目标目录的字典,包含'train', 'val', 'test'对应的图片和标签目录。
- train_ratio: 训练集在总数据中的比例。
- val_ratio: 验证集在总数据中的比例。
"""
test_ratio = 1 - train_ratio - val_ratio
assert test_ratio > 0, "Invalid ratios: Sum of train and val ratios must be less than 1."
# 读取并分割数据
files = [os.path.splitext(file)[0] for file in os.listdir(src_images_dir)]
train_files, test_files = train_test_split(files, test_size=test_ratio, random_state=42)
train_files, val_files = train_test_split(train_files, test_size=val_ratio / (train_ratio + val_ratio),
random_state=42)
# 移动文件到相应的目录
for part, file_list in [('train', train_files), ('val', val_files), ('test', test_files)]:
move_files(src_images_dir, dst_dirs[part]['images'], file_list, '.jpg') #这里使用的是jpg格式,如果是png需要修改
move_files(src_labels_dir, dst_dirs[part]['labels'], file_list, '.txt')
def main():
# 定义源目录
src_images_dir = r'刚才的tain里面的所有图片'
src_labels_dir = r'刚才label所有txt'
# 定义目标目录
dst_dirs = {
'train': {'images': r'训练集图片路径',
'labels': r'训练集标签路径'},
'val': {'images': r'验证集图片路径',
'labels': r'验证集标签路径'},
'test': {'images': r'测试集标签路径',
'labels': r'测试集标签路径'}
}
# 调用函数,进行数据划分和文件移动
split_dataset(src_images_dir, src_labels_dir, dst_dirs, train_ratio=0.8, val_ratio=0.1)
if __name__ == '__main__':
main()
3.3 my_model.yaml
对于这个yaml文件,大家只需要吧第一个参数nc改成自己的实际类别,对于小白其他参数不需要管
# parameters
nc: 1 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
# anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# 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)
]
3.4 my_para.yaml
对于这个yaml文件,下面四个参数都要修改,具体修改下面已经用中文给出大家了
# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Download command: python -c "from yolov5.utils.google_utils import *; gdrive_download('1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', 'coco128.zip')"
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
# /parent_folder
# /coco128
# /yolov5
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: "训练集图片的路径" # 128 images
val: "验证集的图片路径" # 128 images
# 这个地方写类别个数,数字
nc: 1
# class names
names: ['刚才标记的字符串名字','标记的第二个label名字']
4.训练配置
4.1train.py
对于训练,打开train.py 按照我下面截图框进行设置,下面是解释
(1)cfg 这个把刚才的my_mode.yaml路径放进来
(2)data这个把刚才的my_para.yaml路径放进来
(3)epochs:这个设置轮数
(4)device:选择cpu,当然我这里是0 是因为我配置了cuda
(5)img-size:这里注意一下。这里选择672*672 和我一样就行,方便后面转模型,我也没有研究很深入,好像是转模型有点影响
4.2 model/yolo.py
现在我们是训练阶段,我给大家提供的里面就是我现在截图这个样子,
大家需要把红框的这个取消注释
把红框下面的进行注释掉
4.3正式训练
到现在你就可以愉快的进行训练啦
当你看到类似下面这种情况,代表你的模型正常训练了
4.4训练完成
训练时间还是挺漫长的,当你最后训练完成了,你会发现,控制台输出了一个日志,当然不一定和我的一模一样,因为我的是v5-7.0版本
哈哈大同小异,反正根据日志找到位置就好了,然后你会看到,里面有这些,通常我们会使用best.pt
到此为止,windows的模型就训练得到啦
5 ONNX模型转化
这一步就很简单了,大家只需要,先把yolo.py里面的换到处onnx该使用的那一行就好啦
这个时候大家吧expoer.py里面的weights路径改成刚才的pt就行了
运行代码你就会发现,现在onnx成功转化了,然后你就可以使用这个进行下一期的bin模型转化了
6.总结与下期预告
本期,给大家带来了如何在windows训练yolov5,大家可以试试
当然本期介绍确实没啥含金量,网上有关教程也满天飞,我也参考了好多大佬的文章,也就是对于大佬们详细版本的总结一下
下一期讲给大家带来与地平线更为相关的内容,包括bin模型转化环境,以及转化过程的提炼版本