AI学习者 · 2022年08月25日

nnUNet 改进 | 让UNet系列在大赛中雄风依旧!

image.png

语义分割是医学图像计算中最热门的研究领域之一。尽管 nnU-Net 的概念化可以追溯到 2018 年,但它继续为各种分割问题提供具有竞争力的开箱即用解决方案,并经常被用作挑战获胜算法的开发框架。在这里使用 nnU-Net 参加 AMOS2022 挑战赛,该挑战赛附带一组独特的任务:不仅数据集是有史以来最大的数据集之一,拥有 15 个目标结构,而且比赛还需要提交的解决方案来处理 MRICT 扫描。

通过仔细修改 nnU-Net 的超参数、在编码器中添加残差连接以及自定义后处理策略的设计,能够大大改进 nnU-Net 基线。在提供的训练案例的 5 倍交叉验证中,最终集成在任务 1 (CT) 和任务 2 (CT+MRI) 上的 Dice 得分分别为 90.13 和 89.06。

1、简介

医学图像中所有解剖结构和病理的自动描绘是医学图像计算的长期目标。由于需要专家注释器和 3D 注释的时间密集型特性,到目前为止,数据集需要在目标结构的数量和训练案例的数量之间进行仔细的平衡。因此,现有方法要么在许多图像上进行训练,并且可以稳健地分割少数结构,要么在少数图像上进行训练,可以分割许多鲁棒性降低的结构。因此,无论何时需要对患者解剖结构进行整体透视,都必须将来自多个来源的多个专家模型汇集在一起。这不仅增加了推理时间,而且还产生了新的问题,例如潜在的冲突预测。最后,也许最重要的是,每个专家模型都是独立训练的,标签协同作用不能被利用,因此可能会降低模型的标签效率及其鲁棒性。

在此背景下,2022 年腹部多器官分割 (AMOS2022) 挑战赛旨在促进整体分割方法的发展。它带有数量空前的训练图像和带注释的目标结构:在 500 次 CT 和 100 次 MRI 扫描中标记了 15 个感兴趣的器官,分布在 200+40(CT + MRI)训练、100+20 验证和 200+40 测试图像中挑战提出了两项任务:任务 1 是仅针对 CT 图像的经典多器官分割问题,而任务 2 包括 MRI 图像,并希望提交的方法能够处理这两种模式。

在医学图像分割的背景下,nnU-Net 经受住了时间的考验。尽管具有完全自动化的开箱即用特性,但它始终在新的分割数据集发布时提供最先进的结果。此外,nnU-Net 成功地用作特定任务方法优化的基础,不仅使我们,而且使许多其他团队能够赢得竞争激烈的挑战。因此,使用 nnU-Net 来参与 AMOS2022 挑战似乎是再自然不过的事情了。

2、方法

nnU-Net 是一个自动配置和训练基于 U-Net 的分割管道的框架。通过对目标数据集的严格分析,nnU-Net 自动适应patch sizebatch size、预处理、网络拓扑等。

在本节中,作者对 nnU-net 的自动生成管道提出了一些修改,以最大限度地提高 AMOS2022 挑战的分割性能。在整个方法开发过程中,将所有修改应用于这两个任务,唯一的区别是强度归一化方案。

  • 对于任务 1,使用 nnU-Net 的“CT”方案(数据驱动的裁剪和归一化),
  • 对于任务 2,使用所有图像的简单 z-scorennU-Net 的“非CT”设置)。

2.1、对 nnU-Net 优化

nnU-Net 提供的默认“3d fullres”配置开始,探索了多项改进。尝试用残差编码器替换 U-Net 的默认编码器。进一步优化了预处理,特别是patch sizebatch size和目标间距。GPU 内存容量和处理速度的进步允许比标准 nnU-Net 更大的模型和batch size

image.png

通过对提供的训练案例运行 5 折交叉验证来执行方法开发。所有模型都是使用 nnU-Net 框架从头开始训练的。实验产生了任务 1 的三个表现良好的候选者和任务 2 的两个候选者,所有这些都总结在表 1 中。

image.png

图 1 显示了最终配置使用的分割架构。它们共享相同的拓扑,但由于 nnU-Net 自动配置卷积stridekernel-size作为patch size的函数,它们的特征图大小不同。

2.2、推理策略

使用 nnU-Net 默认值(滑动窗口)进行预测。对于验证和测试集预测,使用交叉验证的 5 个模型以及多个不同的配置(任务 1 的 3 个配置:集成中的 3 × 5 = 15 个模型)。集成实现为 softmax 输出的简单平均。

2.3、后处理

nnU-net 提供的后处理旨在涵盖各种用例。作者相信通过分析和定位方法的失败案例,可以在 AMOS2022 上获得额外的性能。为了识别那些,在交叉验证结果上生成了一个混淆矩阵(参见图 2),并对预测的分割图进行了严格的目视检查。

image.png

  • 左右混乱:有时左肾的一部分被归类为右肾,反之亦然(肾上腺也一样)。在这里,去除不连贯的小组件不是最理想的,因为这样肾脏的一部分就会被标记为背景。因此,将最终的肾脏预测汇集到一个联合类中,并使用连接组件分析来确定连接的肾脏区域。对于每个连接区域,计算其在图像中的位置并相应地分配其标签(左/右)。
  • 连接组件过滤:就像默认的 nnU-Net 一样,探索连接组件的过滤。具体来说,确定是否删除除最大连接组件之外的所有组件可以提高 Dice 分数。
  • 器官大小限制:在人体解剖学的驱动下,器官应该有一定的体积。此属性可用于过滤小的假阳性预测,例如在目标器官不存在但假阳性区域存在的图像中。将一个名为“rate”的参数与在训练集中观察到的最小器官大小相结合。该比率旨在作为安全边际:0.75 的比率表示小于最小器官大小的 75% 的组件被移除。
  • 具有大小限制的组件过滤:将前两种技术结合起来,并删除低于阈值的实例,除非它是唯一的连接区域。
  • 后处理步骤的组成。为了找到最佳的后处理方案,创建了前面提到的后处理方案的组合,即:
  1. PP1:左右混乱(仅肾脏和肾上腺),其次是器官大小限制
  2. PP2:左右混乱(仅肾脏和肾上腺),然后是具有大小限制的组件过滤
  3. PP3:左右混乱(仅肾脏和肾上腺),然后是连接成分过滤。
  4. PP4:只有左右混乱(只有肾脏和肾上腺)。

使用对 5 折交叉验证(训练图像)的集成预测来优化后处理方案。每个器官都进行了独立优化,并保留了性能最佳的后处理策略。请注意,我们会在适用的情况下测试“rate”参数的多个可能值。具体可以参考表 4 和 5 了解最终的后处理配置。

image.png
image.png

3、实验结果

image.png

image.png

4、参考

[1].Extending nnU-Net is all you need

作者:ChaucerG
文章来源:集智书童

推荐阅读

更多汽车电子干货请关注 嵌入式AI 专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
18838
内容数
1371
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息