从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文件是空的情况。根据我的经验主要有三个原因图片确实没有标注对象。这种情况是正常的空文件是正确的。类别名称不匹配。比如XML里写的是scratch而classes列表里是scratches。文件路径错误导致程序读取的是空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 划分比例与策略数据集划分不是简单的随机切分那么简单。在工业检测场景中我们需要特别注意保持类别分布均衡。每个缺陷类别在训练集、验证集和测试集中的比例应该接近。考虑样本相关性。同一卷钢材的连续帧可能高度相似这些样本应该被划分到同一个集合中。测试集应该尽可能代表真实场景的分布。对于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_okTrue) os.makedirs(./datasets/defect/labels/train/, exist_okTrue) # 同理创建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 高级划分技巧对于更复杂的场景你可能需要考虑分层抽样确保每个类别在各个集合中的比例一致时间序列划分如果数据是按时间顺序采集的应该按时间划分交叉验证在小数据集上可以使用k折交叉验证我曾经遇到过一个案例随机划分后某个稀有缺陷在测试集中完全没有样本导致测试结果虚高。后来改用分层抽样才解决了这个问题。4. 工程实践中的优化建议4.1 自动化检查脚本在大型项目中建议编写自动化检查脚本定期验证数据质量。以下是我常用的检查项图片与标签文件是否一一对应标注框是否超出图片边界是否有无效或异常标注类别分布是否均衡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训练数据时还要考虑后续可能的数据增强策略。有些增强方式会影响标注信息镜像翻转需要同步调整标注框坐标随机裁剪可能裁剪掉部分目标色彩变换不影响标注框但可能影响模型学习建议在数据准备阶段就规划好增强策略有些团队会在数据转换阶段就生成增强后的样本这样训练时更高效。4.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个样本这种文档在团队协作和后期模型调优时非常有用能快速定位数据相关的问题。