作者：CSDN @ _养乐多_
本文将介绍如何将目标检测中常用的VOC格式数据集转换为YOLO数据集，并进行数据集比例划分，从而方便的进行YOLO目标检测。
如果不想分两步，可以直接看第三节代码。
文章目录
一、将VOC格式数据集转换为YOLO格式数据集二、YOLO格式数据集划分（训练、验证、测试）2.1 版本12.2 版本2三、一步到位
一、将VOC格式数据集转换为YOLO格式数据集
执行以下脚本将VOC格式数据集转换为YOLO格式数据集。
但是需要注意的是：
import osimport shutilimport xml.etree.ElementTree as ET# VOC格式数据集路径voc_data_path = 'E:\\DataSet\\helmet-VOC'voc_annotations_path = os.path.join(voc_data_path, 'Annotations')voc_images_path = os.path.join(voc_data_path, 'JPEGImages')# YOLO格式数据集保存路径yolo_data_path = 'E:\\DataSet\\helmet-YOLO'yolo_images_path = os.path.join(yolo_data_path, 'images')yolo_labels_path = os.path.join(yolo_data_path, 'labels')# 创建YOLO格式数据集目录os.makedirs(yolo_images_path, exist_ok=True)os.makedirs(yolo_labels_path, exist_ok=True)# 类别映射 (可以根据自己的数据集进行调整)class_mapping = { 'head': 0, 'helmet': 1, 'person': 2, # 添加更多类别...}def convert_voc_to_yolo(voc_annotation_file, yolo_label_file): tree = ET.parse(voc_annotation_file) root = tree.getroot() size = root.find('size') width = float(size.find('width').text) height = float(size.find('height').text) with open(yolo_label_file, 'w') as f: for obj in root.findall('object'): cls = obj.find('name').text if cls not in class_mapping: continue cls_id = class_mapping[cls] xmlbox = obj.find('bndbox') xmin = float(xmlbox.find('xmin').text) ymin = float(xmlbox.find('ymin').text) xmax = float(xmlbox.find('xmax').text) ymax = float(xmlbox.find('ymax').text) x_center = (xmin + xmax) / 2.0 / width y_center = (ymin + ymax) / 2.0 / height w = (xmax - xmin) / width h = (ymax - ymin) / height f.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")# 遍历VOC数据集的Annotations目录，进行转换for voc_annotation in os.listdir(voc_annotations_path): if voc_annotation.endswith('.xml'): voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation) image_id = os.path.splitext(voc_annotation)[0] voc_image_file = os.path.join(voc_images_path, f"{image_id}.jpg") yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt") yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.jpg") convert_voc_to_yolo(voc_annotation_file, yolo_label_file) if os.path.exists(voc_image_file): shutil.copy(voc_image_file, yolo_image_file)print("转换完成！")二、YOLO格式数据集划分（训练、验证、测试）
参考：https://docs.ultralytics.com/datasets/detect/#ultralytics-yolo-format
随机将数据集按照0.7-0.2-0.1比例划分为训练、验证、测试数据集。
注意，修改代码中图片的后缀，如果是.jpg，就把.png修改为.jpg。
最终结果，
2.1 版本1
用版本1划分就行，不用版本2了。
import osimport shutilimport randomfrom math import floor# 创建输出目录的函数def create_dirs(output_dir): images_dir = os.path.join(output_dir, 'images') labels_dir = os.path.join(output_dir, 'labels') for split in ['train', 'val', 'test']: os.makedirs(os.path.join(images_dir, split), exist_ok=True) os.makedirs(os.path.join(labels_dir, split), exist_ok=True) return images_dir, labels_dir# 获取图片和对应txt标签的列表def get_files(images_path, labels_path): image_files = [f for f in os.listdir(images_path) if f.endswith(('jpg', 'png', 'jpeg'))] label_files = [f for f in os.listdir(labels_path) if f.endswith('.txt')] # 检查图片和标签是否配对 paired_files = [] for image_file in image_files: base_name = os.path.splitext(image_file)[0] label_file = base_name + '.txt' if label_file in label_files: paired_files.append((image_file, label_file)) return paired_files# 将文件按比例划分并拷贝到相应目录def split_and_copy(paired_files, images_path, labels_path, images_dir, labels_dir, train_ratio, val_ratio): random.shuffle(paired_files) # 随机打乱 total_files = len(paired_files) train_count = floor(total_files * train_ratio) val_count = floor(total_files * val_ratio) test_count = total_files - train_count - val_count splits = { 'train': paired_files[:train_count], 'val': paired_files[train_count:train_count + val_count], 'test': paired_files[train_count + val_count:] } for split, files in splits.items(): for image_file, label_file in files: shutil.copy(os.path.join(images_path, image_file), os.path.join(images_dir, split, image_file)) shutil.copy(os.path.join(labels_path, label_file), os.path.join(labels_dir, split, label_file)) print(f'{split}: {len(files)} files')# 主函数def main(): # 写死的路径 images_path = "E:\\DataSet\\LC\\large_coal_blocked_yolo\\totalImages" # 替换为实际图片文件夹路径 labels_path = "E:\\DataSet\\LC\\large_coal_blocked_yolo\\totalLabels" # 替换为实际txt文件夹路径 output_dir = "E:\\DataSet\\LC\\large_coal_blocked_yolo\\output" # 替换为实际输出主目录路径 # 数据划分比例 train_ratio = 0.7 val_ratio = 0.3 test_ratio = 0 # 容差值用于浮点数比较 epsilon = 1e-6 # 确保比例之和等于1 assert abs(train_ratio + val_ratio + test_ratio - 1) < epsilon, "比例之和必须等于1" # 创建目录 images_dir, labels_dir = create_dirs(output_dir) # 获取文件列表 paired_files = get_files(images_path, labels_path) # 进行划分并拷贝 split_and_copy(paired_files, images_path, labels_path, images_dir, labels_dir, train_ratio, val_ratio)# 调用主函数if __name__ == "__main__": main()2.2 版本2
import osimport shutilimport random# YOLO格式数据集保存路径yolo_images_path1 = 'E:\\DataSet\\helmet-VOC'yolo_labels_path1 = 'E:\\DataSet\\helmet-YOLO'yolo_data_path = yolo_labels_path1yolo_images_path = os.path.join(yolo_images_path1, 'JPEGImages')yolo_labels_path = os.path.join(yolo_labels_path1, 'labels')# 创建划分后的目录结构train_images_path = os.path.join(yolo_data_path, 'train', 'images')train_labels_path = os.path.join(yolo_data_path, 'train', 'labels')val_images_path = os.path.join(yolo_data_path, 'val', 'images')val_labels_path = os.path.join(yolo_data_path, 'val', 'labels')test_images_path = os.path.join(yolo_data_path, 'test', 'images')test_labels_path = os.path.join(yolo_data_path, 'test', 'labels')os.makedirs(train_images_path, exist_ok=True)os.makedirs(train_labels_path, exist_ok=True)os.makedirs(val_images_path, exist_ok=True)os.makedirs(val_labels_path, exist_ok=True)os.makedirs(test_images_path, exist_ok=True)os.makedirs(test_labels_path, exist_ok=True)# 获取所有图片文件名（不包含扩展名）image_files = [f[:-4] for f in os.listdir(yolo_images_path) if f.endswith('.png')]# 随机打乱文件顺序random.shuffle(image_files)# 划分数据集比例train_ratio = 0.7val_ratio = 0.2test_ratio = 0.1train_count = int(train_ratio * len(image_files))val_count = int(val_ratio * len(image_files))test_count = len(image_files) - train_count - val_counttrain_files = image_files[:train_count]val_files = image_files[train_count:train_count + val_count]test_files = image_files[train_count + val_count:]# 移动文件到相应的目录def move_files(files, src_images_path, src_labels_path, dst_images_path, dst_labels_path): for file in files: src_image_file = os.path.join(src_images_path, f"{file}.png") src_label_file = os.path.join(src_labels_path, f"{file}.txt") dst_image_file = os.path.join(dst_images_path, f"{file}.png") dst_label_file = os.path.join(dst_labels_path, f"{file}.txt") if os.path.exists(src_image_file) and os.path.exists(src_label_file): shutil.move(src_image_file, dst_image_file) shutil.move(src_label_file, dst_label_file)# 移动训练集文件move_files(train_files, yolo_images_path, yolo_labels_path, train_images_path, train_labels_path)# 移动验证集文件move_files(val_files, yolo_images_path, yolo_labels_path, val_images_path, val_labels_path)# 移动测试集文件move_files(test_files, yolo_images_path, yolo_labels_path, test_images_path, test_labels_path)print("数据集划分完成！")三、一步到位
如果不想分两步进行格式转换，那么以下脚本结合了以上两步，直接得到最后按比例划分训练、验证、测试的数据集结果。
注意：需要修改 voc_data_path ，yolo_data_path ，class_mapping 以及 ‘.png’ 后缀。
import osimport shutilimport randomimport xml.etree.ElementTree as ETfrom tqdm import tqdm# VOC格式数据集路径voc_data_path = 'E:\\DataSet-VOC'voc_annotations_path = os.path.join(voc_data_path, 'Annotations')voc_images_path = os.path.join(voc_data_path, 'JPEGImages')# YOLO格式数据集保存路径yolo_data_path = 'E:\\DataSet-YOLO'yolo_images_path = os.path.join(yolo_data_path, 'images')yolo_labels_path = os.path.join(yolo_data_path, 'labels')# 创建YOLO格式数据集目录os.makedirs(yolo_images_path, exist_ok=True)os.makedirs(yolo_labels_path, exist_ok=True)# 类别映射 (可以根据自己的数据集进行调整)class_mapping = { 'head': 0, 'helmet': 1, 'person': 2, # 添加更多类别...}def convert_voc_to_yolo(voc_annotation_file, yolo_label_file): tree = ET.parse(voc_annotation_file) root = tree.getroot() size = root.find('size') width = float(size.find('width').text) height = float(size.find('height').text) with open(yolo_label_file, 'w') as f: for obj in root.findall('object'): cls = obj.find('name').text if cls not in class_mapping: continue cls_id = class_mapping[cls] xmlbox = obj.find('bndbox') xmin = float(xmlbox.find('xmin').text) ymin = float(xmlbox.find('ymin').text) xmax = float(xmlbox.find('xmax').text) ymax = float(xmlbox.find('ymax').text) x_center = (xmin + xmax) / 2.0 / width y_center = (ymin + ymax) / 2.0 / height w = (xmax - xmin) / width h = (ymax - ymin) / height f.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")# 遍历VOC数据集的Annotations目录，进行转换print("开始VOC到YOLO格式转换...")for voc_annotation in tqdm(os.listdir(voc_annotations_path)): if voc_annotation.endswith('.xml'): voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation) image_id = os.path.splitext(voc_annotation)[0] voc_image_file = os.path.join(voc_images_path, f"{image_id}.png") yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt") yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.png") convert_voc_to_yolo(voc_annotation_file, yolo_label_file) if os.path.exists(voc_image_file): shutil.copy(voc_image_file, yolo_image_file)print("VOC到YOLO格式转换完成！")# 划分数据集train_images_path = os.path.join(yolo_data_path, 'train', 'images')train_labels_path = os.path.join(yolo_data_path, 'train', 'labels')val_images_path = os.path.join(yolo_data_path, 'val', 'images')val_labels_path = os.path.join(yolo_data_path, 'val', 'labels')test_images_path = os.path.join(yolo_data_path, 'test', 'images')test_labels_path = os.path.join(yolo_data_path, 'test', 'labels')os.makedirs(train_images_path, exist_ok=True)os.makedirs(train_labels_path, exist_ok=True)os.makedirs(val_images_path, exist_ok=True)os.makedirs(val_labels_path, exist_ok=True)os.makedirs(test_images_path, exist_ok=True)os.makedirs(test_labels_path, exist_ok=True)# 获取所有图片文件名（不包含扩展名）image_files = [f[:-4] for f in os.listdir(yolo_images_path) if f.endswith('.png')]# 随机打乱文件顺序random.shuffle(image_files)# 划分数据集比例train_ratio = 0.7val_ratio = 0.2test_ratio = 0.1train_count = int(train_ratio * len(image_files))val_count = int(val_ratio * len(image_files))test_count = len(image_files) - train_count - val_counttrain_files = image_files[:train_count]val_files = image_files[train_count:train_count + val_count]test_files = image_files[train_count + val_count:]# 移动文件到相应的目录def move_files(files, src_images_path, src_labels_path, dst_images_path, dst_labels_path): for file in tqdm(files): src_image_file = os.path.join(src_images_path, f"{file}.png") src_label_file = os.path.join(src_labels_path, f"{file}.txt") dst_image_file = os.path.join(dst_images_path, f"{file}.png") dst_label_file = os.path.join(dst_labels_path, f"{file}.txt") if os.path.exists(src_image_file) and os.path.exists(src_label_file): shutil.move(src_image_file, dst_image_file) shutil.move(src_label_file, dst_label_file)# 移动训练集文件print("移动训练集文件...")move_files(train_files, yolo_images_path, yolo_labels_path, train_images_path, train_labels_path)# 移动验证集文件print("移动验证集文件...")move_files(val_files, yolo_images_path, yolo_labels_path, val_images_path, val_labels_path)# 移动测试集文件print("移动测试集文件...")move_files(test_files, yolo_images_path, yolo_labels_path, test_images_path, test_labels_path)print("数据集划分完成！")# 删除原始的 images 和 labels 文件夹shutil.rmtree(yolo_images_path)shutil.rmtree(yolo_labels_path)print("原始 images 和 labels 文件夹删除完成！")