ARTICLE DETAIL

资讯详情

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

COCO/VOC数据集一键转YOLO格式:脚本实现标注转换与训练集划分

COCO/VOC数据集一键转YOLO格式:脚本实现标注转换与训练集划分 简介这是一份针对计算机视觉与目标检测任务的数据处理工具包专为需要快速完成数据集准备的研究者与开发者设计适合具备一定编程基础的学生、工作1-3年的研发人员以及刚接触人工智能的爱好者。资源聚焦YOLO系列模型训练前的关键环节主要解决两个痛点一是自动划分训练集与测试集二是将COCO、VOC两种常见标注格式批量转换为YOLO系列所需的标签格式避免手工整理带来的繁琐与出错。压缩包体积仅2KB共包含2个Python脚本代码精简、职责清晰一个负责数据集拆分一个负责标注格式转换可直接在本地Python环境中运行调试。目前已有3604人学习使用脚本经过大量实践检验无Bug运行后即可快速得到可训练的数据集显著节省手动处理标注文件的时间。对于正在上手YOLOv5等模型训练、需要统一数据格式的科研与工程场景这份工具能提供直接有效的帮助。1. 一个小脚本为什么值得折腾半天做目标检测的人大概率经历过这个场景手里有一份COCO或者VOC格式的数据集想扔给YOLOv8去训练结果发现YOLO根本不认JSON和XML它只要一个和图片同名的txt文件里面写着归一化后的框坐标。网上各种转换工具不是连不上就是类别ID写死、跑完才发现标签对不上。更头疼的是数据还没分成训练集和测试集模型一训练就跑偏。这个标题里的python脚本解决的就是这整条链路把COCO的JSON、VOC的XML解析成YOLO系列的txt标注顺手按比例切出训练集和测试集。适合正在做yolov5、yolov8训练自己的数据集还没搞定数据预处理的人。读懂这篇文章你可以拿着代码直接改自己的数据集不用再去猜类别映射和坐标公式。2. 三种格式的差异与目录规划转换之前先想清楚这三件事2.1 COCO、VOC、YOLO的标注结构对照转换脚本的核心不是写代码而是理解三种格式各自的存储哲学。COCO把所有标注塞进一个超大JSON文件里images数组记录图片信息annotations数组记录每个框categories数组记录类别名。框的坐标是[x, y, width, height]单位是像素从0开始计数。VOC则是一张图片配一个XML文件框坐标是[xmin, ymin, xmax, ymax]也是像素单位从1开始。YOLO最简单粗暴一个txt文件里每行是“类别ID 中心点x 中心点y 宽度 高度”全部归一化到0到1之间。这三种格式转换时最容易被忽略的坑是坐标系。COCO的bbox是左上角加宽高VOC是左上角加右下角YOLO要的是中心点加宽高。如果拿着COCO的坐标直接套YOLO公式不先把[x, y, w, h]换成中心点坐标出来的框会整体偏移半个身位。下面这个表格把关键字段拆开转换时可以对照着查。格式存储载体框坐标表示坐标系类别ID起点COCO单个JSON[x, y, w, h]0起始绝对像素categories里定义的id不连续VOC每图一个XML[xmin, ymin, xmax, ymax]1起始绝对像素文件夹顺序或自定义从代号写YOLO每图一个txt[cx, cy, w, h]0~1归一化0起始连续整数2.2 目录结构设计一种能直接喂给训练的布局转换脚本的输出目录如果随便乱放后面训练时写data.yaml会非常痛苦。业内常见做法是构造一个images和labels平级、内部再按train和test拆分的结构。我一般这样组织dataset/ ├── images/ │ ├── train/ │ │ ├── 000001.jpg │ │ └── ... │ └── test/ │ ├── 000101.jpg │ └── ... ├── labels/ │ ├── train/ │ │ ├── 000001.txt │ │ └── ... │ └── test/ │ ├── 000101.txt │ └── ... └── data.yaml这样YOLO训练时只需要在data.yaml里写两行路径不需要额外写脚本去配对图片和标签。划分train和test时两条list的顺序要保持一致切出来的图片和标签必须一一对应否则训练集里出现一张没有标签的图模型会拿它当背景学习反向传播一跑起来全是噪声。2.3 类别映射表COCO的80类与VOC的20类要重新编码COCO原始JSON里person的category_id通常是1car是2但这些ID存在跳号。VOC的类别ID没有显式定义靠的是XML里object的name标签。转到YOLO后YOLO只认0、1、2这样的序号所以必须建一张新映射表。常见做法是把categories按名字排序或者按原顺序重新编号然后生成一个字典写进代码里。保留映射表还有一个实际好处训练结束后推理模型输出的class_id要还原成中文名或者原数据集的类别名没有映射表就只能对着数字猜。我一般会把映射表同时导出成一个classes.txt一行一个类别名这样data.yaml直接引用它将来做混淆矩阵可视化也方便。写转换脚本时这一步千万别省。3. COCO转YOLO解析JSON、重建类别ID、生成txt标注3.1 从COCO JSON里提取图片信息和标注框COCO转YOLO的第一步是读入JSON并建立三张索引表image_id到文件名的映射、image_id到宽高的映射、annotation里image_id到bbox列表的映射。注意COCO的annotations里可能存在某些图片没有标注的情况这类图片在目标检测里其实应该默认过滤掉因为训练时它们会被当成纯负样本破坏背景与前景的比例。import json import os def load_coco_annotations(json_path): with open(json_path, r, encodingutf-8) as f: coco json.load(f) # 建立图片索引id - (文件名, 宽, 高) image_info {} for img in coco[images]: image_info[img[id]] (img[file_name], img[width], img[height]) # 建立标注索引image_id - 该图所有框的列表 annos {} for ann in coco[annotations]: img_id ann[image_id] x, y, w, h ann[bbox] cat_id ann[category_id] # 把空框和畸形框直接跳过 if w 0 or h 0: continue annos.setdefault(img_id, []).append((cat_id, x, y, w, h)) return coco[categories], image_info, annos逻辑说明这里没有直接遍历annotations并逐条转YOLO而是先按image_id聚合。原因是一个JSON里同一张图片可能有几十个标注框聚合后可以一次性生成完整的一行行txt也方便判断“这张图到底有没有有效标注”。另一种常见做法是遍历annotations再去找image_id对应的图片信息那样每处理一个框就要在字典里查一次代码倒是能跑数据量一大就慢得明显。参数说明coco[categories]是原始类别列表里面虽然也带id但那个id是COCO原始编号不能直接当成YOLO的class_id后面需要重新映射。3.2 重建类别表并映射坐标到归一化值拿到原始categories后新类别ID从0开始递增。这里有一个很容易翻车的细节COCO的categories数组顺序和annotations里category_id的顺序可能不一致直接对annotations里的category_id做“减去某个偏移量”的操作会出错必须通过字典把原始category_id折算成新ID。def convert_coco_to_yolo(coco_json, output_dir, labels_outputNone): categories, image_info, annos load_coco_annotations(coco_json) # 新类别ID从0开始原category_id - 新ID cat_id_map {} for new_id, cat in enumerate(categories): cat_id_map[cat[id]] new_id if labels_output: labels_output.write(cat[name] \n) for img_id, (filename, width, height) in image_info.items(): txt_name os.path.splitext(filename)[0] .txt lines [] for cat_id, x, y, w, h in annos.get(img_id, []): new_id cat_id_map[cat_id] # COCO的bbox是[x,y,w,h]转成YOLO中心点宽高再归一化 cx (x w / 2) / width cy (y h / 2) / height nw w / width nh h / height # 防止浮点越界裁到安全范围 cx min(max(cx, 0.0), 1.0) cy min(max(cy, 0.0), 1.0) nw min(max(nw, 0.0), 1.0) nh min(max(nh, 0.0), 1.0) lines.append(f{new_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}) with open(os.path.join(output_dir, txt_name), w) as f: f.write(\n.join(lines))逻辑说明坐标换算里最容易错的是混淆了“中心点”和“左上角”。x加上w的一半才是中心点的横坐标再除以图片宽度得到归一化值。YOLO训练时内部会乘以图像尺寸还原真实坐标如果这一步写错标注的框会整体偏移loss曲线看起来正常但mAP惨不忍睹。这个脚本里还对结果做了0到1的裁剪处理那种标注框贴边导致浮点计算越界的情况。参数说明output_dir要传labels/train或者labels/test这种子目录不要直接把所有txt扔到一个目录下否则后面划分训练测试集还要再挪一遍没必要。3.3 常见变体只转部分类别、跳过crowd标注COCO数据里有一类标注iscrowd1代表一个区域里有密集人群或物体一个框可能罩着好几个个体。YOLO系列算法对这种标注的处理并不理想转进去训练会学到错误的“一个框装下所有人”的语义。做行人检测时我通常会跳过iscrowd的标注物体检测任务里crowd标注不多保留也无妨但如果你想做精确的检测建议在循环里加一个判断。还有一个需求是只转部分类别比如只要person和car。在load_coco_annotations里加一个类别过滤参数遍历annotations时直接continue掉不需要的category_id即可。这样转完的txt文件里不会出现陌生类别的框省得训练时还要再改一遍标注。注意过滤之后也要同步更新cat_id_map否则新ID仍然按全部类别编号会让类别数量虚高。4. VOC转YOLOXML解析、difficult样本与坐标越界处理4.1 用xml.etree解析VOC的 结构VOC的XML解析在Python里不需要装额外依赖xml.etree.ElementTree是标准库读取速度对常见的万级数据也够用。解析时需要注意VOC的filename和size在每个XML里都有不需要自己去图片文件里再查一遍宽高直接用XML里的size字段就行。这也是VOC格式相对朴素的地方不需要像COCO那样查JSON做关联。import os import xml.etree.ElementTree as ET def parse_voc_xml(xml_path, class_names): tree ET.parse(xml_path) root tree.getroot() filename root.findtext(filename) size root.find(size) width int(size.findtext(width)) height int(size.findtext(height)) # 类别名到新ID的映射 class_id_map {name: i for i, name in enumerate(class_names)} objects [] for obj in root.findall(object): name obj.findtext(name) if name not in class_id_map: # 不在目标类别里的物体直接跳过 continue difficult int(obj.findtext(difficult, 0)) bndbox obj.find(bndbox) xmin float(bndbox.findtext(xmin)) ymin float(bndbox.findtext(ymin)) xmax float(bndbox.findtext(xmax)) ymax float(bndbox.findtext(ymax)) objects.append({ class_id: class_id_map[name], xmin: xmin, ymin: ymin, xmax: xmax, ymax: ymax, difficult: difficult }) return filename, width, height, objects逻辑说明这里特意把difficult也读出来并存进字典是因为VOC里的difficult标记表示这张图很难识别官方建议不参与模型评估。有的场景下应该把它过滤掉有的场景又需要保留做难例挖掘提前读出来比写死跳过或者写死保留更灵活。返回的objects是一个列表一个XML里可能有多个object每个object都会转成txt的一行。如果你只调用了findtext(filename)而没调用find(./size)那遇到网络下载的图片文件名不含扩展名时生成的txt名会和图片名对不上训练会报找不到label的warning。4.2 坐标换算与difficult样本的取舍策略XML里的box是整数型的像素坐标转成YOLO归一化值之前要先把xmax减xmin得到宽度ymax减ymin得到高度。注意VOC的坐标是从1开始计数的宽高的计算要加1还是不加业界有两种做法常见实现是直接用xmax-xmin不做偏移修正。当标注框很大、几乎占满整张图时归一化结果可能略大于1需要做裁剪。def voc_to_yolo_line(obj, width, height, skip_difficultTrue): if skip_difficult and obj[difficult]: return None xmin, ymin obj[xmin], obj[ymin] xmax, ymax obj[xmax], obj[ymax] # 框的宽高不能为0或负数 box_w max(xmax - xmin, 0.0) box_h max(ymax - ymin, 0.0) if box_w 0 or box_h 0: return None # 中心点与归一化 cx (xmin box_w / 2) / width cy (ymin box_h / 2) / height nw box_w / width nh box_h / height # 越界坐标最终兜底 cx max(0.0, min(1.0, cx)) cy max(0.0, min(1.0, cy)) nw max(0.0, min(1.0, nw)) nh max(0.0, min(1.0, nh)) return f{obj[class_id]} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}参数说明skip_difficultTrue时difficult1的样本会被直接丢弃。做基础检测任务时推荐丢弃因为那些难例通常遮挡严重或者目标太小放进训练集会干扰模型对正常样本的特征学习。如果你收集的是电力红外数据集或者遥感数据集对难例有专门的评估需求再改成False。另外这个函数里用max(xmax-xmin, 0.0)来兜底负值实际项目中出现过标注软件导出坐标不合法的情况直接丢掉总比让训练崩溃好。4.3 批量遍历目录一张图对应一个XML的处理策略VOC数据集往往是一个文件夹里全是JPG另一个文件夹里全是XML。转换时要遍历所有XML文件找到同名的图片把生成的txt写到labels目录。常见做法是把XML目录和图片目录传进脚本用os.listdir去匹配同名文件。需要注意如果图片文件夹里有JPEG、jpg、png多种扩展名在拼接图片路径时要做存在性判断。import glob def convert_voc_dir(xml_dir, image_dir, output_label_dir, class_names): os.makedirs(output_label_dir, exist_okTrue) for xml_path in glob.glob(os.path.join(xml_dir, *.xml)): filename, width, height, objects parse_voc_xml(xml_path, class_names) # 找图片兼容jpg/png/jpeg三种扩展名 img_candidates [ os.path.join(image_dir, filename), os.path.join(image_dir, os.path.splitext(filename)[0] .jpg), os.path.join(image_dir, os.path.splitext(filename)[0] .png), ] if not any(os.path.isfile(p) for p in img_candidates): continue txt_name os.path.splitext(filename)[0] .txt lines [] for obj in objects: line voc_to_yolo_line(obj, width, height) if line: lines.append(line) with open(os.path.join(output_label_dir, txt_name), w) as f: f.write(\n.join(lines))逻辑说明遍历时用glob去拿所有XML文件名通过splitext统一改成txt。这里没有做XML到图片的严格1:1校验如果某个XML对应图片不存在会静默跳过后续校验脚本再统一排查。批量转换另一个常见问题是XML里的filename字段和实际文件名大小写不一致Windows下没问题Linux下就会匹配失败稳妥做法是全部转小写后再比较。参数说明class_names列表的顺序决定了新的类别IDVOC的20类建议按官方顺序写以后对比backbone的预训练结果更方便。5. 划分训练集和测试集随机种子、图片配对、均衡分布的坑5.1 按比例划分前先给所有图片洗牌把转换完的txt和图片一起按比例分成训练集和测试集看起来只是一个random.shuffle加一个列表切片但实际做起来有四个容易翻车的地方。先看基础版脚本import random import shutil import os def split_train_test(image_list, ratio0.8, seed42): random.seed(seed) indices list(range(len(image_list))) random.shuffle(indices) split_point int(len(indices) * ratio) train_idx indices[:split_point] test_idx indices[split_point:] return train_idx, test_idx逻辑说明先用固定seed初始化随机数让每次跑脚本得到完全相同的训练测试划分这样实验之间可以公平对比。如果不设定seed每次运行划分结果都不同模型参数不变但数据分布变了实验结论就失去了可重复性。返回的是索引列表而不是直接把图片路径返回这样方便调用处同时操作images和labels两个目录。参数说明ratio0.8代表80%做训练、20%做测试。如果你想再切一个验证集出来可以把ratio改成0.7然后把test_idx再按5:5拆成测试和验证或者干脆直接从测试集里割一半。5.2 图片和标签必须成对移动不能漏掉一个常见的翻车现场是只移动了图片没移动txt或者反向操作。YOLO训练时发现大量图片没有对应标签会默认按背景处理训练出来模型完全不输出框。更隐蔽的是多移动了一部分txt到test目录导致训练集合里缺了一批框类别分布失衡。正确思路是循环时同时构造图片源路径、图片目标路径、标签源路径、标签目标路径四个路径用同一个文件名拼接。def move_paired_files(img_name, src_img, src_lbl, dst_img, dst_lbl): shutil.copy2(src_img, dst_img) txt_src os.path.splitext(src_lbl)[0] .txt txt_dst os.path.splitext(dst_lbl)[0] .txt if os.path.isfile(txt_src): shutil.copy2(txt_src, txt_dst) else: # 源标签不存在时记录日志不要静默 print(fmissing label: {txt_src})逻辑说明用copy2而不是move的原因是防止脚本中途报错导致原数据被破坏。原始数据集只有一份move出错很难恢复copy至少还能留个后悔药。判断txt文件存在之后再复制避免shutil报FileNotFoundError。有些数据集里存在没有标注的图片这类图在YOLO训练里可以当作负样本如果你希望保留它们就不要在拷贝时报错改为写进一个missing_log.txt。参数说明src_img和src_lbl要传绝对路径因为后面划分脚本会在多个目录间切换工作路径相对路径容易突然失效。5.3 类别不均衡时按类别加权做划分目标检测数据集经常存在类别严重不均衡的问题比如一个数据集里person有一万张helmet只有两百张。如果只按图片比例划分可能测试集里完全没有helmet的样本模型在helmet类别上的mAP就等于0。稍微高级一点的划分方法是先统计每种类别出现在哪些图片里做分层抽样确保训练集和测试集的类别比例接近。from collections import defaultdict def class_stratified_indices(label_dir, ratio0.8, seed42): class_to_imgs defaultdict(set) for txt_file in os.listdir(label_dir): img_name os.path.splitext(txt_file)[0] with open(os.path.join(label_dir, txt_file)) as f: for line in f: if line.strip(): cls_id int(line.split()[0]) class_to_imgs[cls_id].add(img_name) train_set set() test_set set() rng random.Random(seed) for cls_id, img_set in class_to_imgs.items(): img_list list(img_set) rng.shuffle(img_list) split int(len(img_list) * ratio) train_set.update(img_list[:split]) test_set.update(img_list[split:]) return train_set, test_set逻辑说明先读一遍所有txt标注建立类别到图片名的倒排索引再对每个类别单独做一次shuffle和切片。这样即使helmet只出现在几十张图里也有80%进训练集、20%进测试集。缺点是多读了一次所有txt文件耗时增加但对万级数据量来说完全可接受。参数说明rng需要单独创建Random实例而不是使用全局random因为划分时两个策略可能同时存在局部实例不会污染全局随机状态。这个函数返回的是set方便后面判断某张图属于哪个数据集直接查集合比遍历列表快一个数量级。5.4 写一个划分后的完整性检查做错能当场发现数据和标签都移动完之后必须跑一遍体检脚本这是血泪经验里最值钱的一步。之前遇到过移动脚本内部多个线程同时写文件、结果标签和图片配错对的情况训练集里模型学得一团糟功劳全部送给划分脚本的黑匣子问题。检查脚本只做三件事def verify_split(image_dir, label_dir): imgs {os.path.splitext(f)[0] for f in os.listdir(image_dir)} labels {os.path.splitext(f)[0] for f in os.listdir(label_dir)} only_img imgs - labels only_label labels - imgs if only_img: print(图片存在但缺少txt:, sorted(only_img)[:10]) if only_label: print(txt存在但缺少图片:, sorted(only_label)[:10]) if not only_img and not only_label: print(校验通过图片和txt一一对应)逻辑说明用集合差集找出左右两边不对应的文件名一一对应的判断通过后再随机抽10个txt打印前几行人工扫一眼框坐标是否在0~1范围内。这个脚本跑完再丢进yolov8训练自己的数据集能让整个流程里的问题提早暴露不用等到训练出来才发现标注方向错了。如果你把训练集测试集分在不同目录记得对images/train、labels/train、images/test、labels/test四组全部跑一遍。6. 进阶一条命令完成转换加划分顺便统计类别分布6.1 串联整个pipeline的入口函数前面几章的脚本都拆开了实际用的时候可以串成一个入口接收几个路径参数内部依次执行解析、转换、划分、校验最后打印一个类别分布统计。这样每次拿到新数据集只需要改一行参数就能跑完不打断思路。def run_pipeline(dataset_format, src_dir, output_root, class_names, ratio0.8, seed42): if dataset_format coco: # coco的src_dir指向json文件标注和图片在同一根目录 convert_coco_to_yolo(src_dir, output_root /labels_all) elif dataset_format voc: convert_voc_dir(src_dir /Annotations, src_dir /JPEGImages, output_root /labels_all, class_names) # 收集全部图片并切分 all_images [f for f in os.listdir(src_dir) if f.lower().endswith((.jpg, .png))] train_idx, test_idx split_train_test(all_images, ratio, seed) for img_name in train_idx: move_paired_files(img_name, ...) # 移动到train for img_name in test_idx: move_paired_files(img_name, ...) # 移动到test verify_split(output_root /images/train, output_root /labels/train) print_class_distribution(output_root /labels/train, class_names)逻辑说明这个入口把前面的函数按顺序粘起来典型的使用方式是在终端执行python trans_yolo.py --format voc --src VOCdevkit --out dataset。class_names既传给了COCO转换也传给了VOC转换它们内部的映射表保持一致不会出现训练集里用一套ID、测试集里用另一套ID的翻车事故。ratio和seed也在一处定义想调整划分比例或更换随机种子时不用翻到函数内部去改。参数说明如果数据集已经划分好了train和test两个子集直接分别对两个子目录调用run_pipeline或convert函数不需要再走split的步骤。6.2 统计各类别框数量提前发现标签异常划分完并不是终点看一眼每个类别的框数量分布能帮你发现很多潜在问题。比如某个类别框数极少说明原数据集标注不均衡模型在这个类别上大概率欠拟合。某个类别框数为0则可能是类别映射出了问题class_names里写的名字和数据集里实际标注的name对不上。def print_class_distribution(label_dir, class_names): counts [0] * len(class_names) for txt in os.listdir(label_dir): with open(os.path.join(label_dir, txt)) as f: for line in f: if line.strip(): cls_id int(line.split()[0]) if 0 cls_id len(class_names): counts[cls_id] 1 for name, cnt in zip(class_names, counts): print(f{name}: {cnt}) print(f总计框数: {sum(counts)})逻辑说明这里对txt里每行取第一个字段当类别ID统计进对应下标。顺带做了一次范围检查如果class_id超出class_names的长度说明转换时映射表有问题打印出来一眼就能看到。我之前有一次VOC转YOLOclasses.txt里少了两个类名统计结果直接暴露了这个错误否则等训练跑到一半才报shape mismatch来回排查浪费大半天。这个统计函数不只用于训练前也可以用于测试集看看测试集和训练集类别分布是否接近。结尾用一条个人习惯收住我现在每次拿到新数据集都会先跑一遍完整转换和校验再回头检查生成的每个txt文件里坐标是不是都在0到1之间抽查五六个就放心丢进训练脚本。这个习惯帮我挡掉过不少标注数据和代码之间莫名其妙的玄学问题希望帮到你。本文还有配套的精品资源点击获取
返回列表