[논문 리뷰] Conditional Generative Refinement Adversarial Networks for Unbalanced Medical Image Semantic Segmentation

Mina Rezaei, Haojin Yang외 1명|arXiv (Cornell University)|2018. 10. 09.

AI in cancer detection참고 문헌 4인용 수 22

한 줄 요약

이 논문은 의료 영상 의미 분할에서의 클래스 불균형 문제를 해결하기 위해 생성망, 구분망, 정제망을 결합한 조건부 생성 보정 적대망(cG-RAN)을 제안한다. 실제 및 합성 분할 결과를 함께 학습하고, 잠재적 거짓 양성 및 거짓 음성 예측을 명시적으로 모델링함으로써, 간 병변, 뇌 종양, 미세세포 분할 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성하며, 정밀도를 훼손하지 않은 채 민감도를 크게 향상시킨다.

ABSTRACT

We propose a new generative adversarial architecture to mitigate imbalance data problem in medical image semantic segmentation where the majority of pixels belongs to a healthy region and few belong to lesion or non-health region. A model trained with imbalanced data tends to bias toward healthy data which is not desired in clinical applications and predicted outputs by these networks have high precision and low sensitivity. We propose a new conditional generative refinement network with three components: a generative, a discriminative, and a refinement network to mitigate unbalanced data problem through ensemble learning. The generative network learns to a segment at the pixel level by getting feedback from the discriminative network according to the true positive and true negative maps. On the other hand, the refinement network learns to predict the false positive and the false negative masks produced by the generative network that has significant value, especially in medical application. The final semantic segmentation masks are then composed by the output of the three networks. The proposed architecture shows state-of-the-art results on LiTS-2017 for liver lesion segmentation, and two microscopic cell segmentation datasets MDA231, PhC-HeLa. We have achieved competitive results on BraTS-2017 for brain tumour segmentation.

연구 동기 및 목표

• 건강한 조직이 지배하고 병변이 희귀한 의료 영상 의미 분할에서의 클래스 불균형 문제를 해결한다.
• 임상적 맥락에서 병변을 놓치는 데 높은 비용이 수반되므로, 정밀도를 저하시키지 않은 채 모델의 민감도를 향상시킨다.
• 생성망, 구분망, 정제망의 세 네트워크를 통합한 앙상블 학습 기반의 딥 러닝 프레임워크를 개발한다.
• 정제망을 통해 거짓 양성 및 거짓 음성 예측을 명시적으로 모델링함으로써 다수 클래스에 대한 편향을 완화한다.
• 3D CT, MRI 및 미세세포 영상과 같은 다양한 의료 영상 작업에 대해 일반화 능력을 입증한다.

제안 방법

• 생성망 네트워크는 2차원 다중 모odal 영상 슬라이스를 조건으로 하여 픽셀 수준의 분할 마스크를 생성하며, 슬라이스 간 시간적 의존성을 모델링하기 위해 이방향 LSTMs를 사용한다.
• 구분망은 생성망의 출력이 실제 진짜 마스크와 얼마나 유사한지 평가하며, 최소 최대 학습 목표를 사용하여 실제 진짜 마스크와 생성 마스크를 구분한다.
• 정제망은 생성망의 출력에서 거짓 양성 및 거짓 음성 마스크를 예측하도록 학습되어, 잘못 분류된 영역을 보정할 수 있다.
• 최종 분할 마스크는 생성망의 출력과 정제망의 보정을 조합하여 형성되며, 이로 인해 정밀도와 재현율이 모두 향상된다.
• 작은 데이터셋에서 특히 안정적인 학습을 위해 환자 기반 미니배치 정규화를 적용한다. 이는 동일한 환자로부터의 모든 영상에 대해 정규화를 수행한다.
• 모델은 적대적 손실과 분할 품질 지표(예: Dice, Hausdorff 거리)에 중점을 두고 조합된 손실 함수를 사용하여 엔드 투 엔드로 학습된다.

실험 결과

연구 질문

• RQ1 전용 정제망을 갖춘 조건부 GAN이 정밀도를 저하시키지 않은 채 불균형한 의료 영상 분할에서 민감도를 향상시킬 수 있는가?
• RQ2 생성망, 구분망, 정제망의 세 개의 별도 네트워크 간 앙상블 학습이 의료 영상 분할에서 다수 클래스에 대한 편향을 어떻게 완화하는가?
• RQ3 환자 기반 미니배치 정규화가 제한된 의료 데이터셋에서 학습할 때 모델 성능에 어떤 영향을 미치는가?
• RQ4 거짓 양성 및 거짓 음성 예측 헤드의 포함이 희귀한 해부학적 또는 병리학적 구조의 분할 품질에 실제로 유의미한 향상을 가져오는가?
• RQ5 이 아키텍처는 CT, MRI 및 미세세포 영상과 같은 다양한 의료 영상 작업에 얼마나 잘 일반화되는가?

주요 결과

• 간 병변 분할을 위한 LiTS-2017 데이터셋에서, 이 방법은 최신 기술 수준 성능을 달성하며 기준 모델 대비 민감도를 크게 향상시켰다.
• PhC-HeLa 미세세포 데이터셋에서 cGAN+Refinement 모델은 교차율(IoU) 0.951을 기록하여 MISS-GAN(0.943)과 U-Net(0.92)를 능가했다.
• MDA231 데이터셋에서, 이 방법은 IoU 0.93, 민감도 0.92를 달성하여 U-Net 및 KTH-SE 모델을 초월했다.
• 정제망은 거짓 양성 및 거짓 음성 비율을 효과적으로 감소시켰으며, MDA231에서 FPR과 FNR이 각각 0.07과 0.08로 떨어졌다.
• 환자 기반 미니배치 정규화는 MDA231 및 PhC-HeLa와 같은 작은 데이터셋에서 성능을 향상시켰으며, 데이터 증강과 조합할 경우 尤히 유의미했다.
• 학습 중 가우시안 노이즈를 추가하면 작은 데이터셋에서는 성능에 악영향을 미쳤지만, BraTS-2017와 같은 큰 데이터셋에서는 영향이 미미했다.