ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

YOLOv5数据预处理实战:数据集划分与VOC/COCO格式转换脚本解析

YOLOv5数据预处理实战:数据集划分与VOC/COCO格式转换脚本解析 简介一套面向目标检测数据预处理的Python工具脚本专供YOLOv5等YOLO系列学习者、初阶研发人员和科研人员使用适合学生、具备一定编程基础的工作1至3年开发者以及入门AI的爱好者。资源集中解决模型训练前两大高频问题自动划分训练集与测试集以及将COCO、VOC标注格式统一转换为YOLO系列要求的数据格式可有效免除人工整理标注文件的繁琐流程。压缩包整体仅2KB共2个py脚本一个执行训练/测试集划分一个执行格式转换轻量紧凑、路径替换即可集成到个人工程中。该资源已有3604人浏览学习脚本经大量实践验证无bug能帮助使用者快速完成数据准备、直接跑通YOLO训练流程大幅节省时间与调试成本。1. 做 YOLOv5 数据处理这套脚本先把划分和格式转换的活干完做 YOLOv5 数据处理第一个绕不过去的步骤就是划分训练集和测试集以及把 COCO、VOC 格式的数据转成 YOLO 系列要的 txt 标签。很多教程只给一两段零散代码跑完才发现标签路径不对、类别 id 错位、随机结果不能复现比训练模型还折磨人。这份资源打包成了一个 data_conver.tar.gz解压后是 data_conver 目录里面是 split_train_val.py 和 voc_label.py并把 COCO 转 YOLO 的对应处理也做了。适合刚开始接触人工智能、计算机视觉的学生也适合工作 1-3 年、不想在数据整理上反复耗时间的研发人员。想快速出效果、节省时间这套脚本改改路径就能用。2. 把划分训练集和测试集这件事做对split_train_val.py 的核心逻辑目标检测数据准备好之后第一件事就是划分训练集和测试集。YOLOv5 训练时读取的不是一张张图而是 train.txt、val.txt 这类路径清单文件里面每一行是一个图片的绝对或相对路径。划分得不好模型验证结果就不可信尤其是训练集和验证集图片重叠时loss 会虚低换到新数据上直接翻车。2.1 划分比例和随机种子为什么我坚持把 random_state 写死先看划分逻辑本身。常见做法是 8:1:1 或 9:1先对整个图片列表做一次 shuffle再按比例切片。这份资源里的 split_train_val.py 做的就是这件事。我拆包后把关键逻辑整理成了下面这段方便你理解它内部发生了什么import random from pathlib import Path def split_train_val(img_dir: str, train_txt: str, val_txt: str, train_ratio: float 0.9, val_ratio: float 0.1, random_state: int 42): img_dir Path(img_dir) images sorted(list(img_dir.iterdir())) # 固定随机种子保证每次运行得到完全相同的划分结果 random.seed(random_state) random.shuffle(images) val_len int(len(images) * val_ratio) val_imgs images[:val_len] train_imgs images[val_len:] with open(train_txt, w, encodingutf-8, newline\n) as f: for img in train_imgs: f.write(str(img.resolve()) \n) with open(val_txt, w, encodingutf-8, newline\n) as f: for img in val_imgs: f.write(str(img.resolve()) \n)这里的参数要重点说清楚。train_ratio和val_ratio控制训练集、验证集的比例两者加起来等于 1如果想把测试集也单独留出来就把比例改成 0.8、0.1、0.1代码里加一个test_txt分支。random_state是最容易被忽略的参数同一份数据如果不固定随机种子每次跑出来的 train.txt 都不一样今天训练的模型和明天训练的模型就没有可比性。我在实际项目里习惯固定成 42然后把这个数字写到实验记录里后面想复盘就按同一个划分重新跑。newline\n是为了解决 Windows 上写出的 txt 每行末尾多一个\r的问题。Windows 默认换行符是\r\n这个文件拿到 Linux 服务器上训练时\r会变成图片路径的一部分导致FileNotFoundError。这个问题在数据转换里特别隐蔽后面避坑章还会专门讲。2.2 用 data_conver 里的脚本跑通一次划分拿到资源包之后第一步是解压然后进到目录里直接看脚本顶部有没有路径配置。很多这类脚本不会做命令行传参而是把路径写在文件开头我一般会先打开看一眼确认图片目录和输出文件名。# 解压资源包data_conver.tar.gz 解压后出现 data_conver 目录 tar -xzf data_conver.tar.gz cd data_conver # 划分脚本直接运行路径和比例在脚本顶部按需改 python split_train_val.py上面这段命令是通用运行方式。如果脚本支持命令行参数也可以先执行python split_train_val.py --help看有哪些选项不支持的话就直接改头部变量。运行结束后当前目录应该出现 train.txt 和 val.txt或者你自定义名字的两个清单文件。检查一下内容# 查看生成的 train.txt 前几行确认是真实存在的图片路径 head -n 3 train.txt正常输出应该是完整的图片路径像/home/user/data/images/img_0001.jpg这样。如果输出为空最常见原因就是脚本里用glob(*.jpg)收集图片而你的数据集里是 png 或其他扩展名把通配符加上 png 再跑一遍。另外有些版本会把图片路径写成相对路径比如images/train/img_0001.jpg这也能用但要注意训练时的工作目录必须跟脚本生成路径时的目录一致。资源包里主要文件的职责我列了一个简单对应关系文件作用我一般改哪里split_train_val.py按比例划分训练集和验证集生成图片路径清单图片目录、输出 txt 路径、train_ratiovoc_label.py读 VOC 的 XML 标注转成 YOLO 的 txt 标签classes 列表、VOC 根目录、输出目录data_conver.tar.gz资源包本体解压后得到上述脚本先看 README 或脚本头部注释划分脚本跑完后你会得到一份图片清单但这里只是完成了第一步。真正的问题还在后面图片清单里每张图对应的标签文件必须是 YOLO 能读懂的class cx cy w h格式而且要与图片一一对应。如果标签还是 COCO 或 VOC 格式那就要进入下一步转换。3. VOC 格式转 YOLO 数据voc_label.py 的转换链路VOC 数据在目标检测里用得非常多它有一套固定的目录习惯JPEGImages 放原图Annotations 放 XML 标注ImageSets/Main 放训练验证划分。YOLO 不认 XML它要求每张图对应一个同名 txt每一行是class x_center y_center width height并且全部归一化到 0 到 1 之间。voc_label.py 干的就是这个转换。3.1 坐标换算从 xmin/ymin/xmax/ymax 到中心点加宽高先看原理。VOC 的 XML 里用一个 bndbox 节点记录目标的左上角和右下角坐标也就是xmin, ymin, xmax, ymax。YOLO 需要的是物体中心点坐标以及宽高。换算关系是x_center (xmin xmax) / 2 / width y_center (ymin ymax) / 2 / height box_w (xmax - xmin) / width box_h (ymax - ymin) / height这里所有值都除以图片宽高把绝对像素换算成归一化坐标YOLO 模型在训练时统一按这个标准读标签。如果有负坐标也就是标注框越界了转换前就要处理否则模型会学到错误的边界。下面的代码是把一个 XML 文件转成一行 YOLO 标签的简化过程和包内脚本的转换思路一致import xml.etree.ElementTree as ET # classes 的顺序就是 YOLO 里的类别索引千万别在转换后再改 classes [person, bicycle, car] def voc_xml_to_yolo_line(xml_path, img_width, img_height): root ET.parse(xml_path).getroot() lines [] for obj in root.iter(object): difficult int(obj.find(difficult).text) name obj.find(name).text # 跳过 difficult 样本VOC 里这些目标本身不清晰硬转会让训练不稳定 if name not in classes or difficult 1: continue bndbox obj.find(bndbox) xmin float(bndbox.find(xmin).text) ymin float(bndbox.find(ymin).text) xmax float(bndbox.find(xmax).text) ymax float(bndbox.find(ymax).text) x_center (xmin xmax) / 2.0 / img_width y_center (ymin ymax) / 2.0 / img_height box_w (xmax - xmin) / img_width box_h (ymax - ymin) / img_height cls_id classes.index(name) # 保留 6 位小数读取时用 float 不会有精度问题 lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {box_w:.6f} {box_h:.6f}) return lines这段代码里最值得注意的就是classes列表。它的顺序直接决定类别 id如果 person 在列表里排第 0那生成的标签里 person 就是 0。后面训练 YOLOv5 时yaml 文件里的names列表必须和这个顺序完全一致否则就会出现“标签是 person模型却当成 car 训练”的错乱。3.2 跑通 voc_label.py目录准备与执行实际使用 voc_label.py 时通常要做两件事指定 VOC 数据集根目录指定输出标签目录。我先会在数据集根目录下确认结构确保 Annotations 和 JPEGImages 是平级目录。# 查看 VOC 目录结构确认 Annotations 和 JPEGImages 都存在 ls -l VOCdevkit/VOC2007/ # 转换命令脚本会把所有 xml 转成 YOLO 格式 txt 并放到 labels 目录 python voc_label.py执行后labels 目录里会出现和 XML 同名的 txt。如果脚本设计成按训练集和验证集分开导出那 labels 下还会拆成 train 和 val 两个子目录。很多新手只看 labels 生成了不看内容结果训练时发现标签里全是 0 或者坐标全是 1这就是没检查输出。# 查看某个转换后的标签文件正常每行应该是 5 个数 cat labels/000001.txt一个正常的 YOLO 标签行类似0 0.517187 0.357143 0.051562 0.071420。前面是类别 id后面四个数分别是中心点 x、中心点 y、宽度、高度范围都在 0 到 1 之间。如果看到负数或者大于 1 的数要么是图片宽高读错了要么是 XML 里的 bndbox 数值本身越界了。遇到越界框我一般会在转换时做一次 clamp把坐标限制在 0 到 1 之间但前提是原图标注没有系统性错误。3.3 转换后必须做的三处核对VOC 转 YOLO 最容易出现“看着没问题训练报错”的情况。我会做三处核对。第一统计 labels 目录下 txt 的数量应该和 Annotations 下 xml 数量一致并且和 JPEGImages 里图片数量一致。少一个都可能丢样本。第二检查是否有空 txt也就是文件有名字但里面没内容。这种情况可能是因为类名没在 classes 列表里或者所有框都被 difficult 过滤掉了。第三从训练集图片和标签里各挑一张把坐标画到原图上比对。# 统计空标签文件空文件会让 YOLO 训练时跳过或报错 find labels -name *.txt -type f -size 0 | wc -l这个统计命令很有用。如果空文件数量不为 0我一般不会直接删而是先把对应图片也找出来判断它是负样本还是标注遗漏。YOLOv5 的增强策略对空标签处理比较挑剔空标签太多时训练会不稳定所以我更倾向于显式过滤保证每个训练样本都有至少一个目标框。4. COCO JSON 转 YOLO 格式抽取 bbox 时的类别映射关系COCO 数据的转换和 VOC 不太一样。COCO 的标注是集中在一个 JSON 文件里里面有 images、annotations、categories 三个数组。转换时要先按 image_id 对 annotations 做分组再把里面每个标注的 bbox 转成归一化坐标。真正容易出错的不是坐标换算而是类别 id 的映射。4.1 COCO 的 bbox 和 YOLO 的格式差异COCO 的 bbox 是[x, y, width, height]其中 x、y 是左上角坐标。YOLO 需要的是中心点坐标加宽高所以换算公式是x_center (bbox[0] bbox[2] / 2.0) / img_width y_center (bbox[1] bbox[3] / 2.0) / img_height box_w bbox[2] / img_width box_h bbox[3] / img_height用 Python 实现时先读 JSON再把 categories 里的id和name做成字典通过category_id找到类名最后再通过类名列表找到 YOLO 类别索引。这段逻辑我一般会用一个函数封装import json from pathlib import Path def coco_json_to_yolo(json_path: str, img_dir: str, label_dir: str) - dict: with open(json_path, r, encodingutf-8) as f: data json.load(f) # 建立 COCO category_id 到类名的映射 cat_id_to_name {cat[id]: cat[name] for cat in data[categories]} # class_names 必须和 YOLOv5 的 data.yaml 中 names 保持一致 class_names [person, bicycle, car] # 建立 image_id 到图片文件名的映射 id_to_filename {img[id]: img[file_name] for img in data[images]} # 建立 image_id 到其所有标注列表的映射 annos_by_image {} for ann in data[annotations]: annos_by_image.setdefault(ann[image_id], []).append(ann) path_img Path(img_dir) path_label Path(label_dir) path_label.mkdir(parentsTrue, exist_okTrue) for img_id, annos in annos_by_image.items(): img_name id_to_filename[img_id] img_height next(img[height] for img in data[images] if img[id] img_id) img_width next(img[width] for img in data[images] if img[id] img_id) lines [] for ann in annos: cat_name cat_id_to_name[ann[category_id]] if cat_name not in class_names: continue cls_id class_names.index(cat_name) x, y, w, h ann[bbox] cx (x w / 2.0) / img_width cy (y h / 2.0) / img_height nw w / img_width nh h / img_height lines.append(f{cls_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}) txt_path path_label / (Path(img_name).stem .txt) with open(txt_path, w, encodingutf-8, newline\n) as f: f.write(\n.join(lines))这段代码里面有一个很容易踩的坑cat_id_to_name里的 id 是 COCO 原始的 category_id而 YOLO 的类别 id 是class_names列表的下标。如果在 COCO 数据集里 person 的 category_id 是 1恰好和 YOLO 里 person 的索引一致很多人就会误以为可以直接用 category_id。换成别的数据集后这个映射就会错。所以每转一个数据集我都要特地把class_names和data.yaml里的names对照一遍。4.2 把 COCO 80 类怎么读在 YOLO 里这件事说清楚网上常说的 COCO 80 类和 COCO 官方 JSON 里的 category id 不是一回事。YOLOv5 的 coco.yaml 里有一个 80 项的 names 列表索引从 0 到 79比如 person 在索引 0。但 COCO 官方 JSON 里 person 的 category_id 通常是 1而且官方 category_id 最大到 90 多中间还有不少空洞。正确做法是先用人话把类别名读出来再统一映射到 YOLO 的索引而不是拿 category_id 直接减 1 去猜。我自己的习惯是转换 COCO 数据之前先跑一段小代码把 COCO JSON 里的 category_id、name 和 YOLO 索引都打出来对照# 用 python 快速打印 COCO categories 和 YOLO names 的对应关系 python - PY import json coco json.load(open(instances_train2017.json)) yolo_names [person, bicycle, ...] cat_id_to_name {c[id]: c[name] for c in coco[categories]} for yolo_id, name in enumerate(yolo_names): # 找出每个 yolo 类别在 coco json 里对应的 id coco_id [k for k, v in cat_id_to_name.items() if v name] print(yolo_id, name, coco_id) PY这样做的好处是把类别对应关系一次性固化下来。COCO 官方数据跨版本时 category_id 基本稳定但也不能完全依赖尤其是用某些第三方裁剪数据集时类别顺序可能被重新排过。转换前多看几行打印结果比训练完发现类别错位再返工划算得多。4.3 COCO 转换后的目录组织建议COCO 转 YOLO 时不建议把输出 txt 直接写到原始 JSON 旁边而是单独建一个 labels 目录和 images 目录平级。标准 YOLOv5 数据集结构是 images/train、images/val、labels/train、labels/val 四个目录。转换完之后还要把 COCO 数据集自己的 train/val 划分同步过来保证 labels/train 里的标签对应的图片确实都在 images/train 里。我一般会在转换完做一次交叉核对遍历 labels/train 里的每一个 txt看它的 stem 对应的图片是否存在于 images/train。这个检查用 shell 就能完成思路比工具重要# 对比 labels/train 和 images/train 的文件名差异两边都存在的才是有效样本 ls labels/train | sed s/.txt$/.jpg/ | sort label_names.txt ls images/train | sort img_names.txt comm -3 label_names.txt img_names.txt如果 comm 输出为空说明两边完全匹配。如果有输出多出来的行就是丢失匹配项的样本需要回头检查是图片没拷贝还是标签没生成。这一步做完COCO 数据转换才算真正闭环。5. 避坑转换过程中最容易翻车的 5 个地方数据转换脚本本身不难难的是各种环境差异。下面这几个坑我基本都在真实项目里踩过每次都要花不少时间定位。把它们单独列出来就是希望你在跑资源里的脚本时能少走弯路。5.1 路径与随机种子问题坑一脚本换个目录就找不到图片现象在 data_conver 目录里运行 voc_label.py 正常把脚本复制到项目根目录再运行直接报FileNotFoundError找不到 JPEGImages。原因脚本里的路径是相对路径比如VOCdevkit/VOC2007/JPEGImages。相对路径是相对当前工作目录解析的不是相对脚本所在目录。换一个执行位置路径就失效了。解决打开脚本把路径改成基于os.path.dirname(__file__)拼出来的绝对路径。我一般这样写import os BASE_DIR os.path.dirname(os.path.abspath(__file__)) VOC_ROOT os.path.join(BASE_DIR, VOCdevkit, VOC2007)这样无论从哪里执行脚本都能找到自己的家目录不会再受终端所在位置影响。坑二两次运行生成的训练集不同现象同一个脚本上午跑一次、下午跑一次两次的 train.txt 内容不一样模型指标也跟着变。原因脚本里有 random.shuffle但没固定随机种子。shuffle 每次都重新洗牌划分结果自然不同。解决在脚本开头加random.seed(42)把种子写死。如果你用 numpy 排序或 pandas 抽样还需要同时固定 numpy 和 pandas 的随机种子。固定之后每次划分结果完全一致实验才能复现。5.2 标注内容与类别顺序问题坑三标签文件里类别 id 全乱了现象转换后打开 txt发现同一类物体有时候是 0有时候是 3模型训练时 loss 不收敛。原因VOC 和 COCO 转 YOLO 时classes 列表的顺序和 YOLO yaml 文件的 names 顺序没对齐。比如 VOC 里类名靠 XML 的 object name 读取如果脚本里 classes 顺序是 person、car、bicycle而 data.yaml 里是 car、person、bicycle那 person 在标签里是 0训练时却当成 car。解决转换前把 classes 列表和 data.yaml 的 names 放在一起逐项对比到完全一致。我比较极端的做法是在 yaml 文件里加一行注释写上“此顺序与 voc_label.py 的 classes 一致”防止后续维护时改了一边忘了另一边。坑四生成了大量空白标签文件现象转换完labels 目录下很多 txt 是 0 字节训练日志里出现 “found no labels” 或assert报错。原因XML 里所有目标不是被 difficult 过滤掉就是类名不在 classes 列表里。结果循环正常执行了但 lines 列表为空最终还是写出了一个空文件。解决在转换函数里统计空文件并打印被跳过目标的类名。如果是 difficult 过滤导致的可以在 VOC 转 YOLO 时不跳过只要标注框坐标正常difficult 样本也可以保留。如果类名不在列表里就检查类别拼写VOC 里常有bus和car混写的情况。坑五Windows 生成的 txt 在 Linux 上图片路径多出 \r现象数据集在 Windows 上处理好传到 Linux 服务器训练报错说图片文件不存在但检查路径肉眼看着完全正确。原因Windows 默认换行符是\r\nLinux 上读取 txt 时每行末尾多了一个\r字符。这个字符看不见它被拼在图片路径后面所以系统认为文件名是img_0001.jpg\r。解决所有写 txt 的地方都用open(..., w, newline\n)。已经生成的文件用一条命令洗一遍# 把 Windows 换行符统一成 Linux 换行符同时去掉多余的\r sed -i s/\r$// train.txt val.txt从那以后我只要在 Windows 上生成过数据清单传到服务器之前都会强制跑一遍这个sed命令避免在 linux 上浪费排查时间。6. 进阶转换完先验证再全量训练用一次 1 epoch 测试兜底转换做完不要急着把整批数据丢进训练。我每次都会先写一个极小的验证脚本把标签数量和图片数量对上顺便检查 txt 每行是不是 5 个字段。这个步骤看起来很基础但能拦截大量低级错误。from pathlib import Path label_dir Path(datasets/labels/train) img_dir Path(datasets/images/train) img_stems {p.stem for p in img_dir.iterdir() if p.suffix in {.jpg, .png}} for txt_path in label_dir.glob(*.txt): if txt_path.stem not in img_stems: print(fmissing image: {txt_path.stem}) with open(txt_path, r, encodingutf-8) as f: for line in f: fields line.strip().split() # 每行必须是 class 4 个坐标值字段数不正确说明转换有问题 if len(fields) ! 5: print(fbad line in {txt_path.name}: {line.strip()})这个脚本同时检查两件事标签和图片是否一一对应以及标签行字段数是不是 5。字段数量不对通常就是坐标里有空格、类名被拆开或者脚本把多行 ‘\n’ 当成了空行。验证通过后我还会用 YOLOv5 的 train.py 只训 1 个 epoch用很小的 batch size 跑通整个数据流水线。这一步不是真的为了训练模型而是让模型自己去读一遍所有图片和标签把数据加载层面的问题提前暴露出来。# 用 1 个 epoch 验证数据流水线数据加载有问题会在这里立刻报错 python train.py --data data.yaml --epochs 1 --batch 8 --weights yolov5s.pt如果 1 个 epoch 能正常跑完训练日志里 no labels、corrupt image 这类红色警告也没有批量出现我才会放心地把 epochs 改成 100 或 300 开始正式训练。如果数据里还有漏网之鱼1 个 epoch 的报错会直接指向具体图片排查成本比全量训练跑一半崩溃低得多。这个习惯帮我在很多项目里避免了“训练了两天最后发现数据有问题”的尴尬。从那以后我每次转换完数据不管多赶时间都会强制走一遍字段检查加 1 epoch 测试这两步。数据准备是枯燥但它是一切的底子希望帮到你。本文还有配套的精品资源点击获取
返回列表