返回列表

4th Place Solution [Source Codes & Submission Notebook Released!]

639. CZII - CryoET Object Identification | czii-cryo-et-object-identification

开始: 2024-11-06 结束: 2025-02-05 医学影像分析 数据算法赛
第四名解决方案 [源代码 & 提交 Notebook 发布!]
标题:第四名解决方案 [源代码 & 提交 Notebook 发布!]
作者:yu4u & tattaka (Team yu4u & tattaka)
排名:第 4 名
发布时间:2025-02-06

我们首先想感谢竞赛主办方和 Kaggle 工作人员组织了这场出色的竞赛。下面,我们将介绍 yu4u & tattaka 团队的解决方案。

总结

我们采用了一种基于热力图的方法来检测粒子点,这是姿态估计和面部关键点检测中最常用的技术。
由于本竞赛处理的是 3D 图像而不是 2D 图像,我们使用了两种类型的类 UNet 模型(yu4u 的模型和 tattaka 的模型),它们以 3D 体素作为输入并输出 3D 热力图。

我们对本竞赛的方法

首先,我们将解释我们对本竞赛的方法,具体是如何解决 CV(交叉验证)和 LB(排行榜)不相关的问题。我们仅使用 CV 来确认产生的指标是否合理以及用于选择检查点。对于有改进潜力的模型,我们直接提交它们,并依赖 LB 来决定采用或丢弃哪些方法。

创建 Ground Truth 热力图

我们生成模型训练所需的 Ground Truth 热力图。这涉及将 Ground Truth 粒子坐标转换为像素坐标系,并使用高斯函数创建掩膜,其中粒子中心设置为 1.0,对于 yu4u 的模型,sigma 为 6 像素。对于 tattaka 的模型,根据不同粒子的大小使用了不同的 sigma 值。
我们认为在将粒子坐标转换为像素坐标系时应添加 1.0 的偏移量。虽然这个讨论建议添加 0.5,但我们的 Notebook证明 1.0 是正确的值。主要区别在于之前的讨论假设粒子中心位于像素的左上角,而我们认为平均而言,圆应该从像素中心 (0.5, 0.5) 绘制。

Coordinate System Explanation

yu4u 的模型

yu4u Model Architecture

我们采用了 2.5D-UNet,它使用基于 2D 图像的模型作为骨干网络。该骨干网络每个阶段的输出沿深度方向进行池化,从而也能在深度维度上进行分层特征提取。这个想法借鉴自这个优秀的 Notebook。一个有趣的观察是,用步长 3D 卷积替换这个池化操作会降低性能。这是因为池化方法有效地聚合了深度特征,同时尽可能保留了原始 2D 骨干网络的特征图。与许多其他处理 3D 数据的 Kaggle 竞赛类似,使用 2D 骨干网络的 UNet 优于直接使用 3D 骨干网络的 UNet。

我们还受contrails 竞赛第 3 名解决方案的启发,在编码器和解码器之间应用了 3D 卷积。

最初,我们使用普通的 UNet 架构,但处理高分辨率特征图需要大量的内存和计算。为了解决这个问题,我们采用了一个模型,使用 stride 为 4 的特征图通过 pixel shuffle 输出最终热力图。Pixel shuffle(在 TensorFlow 中也称为 depth_to_space)是一种将信息从通道维度重新分配到空间维度的操作。与反卷积相比,它在计算效率和减少伪影方面具有优势。

对于最终提交,我们使用了四个模型,它们使用 ConvNeXt Nano 模型作为骨干网络的不同折叠。

tattaka 的模型

模型架构

tattaka Model Architecture

该模型是一个轻量级的 2.5D UNet,以 ResNetRS-50 作为骨干网络。

模型的输入是大小为 32×128×128 (D×H×W) 的体积,输出相同大小的 3D 热力图。在骨干网络内,前两个阶段使用平均池化将深度 progressively 减半。之后,使用 kernel=3, stride=1, padding=1 的平均池化来维持深度同时计算特征图。因此,骨干网络每个阶段的特征图形状如下:
(bs, ch, 16, 64, 64), (bs, ch, 8, 32, 32), (bs, ch, 8, 16, 16), (bs, ch, 8, 8, 8), (bs, ch, 8, 4, 4)

在解码器中,三个最低分辨率的特征图被送入联合金字塔上采样 (Joint Pyramid Upsampling)。然后使用 3D CNN、SESC 注意力机制和上采样层逐步上采样这些图,直到它们达到与输入体积相同的大小。

损失函数

由于体积内的粒子数量相对较少,训练期间正样本和负样本之间存在显著的类别不平衡。我们尝试调整生成 Ground Truth 热力图的参数,但这并没有带来交叉验证性能的提升。

最终,我们实现了一个简单的基于 MSE 的损失函数来平衡正样本和负样本,这允许更快的收敛:

loss = MeanSquaredError(pred, true)

pos_loss = (loss * true).sum() / (true.sum() + 1e-6)
neg_loss = (loss * (1 - true)).sum() / ((1 - true).sum() + 1e-6)

balanced_loss = pos_loss + neg_loss

推理技巧

最后,我们在最终提交中使用了四个 yu4u 的模型和三个 tattaka 的模型。
为了保持在时间限制内,我们通过将模型转换为 TensorRT 格式来优化模型以进行更快的推理。转换过程基于这个 Notebook
此外,我们选择了具有双 T4 GPU 的 Kaggle Notebook 实例,并利用多进程并行化推理。

后处理

对于最终的热力图,我们首先使用非极大值抑制检测局部最大值,这是通过 kernel 大小为 7 的最大池化实现的。接下来,使用针对不同粒子类型的不同阈值过滤检测到的点。

由于检测到的点位于像素坐标系中,我们需要将它们转换为粒子坐标系。为此,我们按以下步骤操作:

  1. 居中:向像素坐标添加 0.5,从像素的左上角shift到其中心。
  2. 偏移校正:减去在热力图生成期间添加的 1.0 偏移量。
  3. 缩放:乘以 10.012 将调整后的像素坐标转换为粒子坐标系。

对我们无效的方法

  • 两阶段模型:我们构建了一个模型,通过裁剪使用热力图方法检测到的点周围的区域,然后将分类模型应用于这些裁剪区域来细化分数。虽然它在 CV 分数方面表现良好,但并没有提高 LB 性能。
同比赛其他方案