[논문 리뷰] CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection

Jie Liu, Yixiao Zhang외 8명|arXiv (Cornell University)|2023. 01. 02.

Radiomics and Machine Learning in Medical Imaging인용 수 18

한 줄 요약

텍스트 임베딩을 통해 해부학적 관계를 인코딩하는 CLIP 기반의 범용 모델을 도입하여 25개 기관을 분할하고 14개 공개 데이터셋에서 6개 종양 유형을 탐지하며 최첨단 성능과 강한 일반화를 달성한다.

ABSTRACT

An increasing number of public datasets have shown a marked impact on automated organ segmentation and tumor detection. However, due to the small size and partially labeled problem of each dataset, as well as a limited investigation of diverse types of tumors, the resulting models are often limited to segmenting specific organs/tumors and ignore the semantics of anatomical structures, nor can they be extended to novel domains. To address these issues, we propose the CLIP-Driven Universal Model, which incorporates text embedding learned from Contrastive Language-Image Pre-training (CLIP) to segmentation models. This CLIP-based label encoding captures anatomical relationships, enabling the model to learn a structured feature embedding and segment 25 organs and 6 types of tumors. The proposed model is developed from an assembly of 14 datasets, using a total of 3,410 CT scans for training and then evaluated on 6,162 external CT scans from 3 additional datasets. We rank first on the Medical Segmentation Decathlon (MSD) public leaderboard and achieve state-of-the-art results on Beyond The Cranial Vault (BTCV). Additionally, the Universal Model is computationally more efficient (6x faster) compared with dataset-specific models, generalized better to CT scans from varying sites, and shows stronger transfer learning performance on novel tasks.

연구 동기 및 목표

• 다수의 공개 공개 복부 데이터셋에서 부분 라벨링 및 라벨 불일치를 해결한다.
• CLIP 텍스트 임베딩을 활용해 해부학적 관계를 인코딩하여 분할 및 종양 탐지를 향상시킨다.
• 다양한 백본을 지원하고 미세 조정 없이 외부 CT 스캔에 일반화되는 범용 모델 아키텍처를 개발한다.
• MSD와 BTCV 벤치마크에서 최첨단 성능을 입증하고 다운스트림 태스크로의 전이 가능성을 제공한다.

제안 방법

• 25개 부분 주석 기관과 6개 종양을 포함한 총 3,410개의 CT 스캔으로 구성된 14개의 공개 데이터셋을 수집한다.
• 해부학 프롬프트로 CLIP 기반 텍스트 임베딩을 사용하여 세분화를 안내하고 원-핫 레이블을 대체한다.
• CLIP 임베딩 및 글로벌 이미지 특징에 조건부로 MLP를 통해 클래스별 매개변수를 생성하는 텍스트 분기를 구성한다.
• 표준화된 전처리로 CT 스캔을 처리하고 비전 인코더를 적용한 후 1x1x1 conv 기반 텍스트 주도 세그먼트를 적용한다.
• 부분 라벨 데이터에서 사용 가능한 라벨이 있는 클래스에만 학습하도록 마스킹된 역전파를 구현한다.
• 종양 탐지를 위해 Dice 유사도 계수(DSC), 정규화 표면 거리(NSD), 민감도, 특이도 및 조화 평균으로 평가한다.

실험 결과

연구 질문

• RQ1 CLIP 기반 라벨 임베딩이 부분적으로 라벨링된 데이터셋들에서 다기관 분할 및 종양 탐지를 개선하기 위해 해부학적 관계를 포착할 수 있는가?
• RQ2 텍스트 프롬프트를 이용한 범용 모델이 데이터셋별 모델에 비해 외부 CT 데이터와 다양한 스캐너에 대해 더 잘 일반화되는가?
• RQ3 부분 라벨과 데이터 불일치 상황에서 마스킹 역전파가 학습에 어떤 영향을 미치는가?
• RQ4 다양한 의학 프롬프트가 CLIP 임베딩의 분할 성능에 미치는 영향은 무엇인가?
• RQ5 범용 모델이 다운스트림 의학 영상 작업에 효과적인 파운데이션 모델로 작용할 수 있는가?

주요 결과

• CLIP-드리븐 범용 모델이 MSD와 BTCV 벤치마크에서 최고 성능을 달성한다.
• 모델은 6개 기관의 분할 정확도가 높아 동관찰자 인간 가변성과 일치한다.
• 종양 탐지는 높은 조화 평균을 달성하여 데이터셋에 걸쳐 높은 민감도와 낮은 거짓 양성을 나타낸다.
• 데이터셋별 모델에 비해 추론 속도가 약 6배 빨라 우수한 효율성을 보여준다.
• 이 프레임워크는 추가 조정 없이 다양한 병원의 CT 스캔에 일반화되며 백본은 CNN과 Transformer 간에 교체 가능하다.
• 외부 검증에서 3D-IRCADb 및 JHH 데이터셋에서 강력한 일반화를 보여주고 기존 방법을 능가한다.