Y2-使用自己的数据集训练YOLO-v5s
- 零、遇到的问题汇总
- (1)遇到git的`import error`
- (2)`Error:Dataset not found`
- (3)`Error:删除中文后,训练图片路径不存在`
- 一、.xml文件里保存的是什么
- 二、准备好自己的数据
- 三、创建split_train_val.py 文件,运行并生成 train.txt、test.txt、val.txt
- (1)在yolov5-master工程内,新建一个.py文件,并命名为`split_train_val.py`:
- (2)将以下代码写入该文件,设置自己的`.xml`和`.txt`文件路径:
- (3)然后执行该文件,就得到下面的结果:
- (4)创建`voc_label.py`,填充图片路径
- 四、创建 `fruit.yaml `文件
- 五、开始用自己的数据集训练模型
- (1) 输入训练指令
- (2)查看训练结果
本文使用水果数据集、CPU进行训练,包含200张图像,水果类别分为4类,
Banana、Snake fruit、Dragon fruit和Pineapple
。
本文先列出执行过程中遇到的问题以及解决办法,再将执行步骤一一说明,给出的代码是我修改过后的最终的代码。
需要注意的是:!!!文件路径不要包含中文!!!
零、遇到的问题汇总
(1)遇到git的import error
参考【这里】解决:
找到提示报错的路径d:\jupyter notebook\365-DL\.venv\Lib\site-packages\git\cmd.py
,然后在该文件中添加一行:
os.environ['GIT_PYTHON_REFRESH'] = 'quiet'
(2)Error:Dataset not found
路径中含有中文,删除中文。
(3)Error:删除中文后,训练图片路径不存在
将split_train_val.py
和voc_label.py
生成的所有文件删除,重新生成。【注意检查路径】
如下所示:
【注意abs_path的路径,本文的图像路径为D:\jupyter notebook\365-DL\YOLO\Y2\yolov5-master\Y2-fruit_data\images\*.png
,而abs_path=D:\jupyter notebook\365-DL\YOLO\Y2\yolov5-master
】
一、.xml文件里保存的是什么
在annotations/
文件夹里,打开任意一个.xml
文件,这里打开fruit0.xml
,文件内容如下:
注意每个标签组内的信息,后续voc_label.py
文件会提取这些信息。
<annotation>
<folder>images</folder>
<filename>fruit0.png</filename>
<size>
<width>400</width>
<height>300</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>pineapple</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<occluded>0</occluded>
<difficult>0</difficult>
<bndbox>
<xmin>38</xmin>
<ymin>82</ymin>
<xmax>271</xmax>
<ymax>227</ymax>
</bndbox>
</object>
<object>
<name>snake fruit</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<occluded>0</occluded>
<difficult>0</difficult>
<bndbox>
<xmin>244</xmin>
<ymin>174</ymin>
<xmax>280</xmax>
<ymax>207</ymax>
</bndbox>
</object>
<object>
<name>dragon fruit</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<occluded>0</occluded>
<difficult>0</difficult>
<bndbox>
<xmin>254</xmin>
<ymin>228</ymin>
<xmax>351</xmax>
<ymax>300</ymax>
</bndbox>
</object>
</annotation>
二、准备好自己的数据
本次使用水果数据集,数据集包含200张图片,每张图片包含4种不同类别的水果:Banana、Snake fruit、Dragon fruit和Pineapple
。
三、创建split_train_val.py 文件,运行并生成 train.txt、test.txt、val.txt
执行split_train_val.py
前的文件结构:
(1)在yolov5-master工程内,新建一个.py文件,并命名为split_train_val.py
:
(2)将以下代码写入该文件,设置自己的.xml
和.txt
文件路径:
# 导入必要的库
# 导入必要的库
import os
import random
import argparse
# 创建一个参数解析器
parser = argparse.ArgumentParser()
# 添加命令行参数,用于指定XML文件的路径,默认为“Annotations”文件夹
parser.add_argument('--xml_path', default='D:/jupyter notebook/365-DL/YOLO/Y2/yolov5-master/Y2-fruit_data/annotations/', type=str, help='input xml label path')
# 添加命令行参数,用于指定txt标签文件的路径,默认为“ImageSets/Main”文件夹
parser.add_argument('--txt_path', default='D:/jupyter notebook/365-DL/YOLO/Y2/yolov5-master/Y2-fruit_data/ImageSets/Main', type=str, help='output txt label path')
# 解析命令行参数
opt = parser.parse_args()
# 定义训练验证和测试集的划分比例
trainval_percent = 1.0 # 使用全部数据
train_percent = 0.9 # 训练集占验证集的90%
# 设置xml文件的路径,并根据命令行参数指定
xmlfilepath = opt.xml_path
print("xmlfilepath: ", xmlfilepath)
# 设置txt文件的路径,并根据命令行参数指定
txtfilepath = opt.txt_path
# 获取xml文件夹中的所有xml文件列表
total_xml = os.listdir(xmlfilepath)
# 如果输出txt标签文件的文件夹不存在,创建它
if not os.path.exists(txtfilepath):
os.makedirs(txtfilepath)
# 获取xml文件的总数
num = len(total_xml)
# 创建一个包含所有xml文件索引的列表
list_index = range(num)
# 计算训练验证集的数量
tv = int(num*trainval_percent)
# 计算训练集的数量
tr = int(num*train_percent)
# 从所有xml文件索引中随机选择出训练验证集的索引
trainval = random.sample(list_index, tv)
# 从训练验证集的索引中随机选择出训练集的索引
train = random.sample(list_index, tr)
# 打开要写入的训练验证集、测试集、训练集、验证集的txt文件
file_trainval = open(txtfilepath + '/trainval.txt', 'w')
file_test = open(txtfilepath + '/test.txt', 'w')
file_train = open(txtfilepath + '/train.txt', 'w')
file_val = open(txtfilepath + '/val.txt', 'w')
# 遍历所有xml文件的索引
for i in list_index:
name = total_xml[i][:-4] + '\n' # 获取所有文件的名称(去掉后缀.xml),并添加换行符
# 如果该索引在训练验证集中, 写入训练验证集txt文件,否则写入测试集txt文件
if i in trainval:
file_trainval.write(name) #
if i in train: # 如果该索引在训练集中, 写入训练集txt文件,否则写入验证集txt文件
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
# 关闭所有打开的文件
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
(3)然后执行该文件,就得到下面的结果:
打开任意一个文件,查看内容:【此处打开val.txt
,文件内保存的是个文件名】
(4)创建voc_label.py
,填充图片路径
voc_label.py
代码如下:
# 导入必要的库
import xml.etree.ElementTree as ET
import os
from os import getcwd
# 定义数据集的名称
sets = ['train', 'val', 'test']
# 根据所用数据集,填写类别名称,本文使用水果数据集,包含4类,分别如下:
classes = ["banana", "snake fruit", "dragon fruit", "pineapple"]
# 获取当前工作目录的绝对路径
abs_path = os.getcwd() # abs_path: D:\jupyter notebook\365-DL\YOLO\Y2\yolov5-master
print("abs_path: ", abs_path)
# 定义一个函数,将边界框的坐标绝对值转换为相对于图像大小的比例
def convert(size, box):
dw = 1./(size[0]) # 计算图像宽度的倒数
dh = 1./(size[1]) # 计算图像高度的倒数
x = (box[0] + box[1])/ 2.0 - 1 # 计算中心点的x坐标
y = (box[2] + box[3])/ 2.0 - 1 # 计算中心点的y坐标
w = box[1] - box[0] # 计算边界框的宽度
h = box[3] - box[2] # 计算边界框的高度
x = x * dw # 缩放x坐标
w = w * dw # 缩放宽度
y = y * dh # 缩放y坐标
h = h * dh # 缩放高度
return x,y,w,h
# 定义一个函数,将标注文件从xml格式转为YOLO格式
dir = "D:/jupyter notebook/365-DL/YOLO/Y2/yolov5-master/Y2-fruit_data/"
def convert_annotations(image_id):
# 打开xml标注文件
in_file = open(dir + "annotations/%s.xml" % (image_id), encoding='UTF-8')
# 打开要写入的YOLO格式标签文件
out_file = open(dir + "labels/%s.txt" % (image_id), 'w')
# 解析xml文件
tree = ET.parse(in_file)
root = tree.getroot()
# 获取图像文件名
filename = root.find('filename').text
# 获取图像文件格式
filenameFormat = filename.split(".")[1]
# 获取图像尺寸信息
size = root.find('size')
# 获取图像的宽、高
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
# 获取对象的难度标志
difficult = obj.find('difficult').text
# 获取对象的类别名称
cls = obj.find('name').text
if cls not in classes or int(difficult)==1:
continue
# 获取类别索引
cls_id = classes.index(cls)
# 获取对象的边界框信息,包括:左上角x坐标、左上角y坐标、右下角x坐标、右下角y坐标
xmlbox = obj.find('bndbox')
b = ( float(xmlbox.find('xmin').text),
float(xmlbox.find('xmax').text),
float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text) )
b1,b2,b3,b4 = b
# 标注越界修正
if b2 > w:
b2 = w
if b4 > h:
b4 = h
b = (b1,b2,b3,b4)
# 调用convert()函数,将边界框坐标转换为YOLO格式
bb = convert((w,h), b)
out_file.write(str(cls_id)+" " + " ".join([str(a) for a in bb]) + "\n")
return filenameFormat
# 获取当前工作目录
wd = getcwd()
# 遍历每个数据集(train、val、test)
for image_set in sets:
# 如果labels目录不存在,就创建它
if not os.path.exists(dir + "labels/"):
os.makedirs(dir + "labels/")
# 从数据集文件中获取图像id列表
image_ids = open(dir + "ImageSets/Main/%s.txt" % (image_set)).read().strip().split()
# 打开要写入的文件,写入图像文件路径和格式
list_file = open(dir + "ImageSets/Main/%s.txt" % (image_set), 'w')
for image_id in image_ids:
filenameFormat = convert_annotations(image_id)
list_file.write(abs_path + '/Y2-fruit_data/images/%s.%s\n' % (image_id,filenameFormat))
list_file.close()
执行后得到结果:
四、创建 fruit.yaml
文件
新建fruit.yaml
文件:
train: D:/jupyter notebook/365-DL/YOLO/Y2/yolov5-master/Y2-fruit_data/ImageSets/Main/train.txt
val: D:/jupyter notebook/365-DL/YOLO/Y2/yolov5-master/Y2-fruit_data/ImageSets/Main/val.txt
# number of classes
nc: 4
# class names
names: ["banana", "snake fruit", "dragon fruit", "pineapple"]
五、开始用自己的数据集训练模型
(1) 输入训练指令
由于本机没有GPU,所以执行:python .\train.py --img 900 --batch 2 --epoch 100 --data .\fruit.yaml --cfg .\models\yolov5s.yaml --weights .\yolov5s.pt --device cpu
如果有GPU,则执行:
python .\train.py --img 900 --batch 2 --epoch 100 --data .\fruit.yaml --cfg .\models\yolov5s.yaml --weights .\yolov5s.pt --device '0'
执行命令后,出现如下提示,表明训练进行中,等待训练完成,查看训练结果。
(2)查看训练结果
如图中所示,使用YOLO-v5s训练本文的数据集:
- a. 100个epoch需要的时间是3.382小时
- b. YOLOv5s 网络结构: 157 层, 参数量是7020913 , 梯度是0 , GFLOPs是15.8
- c. 还显示了类别的训练结果,包括P-R值、mAP50的值
- d. 训练结果保存在
runs\train\exp7
,在该路径下生成了许多文件:
打开其中一张图片,如val_batch1_labels.jpg
,如下图,显示了各水果的标签:
打开val_batch1_labels.jpg
,则显示了带预测值的标签: