[논문 리뷰] Universal Segmentation of 33 Anatomies

Pengbo Liu, Yang Deng외 13명|arXiv (Cornell University)|2022. 03. 04.

Medical Imaging and Analysis인용 수 6

한 줄 요약

이 논문은 7개의 부분적으로 레이블이 붙은 데이터셋, 특히 1,003개의 3D CT 영상과 11,000개 이상의 추정 척추 구조물 레이블을 포함한 새로운 대규모 CTSpine1K 데이터셋을 기반으로 훈련된 유니버설 33해부학 분할 모델을 제안한다. 패치 기반 훈련에서 제한된 맥락을 보완하기 위해 장거리 특징 집합을 향상시키는 크로스패치 트랜스포머 모듈(CPTM)을 도입하여, 미세조정 후 Verse19 및 Verse20과 같은 오픈 벤치마크에서 최신 기술 수준(SOTA)의 일반화 성능을 달성한다.

ABSTRACT

In the paper, we present an approach for learning a single model that universally segments 33 anatomical structures, including vertebrae, pelvic bones, and abdominal organs. Our model building has to address the following challenges. Firstly, while it is ideal to learn such a model from a large-scale, fully-annotated dataset, it is practically hard to curate such a dataset. Thus, we resort to learn from a union of multiple datasets, with each dataset containing the images that are partially labeled. Secondly, along the line of partial labelling, we contribute an open-source, large-scale vertebra segmentation dataset for the benefit of spine analysis community, CTSpine1K, boasting over 1,000 3D volumes and over 11K annotated vertebrae. Thirdly, in a 3D medical image segmentation task, due to the limitation of GPU memory, we always train a model using cropped patches as inputs instead a whole 3D volume, which limits the amount of contextual information to be learned. To this, we propose a cross-patch transformer module to fuse more information in adjacent patches, which enlarges the aggregated receptive field for improved segmentation performance. This is especially important for segmenting, say, the elongated spine. Based on 7 partially labeled datasets that collectively contain about 2,800 3D volumes, we successfully learn such a universal model. Finally, we evaluate the universal model on multiple open-source datasets, proving that our model has a good generalization performance and can potentially serve as a solid foundation for downstream tasks.

연구 동기 및 목표

• 다양한 해부학적 기관의 분할을 위한 대규모로 완전히 레이블링된 3D 의료 영상 데이터셋의 부족 문제를 해결하기 위해.
• 다양한 데이터셋에서 척추, 골반 뼈, Abdominal 기관 등 다양한 해부학적 기관을 동시에 분할할 수 있는 단일 유니버설 모델을 개발하기 위해.
• CTSpine1K를 구축하고 공개하여 척추 분할 분야의 데이터 격차를 해소하고, 1,000개 이상의 3D CT 영상과 11,000개 이상의 척추 레이블을 포함한 대규모 오픈소스 데이터셋을 제공하기 위해.
• 특히 기립형 구조물인 척추와 같은 기관의 분할에서 패치 기반 접근 방식의 국한된 맥락을 보완하기 위해, 새로운 크로스패치 트랜스포머 모듈(CPTM)을 통해 장거리 맥락 모델링을 향상시키기 위해.
• 미리 훈련된 모델이 새로운 오픈 데이터셋에서 강력한 일반화 성능을 보이며, 향후 작업을 위한 기초 모델로서의 잠재력을 입증하기 위해.

제안 방법

• 7개의 부분적으로 레이블이 붙은 데이터셋에서 다양한 출처의 데이터를 조합하여, 가짜 레이블링 전략을 사용해 유니버설 모델을 훈련한다.
• 이웃한 3D 패치 간의 특징을 융합하기 위해 크로스패치 트랜스포머 모듈(CPTM)을 도입하여, 효과적인 수신 영역을 확장하고 장거리 맥락 학습을 향상시킨다.
• CPTM는 세 개의 겹치는 3D 패치에서 작동하며, 자기주의 기반 메커니즘을 활용해 패치 경계를 넘어서 정보를 집합적으로 처리함으로써 기립형 기관의 특징 표현을 강화한다.
• 모델는 nnU-Net 프레임워크를 기반으로 구축되어, 3D 의료 영상 분할 작업에 있어 뛰어난 안정성과 적응성을 확보한다.
• GPU 메모리 제약으로 인해 훈련 과정에서 3D 패치를 잘라서 사용하며, CPTM는 전역 맥락의 손실을 보완하도록 특별히 설계되었다.
• 유니버설 모델는 재학습 없이 여러 오픈 데이터셋에서 평가되며, 목표 데이터셋에서의 미세조정을 통해 이식 가능성과 성능 향상을 평가한다.

실험 결과

연구 질문

• RQ1 단일 딥러닝 모델이 여러 부분적으로 레이블이 붙은 데이터셋에서 33개의 다양한 해부학적 기관에 대해 강력하고 유니버설한 분할 성능을 달성할 수 있는가?
• RQ2 기존의 표준 패치 기반 접근 방식에 비해 크로스패치 트랜스포머 모듈(CPTM)이 척추와 같은 기립형 해부학적 기관의 분할 정확도 향상에 얼마나 효과적인가?
• RQ3 부분적으로 레이블이 붙은 데이터셋의 결합에서 미리 훈련된 모델이 Amos, CLINIC, Verse19, Verse20와 같은 오픈 벤치마크로의 일반화 능력은 어느 정도인가?
• RQ4 최근에 구축된 CTSpine1K 데이터셋이 기존의 소규모 척추 데이터셋에 비해 척추 분할 모델의 성능과 일반화 능력을 얼마나 향상시키는가?
• RQ5 유니버설 모델은 목표 데이터셋에서의 미세조정 후에도 강력한 기초 모델로 기능할 수 있는가, 특히 후속 작업에 대해?

주요 결과

• 모델은 미세조정 후 Verse19 테스트 세트에서 딱시 스코어 0.904와 HD95 10.28을 기록하며, 복잡한 3단계 설계를 가진 SOTA 방법을 초월한다.
• Verse20 테스트 세트에서는 딱시 스코어 0.866과 HD95 18.03을 기록하여, 다른 척추 분할 벤치마크에서도 강력한 일반화 성능을 입증한다.
• 기본 모델 대비 Amos 데이터셋에서 딱시 스코어 2.5점 향상 및 HD95 35.7% 향상되었고, CLINIC 데이터셋에서는 각각 2.0점과 23.5% 향상되었다.
• 크로스패치 트랜스포머 모듈(CPTM)은 척추 구조물 식별률을 크게 향상시켰으며, 미세조정 후 Verse19 검증 세트에서 99.7%의 식별률을 기록하여 구조적 인식 능력 향상을 시사한다.
• Amos 데이터셋에서는 93.4%의 척추 식별률, Verse19 테스트 세트에서는 96.8%의 식별률을 기록하여 다양한 해부학적 구성에서도 뛰어난 성능을 보였다.
• 1,003개의 3D CT 영상과 11,000개 이상의 레이블을 포함한 CTSpine1K 데이터셋은 척추 분할 분야의 핵심 데이터 격차를 메우며, 커뮤니티 지원을 위해 오픈소스로 공개되었다.