[논문 리뷰] A Probabilistic U-Net for Segmentation of Ambiguous Images

Simon Köhl, Bernardino Romera‐Paredes외 7명|arXiv (Cornell University)|2018. 06. 13.

Generative Adversarial Networks and Image Synthesis참고 문헌 36인용 수 151

한 줄 요약

논문은 U-Net과 통합된 CVAE를 사용하여 조건부 분포를 모델링하는 Probabilistic U-Net을 도입합니다. 이는 모호한 이미지에 대해 여러 개의 그럴듯한 분할 가설을 가능하게 하고 모드 빈도수를 보정합니다.

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과 조건부 변분 오토인코더를 결합합니다.
• 저차원 잠재 공간을 사용하여 분할 변이를 인코딩하고 잠재 샘플을 U-Net에 주입하여 다중 분할 맵을 생성합니다.
• Ground-truth 변이를 잠재 공간 위치로 매핑하는 후방 네트워크와 KL 발산 항이 포함된 변분 하한을 최적화하여 학습합니다.
• 입력 이미지에 조건을 걸어 학습된 그럴듯한 분할 분포에서 샘플을 생성합니다.
• IoU를 기반으로 한 일반화된 에너지 거리를 사용하여 예측 분포와 실제 분할 분포의 질을 평가합니다.

실험 결과

연구 질문

• RQ1 조건부 생성 모델이 모호한 이미지에 대해 다양하고 일관된 분할 가설을 생성할 수 있는가?
• RQ2 제안한 모델이 일반적인 모드와 희귀한 모드를 모두 정확한 빈도로 포착하는가?
• RQ3 Probabilistic U-Net가 분포 정확성과 보정 측면에서 기준선 다중 가설 접근법과 어떻게 비교되는가?
• RQ4 임상 또는 도시 현장 이해를 위한 실용적 사용을 위해 모델에서 샘플링이 계산적으로 충분히 효율적인가?

주요 결과

• Probabilistic U-Net가 의료 및 Cityscapes 유사 데이터에서 다중 모드 분할 분포를 재현하는 데 기준선을 능가합니다.
• 모델은 일반적인 분할뿐만 아니라 희귀 모드도 적절한 빈도로 학습하고 재현할 수 있으며 보정된 모드 확률을 보여줍니다.
• 샘플당 네트워크의 일부분만 재평가하면 되므로 다수의 분할 가설 샘플링은 계산적으로 비용이 적습니다.
• 잠재 공간을 통해 분할 변형과 그 구조를 시각화하고 분석할 수 있습니다.
• 제약 조건에 따라 prior를 조정하고 잠재 특징을 적절한 네트워크 단계에 주입하는 것이 성능에 중요하다는 변형 실험 결과가 나옵니다.