logo

开发者热搜

Python如何生成图片姿态估计的数据集

作者：KAKAKA2025.09.26 22:11浏览量：0

简介：本文详细介绍了如何使用Python生成图片姿态估计数据集的方法，包括数据采集、标注工具选择、标注文件生成、数据增强及数据集组织等关键步骤，为开发者提供实用指导。

Python如何生成图片姿态估计的数据集

姿态估计（Pose Estimation）是计算机视觉领域的重要任务，旨在从图像或视频中识别并定位人体或物体的关键点（如关节、面部特征等）。生成高质量的姿态估计数据集是训练和评估模型的关键环节。本文将详细介绍如何使用Python生成图片姿态估计的数据集，涵盖数据采集、标注工具选择、标注文件生成、数据增强及数据集组织等关键步骤。

一、数据采集与预处理

1. 数据来源选择

姿态估计数据集的数据来源多样，包括公开数据集（如COCO、MPII、Human3.6M等）和自定义采集的数据。自定义采集的数据能更好地满足特定场景的需求，但需投入更多资源。

自定义采集步骤：

  • 设备准备：使用高清摄像头或智能手机进行视频录制，确保光线充足、背景简洁。
  • 场景设计：根据目标应用场景（如运动分析、医疗康复等）设计动作序列。
  • 多视角采集：从不同角度录制同一动作，增加数据多样性。

2. 数据预处理

采集到的原始视频需进行预处理，包括帧提取、尺寸统一和格式转换。

Python实现示例：

  1. import cv2
  2. import os
  4. def extract_frames(video_path, output_folder, frame_interval=1):
  5.     """
  6.     从视频中提取帧并保存为图片
  7.     :param video_path: 视频文件路径
  8.     :param output_folder: 输出文件夹
  9.     :param frame_interval: 每隔多少帧提取一帧
  10.     """
  11.     cap = cv2.VideoCapture(video_path)
  12.     frame_count = 0
  13.     saved_count = 0
  15.     if not os.path.exists(output_folder):
  16.         os.makedirs(output_folder)
  18.     while True:
  19.         ret, frame = cap.read()
  20.         if not ret:
  21.             break
  23.         if frame_count % frame_interval == 0:
  24.             output_path = os.path.join(output_folder, f"frame_{saved_count:04d}.jpg")
  25.             cv2.imwrite(output_path, frame)
  26.             saved_count += 1
  28.         frame_count += 1
  30.     cap.release()
  31.     print(f"共提取 {saved_count} 帧")
  33. # 使用示例
  34. extract_frames("input_video.mp4", "output_frames")

二、标注工具选择与标注文件生成

1. 标注工具选择

姿态估计数据集的标注需标记关键点位置及连接关系。常用标注工具包括：

  • LabelImg：支持矩形框标注，但需扩展以支持关键点标注。
  • VGG Image Annotator (VIA)：支持自定义标注类型，包括关键点。
  • COCO Annotator：专为COCO格式设计，支持关键点标注。
  • SageMaker Ground Truth（AWS服务，非纯Python，但可集成）：提供托管标注服务。

推荐工具：对于Python开发者，VIA是轻量级且灵活的选择，支持JSON格式输出，便于后续处理。

2. 标注文件生成

使用VIA标注后，需将标注文件转换为模型训练所需的格式（如COCO、OpenPose等）。

COCO格式标注文件结构：

  1. {
  2.     "images": [
  3.         {
  4.             "id": 1,
  5.             "file_name": "frame_0000.jpg",
  6.             "width": 640,
  7.             "height": 480
  8.         }
  9.     ],
  10.     "annotations": [
  11.         {
  12.             "id": 1,
  13.             "image_id": 1,
  14.             "category_id": 1,
  15.             "keypoints": [x1, y1, v1, x2, y2, v2, ...],  # x,y坐标及可见性标记
  16.             "num_keypoints": 17,
  17.             "bbox": [x, y, width, height]
  18.         }
  19.     ],
  20.     "categories": [
  21.         {
  22.             "id": 1,
  23.             "name": "person",
  24.             "keypoints": ["nose", "neck", ...],  # 关键点名称列表
  25.             "skeleton": [[16, 14], [14, 12], ...]  # 关键点连接关系
  26.         }
  27.     ]
  28. }

Python转换示例：

  1. import json
  3. def via_to_coco(via_annotations, output_path):
  4.     """
  5.     将VIA标注转换为COCO格式
  6.     :param via_annotations: VIA导出的JSON标注
  7.     :param output_path: COCO格式输出文件路径
  8.     """
  9.     coco_data = {
  10.         "images": [],
  11.         "annotations": [],
  12.         "categories": [{
  13.             "id": 1,
  14.             "name": "person",
  15.             "keypoints": ["nose", "neck", "r_shoulder", "r_elbow", "r_wrist",
  16.                           "l_shoulder", "l_elbow", "l_wrist", "r_hip", "r_knee",
  17.                           "r_ankle", "l_hip", "l_knee", "l_ankle", "r_eye",
  18.                           "l_eye", "r_ear", "l_ear"],
  19.             "skeleton": [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13],
  20.                           [6, 12], [7, 13], [6, 7], [6, 8], [7, 9],
  21.                           [8, 10], [9, 11]]
  22.         }]
  23.     }
  25.     image_id = 1
  26.     annotation_id = 1
  28.     for img_ann in via_annotations["annotations"]:
  29.         # 假设via_annotations中包含关键点坐标和可见性
  30.         keypoints = []
  31.         for kp in img_ann["keypoints"]:
  32.             x, y, v = kp["x"], kp["y"], 1 if kp["visible"] else 0
  33.             keypoints.extend([x, y, v])
  35.         coco_data["images"].append({
  36.             "id": image_id,
  37.             "file_name": img_ann["filename"],
  38.             "width": img_ann["width"],
  39.             "height": img_ann["height"]
  40.         })
  42.         coco_data["annotations"].append({
  43.             "id": annotation_id,
  44.             "image_id": image_id,
  45.             "category_id": 1,
  46.             "keypoints": keypoints,
  47.             "num_keypoints": len(keypoints) // 3,
  48.             "bbox": img_ann["bbox"]  # 假设VIA标注中包含bbox
  49.         })
  51.         image_id += 1
  52.         annotation_id += 1
  54.     with open(output_path, "w") as f:
  55.         json.dump(coco_data, f, indent=4)
  57. # 使用示例（需先通过VIA标注并导出JSON）
  58. via_annotations = json.load(open("via_annotations.json"))
  59. via_to_coco(via_annotations, "coco_annotations.json")

三、数据增强与扩充

为提升模型泛化能力，需对数据集进行增强，包括几何变换、颜色空间调整和模拟遮挡等。

Python实现示例（使用Albumentations库）：

  1. import albumentations as A
  2. from albumentations.core.transforms_interface import ImageOnlyTransform
  4. class PoseKeypointAugmentation(A.DualTransform):
  5.     """自定义关键点增强，确保关键点与图像同步变换"""
  6.     def __init__(self, rotate_limit=30, scale_limit=0.2, always_apply=False, p=0.5):
  7.         super().__init__(always_apply, p)
  8.         self.rotate_limit = rotate_limit
  9.         self.scale_limit = scale_limit
  11.     def get_params(self):
  12.         return {
  13.             "angle": self.rotate_limit * (2 * self.random.random() - 1),
  14.             "scale": 1 + self.scale_limit * (2 * self.random.random() - 1)
  15.         }
  17.     def apply(self, img, angle=0, scale=1, **params):
  18.         # 图像旋转与缩放
  19.         h, w = img.shape[:2]
  20.         center = (w // 2, h // 2)
  21.         M = cv2.getRotationMatrix2D(center, angle, scale)
  22.         img_rotated = cv2.warpAffine(img, M, (w, h))
  23.         return img_rotated
  25.     def apply_to_keypoints(self, keypoints, angle=0, scale=1, **params):
  26.         # 关键点同步变换
  27.         transformed_kps = []
  28.         center = (params["columns"] // 2, params["rows"] // 2)  # 假设params包含图像尺寸
  29.         for kp in keypoints:
  30.             x, y, v = kp
  31.             # 旋转
  32.             x_rot = (x - center[0]) * math.cos(math.radians(angle)) - (y - center[1]) * math.sin(math.radians(angle)) + center[0]
  33.             y_rot = (x - center[0]) * math.sin(math.radians(angle)) + (y - center[1]) * math.cos(math.radians(angle)) + center[1]
  34.             # 缩放
  35.             x_scaled = x_rot * scale
  36.             y_scaled = y_rot * scale
  37.             transformed_kps.append([x_scaled, y_scaled, v])
  38.         return transformed_kps
  40. # 定义增强管道
  41. aug = A.Compose([
  42.     A.HorizontalFlip(p=0.5),
  43.     PoseKeypointAugmentation(p=0.8),
  44.     A.RandomBrightnessContrast(p=0.2),
  45. ], keypoint_params=A.KeypointParams(format="xyas"))  # xyas: x,y,angle,scale（此处简化为xyv）
  47. # 应用增强
  48. def augment_sample(image, keypoints):
  49.     augmented = aug(image=image, keypoints=keypoints)
  50.     return augmented["image"], augmented["keypoints"]

四、数据集组织与验证

1. 数据集组织

遵循标准目录结构，便于模型加载：

  1. dataset/
  2. ├── train/
  3. │   ├── images/
  4. │   └── annotations/
  5. ├── val/
  6. │   ├── images/
  7. │   └── annotations/
  8. └── test/
  9.     ├── images/
  10.     └── annotations/

2. 数据集验证

检查标注文件与图像的匹配性，确保无缺失或错误。

Python验证示例：

  1. import os
  2. import json
  4. def validate_dataset(images_dir, annotations_path):
  5.     annotations = json.load(open(annotations_path))
  6.     image_files = set(os.listdir(images_dir))
  7.     annotated_files = {ann["file_name"] for ann in annotations["images"]}
  9.     missing_files = annotated_files - set(image_files)
  10.     extra_files = set(image_files) - annotated_files
  12.     print(f"缺失的文件: {missing_files}")
  13.     print(f"多余的文件: {extra_files}")
  14.     assert not missing_files and not extra_files, "数据集验证失败"
  16. # 使用示例
  17. validate_dataset("dataset/train/images", "dataset/train/annotations/coco_annotations.json")

五、总结与建议

生成姿态估计数据集需经历数据采集、标注、格式转换、增强及验证等环节。Python生态提供了丰富的工具（如OpenCV、Albumentations、JSON处理库）以支持高效开发。

实践建议：

  • 标注质量控制：采用多人标注与交叉验证，减少误差。
  • 数据平衡：确保不同动作、视角的数据分布均匀。
  • 持续更新：根据模型表现反馈，迭代扩充数据集。

通过系统化的数据集生成流程，开发者能构建出高质量、多样化的姿态估计数据集，为模型训练提供坚实基础。

相关文章推荐

发表评论

最热文章

关于作者

  • 被阅读数
  • 被赞数
  • 被收藏数
活动
咨询