[논문 리뷰] Bottleneck Supervised U-Net for Pixel-wise Liver and Tumor Segmentation

Song Li, Kwok Fai Tso|arXiv (Cornell University)|2018. 10. 16.

Advanced Neural Network Applications참고 문헌 32인용 수 14

한 줄 요약

이 논문은 CT 영상에서 간 및 종양의 정확한 분할을 위한 Bottleneck Supervised U-Net (BS U-Net)을 제안한다. 지도 학습을 통해 진단 레이블의 지도를 이용해 해부학적 특징을 추출하는 인코딩 U-Net을 훈련하고, 이 특징을 가중치가 부여된 덱스 손실과 MSE 손실을 통해 분할 U-Net을 지도하는 방식으로, 형태 왜곡과 가짜 양성/음성 반응을 감소시켜 LiTS 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

In this paper, we propose a bottleneck supervised (BS) U-Net model for liver and tumor segmentation. Our main contributions are: first, we propose a variation of the original U-Net that incorporates dense modules, inception modules and dilated convolution in the encoding path; second, we propose a bottleneck supervised (BS) U-Net that contains an encoding U-Net and a segmentation U-Net. To train the BS U-Net, the encoding U-Net is first trained to get encodings of the label maps that contain the anatomical information (shape and location). Subsequently, this information is used to guide the training of the segmentation U-Net so as to reserve the anatomical features of the target objects. More specifically, the loss function for segmentation U-Net is set to be the weighted average of the dice loss and the MSE loss between the encodings and the bottleneck feature vectors. The model is applied to a public liver and tumor CT scan dataset. Experimental results show that besides achieving excellent overall segmentation performance, BS U-Net also works great in controlling shape distortion, reducing false positive and false negative cases.

연구 동기 및 목표

• 형태 왜곡과 높은 가짜 양성/음성 비율로 인해 CT 영상에서 간 및 종양 분할이 정확하지 않은 문제를 해결하기 위해.
• 레이블 맵에서 유도된 해부학적 사전 지식을 훈련 과정에 통합하여 분할 성능을 향상시키기 위해.
• 특징 학습과 분할을 분리하여 보다 우수한 일반화 성능을 달성하기 위한 이중 단계 딥 러닝 프레임워크를 개발하기 위해.
• 데이터 증강 및 수동 후처리에 대한 의존도를 줄이기 위해 손실 함수에 직접적으로 구조적 사전 지식을 통합하기 위해.

제안 방법

• 인코딩 경로에 수정된 U-Net 아키텍처를 도입하여 밀집 모듈, 인셉션 모듈 및 확장 컨볼루션을 통합해 특징 추출 능력을 향상시켰다.
• 이중 브랜치 네트워크 구조를 제안한다: 먼저 진단 레이블 맵으로 훈련된 인코딩 U-Net을 통해 해부학적 특징을 학습하고, 그 다음 분할 U-Net을 훈련한다.
• 분할 U-Net은 예측된 블로킹 특징과 사전 훈련된 인코딩 U-Net의 특징 간의 덱스 손실과 평균 제곱 오차(MSE) 손실을 조합한 하이브리드 손실 함수를 사용하여 훈련된다.
• 인코딩 U-Net의 블로킹 특징 벡터는 분할 출력의 형태와 공간 일관성을 유지하기 위한 지도 신호 역할을 한다.
• 데이터 전처리로는 자르기, 정사각형으로 팞딩, 224×224로 크기 조정을 통해 입력 표준화를 수행하며, 랜덤 스케일링 및 회전을 통한 데이터 증강을 실시한다.
• 훈련 파이프라인은 코사인 감쇠 학습률 스케줄과 50 에포크 동안 배치 크기 20를 사용하며, 검증 손실 기반 조기 정지 전략을 적용한다.

실험 결과

연구 질문

• RQ1 레이블 맵 인코딩에서 유도된 해부학적 사전 지식을 통합함으로써 간 및 종양 분할의 정확도를 향상시키고 형태 왜곡을 줄일 수 있는가?
• RQ2 먼저 인코딩 U-Net을 사전 훈련하고, 그 블로킹 특징을 사용해 분할 U-Net을 지도하는 이중 단계 훈련 전략이 기존 U-Net 훈련 방식보다 더 우수한 성능을 낼 수 있는가?
• RQ3 하이브리드 손실 함수(블로킹 특징에 대한 덱스 + MSE)는 기존의 덱스 손실에 비해 가짜 양성 및 음성 사례를 얼마나 효과적으로 줄이는가?
• RQ4 제안된 BS U-Net은 기준 U-Net 모델에 비해 다양한 네트워크 아키텍처와 데이터셋에 대해 얼마나 잘 일반화되는가?
• RQ5 광범위한 데이터 증강이나 후처리 없이도 공개 벤치마크인 LiTS 챌린지에서 기존 모델보다 뛰어난 성능을 달성할 수 있는가?

주요 결과

• BS U-Net은 LiTS 테스트 세트에서 케이스 기반 덱스(DPC) 0.9610과 글로벌 덱스(DG) 0.9640을 기록하여 원본 U-Net(DPC: 0.957, DG: 0.960)을 능가했다.
• 모델은 볼륨 오버랩 오차(VOE)를 0.075로 줄였고, 상대 볼륨 차이(RVD)는 0.018로 줄여 기존 모델 대비 뛰어난 형태 유지 능력을 보였다.
• 1C BS U-Net의 평균 표면 거리(MSD)는 47.217 mm로 기존 U-Net의 61.235 mm보다 유의미하게 낮아 표면 정확도 향상을 확인했다.
• BS U-Net은 테스트된 모든 모델 중에서 가장 낮은 3.831 RSSD를 기록하여 케이스 간 분할 일관성에서 뛰어난 성능을 보였다.
• 시각적 비교(그림 7)는 BS U-Net이 더 매끄럽고 해부학적으로 타당한 분할 결과를 도출하며, 가짜 양성 및 음성 반응이 적다는 것을 확인했다.
• 훈련 손실 곡선(그림 9)은 안정적인 수렴을 보였으며, BS U-Net의 덱스 손실은 점진적으로 감소하고, 블로킹 특징 간 MSE 손실은 안정화되어 효과적인 특징 지도 학습이 이루어졌음을 확인했다.