企业官网建设流程全解析

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-windows

1.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', f'person_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_ok=True) os.makedirs(txt_save_dir, exist_ok=True) # 加载COCO标注 coco = COCO(annotation_file) catIds = coco.getCatIds(catNms=['person']) imgIds = coco.getImgIds(catIds=catIds)

2.2 关键点转换逻辑

for imgId in tqdm.tqdm(imgIds, ncols=100): img = coco.loadImgs(imgId)[0] annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None) 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'], dtype=float) 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'], dtype=float) keypoints[::3] /= img_width # x坐标归一化 keypoints[1::3] /= img_height # y坐标归一化 # 构建YOLO格式行 line = f"0 {' '.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", type=str, help="COCO数据集根目录") parser.add_argument("--split", default="train2017", type=str, help="train2017或val2017") parser.add_argument("--output_dir", default="./converted", type=str, 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 处理特殊情况的改进方案

原始代码可能需要针对以下场景增强：

1. 多人物处理：确保每个实例单独一行
2. 关键点过滤：只保留可见点(v>0)
3. 数据校验：检查图像尺寸与标注一致性

改进后的关键点处理逻辑：

# 在关键点处理部分添加过滤 visible_keypoints = [] for i in range(0, len(keypoints), 3): x, y, v = keypoints[i], keypoints[i+1], keypoints[i+2] if v > 0: # 只保留可见点 visible_keypoints.extend([x/img_width, y/img_height, v]) if len(visible_keypoints) >= 6: # 至少需要2个点 line = f"0 {' '.join(map(str, bbox))} {' '.join(map(str, visible_keypoints))}" lines.append(line)

4.2 常见错误解决方案

4.3 性能优化建议

1. 批量处理：使用多进程加速
2. 内存映射：对于超大JSON文件使用mmap
3. 进度显示：添加更详细的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(catIds=coco.getCatIds(catNms=['person'])) with Pool(processes=4) as pool: pool.map(process_image, [(imgId, coco, args) for imgId in imgIds])

5. 结果验证与YOLO训练准备

5.1 转换结果检查

转换完成后，建议检查：

1. 标注文件与图像是否一一对应
2. 关键点坐标是否在[0,1]范围内
3. 可见性标志是否正确保留

验证脚本示例：

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[i+1]*h if keypoints[i+2] > 0: # 可见点 plt.scatter(x, y, c='r', s=10) 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一致，无需额外调整