从NEU-DET到YOLOv7：实战数据集格式转换与划分全流程解析

1. 认识NEU-DET数据集与YOLOv7格式要求

第一次接触NEU-DET数据集的朋友可能会有点懵，这个数据集是东北大学发布的钢材表面缺陷检测数据集，包含六种常见缺陷：裂纹（crazing）、夹杂物（inclusion）、斑块（patches）、点蚀（pitted_surface）、轧入氧化皮（rolled-in_scale）和划痕（scratches）。原始数据采用Pascal VOC格式，也就是每张图片对应一个XML标注文件。

而YOLOv7需要的标注格式则完全不同，它要求的是简单的TXT文本文件，每行表示一个目标物体，格式为：类别编号 x_center y_center width height。这里的坐标都是相对于图片宽高的归一化值（0到1之间）。这种格式转换看似简单，但实际操作中会遇到各种坑，比如路径设置错误、类别映射不对应、归一化计算错误等。

我去年在帮一个钢厂做缺陷检测项目时，就遇到过因为坐标归一化计算错误导致模型完全学不会的情况。后来发现是XML解析时把宽高读反了，导致所有标注框位置都错了。所以特别提醒大家，在转换格式时一定要仔细检查几个关键点：XML解析是否正确、坐标转换公式是否写对、文件路径是否配置正确。

2. XML到TXT格式转换实战

2.1 基础转换代码解析

先来看核心的转换代码，我把它拆解成几个关键部分：

import xml.etree.ElementTree as ET import os import glob classes = ["crazing","inclusion","patches","pitted_surface","rolled-in_scale","scratches"] def convert(size, box): # 将Pascal VOC的(xmin,ymin,xmax,ymax)转换为YOLO的(x_center,y_center,width,height) dw = 1.0 / size[0] dh = 1.0 / 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)

这个convert函数是格式转换的核心，它完成了两个关键操作：一是将矩形框表示方法从两点式转为中点+宽高式，二是进行了归一化处理。这里特别容易出错的是归一化顺序，一定要先计算原始坐标，最后再做归一化。

2.2 路径配置与文件处理

实际项目中，文件路径配置不当是最常见的错误之一。建议采用os.path.join来构建路径，这样能避免不同操作系统下的路径分隔符问题：

def convert_annotation(image_name): # 输入输出路径配置 xml_dir = './mydata/neu/label/' txt_dir = './mydata/neu/labels/' in_file = open(os.path.join(xml_dir, image_name[:-3] + 'xml')) out_file = open(os.path.join(txt_dir, image_name[:-3] + 'txt'), 'w') # 解析XML文件 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'): cls = obj.find('name').text if cls not in classes: print(f"警告：发现未定义类别 {cls}，在文件 {image_name} 中") 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) out_file.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n")

2.3 常见问题排查

在实际运行中，经常会遇到生成的TXT文件是空的情况。根据我的经验，主要有三个原因：

1. 图片确实没有标注对象。这种情况是正常的，空文件是正确的。
2. 类别名称不匹配。比如XML里写的是"scratch"而classes列表里是"scratches"。
3. 文件路径错误，导致程序读取的是空XML文件。

建议在转换完成后，运行一个简单的检查脚本：

import os label_dir = './mydata/neu/labels/' empty_files = [] for file in os.listdir(label_dir): if os.path.getsize(os.path.join(label_dir, file)) == 0: empty_files.append(file) if empty_files: print(f"发现 {len(empty_files)} 个空标签文件，建议检查：") print("\n".join(empty_files[:5])) # 只打印前5个避免输出太多 else: print("未发现空标签文件，转换成功！")

3. 数据集划分的科学方法

3.1 划分比例与策略

数据集划分不是简单的随机切分那么简单。在工业检测场景中，我们需要特别注意：

1. 保持类别分布均衡。每个缺陷类别在训练集、验证集和测试集中的比例应该接近。
2. 考虑样本相关性。同一卷钢材的连续帧可能高度相似，这些样本应该被划分到同一个集合中。
3. 测试集应该尽可能代表真实场景的分布。

对于NEU-DET这样相对均衡的数据集，常用的比例是训练集70-80%，验证集10-15%，测试集10-15%。我个人的经验是8:1:1的比例在大多数情况下效果不错。

3.2 代码实现与路径管理

下面是一个更健壮的划分实现，增加了类别平衡检查：

import shutil import random import os from collections import defaultdict def split_dataset(): # 路径配置 image_original_path = "./mydata/neu/images/" label_original_path = "./mydata/neu/labels/" # 创建输出目录 os.makedirs("./datasets/defect/images/train/", exist_ok=True) os.makedirs("./datasets/defect/labels/train/", exist_ok=True) # 同理创建val和test目录... # 统计每个类别的样本数 class_counts = defaultdict(int) label_files = [f for f in os.listdir(label_original_path) if f.endswith('.txt')] for file in label_files: with open(os.path.join(label_original_path, file)) as f: for line in f: class_id = int(line.split()[0]) class_counts[class_id] += 1 # 打印类别分布 print("原始数据集类别分布：") for cls_id, count in class_counts.items(): print(f"{classes[cls_id]}: {count}个样本") # 随机打乱并划分 random.shuffle(label_files) total = len(label_files) train_end = int(total * 0.8) val_end = train_end + int(total * 0.1) # 复制文件并统计新分布 train_counts = defaultdict(int) for i, file in enumerate(label_files): name = file[:-4] src_img = os.path.join(image_original_path, name + '.jpg') src_lbl = os.path.join(label_original_path, file) if i < train_end: dst_dir = "train" elif i < val_end: dst_dir = "val" else: dst_dir = "test" # 复制文件 shutil.copy(src_img, f"./datasets/defect/images/{dst_dir}/") shutil.copy(src_lbl, f"./datasets/defect/labels/{dst_dir}/") # 统计新分布 with open(src_lbl) as f: for line in f: cls_id = int(line.split()[0]) train_counts[(dst_dir, cls_id)] += 1 # 打印划分后的分布 print("\n划分后类别分布：") for (split, cls_id), count in train_counts.items(): print(f"{split}集 {classes[cls_id]}: {count}个样本")

3.3 高级划分技巧

对于更复杂的场景，你可能需要考虑：

1. 分层抽样：确保每个类别在各个集合中的比例一致
2. 时间序列划分：如果数据是按时间顺序采集的，应该按时间划分
3. 交叉验证：在小数据集上可以使用k折交叉验证

我曾经遇到过一个案例，随机划分后某个稀有缺陷在测试集中完全没有样本，导致测试结果虚高。后来改用分层抽样才解决了这个问题。

4. 工程实践中的优化建议

4.1 自动化检查脚本

在大型项目中，建议编写自动化检查脚本，定期验证数据质量。以下是我常用的检查项：

1. 图片与标签文件是否一一对应
2. 标注框是否超出图片边界
3. 是否有无效或异常标注
4. 类别分布是否均衡

def validate_dataset(image_dir, label_dir): # 检查文件对应关系 image_files = set(f[:-4] for f in os.listdir(image_dir) if f.endswith('.jpg')) label_files = set(f[:-4] for f in os.listdir(label_dir) if f.endswith('.txt')) missing_images = label_files - image_files missing_labels = image_files - label_files if missing_images: print(f"警告：{len(missing_images)}个标签文件没有对应的图片") if missing_labels: print(f"警告：{len(missing_labels)}张图片没有对应的标签文件") # 检查标注有效性 for label_file in label_files: with open(os.path.join(label_dir, label_file + '.txt')) as f: for line in f: parts = line.strip().split() if len(parts) != 5: print(f"无效行格式：{label_file}.txt -> {line}") continue cls_id, x, y, w, h = map(float, parts) if not (0 <= x <= 1 and 0 <= y <= 1 and 0 <= w <= 1 and 0 <= h <= 1): print(f"异常坐标值：{label_file}.txt -> {line}")

4.2 数据增强考虑

在准备YOLOv7训练数据时，还要考虑后续可能的数据增强策略。有些增强方式会影响标注信息：

1. 镜像翻转：需要同步调整标注框坐标
2. 随机裁剪：可能裁剪掉部分目标
3. 色彩变换：不影响标注框但可能影响模型学习

建议在数据准备阶段就规划好增强策略，有些团队会在数据转换阶段就生成增强后的样本，这样训练时更高效。

4.3 版本控制与文档

最后提醒一个容易被忽视的点：数据版本控制。每次数据变更都应该记录：

1. 数据来源和版本
2. 转换和划分的参数
3. 发现的特殊情况和处理方式

可以用一个简单的README文件记录这些信息：

# NEU-DET数据集准备日志 - 数据来源：NEU-DET官方下载，版本1.0 - 获取日期：2023-08-15 - 转换参数： - 类别列表：['crazing','inclusion','patches','pitted_surface','rolled-in_scale','scratches'] - 划分比例：train 80%, val 10%, test 10% - 特殊处理： - 发现5张图片没有对应标注，已排除 - 发现2个标注文件为空，检查后确认对应图片确实无缺陷 - 校验结果： - 训练集：1200个样本 - 验证集：150个样本 - 测试集：150个样本

这种文档在团队协作和后期模型调优时非常有用，能快速定位数据相关的问题。