[论文解读] A Probabilistic U-Net for Segmentation of Ambiguous Images

Simon Köhl, Bernardino Romera‐Paredes+7 more|arXiv (Cornell University)|Jun 13, 2018

Generative Adversarial Networks and Image Synthesis参考文献 36被引用 151

一句话总结

本文介绍了一个 Probabilistic U-Net，它使用一个与 U-Net 集成的 CVAE 来对分割图的条件分布进行建模，从而在模糊图像中实现多种可信的分割假设，并对模式频率进行校准。

ABSTRACT

Many real-world vision problems suffer from inherent ambiguities. In clinical applications for example, it might not be clear from a CT scan alone which particular region is cancer tissue. Therefore a group of graders typically produces a set of diverse but plausible segmentations. We consider the task of learning a distribution over segmentations given an input. To this end we propose a generative segmentation model based on a combination of a U-Net with a conditional variational autoencoder that is capable of efficiently producing an unlimited number of plausible hypotheses. We show on a lung abnormalities segmentation task and on a Cityscapes segmentation task that our model reproduces the possible segmentation variants as well as the frequencies with which they occur, doing so significantly better than published approaches. These models could have a high impact in real-world applications, such as being used as clinical decision-making algorithms accounting for multiple plausible semantic segmentation hypotheses to provide possible diagnoses and recommend further actions to resolve the present ambiguities.

研究动机与目标

• 在模糊成像情景（如医疗 CT）中体现多假设分割的必要性，以避免仅凭单一分割而导致的误诊。
• 提出一个生成模型，学习在输入图像条件下的分割分布。
• 实现高效生成多种分割假设并对其频率进行量化。
• 在医学和城市数据集上进行评估，以展示对变体分布和罕见模态的准确再现。

提出的方法

• 将 U-Net 与条件变分自编码器（conditional variational autoencoder）结合起来，以建模给定图像的分割条件密度。
• 使用低维潜在空间来编码分割变体，并将潜在样本注入到 U-Net 中以生成多张分割图。
• 用后验网络对齐，将真实变体映射到潜在空间的位置，优化带有 KL 散度项的变分下界。
• 将先验条件化在输入图像上，以从已学习的可信分割分布中产生样本。
• 使用基于 IoU 的广义能量距离来评估分布质量，以比较预测分布与真实分割分布。

实验结果

研究问题

• RQ1 一个条件生成模型是否能够为模糊图像产生多样且连贯的分割假设？
• RQ2 所提出的模型是否能正确捕捉常见和罕见分割模态及其频率？
• RQ3 在分布精度和校准方面，Probabilistic U-Net 与基线多假设方法相比如何？
• RQ4 从模型进行采样在计算上是否足够高效，适用于临床或城市场景理解的实际应用？

主要发现

• Probabilistic U-Net 在再现多模态分割分布方面优于基线，在医学数据和类似 Cityscapes 的数据上均如此。
• 该模型不仅能学习并再现常见分割，还能以适当的频率再现罕见模态，显示出经校准的模态概率。
• 多次采样分割假设的计算成本较低，因为每个样本只需重新评估网络的一小部分。
• 潜在空间使得分割变体及其结构的可视化和分析成为可能。
• 消融研究表明，先验条件化以及在网络的合适阶段注入潜在特征对性能重要。