更多有用项目见----------->>>>>> 深度学必学项目!!
数据集划分与文件列表生成脚本总结
此Python脚本用于将标注文件(以XML格式存储)划分为训练集、验证集和测试集,并生成相应的文本文件列表,以便于机器学习或深度学习模型的训练和评估。以下是脚本的工作流程及其关键功能的详细说明:
1. 模块导入
首先,脚本导入了三个Python标准库模块:
os:提供与操作系统交互的功能,如创建目录、检查路径存在性等。random:提供随机数生成功能,用于随机选择数据样本。argparse:支持命令行参数解析,使用户可以通过命令行指定输入输出路径和其他配置选项。
2. 命令行参数解析
通过argparse定义并解析两个命令行参数:
--xml_path:指定包含XML标注文件的目录路径,默认值为D:\Desktop\YOLOv9\data\Annotations。--txt_path:指定保存生成的TXT文件列表的目录路径,默认值为D:\Desktop\YOLOv9\data\images。
3. 辅助函数
为了提高代码复用性和可读性,脚本定义了几个辅助函数:
create_dir_if_not_exists(path):如果给定的路径不存在,则递归地创建该路径所指向的目录。write_to_file(file, filenames):将一个字符串列表写入指定文件中,每个元素占一行。
4. 主逻辑 (main函数)
主逻辑包括以下几个步骤:
- 解析命令行参数:调用
parse_args()获取用户提供的参数。 - 设置数据划分比例:定义训练+验证集占总数据量的比例(
trainval_percent)以及训练集中训练数据的比例(train_percent)。 - 获取所有XML文件名:从指定的XML路径中提取所有
.xml文件的名字(去掉扩展名),形成一个总的文件名列表。 - 确保输出目录存在:使用
create_dir_if_not_exists(txtsavepath)保证输出TXT文件的目录已创建。 - 计算各集合大小:根据设定的比例计算出训练+验证集和训练集的具体数量。
- 随机分配文件到不同集合:利用
random.sample()方法,按照预定的比例随机选取文件索引,分别构建训练+验证集、训练集、验证集和测试集。 - 生成并保存文件列表:对于每个集合,构造对应的TXT文件路径,并调用
write_to_file()将相关文件名写入文件中。
5. 执行入口
最后,在if __name__ == "__main__":条件下调用main()函数启动整个流程,确保当脚本被直接执行时会自动运行核心处理逻辑。
通过这种方式,该脚本能够自动化地完成数据集的划分任务,并准备好用于后续训练过程所需的文件列表,简化了数据预处理的工作量,提高了工作效率。
划分训练测试验证数据集脚本
import os # 导入操作系统相关模块 import random # 导入随机数生成模块 import argparse # 导入命令行参数解析模块 def parse_args(): """解析命令行参数""" parser = argparse.ArgumentParser(description="划分数据集并生成训练、验证和测试集的文件列表") parser.add_argument('--xml_path', default=r'D:\Desktop\YOLOv9\data\Annotations', type=str, help='输入 XML 标签路径') parser.add_argument('--txt_path', default=r'D:\Desktop\YOLOv9\data\images', type=str, help='输出 TXT 标签路径') return parser.parse_args() # 返回解析后的命令行参数 def create_dir_if_not_exists(path): """如果路径不存在,则创建它""" if not os.path.exists(path): # 检查路径是否存在 os.makedirs(path) # 创建目录及其所有父目录 def write_to_file(file, filenames): """将文件名列表写入文件""" with open(file, 'w') as f: # 以写模式打开文件 for name in filenames: # 遍历文件名列表 f.write(name + '\n') # 将文件名写入文件,并在末尾添加换行符 def main(): # 解析命令行参数 opt = parse_args() # 数据划分比例 trainval_percent = 0.9 # 训练集和验证集的总比例 train_percent = 0.8 # 训练集中训练数据的比例 # 获取 XML 文件路径和输出 TXT 文件路径 xmlfilepath = opt.xml_path txtsavepath = opt.txt_path # 获取 XML 文件夹下的所有文件名(不包括扩展名) total_xml = [f[:-4] for f in os.listdir(xmlfilepath) if f.endswith('.xml')] # 如果输出目录不存在,则创建该目录 create_dir_if_not_exists(txtsavepath) # 总文件数量 num = len(total_xml) # 计算 XML 文件的总数量 list_index = list(range(num)) # 生成文件索引列表 # 计算训练+验证集数量和训练集数量 tv = int(num * trainval_percent) # 训练+验证集的数量 tr = int(tv * train_percent) # 训练集的数量 # 随机抽取训练+验证集和训练集 trainval = random.sample(list_index, tv) # 随机选择 tv 个文件作为训练+验证集 train = random.sample(trainval, tr) # 从训练+验证集中随机选择 tr 个文件作为训练集 # 分配文件名到对应的集合 trainval_files = [total_xml[i] for i in trainval] # 获取训练+验证集文件名 test_files = [total_xml[i] for i in list_index if i not in trainval] # 获取测试集文件名 train_files = [total_xml[i] for i in train] # 获取训练集文件名 val_files = [total_xml[i] for i in trainval if i not in train] # 获取验证集文件名 # 定义文件路径 trainval_path = os.path.join(txtsavepath, 'trainval.txt') # 训练+验证集文件路径 test_path = os.path.join(txtsavepath, 'test.txt') # 测试集文件路径 train_path = os.path.join(txtsavepath, 'train.txt') # 训练集文件路径 val_path = os.path.join(txtsavepath, 'val.txt') # 验证集文件路径 # 将文件名写入对应的 TXT 文件中 write_to_file(trainval_path, trainval_files) # 写入训练+验证集文件名 write_to_file(test_path, test_files) # 写入测试集文件名 write_to_file(train_path, train_files) # 写入训练集文件名 write_to_file(val_path, val_files) # 写入验证集文件名 if __name__ == "__main__": main() # 运行主函数xml2txt代码
数据集转换脚本总结与阐述
此Python脚本用于将Pascal VOC格式的XML标注文件转换为YOLO格式的TXT文件,适用于目标检测任务中的数据预处理。以下是该脚本的工作流程及关键功能的详细说明:
1. 环境设置与导入模块
首先,脚本设置了字符编码为UTF-8,并导入了必要的模块:
xml.etree.ElementTree(简称ET):用于解析和操作XML文件。os:提供操作系统相关的功能,如路径操作、创建目录等。getcwd:获取当前工作目录。
2. 配置信息
定义了一些全局变量来配置数据集的基本信息:
- 子集划分:
sets = ['train', 'val'],表示将数据集划分为训练集和验证集两个部分。 - 类别列表:
classes = ["fry"],指定了数据集中包含的目标类别。这里仅有一个类别"fry"。 - 绝对路径:
abs_path = os.getcwd(),打印当前工作目录以确认脚本运行的位置。
3. 辅助函数
为了提高代码复用性和可读性,脚本定义了几个辅助函数:
坐标转换 (
convert)- 功能:将XML文件中给出的边界框坐标(xmin, xmax, ymin, ymax)转换为YOLO格式(x_center, y_center, width, height),并归一化到[0,1]区间内。
- 参数:
size:图像的宽度和高度。box:原始边界框的坐标。
- 返回值:转换后的YOLO格式坐标。
单个文件转换 (
convert_annotation)- 功能:读取一个XML文件,提取其中的对象信息,并根据
convert函数的结果生成相应的TXT文件。 - 参数:
image_id:图像的唯一标识符。in_dir:输入XML文件所在的目录。out_dir:输出TXT文件保存的目录。
- 特点:处理时会跳过困难样本(
difficult == 1),并且对越界的边界框进行修正,确保其不超过图像尺寸。
- 功能:读取一个XML文件,提取其中的对象信息,并根据
4. 主逻辑 (process_dataset)
主逻辑包括以下几个步骤:
- 创建标签目录:如果指定的输出标签目录不存在,则创建它。
- 遍历数据子集:对于每个子集(如
train或val),读取对应的.txt文件以获取图像ID列表。 - 构建文件列表:为每个子集创建一个文本文件,列出所有参与训练或验证的图像路径。
- 调用转换函数:针对每一个图像ID,调用
convert_annotation完成从XML到TXT的转换过程。
5. 执行入口
最后,通过设置具体的目录路径(xml_dir,labels_dir,images_dir),并调用process_dataset函数启动整个流程。这一步骤确保脚本能够自动化地完成数据集的转换任务,准备用于后续的目标检测模型训练。
进一步阐述
在实际应用中,这个脚本可以大大简化数据预处理的工作量,尤其是在需要频繁更新数据集或者调整类别标签的情况下。通过简单的修改,比如更改classes列表中的内容或调整sets数组,用户可以轻松适应不同的项目需求。
此外,该脚本还具备良好的错误处理机制,例如检查文件是否存在以及处理边界框越界的问题,保证了程序的稳定性和鲁棒性。同时,由于使用了标准库中的xml.etree.ElementTree来进行XML解析,使得代码易于理解和维护。
总之,这是一个非常实用的数据集转换工具,特别适合那些采用YOLO框架进行目标检测的研究人员或开发者。它不仅实现了从一种常见标注格式到另一种格式的高效转换,而且提供了足够的灵活性以应对各种特殊情况。
代码
# -*- coding: utf-8 -*- import xml.etree.ElementTree as ET import os from os import getcwd # 数据集的子集 sets = ['train', 'val'] # 自定义类别 classes = ["fry"] # 获取当前工作目录 abs_path = os.getcwd() print(abs_path) def convert(size, box): """ 将标注的边界框坐标转换为 YOLO 格式 :param size: 图片的宽和高 :param box: 边界框的坐标 (xmin, xmax, ymin, ymax) :return: 转换后的坐标 (x, y, w, h) """ dw = 1. / size[0] dh = 1. / size[1] x = (box[0] + box[1]) / 2.0 - 1 y = (box[2] + box[3]) / 2.0 - 1 w = box[1] - box[0] h = box[3] - box[2] x = x * dw w = w * dw y = y * dh h = h * dh return x, y, w, h def convert_annotation(image_id, in_dir, out_dir): """ 将单个 XML 文件转换为 YOLO 格式的 TXT 文件 :param image_id: 图片的 ID :param in_dir: 输入 XML 文件的目录 :param out_dir: 输出 TXT 文件的目录 """ in_file_path = os.path.join(in_dir, f'{image_id}.xml') out_file_path = os.path.join(out_dir, f'{image_id}.txt') try: in_file = open(in_file_path, encoding='UTF-8') except FileNotFoundError: print(f"文件 {in_file_path} 不存在,请检查路径和文件名。") return with open(out_file_path, 'w') as out_file: tree = ET.parse(in_file) root = tree.getroot() 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) 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) bb = convert((w, h), b) out_file.write(f"{cls_id} {' '.join(map(str, bb))}\n") def process_dataset(sets, in_dir, out_dir, img_dir): """ 处理整个数据集,将 XML 文件转换为 YOLO 格式 :param sets: 数据集的子集列表 :param in_dir: 输入 XML 文件的目录 :param out_dir: 输出 TXT 文件的目录 :param img_dir: 图像文件的目录 """ # 创建标签目录(如果不存在) if not os.path.exists(out_dir): os.makedirs(out_dir) for image_set in sets: image_ids_path = os.path.join(img_dir, f'{image_set}.txt') if not os.path.exists(image_ids_path): print(f"文件 {image_ids_path} 不存在,请检查路径和文件名。") continue with open(image_ids_path) as f: image_ids = f.read().strip().split() list_file_path = os.path.join(getcwd(), f'{image_set}.txt') with open(list_file_path, 'w') as list_file: for image_id in image_ids: list_file.write(f'{os.path.join(img_dir, image_id)}.jpg\n') convert_annotation(image_id, in_dir, out_dir) # 设置目录路径 xml_dir = r'D:\Desktop\YOLOv9\data\Annotations' labels_dir = r'D:\Desktop\YOLOv9\data\labels' images_dir = r'D:\Desktop\YOLOv9\data\images' # 处理数据集 process_dataset(sets, xml_dir, labels_dir, images_dir)