AggPose: Deep Aggregation Vision Transformer for Infant Pose Estimation 论文笔记

IJCAI 2022- AggPose：用于婴儿姿态估计的深度聚合 vision transformer

论文链接

To appear in the 31th International Joint Conference on Artificial Intelligence (IJCAI 2022)

---

摘要： 通过估计新生儿的运动和姿态，有相关经验的儿科医生可以对其进行神经发育障碍的预测，从而对相关疾病进行早期干预。然而，大多最新人体姿态估计 AI 方法都集中在成年人身上，且缺乏用于婴儿姿态估计的公开基准。 本文提出了婴儿姿态数据集和深度聚合 vision transformer 人体姿态估计方法来填补这一研究领域的空白。深度聚合VIT 引入了一个快速训练的 full transformer 框架，无需在早期阶段使用卷积运算来提取特征。深度聚合 VIT 将 Transformer+MLP 推广到特征图内高分辨率的深层聚合中，从而实现不同 vision 级别间的信息融合。我们在 COCO 姿态数据集上预训练 AggPose，并将其应用于我们新发布的大规模婴儿姿态估计数据集中，结果表明：AggPose 能够有效学习不同分辨率间的多尺度特征，显著提高婴儿姿态估计的性能。在婴儿姿态估计数据集上的结果表明：AggPose 优于混合模型 HRFormer 和 TokenPose。此外，在 COCO val 姿态估计数据集上，AggPose 比 HRFormer 平均高出 0.7%AP。代码链接

---

---

1 Introduction

每年，全世界约有500万新生儿患有神经发育障碍。由于缺乏早期诊断和干预，许多婴儿严重致残并被父母遗弃，尤其在那些缺乏经验丰富的神经发育障碍儿科医生的国家，这已成为困扰世界各地许多家庭的难题。
最新发展的基于深度学习方法为开发用于神经发育障碍早期干预的计算机辅助运动评估工具提供了可能。一般运动评估（general movements assessment: GMA）是早期脑瘫诊断最具预测性的工具之一，它需要在许多小幅度运动中区分烦躁和非烦躁运动，计算机对检测此类运动更为敏感。研究人员使用 OpenPose 等人体姿势估计方法捕捉婴儿姿态，然后生成婴儿运动序列以检测脑瘫。与手工的 GMA 检测相比，基于计算机的方法速度更快，成本更低。 然而，考虑到婴儿姿态的复杂场景，且世界各地缺乏大规模的公共婴儿姿态数据集，这项任务在实际应用中具有挑战性。此外，由于缺乏临床意义和可操作性，COCO 数据集定义的 17 个成人关键点不能很好地支持婴儿运动检测。
另一个问题是姿态估计方法的性能。尽管具有强大的表示学习和语义理解能力的基于CNN的方法将人体姿态估计推到了一个新的水平，但它在理解身体部位之间的全局约束关系方面仍然表现不佳。研究人员将Vision Transformer与CNN结合用在混合模型中，让ViT扩展感受野，并增强模型捕捉身体部位之间约束关系的能力来解决此问题。在最近的进展中，Swin-Transformer 的 local-window self-attention 结构和 SegFormer 的混合前馈网络（Mix Feed Forward Network：Mix-FNN）在多尺度特征表示学习方面显示出巨大的潜力。
然而，将Transformer应用于人体姿态估计仍有一些问题：(1) 混合模型的第一阶段高度依赖于预训练的 HRNet 卷积层，这将无法利用具有最新自监督掩码自动编码器的大规模未标记数据；(2) 在训练过程中难以融合；(3) 很难将模型从一个 domain (领域) 转移到另一个 domain 。
本文提出了一种将多尺度 transformer 结构推广到深度聚合网络的方法：AggPose。不同于 HRFormer，AggPose 不使用卷积层作为初始特征提取和融合模块，相反，它使用逐层 Mix transformer 和交叉分辨率 MLP 融合模块。Transformer 接收前一层的输入，应用 self-attention 操作和 Mix-FNN，并将消息发送到下一层，MLP 融合模块集成了来自不同分辨率级别的丰富空间信息，并将结果发送到下一阶段。
我们在 COCO 人体姿态估计数据集上进行实验，然后在本文提出的大规模有标记的婴儿姿态估计数据集上对模型进行微调。AggPose 在这两个基准上都取得了有竞争力的性能，例如，在 COCO val set 上，AggPose-L 比 HRFormer 和 TokenPose 分别高出 0.7 AP 和 0.5 AP。AggPose-L的鲁棒性和快速收敛性使其能够轻松从COCO数据集转到我们的婴儿姿态数据集上，并达到最高 95.0 AP。
贡献总结如下：

• 本文提出了一种新的基于 Transformer+MLP 的聚合架构，不使用 HRNet’s CNN backbone 和基于 CNN 的多尺度融合模块。
• 为了提高深层聚合的效率，本文设计了一种特殊的深层聚合 MLP 结构来融合不同分辨率的信息。
• 为了促进神经发育障碍的早期干预研究，本文提供了一个大规模的婴儿挑战性数据集，包括 20748 张姿态标记图像。实验结果表明，我们的框架在COCO和这个新数据集上都具有鲁棒性。据我们所知，这是为临床应用构建的最大的婴儿姿态估计数据集。

2 Related Works

2.1 Infant Pose Estimation 婴儿姿态估计

婴儿姿势估计在临床研究中具有很高的应用价值。最常用的脑瘫评估工具，如 GMA 和 CPVC 已经使用自动姿态估计方法来辅助诊断。然而，现有的大多数自动婴儿姿态算法都基于传统的机器学习方法，限制了它们处理复杂情况的能力。与此同时，人工智能界很少尝试处理婴儿图像，只有 [McCayet al., 2019; Reichet al., 2021] 首次尝试采用 OpenPose 提取婴儿关键点，但缺乏大型通用数据集。MINIRGBD 是该领域最著名的开源数据集，但它仅包含700幅真实的婴儿图像和一小部分合成的婴儿图像。以上这些问题使婴儿姿态自动评估方法在实际的临床中不可靠。

2.2 Vision Transformers for Pose Estimation 姿态估计 VIT

多年来，深度卷积神经网络已被应用于人体姿态估计。在所有基于CNN的姿态估计算法中，在整个网络中保持高分辨率表示的方法取得了巨大的成功。 最具代表性的模型是 HRNet、HigherNet、UDP、DARK。然而，由于CNN的感受野限制了其理解全局空间关系的能力，使得 CNN 难以捕捉人类关键点之间的约束关系。
最近一些工作将 transformer 用于人体姿态估计。TokenPose (2021) 引入Transformer，将关键点表示为人类姿态估计的 token embeddings。HRFormer (2021) 将HRNet与Swin Transformer集成，充分利用不同非重叠图像 window 上的多分辨率并行信息。然而，HRFormer和TokenPose的第一阶段在 HRNet’s CNN backbone 上进行了微调，所以两者都未摒弃卷积运算来获得初始特征。 本文提出聚合视觉转换器（Aggregation Vision Transformers: AViT），提供了一种不同的方法来解决ViT的低分辨率问题，并用 Overlapped patch embedding 代替卷积 操作在 early stages 提取特征。

3 Proposed Method 提出的方法

我们的目标是提出一个新的婴儿姿态数据集，并建立一个新的基准，通过 VIT 快速提取婴儿姿态。图2所示为模型的 pipeline 。我们的婴儿姿态估计研究已通过深圳市宝安妇幼医院的伦理检查。

图2：AggPose pipiline。每个 module 由多个连续的 Mix Transformer blocks 组成。不同分辨率的特征通过MLP层（图中的蓝色方块）连接。

3.1 Infant Pose Detection Dataset 婴儿姿态检测数据集

本文提出了一个具有挑战性的大规模新生儿姿态提取和检测数据集，它可以用于预测婴儿的运动顺序，并设计像 automatic GMA 那样的自动临床工具。婴儿姿态检测很重要也很困难，现有的数据集要么太小，要么太简单，需要一个大型的公共标注基准来比较不同的方法。此外，没有数据集为婴儿图像提出合适的关键点标注，它们采用了 COCO 的17个关键点格式，这会丢失婴儿的许多重要的精细姿态和运动特征。
受 [Silvaet al., 2021; Huanget al., 2021] 的启发，我们发布了新的开源婴儿姿态数据集和新的婴儿关键点格式。为收集数据，我们采用 GMA 设备记录2013年至今的婴儿运动视频，收集了 216 个多小时的视频，提取了1500万帧。我们的数据集的大小和可伸缩性都比 MINI-RGBD 数据集好得多。我们从视频中随机抽取 20748 帧，让专业临床医生标注婴儿关键点。然后，我们将 11756 张图片分给训练集和验证集，8992 张图片分给测试集。婴儿姿态的 21个关键点格式是由研究神经发育障碍超过30年的经验丰富的临床医生提出的。图1所示为数据集的示例，考虑到临床应用要求和保护患者隐私，我们的数据集减少了婴儿头部的关键点，并包括更精细的身体关键点，如手指、脚趾和肚脐。对于公共版本，我们将重新格式化数据集以解决伦理问题：在最终发布的关键点数据集中，所有婴儿的头部都将覆盖马赛克，以保护患者的隐私。禁止商用婴儿姿态数据集。

图1: (a) 我们提出的婴儿姿态数据集示例。(b) 21个婴儿身体关键点

本文重点关注人体/婴儿仰卧位姿态检测，这是新数据集最直接的应用。 该数据集也可用于其他临床领域，因为它包含超过200小时的婴儿运动序列，并与脑瘫和其他神经发育障碍的自动预测有高度联系。我们希望利用我们的数据集对疾病进行早期诊断和干预，造福所有年龄段的人，尤其是儿童。未来，我们还将发布由 AggPose 生成的200多个小时的新婴儿姿态序列以及相应的的GMA标签。回顾性研究得到了我们的机构审查委员会的批准。

3.2 Deep Aggregation Vision Transformers 深度聚合 VIT

Overlapped Patch Embedding
Transformers 在 early stage 将输入视为1D向量，专门关注全局上下文建模，会丢失细节的局部信息，因此 Early convolutions 被认为是为 hybrid transformer 架构提取低级特征的实用工具。 HRNet及其预训练的CNN参数是几乎所有最新人体姿态估计模型的基石。 受 SegFormer 的启发，我们采用具有 Overlapped Patch Embedding 的 full transformer 来取代 HRNet’s CNN 特征提取器和每个 stage 的下采样 stem。与HRNet、HRFormer和TokenPose中 early stage 中的卷积相比，Overlapped Patch Embedding 可以获得更好的 low-level 特征，增强高分辨率 transformer 的特征表示，降低计算复杂度。
Aggregation Vision Transformers (AViTs) Architecture 聚合VIT架构
遵循 Mix transformer 的 transformer 模块设计，从第一阶段 patch size =7，stride =4，padding size=3 的 Overlapped Patch Embedding 生成的高分辨率特征图开始，然后通过 Overlapped Patch Embedding 逐个添加 high-to-low 分辨率流，对每个分辨率流使用多个 mutil-head self-attention 块来更新特征表示。为了构造不同深度的模型，我们分别提出了小（AggPose-S）和大（AggPose-L）模型。表2显示了AggPose-L中每个 stage 的 transformer 层数

表2: 每个 stage 的 transformer 层数。

与Swin-Transformer 和 HRFormer相比，考虑到 overlapped patches 的使用，我们没有使用 local-window self-attention 来增强局部信息理解，而是使用序列缩减过程，这大大减少了 transformer 内部的计算量，并加快了模型训练期间的收敛。对于每个 transformer block，self-attention 被估计为：

其中，K是初始 shape 为 N×C 的 token 表示。γ 是将 K 的维数从 N×C 缩减到 N/γ×C 的比率。

MLP Cross-Layer Aggregation
ViT 和 Swin Transformer 都使用 positional embedding 来引入跨层位置信息，但 positional embedding 的分辨率是固定的。local-window Transformer 缺少跨 windows 的信息交换，因此，SegFormer 和 HRFormer 都在前馈网络（feed-forward network FFN）中引入 3×3 depth-wise 卷积来扩大感受野，减少 positional embedding 带来的坏处。具有 depth-wise 卷积的FFN（HRFormer） 和 MixFFN（SegFormer）使用了非常相似的计算。

DW Conv 指 3×3 depth-wise卷积。
在AggPose中，将 Mix-FFN 的使用扩展到跨不同分辨率层的深度聚合方法中。对于CNN中的深层聚合，如 CAggNet 和 HRNet，聚合从最浅的高分辨率层开始，然后迭代合并更深的低分辨率层，通过这种方式，浅层特征在通过聚合的不同 stage 传播时得到细化。相关研究表明，深度聚合结构并非仅传播前一 block 而是传播所有分辨率的聚合，来更好地保留特征，它被广泛用于语义分割任务中来获得有竞争力的性能。 本文提出的跨层聚合模块由每个分辨率级别中的两个 main step 组成：① 首先，来自不同分辨率的 multi-level 特征通过一个具有 3×3 depth-wise 卷积的混合前馈网络(mixed-ffn)，将通道维度和上采样或下采样 (overlapped patch embedding) 特征图统一到同一 shape 。② 然后将相邻 level 的特征向量 concatenate ，并采用额外的FFN层来融合跨层信息。与 HRFormer 和 HRNet 中的卷积多尺度融合模块相比，MLP 融合模块在提高模型性能的同时加快了收敛速度。

其中，$x_{i,j}$ 是聚合 MLP 层的输入，$x_i$ 表示来自相邻分辨率的特征图。跨层聚合模块定义为：

MixFFN 表示 Eq.(3) 中的 Mix feed forward block。

3.3 Analysis

AggPose 和 我们的大规模婴儿姿态数据集与其他CNN或 hybrid CNN Trans-
former 方法（如HRFormer和TokenPose）以及其他用于婴儿姿态估计的小型数据集相比，有两个主要优点：
(1) Potential of using self-supervised learning 使用自监督学习的潜力. 近年来，通过自监督学习进行预训练的 VIT 受到广泛关注，MAE 构建了一个 inpainting masked 自动编码器任务，从未标记数据中学习表示，并在任何监督任务上微调模型，其结果证明，full ransformers 可以从大规模的未标记数据集中学习合理的语义。如表1所示，我们有 5187 个视频中大量未标记的婴儿运动帧，所有这些数据将有助于通过自监督学习对 transformer-based autoencoder 进行预训练。

表1：其他婴儿姿态数据集间的比较。

(2) Faster convergence 更快的收敛. HRFormer 通过跨层卷积运算实现特征传递，这难以收敛。AggPose 框架由 MLP 跨不同层传播消息，它可以看作是对深层聚合模型的一种修正。我们的实验将表明，这种消息传递方案比基于 hybrid CNN-Transformer 的方法取得更好的结果。

4 Experiment

4.1 Model Variants 模型变化

考虑到大多数基于 transformer 的姿态估计模型训练过程比较复杂，本文提出了一种有效的训练策略。首先，我们 load 在 ImageNet 上预训练的 Mix
Transformer，在COCO关键点训练集上训练 Mix Transformer 。Mix Transformer 收敛后，我们将 Mix Transformer 的参数 load 到 AggPose 的每一层中。然后，我们在不同分辨率水平上逐层固定 AggPose 的参数，并在COCO 和 婴儿姿态数据集上对模型进行微调。
overlapped patch embedding 的 size 和 transformer 层的数量的配置细节如表2所示。注意，表2只提供了使用 MiT-B5 作为 backbone 的 AggPose-L的配置。对于使用MiT-B2作为 backbone 的 AggPose-S 的 transformer 层的数量为 [[3,3,3,3],[4,3,3],[6,3],[3]]。

4.2 Comparing with SOTA Methods 与sota方法对比

数据集： 我们在COCO人体姿态估计数据集 （包含25万多个标记有17个关键点的人体实例）和新的婴儿姿态估计数据集（包含2万个标记有21个关键点的婴儿实例）上研究AggPose的性能。

训练设置： 遵循 HRNet 和HRFormer 的默认训练和评估设置，使用AdamW优化器对模型进行训练，学习率初始值为 0.001。训练的 batch size=32，采用4×48G-RTX8000 GPU。

评估指标： 采用通过 OKS 计算的AP评估。

COCO姿态估计中的关键点检测。

表3: COCO验证集的比较结果，提供与HRNet相同的 detected human boxes

婴儿姿势估计的关键点检测。 表4 为婴儿姿态测试集的比较结果。我们将AggPose与最具代表性的 bottom-up 方法 OpenPose 进行比较，因为几乎所有最新的婴儿姿态估计框架都使用它。 我们还将 AggPose 与最近的几个CNN和混合transformer模型（如HRNet、TokenPose和HRFormer）进行了比较。

表4：婴儿姿态测试集比较，配备相同的 object detection boxes（OpenPose是bottom-up法，不需要object detection。我们之所以选择OpenPose进行比较，是因为大多数最新提出的婴儿姿态估计框架都选择OpenPose作为backbone。）

4.3 Ablation Experiments

Full Transformer backbone 的作用： 考虑到所有其他新提出的混合方法都使用 HRNet’s CNN encoder 作为 backbone，我们将我们的方法（第一阶段不使用卷积层）与表3 中 HRNet’s CNN 方案进行比较。Backbone 列显示模型内第一阶段的差异。AggPose-S 和 AggPose-L都使用基于Transformer的 SegFormer 作为第一 stage 层。虽然其他作者声称在 early stage 的 Transformers 将导致缺乏局部细节信息，但我们观察到 AggPose 的 full Transformer-based early stage 仍然可以实现更好的性能。

Deep Aggregation Framework 的影响： 表5为基于两个著名的full transformer models：Swin Transformer 和 SegFormer 在COCO数据集上的姿态估计结果。事实上，AggPose-L可以被视为具有 MLP skip-connect 的SegFormer-B5 的深层聚合结构。表5的结果表明：我们提出的多分辨率聚合框架（AggPose）比Swin Transformer和SegFormer 性能优越。

表5：在 COCO pose estimation val 上与非聚合框架（Swin-Transformer，SegFormer）的比较

4.4 Visualization Analysis 可视化分析

图3：infant pose test set 上显示基于 AggPose-L 的姿态估计热图结果。

表4：COCO val 上基于AggPose-L 姿态估计结果可视化

表5：infant pose test 上基于AggPose-L 姿态估计结果可视化

5 Conclusion

通过利用带有姿态标签和临床标签的新数据集，我们构建了一个基于 transformer 的婴儿姿态估计框架，该框架可以从视频的运动帧中准确地检测婴儿仰卧姿态。AggPose模型的关键是具有cross-layer MLP connection 的深度聚合 Transformer 。 我们模型生成的姿态序列已用于新生儿神经发育障碍预测和相关疾病的早期评估。此外，我们的方法在未来将部署到移动设备上，以期解决医疗资源不平等和患者隐私保护方面的问题。虽然我们的结果很有希望，但我们承认，要将我们的模型与当前可用的硬件完全端到端地应用，还有很长的路要走。
除了测试 AggPose 在实时婴儿姿态提取和评估中的效用外，在未来，我们下一步的明确目标是通过姿态序列理解模型预测脑瘫——甚至被经验丰富的儿科医生看到之前。对于儿科医生有限的国家，这将大大降低儿童严重残疾的风险。