使用的是YOLOv8.0.157,使用过程中想训练自己的数据集。发现自己的数据大多是使用VOC格式(标签文件是xml),记录一下自己进行数据集转换的过程。
目录
数据格式对比
根据示例(coco128),使用的yolov8的数据集yaml文件格式为:
path: ./datasets # dataset root dir
train: images/train2017 # train images (relative to 'path') 128 images
val: images/train2017 # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
names:
0: person
1: bicycle
···
download: https://ultralytics.com/assets/coco128.zip
其中标签文件为txt,格式为:
45 0.479492 0.688771 0.955609 0.5955
45 0.736516 0.247188 0.498875 0.476417
50 0.637063 0.732938 0.494125 0.510583
45 0.339438 0.418896 0.678875 0.7815
49 0.646836 0.132552 0.118047 0.0969375
49 0.773148 0.129802 0.0907344 0.0972292
49 0.668297 0.226906 0.131281 0.146896
49 0.642859 0.0792187 0.148063 0.148062
也就是(类别 x中心点 y中心点 宽度 高度),且对坐标完成了归一化。
使用的xml文件格式为:
<annotation>
<folder>2023-02-07 000000</folder>
<filename>图片名称</filename>
<path>路径</path>
<source>
<database>Unknown</database>
</source>
<size>
<width>宽度</width>
<height>高度</height>
<depth>深度</depth>
</size>
<segmented>0</segmented>
<object>
<name>3</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>x轴最小值</xmin>
<ymin>y轴最小值</ymin>
<xmax>x轴最大值</xmax>
<ymax>y轴最小值</ymax>
</bndbox>
</object>
</annotation>
完整代码
不多说,先看完整代码,解析在后面:
import os, shutil
from tqdm import tqdm
from collections import Counter
import xml.etree.ElementTree as ET
from PIL import Image
import yaml
import random
class Dataset_Transforme_Yolov8:
def __init__(self, jpg_path: str, xml_path: str, save_path: str, divide=False):
self.jpg_path = jpg_path
self.xml_path = xml_path
self.save_path = save_path
self.divide = divide
def get_classes(self):
'''
统计路径下xml里的各类别标签数量
'''
names = []
files_1 = []
for root, dirs, files in os.walk(self.xml_path):
for file in files:
if ".xml" in file:
file = os.path.join(root, file)
subs = ET.parse(file).getroot().findall("object")
if len(subs) != 0:
files_1.append(file)
for sub in subs:
name = sub.find("name").text
names.append(name)
result = Counter(names)
return result, files_1
def xml2txt(self, classes, file_path, txt_save_path, image_width, image_height):
'''
根据xml文件生成txt标签文件
:param classes:类别列表
:param file_path:xml文件路径
:param image_width:图片宽度
:param image_height:图片高度
'''
tree = ET.parse(file_path)
root = tree.getroot()
objects = root.findall('object')
bboxes = []
class_names = []
for obj in objects:
bbox = obj.find('bndbox')
xmin = int(bbox.find('xmin').text)
ymin = int(bbox.find('ymin').text)
xmax = int(bbox.find('xmax').text)
ymax = int(bbox.find('ymax').text)
class_name = obj.find('name').text
c1 = round((xmin + xmax) / (image_width * 2), 6)
c2 = round((ymin + ymax) / (image_height * 2), 6)
c3 = round((xmax - xmin) / image_width, 6)
c4 = round((ymax - ymin) / image_height, 6)
if class_name in classes:
# print(class_name)
bboxes.append([c1, c2, c3, c4])
class_names.append(class_name)
# 将数据写入到YOLO的TXT文件中
with open(txt_save_path, 'w') as file:
for bbox, class_name in zip(bboxes, class_names):
file.write(
f"{classes.index(class_name)} {bbox[0]} {bbox[1]} {bbox[2]} {bbox[3]}\n")
def data_split(self, full_list, train, val, shuffle, seed):
"""
数据集拆分: 将列表full_list按比例划分为3个子列表
:param full_list: 数据列表
:param ratio: 比例
:param shuffle: 是否打乱顺序
:param seed: 随机种子
:return:
"""
n_total = len(full_list)
offset_train = int(n_total * train)
offset_val = int(n_total * val)
offset_trainval = offset_train + offset_val
if n_total == 0 or offset_train < 1 or train + val > 1.0:
raise ValueError("划分出错,请检查列表与划分比例!参考格式train,val->0.7,0.1")
if shuffle:
random.seed(seed)
random.shuffle(full_list)
train_list = full_list[:offset_train]
val_list = full_list[offset_train:offset_trainval]
test_list = full_list[offset_trainval:]
return train_list, val_list, test_list
def voc2yolov8(self, train, val, shuffle=True, seed=10):
"""
:param train: 训练集比例
:param val: 验证集比例
"""
obj_classes, files = self.get_classes()
os.makedirs(os.path.join(self.save_path, "images/train"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "images/val"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "images/test"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "labels/train"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "labels/val"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "labels/test"), exist_ok=True)
classes = list(obj_classes.keys())
train_list, val_list, test_list = self.data_split(files, train, val, shuffle, seed)
for file in tqdm(files):
if ".xml" in file:
name = file.replace("\\", "/").split("/")[-1].split(".")[0]
xml_file = file.replace("\\", "/")
jpg_file = os.path.join(self.jpg_path, name + ".jpg").replace("\\", "/")
img_w, img_h = Image.open(jpg_file).size
if file in val_list:
txt_save_path=os.path.join(self.save_path, "labels/val",name+".txt")
self.xml2txt(classes, xml_file, txt_save_path, img_w, img_h)
elif file in test_list:
txt_save_path = os.path.join(self.save_path, "labels/test", name + ".txt")
self.xml2txt(classes, xml_file, txt_save_path, img_w, img_h)
else:
txt_save_path = os.path.join(self.save_path, "labels/train", name + ".txt")
self.xml2txt(classes, xml_file, txt_save_path, img_w, img_h)
if self.divide:
if file in val_list:
shutil.copy(jpg_file, os.path.join(self.save_path, "images/val", name + ".jpg"))
elif file in test_list:
shutil.copy(jpg_file, os.path.join(self.save_path, "images/test", name + ".jpg"))
else:
shutil.copy(jpg_file, os.path.join(self.save_path, "images/train", name + ".jpg"))
# 编写yaml文件
classes_txt = {i: classes[i] for i in range(len(classes))} # 标签类别
data = {
'path': os.path.join(os.getcwd(),self.save_path),
'train': "images/train",
'val': "images/val",
'test': "images/test",
'names': classes_txt,
'download': ''
}
with open(self.save_path + '/dataset.yaml', 'w', encoding="utf-8") as file:
yaml.dump(data, file, allow_unicode=True)
print("标签:", dict(obj_classes))
print("有标签文件数量:", len(files))
if __name__ == '__main__':
jpg_path = r"" # 你的图片路径
xml_path = r"" # 你的xml路径
save_path = r"dataset_yolov8" # 保存数据路径
deals = Dataset_Transforme_Yolov8(jpg_path, xml_path, save_path, divide=True)
deals.voc2yolov8(0.7, 0.1)#训练集与验证集占比
get_classes
获取路径下xml内各类别标签的数量及标签文件名称
def get_classes(self):
'''
统计路径下xml里的各类别标签数量
'''
names = []
files_1 = []
for root, dirs, files in os.walk(self.xml_path):
for file in files:
if ".xml" in file:
file = os.path.join(root, file)
subs = ET.parse(file).getroot().findall("object")
if len(subs) != 0:
files_1.append(file)
for sub in subs:
name = sub.find("name").text
names.append(name)
result = Counter(names)
return result, files_1
xml2txt
将xml内的标签转化为txt格式,c1、c2、c3、c4分别对应x轴中心,y轴中心,宽,高
def xml2txt(self, classes, file_path, txt_save_path, image_width, image_height):
'''
根据xml文件生成txt标签文件
:param classes:类别列表
:param file_path:xml文件路径
:param image_width:图片宽度
:param image_height:图片高度
'''
tree = ET.parse(file_path)
root = tree.getroot()
objects = root.findall('object')
bboxes = []
class_names = []
for obj in objects:
bbox = obj.find('bndbox')
xmin = int(bbox.find('xmin').text)
ymin = int(bbox.find('ymin').text)
xmax = int(bbox.find('xmax').text)
ymax = int(bbox.find('ymax').text)
class_name = obj.find('name').text
c1 = round((xmin + xmax) / (image_width * 2), 6)
c2 = round((ymin + ymax) / (image_height * 2), 6)
c3 = round((xmax - xmin) / image_width, 6)
c4 = round((ymax - ymin) / image_height, 6)
if class_name in classes:
# print(class_name)
bboxes.append([c1, c2, c3, c4])
class_names.append(class_name)
# 将数据写入到YOLO的TXT文件中
with open(txt_save_path, 'w') as file:
for bbox, class_name in zip(bboxes, class_names):
file.write(
f"{classes.index(class_name)} {bbox[0]} {bbox[1]} {bbox[2]} {bbox[3]}\n")
data_split
将列表full_list按比例随机划分为3个子列表
def data_split(self, full_list, train, val, shuffle, seed):
"""
数据集拆分: 将列表full_list按比例随机划分为3个子列表
:param full_list: 数据列表
:param ratio: 比例
:param shuffle: 是否打乱顺序
:param seed: 随机种子
:return:
"""
n_total = len(full_list)
offset_train = int(n_total * train)
offset_val = int(n_total * val)
offset_trainval = offset_train + offset_val
if n_total == 0 or offset_train < 1 or train + val > 1.0:
raise ValueError("划分出错,请检查列表与划分比例!参考格式train,val->0.7,0.1")
if shuffle:
random.seed(seed)
random.shuffle(full_list)
train_list = full_list[:offset_train]
val_list = full_list[offset_train:offset_trainval]
test_list = full_list[offset_trainval:]
return train_list, val_list, test_list
voc2yolov8
voc转yolov8主要运行逻辑(万一以后还有其他拓展呢)。
创建出来的文件夹格式为:
dataset_yolov8
├─images
│ ├─test
│ ├─train
│ └─val
└─labels
├─test
├─train
└─val
通过yaml包编写yaml文件,这里生成的“path”为绝对路径。可以通过输入训练集与验证集的比例来控制数据集的划分,使用格式参考完整代码中的示例(训练与验证比例之和不超过1.0)。
def voc2yolov8(self, train, val, shuffle=True, seed=10):
"""
:param train: 训练集比例
:param val: 验证集比例
"""
obj_classes, files = self.get_classes()
os.makedirs(os.path.join(self.save_path, "images/train"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "images/val"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "images/test"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "labels/train"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "labels/val"), exist_ok=True)
os.makedirs(os.path.join(self.save_path, "labels/test"), exist_ok=True)
classes = list(obj_classes.keys())
train_list, val_list, test_list = self.data_split(files, train, val, shuffle, seed)
for file in tqdm(files):
if ".xml" in file:
name = file.replace("\\", "/").split("/")[-1].split(".")[0]
xml_file = file.replace("\\", "/")
jpg_file = os.path.join(self.jpg_path, name + ".jpg").replace("\\", "/")
img_w, img_h = Image.open(jpg_file).size
if file in val_list:
txt_save_path=os.path.join(self.save_path, "labels/val",name+".txt")
self.xml2txt(classes, xml_file, txt_save_path, img_w, img_h)
elif file in test_list:
txt_save_path = os.path.join(self.save_path, "labels/test", name + ".txt")
self.xml2txt(classes, xml_file, txt_save_path, img_w, img_h)
else:
txt_save_path = os.path.join(self.save_path, "labels/train", name + ".txt")
self.xml2txt(classes, xml_file, txt_save_path, img_w, img_h)
if self.divide:
if file in val_list:
shutil.copy(jpg_file, os.path.join(self.save_path, "images/val", name + ".jpg"))
elif file in test_list:
shutil.copy(jpg_file, os.path.join(self.save_path, "images/test", name + ".jpg"))
else:
shutil.copy(jpg_file, os.path.join(self.save_path, "images/train", name + ".jpg"))
# 编写yaml文件
classes_txt = {i: classes[i] for i in range(len(classes))} # 标签类别
data = {
'path': os.path.join(os.getcwd(),self.save_path),
'train': "images/train",
'val': "images/val",
'test': "images/test",
'names': classes_txt,
'download': ''
}
with open(self.save_path + '/dataset.yaml', 'w', encoding="utf-8") as file:
yaml.dump(data, file, allow_unicode=True)
print("标签:", dict(obj_classes))
print("有标签文件数量:", len(files))
大概就是这样,都是一些较为基础的代码,也不再过多赘述。经过测试生成的数据集是可以直接通过使用生成的dataset.yaml进行训练的,希望对你有所帮助。
对了,如果生成完成后移动过文件夹,记得把yaml文件中的路径也改了,不然会报错。