使用自己的数据集训练YOLOV5

2023-02-19 05:16| 来源: 网络整理| 查看: 265

一、代码：

yolov5的代码需要大家自己去github去下载了。

二、准备存放数据的文件夹：

首先我们需要创建三个文件夹来存放数据（文件夹的名字可以自己定义）：

1.存放所有数据的文件

其中 image-data 存放的是我们所有的图片；image-txt中存放的是由xml转换以后的txt文件（下面会附上转换的txt文件的代码）；xml-data 中存放的是我们的标签文件也就是xml文件。

2.存放图片的文件

这个文件夹中 train存放我们的训练集， val 存放我们的验证集。

3.存放转换以后的txt文件

这里跟上面那个文件夹是同一个道理只不过我们这里存放的是转换后的txt文件。

三、将我们标好的xml文件转换成txt文件：

我们知道使用labelimg标注好的文件格式是.xml文件如果我们要使用yolov5就需要把它转换成yolov5需要的文件格式 .txt （这是重点）代码附下：

# -*- coding: utf-8 -*- import os import xml.etree.ElementTree as ET dirpath = r'E:\yolov5-5.0\face - data\face\xml-data' # 原来存放xml文件的目录 newdir = r'E:\yolov5-5.0\face - data\face\image-txt' # 修改label后形成的txt目录 if not os.path.exists(newdir): os.makedirs(newdir) dict_info = {'eyes': 0, 'nose': 1, 'mouth': 2} # 有几个属性填写几个label names for fp in os.listdir(dirpath): if fp.endswith('.xml'): root = ET.parse(os.path.join(dirpath, fp)).getroot() xmin, ymin, xmax, ymax = 0, 0, 0, 0 sz = root.find('size') width = float(sz[0].text) height = float(sz[1].text) filename = root.find('filename').text for child in root.findall('object'): # 找到图片中的所有框 sub = child.find('bndbox') # 找到框的标注值并进行读取 label = child.find('name').text label_ = dict_info.get(label) if label_: label_ = label_ else: label_ = 0 xmin = float(sub[0].text) ymin = float(sub[1].text) xmax = float(sub[2].text) ymax = float(sub[3].text) try: # 转换成yolov3的标签格式，需要归一化到（0-1）的范围内 x_center = (xmin + xmax) / (2 * width) x_center = '%.6f' % x_center y_center = (ymin + ymax) / (2 * height) y_center = '%.6f' % y_center w = (xmax - xmin) / width w = '%.6f' % w h = (ymax - ymin) / height h = '%.6f' % h except ZeroDivisionError: print(filename, '的 width有问题') with open(os.path.join(newdir, fp.split('.xml')[0] + '.txt'), 'a+') as f: f.write(' '.join([str(label_), str(x_center), str(y_center), str(w), str(h) + '\n'])) print('ok')

四、构造yaml文件：

因为yolov5训练数据调用的都是.yaml文件，所以我们需要在date目录下创建一个yaml文件来存放我们已经整理好的数据，其中train是训练集的txt文件，需要存放其绝对路径，相对路径也可以但是容易出错。val是验证集.txt文件，路径同上。

nc：类别个数

names：类别名称

五、运行train.py文件

我们只需要将.yaml文件换成我们刚才新建的就可以成功运行了（看下图）。

【本文地址】

公司简介

联系我们