两阶段系统

摘要

与其他参赛者类似，我们实现了一个由三个核心模型（加上一个用于预训练的额外模型）组成的两阶段系统：

• 关键点 (Keypoints)： 定位椎间点。
• 级别 (Levels)： 将轴向图像分类为脊柱级别。
• 预训练 (Patch)： 分类单个补丁。
• 序列 (Sequence)： 结合 Transformer 的 2D 模型用于最终预测。

有关详细的训练 specifics，请参考附带的代码。

模型

1. 关键点 (Keypoints)

我们最初尝试了边界框检测器和分割模型，但关键点提供了最好的结果和灵活性。模型输出 6 对 XY 坐标（共 12 个），其中 5 对用于矢状面图像，1 对用于轴向图像。我们采用了 修正的数据集标签，由 @brendanartley 分享。这允许我们使用倾斜裁剪，尽管性能提升微乎其微。

损失函数基于欧几里得距离，并应用掩码以避免在矢状面图像中传播轴向关键点误差，反之亦然。我们还在池化 (pooling) 之前将位置编码引入骨干特征，以增强模型性能。

2. 级别 (Levels)

对于轴向图像，我们开发了一个简单的图像模型来预测相应的级别。一些元数据与骨干特征连接以增强模型性能。

我们观察到偶尔的预测不一致（例如，在同一系列中 L5-S1 预测在 L1-L2 旁边）。为了解决这个问题，我们将任务重新定义为回归问题，预测 0.0 到 1.0 之间的值，间隔为 0.25（例如，L1-L2 为 0.0，L2-L3 为 0.25 等），使用均方误差 (MSE) 作为损失函数。这种方法更严厉地惩罚远离真实值的预测。

3. 序列 (Sequence)

利用解剖对称性，我们使用研究 - 级别 - 侧输入训练了一个模型。

• 输入： 例如，一行是 L1-L2 右侧的所有补丁，另一行是 L4-L5 左侧的所有补丁。
• 输出： neural_foraminal_narrowing（神经椎间孔狭窄）, subarticular_stenosis（关节下狭窄）, spinal_canal_stenosis（椎管狭窄）。

详情如下所述。

输入

我们使用关键点和级别模型创建补丁，每个轴向补丁对应一个或多个脊柱级别。由于级别模型是回归模型，我们应用了带有公差范围的阈值来捕获额外的上下文。例如，对应于 L2-L3（0.25 输出）的补丁可以跨越 0.10 到 0.40。

为了确保一致性，我们在每个平面保持了统一的像素间距。我们使用了 96x96 的补丁。矢状面补丁在顶部和底部偏移以防止级别泄漏。

示例：

预训练：补丁模型

在训练序列模型之前，我们使用相同的损失函数（稍后定义）使用单个补丁预训练了一个补丁模型。这一步显著提高了序列模型的收敛性。

架构

2D 模型 + 编码器 - 解码器 Transformer。

该架构结合了 2D 图像模型与编码器 - 解码器 Transformer。我们发现最好的骨干网络是 regnetz_b16.ra3_in1k 和 hgnet_tiny.ssld_in1k。

元数据（修改后的 XYZ 世界坐标）被馈送到编码器，而图像特征被传递到解码器。我们在解码器中添加了三个可学习向量，每个代表一种疾病分类（类似于 CLS 令牌）。测试了各种池化策略，包括 max/mean pooling 和 attention pooling。

可重塑模型

该模型设计为灵活的，允许它处理整个研究或单个级别 - 侧。这种适应性使我们能够在级别 - 侧训练模型，而在整个研究上评估它。

如何处理脊柱（狭窄）？

由于脊柱狭窄影响中心而不是特定侧，我们在训练期间为两侧复制了标签。在推理时，我们在将可学习向量传递到分类头之前聚合了两侧的向量（解码器中添加的那些）。我们测试了多种聚合方法，最终选择了 max pooling。

损失

我们使用了 PyTorch 中实现的特定于竞赛的损失函数。虽然我们探索了几种变体，例如竞赛损失的 focal 版本，但都没有产生显著更好的结果。

更多信息

选定的两个提交分别是本地验证最好的和公共榜单最好的。公共榜单最好的提交在私人榜单中排名第 3。

代码