MMPose项目中修改COCO数据集关键点个数时的解决方案

问题背景

在使用MMPose进行人体关键点检测时，研究人员有时需要针对特定任务修改COCO数据集中的关键点数量。本文以COCO WholeBody Face数据集为例，探讨如何将原本68个面部关键点修改为17个关键点，并解决在此过程中遇到的技术问题。

关键问题分析

当尝试减少COCO WholeBody Face数据集的关键点数量时，主要会遇到以下两个技术难点：

1. 维度不匹配错误：在训练过程中出现"stack expects each tensor to be equal size"错误，提示热图维度不一致（部分为17x64x64，部分为68x64x64）。

2. 关键点索引问题：原始数据集的关键点ID与自定义关键点ID之间的映射关系需要正确处理。

解决方案详解

1. 数据集类修改

首先需要创建一个新的数据集类CocoWholeBodyFace17PointDataset，继承自基础数据集类。这个类需要正确指定只包含17个关键点的元信息：

class CocoWholeBodyFace17PointDataset(BaseCocoStyleDataset):
    METAINFO: dict = dict(
        from_file='configs/_base_/datasets/coco_wholebody_face_17point.py')

2. 元信息配置文件

在coco_wholebody_face_17point.py配置文件中，需要正确定义17个关键点的信息。特别注意以下几点：

• 每个关键点的ID必须与其在字典中的键一致
• 需要正确定义关键点之间的对称关系（swap属性）
• 需要设置正确的关节权重和sigmas值

keypoint_info={
    30: dict(name='face-30', id=30, color=[255, 0, 0], type='', swap=''),
    36: dict(name='face-36', id=36, color=[255, 0, 0], type='', swap='face-45'),
    # 其他关键点定义...
    57: dict(name='face-57', id=57, color=[255, 0, 0], type='', swap=''),
},
skeleton_info={},
joint_weights=[1.] * 17,
sigmas=[
    0.007, 0.017, 0.011, 0.009, 0.011, 
    0.009, 0.007, 0.013, 0.008, 0.011, 
    0.012, 0.010, 0.034, 0.008, 0.008, 
    0.010, 0.009]

3. 模型头部调整

在模型配置文件中，需要将HeatmapHead的输出通道数改为17，以匹配新的关键点数量：

head=dict(
    type='HeatmapHead',
    in_channels=1280,
    out_channels=17,  # 修改为17个关键点
    loss=dict(type='KeypointMSELoss', use_target_weight=True),
    decoder=codec),

4. 关键实现细节

1. 关键点提取逻辑：在自定义数据集类中，需要重写关键点提取逻辑，确保只从原始数据中提取所需的17个关键点。

2. 评估指标调整：需要修改评估过程中使用的关键点索引，特别是对于需要计算距离的关键点对（如左右眼外眼角）。

3. 数据增强处理：对于随机翻转等数据增强操作，需要确保关键点的对称关系正确映射到新的关键点集合。

常见问题排查

如果在实施上述修改后仍然遇到问题，可以检查以下几个方面：

1. 确保数据集注册正确，新的数据集类被正确导入和使用
2. 检查关键点ID是否连续且无冲突
3. 验证数据加载过程中是否所有样本都被正确处理
4. 确认评估指标计算使用的关键点索引是否正确

总结

在MMPose项目中修改COCO数据集的关键点数量是一个需要谨慎处理的过程。通过创建自定义数据集类、正确定义元信息、调整模型头部输出以及正确处理评估逻辑，可以成功实现对数据集关键点数量的修改。本文提供的解决方案不仅适用于面部关键点检测，其思路也可以推广到其他关键点检测任务的定制化需求中。