告别混乱标注！用Python脚本一键搞定YOLOv4训练所需的VOC格式数据集（附完整代码）-程序员充电站

告别混乱标注！用Python脚本一键搞定YOLOv4训练所需的VOC格式数据集（附完整代码）

当你准备开始YOLOv4的训练时，最令人头疼的往往不是模型配置，而是数据集的准备工作。那些散落在不同文件夹中的图片、格式各异的标注文件、需要手动划分的训练测试集，常常让开发者陷入混乱。本文将带你用Python脚本实现从原始数据到YOLOv4可用的VOC格式数据集的自动化转换，彻底告别手工操作的繁琐与错误。

1. 理解YOLOv4的数据需求

在开始编写脚本之前，我们需要清楚YOLOv4对训练数据的具体要求。与常见的PASCAL VOC格式不同，YOLOv4有其特定的标注格式和目录结构需求。

YOLOv4训练需要以下几种关键文件：

图片文件：通常存放在JPEGImages目录下，支持JPG、PNG等常见格式
标注文件：每个图片对应一个.txt文件，包含归一化后的边界框坐标
数据集划分文件：train.txt和test.txt，列出用于训练和测试的图片路径
类别定义文件：obj.names，列出所有类别名称

典型的VOC格式数据集目录结构如下：

VOCdevkit/ └── VOC2020/ ├── Annotations/ # 存放原始XML标注文件 ├── ImageSets/ │ └── Main/ # 存放train.txt和test.txt ├── JPEGImages/ # 存放所有图片 └── labels/ # 存放转换后的YOLO格式标注

2. 自动化转换脚本设计

我们将开发一个完整的Python脚本工具链，实现从原始数据到YOLOv4可用数据的一键转换。这个工具链包含三个核心脚本：

数据集划分脚本：自动将图片划分为训练集和测试集
标注格式转换脚本：将XML标注转换为YOLO格式的TXT文件
路径生成脚本：生成Darknet训练所需的绝对路径文件

2.1 数据集划分脚本

首先创建一个split_dataset.py脚本，用于自动划分训练集和测试集：

import os import random from shutil import copyfile def split_dataset(image_dir, output_dir, train_ratio=0.8): """ 自动划分数据集为训练集和测试集 :param image_dir: 原始图片目录 :param output_dir: 输出目录(VOC格式) :param train_ratio: 训练集比例 """ # 创建VOC标准目录结构 os.makedirs(os.path.join(output_dir, 'Annotations'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'ImageSets/Main'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'JPEGImages'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True) # 获取所有图片文件 image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))] random.shuffle(image_files) # 随机打乱 # 计算划分点 split_point = int(len(image_files) * train_ratio) train_files = image_files[:split_point] test_files = image_files[split_point:] # 生成train.txt和test.txt with open(os.path.join(output_dir, 'ImageSets/Main/train.txt'), 'w') as f: f.write('\n'.join([os.path.splitext(f)[0] for f in train_files])) with open(os.path.join(output_dir, 'ImageSets/Main/test.txt'), 'w') as f: f.write('\n'.join([os.path.splitext(f)[0] for f in test_files])) # 拷贝图片到JPEGImages目录 for img_file in image_files: src = os.path.join(image_dir, img_file) dst = os.path.join(output_dir, 'JPEGImages', img_file) copyfile(src, dst) print(f"数据集划分完成: 共{len(image_files)}张图片, 训练集{len(train_files)}张, 测试集{len(test_files)}张") if __name__ == '__main__': split_dataset('raw_images', 'VOCdevkit/VOC2020')

提示：在实际项目中，建议保持训练集和测试集的比例在7:3到8:2之间，确保模型有足够的训练数据同时也能充分评估性能。

2.2 标注格式转换脚本

接下来是核心的标注格式转换脚本voc_to_yolo.py，它将XML格式的标注转换为YOLO需要的TXT格式：

import xml.etree.ElementTree as ET import os def convert_annotation(xml_file, output_dir, classes): """ 将单个XML标注文件转换为YOLO格式 :param xml_file: XML标注文件路径 :param output_dir: 输出目录 :param classes: 类别列表 """ tree = ET.parse(xml_file) root = tree.getroot() # 获取图片尺寸 size = root.find('size') w = int(size.find('width').text) h = int(size.find('height').text) # 准备输出文件 image_id = os.path.splitext(os.path.basename(xml_file))[0] out_file = os.path.join(output_dir, f"{image_id}.txt") with open(out_file, 'w') as f: for obj in root.iter('object'): cls = obj.find('name').text if cls not in classes: 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) ) bb = convert((w, h), b) f.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n") def convert(size, box): """ 坐标转换函数：将VOC格式的绝对坐标转换为YOLO格式的相对坐标 :param size: 图片尺寸 (w, h) :param box: 边界框坐标 (xmin, xmax, ymin, ymax) :return: YOLO格式坐标 (x_center, y_center, width, height) """ dw = 1. / size[0] dh = 1. / size[1] x = (box[0] + box[1]) / 2.0 y = (box[2] + box[3]) / 2.0 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 batch_convert(voc_dir, classes): """ 批量转换VOC标注为YOLO格式 :param voc_dir: VOC格式数据集目录 :param classes: 类别列表 """ ann_dir = os.path.join(voc_dir, 'Annotations') out_dir = os.path.join(voc_dir, 'labels') os.makedirs(out_dir, exist_ok=True) for xml_file in os.listdir(ann_dir): if xml_file.endswith('.xml'): convert_annotation( os.path.join(ann_dir, xml_file), out_dir, classes ) print(f"标注转换完成: 共转换{len(os.listdir(ann_dir))}个XML文件") if __name__ == '__main__': # 定义你的类别列表 CLASSES = ['person', 'car', 'dog', 'cat'] # 替换为你的实际类别 batch_convert('VOCdevkit/VOC2020', CLASSES)

2.3 路径生成脚本

最后，我们需要一个脚本generate_paths.py来生成Darknet训练所需的路径文件：

import os def generate_path_file(voc_dir, output_file, set_type='train'): """ 生成Darknet训练所需的路径文件 :param voc_dir: VOC格式数据集目录 :param output_file: 输出文件路径 :param set_type: 数据集类型 (train/test) """ image_dir = os.path.join(voc_dir, 'JPEGImages') set_file = os.path.join(voc_dir, 'ImageSets/Main', f'{set_type}.txt') with open(set_file, 'r') as f: image_ids = f.read().strip().split('\n') with open(output_file, 'w') as f: for image_id in image_ids: image_path = os.path.abspath(os.path.join(image_dir, f"{image_id}.jpg")) f.write(f"{image_path}\n") print(f"路径文件生成完成: {output_file}") if __name__ == '__main__': # 生成训练集路径文件 generate_path_file( 'VOCdevkit/VOC2020', '2020_train.txt', 'train' ) # 生成测试集路径文件 generate_path_file( 'VOCdevkit/VOC2020', '2020_test.txt', 'test' )

3. 完整工作流整合

现在我们将上述脚本整合成一个完整的自动化工作流。创建一个run_pipeline.py脚本作为入口点：

import os from split_dataset import split_dataset from voc_to_yolo import batch_convert from generate_paths import generate_path_file def main(): # 配置参数 RAW_IMAGE_DIR = 'raw_images' # 原始图片目录 RAW_ANNOTATION_DIR = 'raw_annotations' # 原始XML标注目录 VOC_DIR = 'VOCdevkit/VOC2020' # VOC格式输出目录 CLASSES = ['person', 'car', 'dog', 'cat'] # 你的类别列表 # 1. 数据集划分 print("=== 开始数据集划分 ===") split_dataset(RAW_IMAGE_DIR, VOC_DIR) # 2. 拷贝原始标注到Annotations目录 print("\n=== 准备标注文件 ===") os.makedirs(os.path.join(VOC_DIR, 'Annotations'), exist_ok=True) for ann_file in os.listdir(RAW_ANNOTATION_DIR): if ann_file.endswith('.xml'): src = os.path.join(RAW_ANNOTATION_DIR, ann_file) dst = os.path.join(VOC_DIR, 'Annotations', ann_file) os.replace(src, dst) print(f"已拷贝{len(os.listdir(os.path.join(VOC_DIR, 'Annotations')))}个标注文件") # 3. 标注格式转换 print("\n=== 开始标注格式转换 ===") batch_convert(VOC_DIR, CLASSES) # 4. 生成路径文件 print("\n=== 生成路径文件 ===") generate_path_file(VOC_DIR, '2020_train.txt', 'train') generate_path_file(VOC_DIR, '2020_test.txt', 'test') # 5. 生成obj.names文件 print("\n=== 生成类别文件 ===") with open('obj.names', 'w') as f: f.write('\n'.join(CLASSES)) print(f"已生成类别文件，共{len(CLASSES)}个类别") # 6. 生成obj.data文件 print("\n=== 生成配置文件 ===") with open('obj.data', 'w') as f: content = f"""classes = {len(CLASSES)} train = 2020_train.txt valid = 2020_test.txt names = obj.names backup = backup/ """ f.write(content) print("配置文件生成完成") if __name__ == '__main__': main()

4. 常见问题与解决方案

在实际使用这套自动化工具时，可能会遇到一些典型问题。以下是常见问题及其解决方案：

4.1 标注文件与图片不匹配

症状：运行脚本时报错，提示找不到对应的图片或标注文件。

解决方案：

确保所有图片都有对应的标注文件，且文件名（不含扩展名）完全一致
使用以下脚本检查匹配情况：

import os def check_matching(image_dir, ann_dir): images = {os.path.splitext(f)[0] for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))} annotations = {os.path.splitext(f)[0] for f in os.listdir(ann_dir) if f.endswith('.xml')} missing_ann = images - annotations missing_img = annotations - images if missing_ann: print(f"警告: {len(missing_ann)}张图片缺少标注文件") for name in sorted(missing_ann)[:5]: # 只显示前5个 print(f" - {name}") if missing_img: print(f"警告: {len(missing_img)}个标注文件缺少对应图片") for name in sorted(missing_img)[:5]: print(f" - {name}") if not missing_ann and not missing_img: print("所有图片和标注文件匹配正确") check_matching('VOCdevkit/VOC2020/JPEGImages', 'VOCdevkit/VOC2020/Annotations')

4.2 类别ID不连续

症状：训练时报错，提示类别ID超出范围。

解决方案：

确保CLASSES列表中的类别顺序与标注文件中的类别名称完全一致
使用以下脚本验证类别一致性：

def check_classes(voc_dir, classes): ann_dir = os.path.join(voc_dir, 'Annotations') class_set = set(classes) found_classes = set() for xml_file in os.listdir(ann_dir): if not xml_file.endswith('.xml'): continue tree = ET.parse(os.path.join(ann_dir, xml_file)) root = tree.getroot() for obj in root.iter('object'): cls = obj.find('name').text found_classes.add(cls) unknown_classes = found_classes - class_set if unknown_classes: print(f"错误: 发现未定义的类别: {unknown_classes}") print("请在CLASSES列表中添加这些类别或检查标注文件") else: print("所有标注类别都已定义") check_classes('VOCdevkit/VOC2020', CLASSES)

4.3 坐标转换错误

症状：训练时损失不下降或预测框位置明显错误。

解决方案：

验证坐标转换是否正确，使用以下脚本抽样检查：

import cv2 import random def visualize_annotations(image_dir, label_dir, classes, sample_size=3): image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))] samples = random.sample(image_files, min(sample_size, len(image_files))) for img_file in samples: img_path = os.path.join(image_dir, img_file) label_path = os.path.join(label_dir, os.path.splitext(img_file)[0] + '.txt') img = cv2.imread(img_path) h, w = img.shape[:2] if os.path.exists(label_path): with open(label_path, 'r') as f: lines = f.readlines() for line in lines: parts = line.strip().split() cls_id = int(parts[0]) x, y, box_w, box_h = map(float, parts[1:]) # 转换回绝对坐标 x = x * w y = y * h box_w = box_w * w box_h = box_h * h x1 = int(x - box_w / 2) y1 = int(y - box_h / 2) x2 = int(x + box_w / 2) y2 = int(y + box_h / 2) cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(img, classes[cls_id], (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) cv2.imshow(f"验证: {img_file}", img) cv2.waitKey(0) cv2.destroyAllWindows() visualize_annotations('VOCdevkit/VOC2020/JPEGImages', 'VOCdevkit/VOC2020/labels', CLASSES)