[논문 리뷰] A two-stage 3D Unet framework for multi-class segmentation on full resolution image

Chengjia Wang, Tom MacGillivray외 3명|arXiv (Cornell University)|2018. 04. 12.

Advanced Image Processing Techniques참고 문헌 16인용 수 60

한 줄 요약

이 논문은 ROI를 먼저 국지화한 다음 재샘플링 없이 전체 해상도 3D 의료 영상에서 다중 클래스 분할을 수행하는 두 단계 연결된 U-Net 프레임워크를 제시합니다. MM-WHS 데이터에서 최신 U-Net보다 우수한 성능을 달성합니다.

ABSTRACT

Deep convolutional neural networks (CNNs) have been intensively used for multi-class segmentation of data from different modalities and achieved state-of-the-art performances. However, a common problem when dealing with large, high resolution 3D data is that the volumes input into the deep CNNs has to be either cropped or downsampled due to limited memory capacity of computing devices. These operations lead to loss of resolution and increment of class imbalance in the input data batches, which can downgrade the performances of segmentation algorithms. Inspired by the architecture of image super-resolution CNN (SRCNN) and self-normalization network (SNN), we developed a two-stage modified Unet framework that simultaneously learns to detect a ROI within the full volume and to classify voxels without losing the original resolution. Experiments on a variety of multi-modal volumes demonstrated that, when trained with a simply weighted dice coefficients and our customized learning procedure, this framework shows better segmentation performances than state-of-the-art Deep CNNs with advanced similarity metrics.

연구 동기 및 목표

• 다운샘플링이나 후처리 없이 높은 해상도 3D 전심(segmentation) 전체를 목표로 한다.
• 먼저 ROI를 국지화하고 그 후 전체 해상도에서 볼륨 단위 라벨을 정제하는 두 단계 DCNN을 개발한다.
• Skip 연결이 있는 근접 SRCNN에서 영감 받은 정제 기능을 활용하여 분할 정확도를 향상시킨다.
• MM-WHS 챌린지의 다중 모달 심장 CT/MR 데이터에서 효과를 입증한다.
• 최신 3D U-Net과 비교하고 제한된 데이터로도 좋은 성능을 달성하기 위한 학습 전략을 평가한다.

제안 방법

• Net1이 다운샘플링된 부피에서 거칠고 ROI 중심의 예측을 제공하고 Net2가 전체 해상도에서 라벨을 정제하는 두 개의 연결된 U-Net 아키텍처.
• Net1은 수용 영역을 확장하기 위해 수축 경로에서 증가하는 확장을 갖는 5x5x5 컨볼루션을 사용한다.
• Net2는 Skip 연결과 재귀적 구조를 갖춘 SRCNN 영감 설계를 따르며 입력으로 Net1 출력과 원본 데이터를 모두 사용한다.
• 4단계 학습 절차에는 (1) 전체 부피를 이용한 ROI 국지화, (2) 부분 부피를 사용한 거친 다중 클래스 분할, (3) 끝에서 끝까지 학습하는 공동 거친+정밀 분할, (4) 축 슬라이스의 하위 볼륨(K 슬라이스)에서 Net2를 미세 조정이 포함됩니다.
• (Eq. 2)에 따른 가중 다중 클래스 Dice 손실과 ROI 국지화를 위한 전경 초점 Dice 손실(Eq. 3–4)을 사용하여 학습을 안내하며, 학습률 0.0001의 Adam 옵티마이저를 사용합니다.
• 데이터 증강과 계층적/서브볼륨 샘플링은 데이터 불균형과 메모리 제약을 해결하고 배치 설계는 저주파 클래스의 균형을 맞추기 위해 중앙 볼륨 클래스에 중점을 둡니다.

실험 결과

연구 질문

• RQ1 이것은 후처리나 재샘플링 없이 단일 샷 3D U-Net보다 성능이 높은 전체 해상도 3D 분할 프레임워크인가?
• RQ2 ROI 기반 로컬라이제이션과 정제 단계의 결합이 고해상도 심장 CT/MR 데이터에서 다중 클래스 분할을 향상시키는가?
• RQ3 단계적 학습 절차와 특수한 손실 함수가 MM-WHS 데이터에서 수렴성과 정확도에 어떤 영향을 미치는가?
• RQ4 제안된 프레임워크가 CT와 MR 모달리티 모두에 대해 MM-WHS 벤치마크에서 최신 U-Net과 경쟁력이 있는가?

주요 결과

• Net2가 CT 구조들에 대해 기본 3D U-Net보다 Dice 및 Jaccard 점수를 더 높게 달성했습니다.
• Net1은 상대적으로 다운샘플링된 U-Net에 비해 여러 구조에서 개선을 제공했으며 Net2와 함께 사용될 때 특히 그렇습니다.
• MM-WHS 테스트 데이터에서 프레임워크는 매우 경쟁력 있는 MR 결과를 보였고(평균 Dice 약 0.83), CT 결과는 베이스라인 대비 현저한 개선을 보였으며 추가 후처리 없이도 대회 우승자에 근접한 성능을 더 적은 학습 반복으로 달성했습니다.
• 전체 해상도에서 효과적인 분할을 보여 주며, 고해상도 3D 의료 영상 분할을 위한 ROI-주도 이단계 접근 방식을 검증합니다.
• 거칠은 손실과 정밀 손실을 결합한 엔드투엔드 학습은 수렴을 가속하고 개별 네트워크를 학습하는 것보다 성능을 향상시킵니다.