第 7 名解决方案 [3D-UNet 与高斯热力图]

首先，我要衷心感谢比赛组织者、数据集提供者以及竞争对手们举办这场有趣的比赛。尝试不同的模型非常有趣，我在本次比赛中收获良多。这是我获得第 7 名提交的解决方案。

总结

我的方法利用了带有每类高斯热力图预测的 3D 分割模型。模型首先在模拟数据集上进行预训练（1 折，wbp 作为验证），然后在实验数据集上进行微调（4 折）。我使用带有非常高 pos_weight 的加权 BCE 训练所有模型，以使模型生成更多的预测。训练期间进行了大量增强，包括 Mixup、Cutmix、RandomFlip、Affine（仅在 xy 平面旋转，仅用于预训练）、rot90（仅在 xy 平面）以及其他像素值增强。最终提交由三个模型 soup 的集成组成（每个结合四折），应用了 4 倍测试时增强（TTA）和 4 倍滑动窗口方法。

预处理与数据增强

为了减少预训练时间，我使用整个模拟数据集进行训练，并使用实验数据集的 wbp 版本进行验证。

预处理是最小的，涉及百分位 clipping (0.1–99.9%) 和数据集特定的缩放因子：

• 模拟数据：×1.0
• WBP: ×1e4
• 去噪数据：×1e5

我使用步长为 64 的滑动窗口方法，排除最后一个窗口，为每个实验创建 1x8x8=64 个窗口（x 和 y 轴为 0, 64, 128, …, 384, 448，z 轴为 0），然后在生成训练数据时随机移动裁剪坐标。

我使用的增强方法包括：

• shift（从预定裁剪窗口最多移动 64）
• cutmix
• mixup
• randomflip（所有轴）
• rot90（仅在 xy 平面）
• affine（每个轴缩放，仅在 xy 平面旋转）*仅在预训练时
• contrast (0.7 ~ 1.5)
• gamma (0.8 ~ 1.2)
• gaussian noise (0 ~ 0.05)*仅在微调时

模型架构

我在最终提交中使用了 3 个模型，2 个基于 U-Net 的模型和 1 个基于 DeepLab 的模型。在两个基于 U-Net 的模型中，一个具有 ResNet50d 骨干网络，另一个具有 EfficientNetV2-M 骨干网络。基于 DeepLab 的模型具有 ResNet50d 骨干网络。我尝试训练一个带有偏移预测头的模型以提高定位性能，但在我的情况下没有起作用，所以最终提交只使用简单的分割头。除模型 207 外，所有模型均使用 decay 为 0.995 的 ModelEMA。

在 5 个特征阶段中，仅使用前 4 个阶段产生了相当或更好的性能，所以集成中的一些模型（resnet50d）就是这样做的。我猜测这是因为高分辨率特征图包含了定位粒子所需的更好信息。

在最终集成中，使用了 [resnet50d, resnet50d, efficientnetv2-m] 的骨干网络。

虽然未在最终提交中使用，但 DeepLabV3+ 也能够取得相对较高的分数，同时仅需 U-Net 模型一半的推理时间。通过使用 DeepLabV3+ 并进行更多集成，可能会创建更好的提交。

损失函数

损失函数是加权 BCE 损失，其中给予“困难”粒子和目标为正的区域更多的权重。损失函数的公式如下：

其中 \(L_{c, i}\) 是类别 \(c\) 和像素 \(i\) 的损失，\(w_c\) 是类别 \(c\) 的权重，\(p_c\) 是类别 \(c\) 中正像素的权重，\(y_{c, i}\) 是类别 \(c\) 和像素 \(i\) 的预测，\(t_{c, i}\) 是类别 \(c\) 和像素 \(i\) 的目标。\(N\) 是图像中的像素数量，总和取自所有像素和类别。

类别权重和正样本权重如下所示。
类别顺序为 [apo-ferritin, beta-galactosidase, ribosome, thyroglobulin, virus-like-particle]。

推理

推理是在大小为 (192, 128, 128) 的填充体积上进行的。我从 [184, 630, 630] 的断层扫描图中创建了 81 个窗口，x 和 y 轴步长为 64， resulting in 4x TTA。当将结果组合成单个热力图时，平均 logit 比平均 sigmoid 后的热力图更好。
为了减少边缘伪影，我还应用了一个在边缘附近下降的斜权重函数。

翻转增强作为 TTA 完成。根据模型数量，应用了 8 倍到 3 倍（正常，flipx, flipy）的 TTA。

后处理

最初我使用 CCL + 质心来确定坐标，但我 switched 到使用 maxpool 来检测局部最大值。为了减少噪声，在 maxpool 之前应用了 kernel size 5 和 sigma 1.0 的高斯模糊，并在局部最大值检测后应用了 radius=0.5*particle_radius 的加权框融合 (WBF)。使用 maxpool 方法显著提高了推理速度，性能损失最小（~0.002 减少），使我能够进行更多的 TTA 和集成。

由于非常高的 pos_weight，0.5 左右的阈值足以使模型预测许多 FP。最终提交在 logit 空间 中使用了 [0.3, 0.2, -0.2, -0.2, -0.2] 的阈值。

集成策略

在集成不同模型的输出时，平均热力图的 logit 比平均 sigmoid 后的热力图或对最终预测应用 WBF 更好。3 个模型配合 4 倍翻转 TTA 和 4 倍滑动窗口似乎是 TTA 和集成之间的最佳平衡。

探测 LB (Leaderboard)

通过仅针对目标类别进行 3 次提交：

• 正常
• 添加 FP（例如将 (0, 0, 0) * 20 附加到每个预测）
• 双倍预测

可以通过求解方程组来计算 TP, FP, FN 的数量：

mat = np.array([
    [17*(original_score - 1), original_score, 16*original_score],
    [17*(more_fp_score - 1), more_fp_score, 16*more_fp_score],
    [18*double_pred_score - 17, 2*double_pred_score, 16*double_pred_score]
])
TP, FP, FN = np.linalg.solve(mat, [0, -added_fp*more_fp_score, 0])

获得粒子总数后，可以通过 2 次提交计算 TP, FP, FN。（通过将系数矩阵的最后一行更改为 1, 0, 1 并将目标更改为 0, -added_fp*more_fp_score, num_particles）

起初我害怕使用这些信息来调整我的模型，但探测到的粒子计数 (TP + FN) 与论文中陈述的平均计数匹配 (https://www.biorxiv.org/content/10.1101/2024.11.04.621686v1)，所以我假设私有测试集具有与公共测试集相同的粒子分布。在比赛的前半部分，我试图保持提交数量较低，以免过拟合 LB，但在此之后，我决定相信 LB 而不是我的 CV。

分数