)
Python实战COCO人体关键点数据集转YOLO格式全流程解析在计算机视觉领域姿态估计一直是热门研究方向。许多开发者在尝试使用YOLOv5或YOLOv8进行人体姿态估计时遇到的第一个难题就是数据格式转换——COCO数据集提供的JSON标注格式与YOLO训练所需的TXT格式存在显著差异。本文将提供一套完整的解决方案从环境配置到代码解析手把手教你实现格式转换。1. 环境准备与数据理解1.1 必备工具安装转换工作开始前需要确保环境配置正确。以下是核心依赖项pip install pycocotools numpy tqdm常见问题在Windows系统安装pycocotools可能会遇到Microsoft Visual C 14.0报错。解决方案是pip install pycocotools-windows1.2 COCO关键点数据结构解析COCO数据集的人体关键点标注采用JSON格式每个标注包含17个关键点每个点由3个值表示keypoints: [ x1, y1, v1, # 鼻子 x2, y2, v2, # 左眼 ... # 共17个点 ]其中v表示可见性0未标注1标注但被遮挡2标注且可见1.3 YOLO格式要求YOLO格式的标注文件为TXT每行对应一个实例格式为class x_center y_center width height px1 py1 pv1 ... px17 py17 pv17所有坐标需要归一化到[0,1]区间。2. 完整转换代码实现2.1 主程序框架import argparse import os import numpy as np import tqdm import shutil from pycocotools.coco import COCO def main(args): # 初始化路径 annotation_file os.path.join(args.input_dir, annotations, fperson_keypoints_{args.split}.json) img_save_dir os.path.join(args.output_dir, f{args.split[:-4]}_coco2, images) txt_save_dir os.path.join(args.output_dir, f{args.split[:-4]}_coco2, labels) os.makedirs(img_save_dir, exist_okTrue) os.makedirs(txt_save_dir, exist_okTrue) # 加载COCO标注 coco COCO(annotation_file) catIds coco.getCatIds(catNms[person]) imgIds coco.getImgIds(catIdscatIds)2.2 关键点转换逻辑for imgId in tqdm.tqdm(imgIds, ncols100): img coco.loadImgs(imgId)[0] annIds coco.getAnnIds(imgIdsimg[id], catIdscatIds, iscrowdNone) anns coco.loadAnns(annIds) if len(annIds) 0: img_height, img_width img[height], img[width] lines [] for ann in anns: # 跳过非人体或群体标注 if ann[iscrowd] ! 0 or ann[category_id] ! 1: continue # 边界框处理 bbox np.asarray(ann[bbox], dtypefloat) bbox[::2] / img_width # x坐标归一化 bbox[1::2] / img_height # y坐标归一化 # 转换为YOLO格式(x_center, y_center, width, height) bbox[0] bbox[2] / 2 bbox[1] bbox[3] / 2 # 关键点处理 keypoints np.asarray(ann[keypoints], dtypefloat) keypoints[::3] / img_width # x坐标归一化 keypoints[1::3] / img_height # y坐标归一化 # 构建YOLO格式行 line f0 { .join(map(str, bbox))} { .join(map(str, keypoints))} lines.append(line)2.3 结果保存与图像复制if lines: txt_output_path os.path.join( txt_save_dir, os.path.splitext(img[file_name])[0] .txt ) with open(txt_output_path, w) as f: f.write(\n.join(lines)) # 复制图像文件 img_origin_path os.path.join(args.input_dir, args.split, img[file_name]) img_output_path os.path.join(img_save_dir, img[file_name]) shutil.copy(img_origin_path, img_output_path)3. 参数配置与脚本使用3.1 参数解析设置def get_args(): parser argparse.ArgumentParser() parser.add_argument(--input_dir, default/path/to/coco, typestr, helpCOCO数据集根目录) parser.add_argument(--split, defaulttrain2017, typestr, helptrain2017或val2017) parser.add_argument(--output_dir, default./converted, typestr, help输出目录) return parser.parse_args() if __name__ __main__: args get_args() main(args)3.2 运行示例python convert_coco2yolo.py \ --input_dir /data/coco \ --split train2017 \ --output_dir ./yolo_pose_data提示对于大型数据集建议使用SSD硬盘存储以加快IO速度4. 高级技巧与问题排查4.1 处理特殊情况的改进方案原始代码可能需要针对以下场景增强多人物处理确保每个实例单独一行关键点过滤只保留可见点(v0)数据校验检查图像尺寸与标注一致性改进后的关键点处理逻辑# 在关键点处理部分添加过滤 visible_keypoints [] for i in range(0, len(keypoints), 3): x, y, v keypoints[i], keypoints[i1], keypoints[i2] if v 0: # 只保留可见点 visible_keypoints.extend([x/img_width, y/img_height, v]) if len(visible_keypoints) 6: # 至少需要2个点 line f0 { .join(map(str, bbox))} { .join(map(str, visible_keypoints))} lines.append(line)4.2 常见错误解决方案错误类型可能原因解决方案ImportErrorpycocotools安装失败使用pip install pycocotools-windowsKeyErrorJSON文件路径错误检查annotations目录是否存在ValueError图像尺寸为0检查损坏的图像文件IOError权限问题确保输出目录可写4.3 性能优化建议批量处理使用多进程加速内存映射对于超大JSON文件使用mmap进度显示添加更详细的tqdm进度条改进后的主循环结构from multiprocessing import Pool def process_image(args): imgId, coco, args args # 处理逻辑... if __name__ __main__: args get_args() coco COCO(annotation_file) imgIds coco.getImgIds(catIdscoco.getCatIds(catNms[person])) with Pool(processes4) as pool: pool.map(process_image, [(imgId, coco, args) for imgId in imgIds])5. 结果验证与YOLO训练准备5.1 转换结果检查转换完成后建议检查标注文件与图像是否一一对应关键点坐标是否在[0,1]范围内可见性标志是否正确保留验证脚本示例import matplotlib.pyplot as plt def visualize_annotation(img_path, txt_path): img plt.imread(img_path) h, w img.shape[:2] with open(txt_path) as f: for line in f: parts list(map(float, line.strip().split())) bbox parts[1:5] keypoints parts[5:] # 还原像素坐标 bbox[0] * w; bbox[1] * h bbox[2] * w; bbox[3] * h for i in range(0, len(keypoints), 3): x, y keypoints[i]*w, keypoints[i1]*h if keypoints[i2] 0: # 可见点 plt.scatter(x, y, cr, s10) plt.imshow(img) plt.show()5.2 YOLO训练配置转换后的数据可用于YOLO训配置示例# yolov8-pose.yaml train: ./yolo_pose_data/train_coco2/images val: ./yolo_pose_data/val_coco2/images # 关键点配置 kpt_shape: [17, 3] # 17个点每个点(x,y,v) flip_idx: [1,0,3,2,5,4,7,6,9,8,11,10,13,12,15,14,16] # 左右翻转对应点注意YOLOv8要求关键点顺序与COCO一致无需额外调整