YOLOv5 v6.0 自定义数据集训练:从 VOC 到 YOLO 格式的 3 步转换与常见错误修复
在计算机视觉领域,目标检测是一项基础而重要的任务。YOLOv5作为当前最先进的目标检测框架之一,以其高效、准确和易用性受到广泛欢迎。然而,在实际应用中,许多开发者面临的首要挑战是如何将自己的数据集转换为YOLOv5可识别的格式。本文将深入探讨从VOC格式到YOLO格式的转换过程,并提供完整的Python脚本和常见错误解决方案。
1. 数据准备与环境配置
在开始转换之前,我们需要确保环境配置正确并准备好原始数据。YOLOv5 v6.0要求Python 3.8+和PyTorch 1.6+环境。建议使用conda创建虚拟环境:
conda create -n yolov5 python=3.8 conda activate yolov5 pip install torch torchvision数据集目录结构应按照VOC格式组织:
paper_data/ ├── Annotations/ # 存放XML标注文件 ├── images/ # 存放原始图像 └── ImageSets/ └── Main/ # 存放数据集划分文件(train.txt, val.txt等)关键检查点:
- 确保每个XML文件都有对应的图像文件
- 验证XML文件中width和height不为0
- 检查类别名称是否一致且不含特殊字符
2. VOC转YOLO格式的三步转换法
2.1 数据集划分文件生成
首先需要生成训练集、验证集和测试集的划分文件。以下Python脚本可以自动完成这一过程:
import os import random import argparse def split_dataset(xml_path, txt_path, trainval_percent=0.9, train_percent=0.8): xml_files = [f for f in os.listdir(xml_path) if f.endswith('.xml')] total_size = len(xml_files) # 计算各集合大小 trainval_size = int(total_size * trainval_percent) train_size = int(trainval_size * train_percent) # 随机划分 indices = list(range(total_size)) random.shuffle(indices) # 写入文件 os.makedirs(txt_path, exist_ok=True) with (open(f"{txt_path}/trainval.txt", 'w') as f_trainval, open(f"{txt_path}/train.txt", 'w') as f_train, open(f"{txt_path}/val.txt", 'w') as f_val, open(f"{txt_path}/test.txt", 'w') as f_test): for i in indices: name = os.path.splitext(xml_files[i])[0] + '\n' if i < trainval_size: f_trainval.write(name) if i < train_size: f_train.write(name) else: f_val.write(name) else: f_test.write(name) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--xml_path', default='Annotations', help='输入XML标注路径') parser.add_argument('--txt_path', default='ImageSets/Main', help='输出划分文件路径') opt = parser.parse_args() split_dataset(opt.xml_path, opt.txt_path)2.2 核心格式转换脚本
以下完整转换脚本包含边界框越界修正和类别ID映射功能:
import xml.etree.ElementTree as ET import os from os import getcwd sets = ['train', 'val', 'test'] classes = ["cat", "dog"] # 替换为你的类别 def convert(size, box): """将VOC格式的bbox转换为YOLO格式""" 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_file = open(f'Annotations/{image_id}.xml') out_file = open(f'labels/{image_id}.txt', 'w') 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)) # 边界框越界修正 b = list(b) if b[1] > w: b[1] = w if b[3] > h: b[3] = h if b[0] < 0: b[0] = 0 if b[2] < 0: b[2] = 0 bb = convert((w, h), b) out_file.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n") # 主程序 wd = getcwd() for image_set in sets: os.makedirs('labels', exist_ok=True) image_ids = open(f'ImageSets/Main/{image_set}.txt').read().strip().split() with open(f'{image_set}.txt', 'w') as list_file: for image_id in image_ids: list_file.write(f'{wd}/images/{image_id}.jpg\n') convert_annotation(image_id)2.3 生成YOLO配置文件
创建数据集配置文件data/custom.yaml:
# 训练和验证图像路径 train: ./train.txt val: ./val.txt # 类别数量 nc: 2 # 类别名称 names: ['cat', 'dog']3. 常见错误与解决方案
3.1 ZeroDivisionError: float division by zero
错误原因:XML文件中width或height为0
解决方案:
- 检查所有XML文件中的size字段
- 使用以下代码自动过滤无效标注:
def validate_xml(xml_file): try: tree = ET.parse(xml_file) root = tree.getroot() size = root.find('size') w = int(size.find('width').text) h = int(size.find('height').text) return w > 0 and h > 0 except: return False3.2 边界框越界问题
现象:训练时出现loss异常或检测框错位
修复方法: 在转换脚本中加入边界框修正逻辑:
# 在convert_annotation函数中添加 b = list(b) b[0] = max(0, min(w, b[0])) # xmin b[1] = max(0, min(w, b[1])) # xmax b[2] = max(0, min(h, b[2])) # ymin b[3] = max(0, min(h, b[3])) # ymax3.3 类别ID不匹配
错误表现:训练时出现类别预测混乱
解决方法:
- 确保classes列表与XML中的name一致
- 添加类别验证代码:
unknown_classes = set() for obj in root.iter('object'): cls = obj.find('name').text if cls not in classes: unknown_classes.add(cls) if unknown_classes: print(f"警告:发现未定义类别: {unknown_classes}")4. 高级技巧与优化建议
4.1 自动计算Anchor
使用k-means算法为自定义数据集生成优化的anchor尺寸:
import numpy as np from sklearn.cluster import KMeans def calculate_anchors(xml_path, n_clusters=9): boxes = [] for xml_file in os.listdir(xml_path): tree = ET.parse(os.path.join(xml_path, xml_file)) root = tree.getroot() w = int(root.find('size/width').text) h = int(root.find('size/height').text) for obj in root.iter('object'): box = obj.find('bndbox') xmin = float(box.find('xmin').text) / w ymin = float(box.find('ymin').text) / h xmax = float(box.find('xmax').text) / w ymax = float(box.find('ymax').text) / h boxes.append([xmax-xmin, ymax-ymin]) boxes = np.array(boxes) kmeans = KMeans(n_clusters=n_clusters, random_state=42) kmeans.fit(boxes) anchors = kmeans.cluster_centers_ * 640 # 乘以默认输入尺寸 anchors = anchors[np.argsort(anchors[:, 0])] # 按宽度排序 return anchors.round(2)4.2 数据增强配置
在data/hyps/hyp.scratch.yaml中调整数据增强参数:
# 颜色空间增强 hsv_h: 0.015 # 色调增强幅度 hsv_s: 0.7 # 饱和度增强幅度 hsv_v: 0.4 # 明度增强幅度 # 几何变换 degrees: 10.0 # 旋转角度范围 translate: 0.1 # 平移比例 scale: 0.5 # 缩放比例 shear: 0.0 # 剪切角度5. 验证转换结果
转换完成后,使用以下代码验证YOLO格式的正确性:
import cv2 import random def visualize_labels(img_path, label_path, classes): img = cv2.imread(img_path) h, w = img.shape[:2] with open(label_path) as f: for line in f.readlines(): cls, x, y, w_, h_ = map(float, line.split()) # 转换回像素坐标 x1 = int((x - w_/2) * w) y1 = int((y - h_/2) * h) x2 = int((x + w_/2) * w) y2 = int((y + h_/2) * h) # 随机颜色 color = (random.randint(0,255), random.randint(0,255), random.randint(0,255)) cv2.rectangle(img, (x1,y1), (x2,y2), color, 2) cv2.putText(img, classes[int(cls)], (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2) cv2.imshow('Preview', img) cv2.waitKey(0) # 示例使用 visualize_labels('images/000001.jpg', 'labels/000001.txt', classes)通过以上步骤,我们完成了从VOC到YOLO格式的高效转换,并解决了常见问题。在实际项目中,建议将转换过程封装成可复用的工具类,方便不同项目间的迁移使用。