[논문 리뷰] Multi-organ Segmentation over Partially Labeled Datasets with Multi-scale Feature Abstraction

Xi Fang, Pingkun Yan|arXiv (Cornell University)|2020. 01. 01.

Advanced Neural Network Applications참고 문헌 54인용 수 8

한 줄 요약

이 논문은 부분적으로 레이블이 붙은 데이터셋을 사용한 다기관 분할을 위해 다중 척도 특징 추출 기반의 새로운 2D U-Net 기반 네트워크인 PIPO-FAN을 제안한다. 동일 깊이의 컨볼루션과 적응형 가중치를 통해 피라미드 입력 및 출력 특징을 통합하여 단일 스텝 추론과 함께 1개의 GPU에서 슬라이스당 0.04초의 상태최고 성능을 달성한다.

ABSTRACT

Shortage of fully annotated datasets has been a limiting factor in developing deep learning based image segmentation algorithms and the problem becomes more pronounced in multi-organ segmentation. In this paper, we propose a unified training strategy that enables a novel multi-scale deep neural network to be trained on multiple partially labeled datasets for multi-organ segmentation. In addition, a new network architecture for multi-scale feature abstraction is proposed to integrate pyramid input and feature analysis into a U-shape pyramid structure. To bridge the semantic gap caused by directly merging features from different scales, an equal convolutional depth mechanism is introduced. Furthermore, we employ a deep supervision mechanism to refine the outputs in different scales. To fully leverage the segmentation features from all the scales, we design an adaptive weighting layer to fuse the outputs in an automatic fashion. All these mechanisms together are integrated into a Pyramid Input Pyramid Output Feature Abstraction Network (PIPO-FAN). Our proposed method was evaluated on four publicly available datasets, including BTCV, LiTS, KiTS and Spleen, where very promising performance has been achieved. The source code of this work is publicly shared at https://github.com/DIAL-RPI/PIPO-FAN for others to easily reproduce the work and build their own models with the introduced mechanisms.

연구 동기 및 목표

• 다기관 분할에서 완전히 레이블이 붙은 의료 영상 데이터셋의 제한 문제를 해결한다.
• 다양한 부분적으로 레이블이 붙은 데이터셋을 활용해 일반화 능력을 향상시키기 위한 통합 학습 전략을 개발한다.
• 계층적으로 다중 척도 입력 및 출력 특징을 통합하여 딥 네트워크의 특징 추출 능력을 향상시킨다.
• 다중 스텝 3D 방법과 비교해 학습 및 추론 비용을 줄이면서도 분할 정확도를 유지하거나 향상시킨다.
• 다양한 레이블 커버리지로 인해 이질적인 데이터셋에서도 엔드 투 엔드 학습이 가능하게 하여 더 넓은 임상 적용 가능성을 확보한다.

제안 방법

• 다기관 분할을 위한 U-형 아키텍처를 가진 피라미드 입력 및 피라미드 출력(PIPO) 특징 추출 네트워크(PIPO-FAN)를 제안한다.
• 입력 이미지의 다양한 척도 수준에 대해 병렬 컨볼루션을 적용하는 다중 척도 입력 처리 모듈을 도입한다.
• 특징 맵을 척도 간에 정렬하고 스킵 연결에서의 의미적 갭을 줄이기 위해 동일 깊이의 컨볼루션 메커니즘을 활용한다.
• 다양한 척도에서 예측을 정밀하게 보정하기 위해 디코더 단계의 다중 단계에서 심층 감독을 적용한다.
• 특징 중요도에 따라 다중 척도 출력을 동적으로 융합하기 위해 주의 메커니즘을 사용한 적응형 가중치 레이어를 설계한다.
• BTCV, LiTS, KiTS, Spleen의 네 개의 공개 데이터셋에서 부분 레이블이 붙은 데이터를 동일한 전략으로 엔드 투 엔드로 모델을 학습시킨다.

실험 결과

연구 질문

• RQ1 U-Net 아키텍처에서 다중 척도 특징 추출이 부분적으로 레이블이 붙은 의료 영상 데이터셋에서의 분할 정확도 향상에 기여하는가?
• RQ2 동일 깊이의 컨볼루션을 사용한 피라미드 입력 및 출력 특징 통합이 스킵 연결에서의 의미적 갭을 어떻게 줄이는가?
• RQ3 주의 메커니즘을 활용한 적응형 특징 융합이 여러 기관에 걸쳐 분할 성능에 어떤 영향을 미치는가?
• RQ4 단일 2D 네트워크가 다중 스텝 3D 네트워크에 비해 더 낮은 계산 비용으로 경쟁 가능한 성능을 달성할 수 있는가?
• RQ5 다양한 부분적으로 레이블이 붙은 데이터셋에서 학습함으로써 다기관 분할에서 모델의 강건성과 일반화 능력은 어떻게 향상되는가?

주요 결과

• PIPO-FAN은 LiTS 챌린지 테스트 세트에서 96.1%의 Dice 스코어를 기록하여 최상위 성능을 보인 메서드(DeepX)와 동일한 성능을 내지만 단일 추론 스텝만을 사용한다.
• 1개의 Titan Xp GPU에서 CT 슬라이스를 0.04초 내로 분할하여, 비교적 높은 성능을 내는 동시에 보고된 바 중 가장 빠른 속도를 기록한다.
• 기존 2D 및 3D 네트워크보다 대칭 표면 거리(SSD: 1.413 vs. 1.450)와 평균 SSD(24.408 vs. 27.118)에서 뛰어난 성능을 보인다.
• 5개의 입력 및 출력 척도를 사용할 경우 최고의 성능를 기록하였으며, 3 척도(94.5%)에서 5 척도(95.6%)로 증가함에 따라 Dice 스코어가 상승하여 더 향상된 문맥 기반 특징 추출 능력을 입증한다.
• 여러 부분적으로 레이블이 붙은 데이터셋에서 통합된 학습 전략을 통해 다양한 해부학적 영역에서 모델의 강건성과 일반화 능력이 크게 향상된다.
• 소스 코드는 GitHub에 공개되어 있어 재현 가능성이 보장되며, 다른 영상 모odal리티 및 데이터셋으로의 확장도 가능하다.