[논문 리뷰] Automatic Segmentation of Head and Neck Tumor: How Powerful Transformers Are?

Ikboljon Sobirov, Otabek Nazarov외 2명|arXiv (Cornell University)|2022. 01. 17.

Radiomics and Machine Learning in Medical Imaging인용 수 8

한 줄 요약

이 논문은 다중 모odal CT/PET 스캔에서 자동으로 두경부 종양을 분할하기 위한 비전 트랜스포머 기반 모델(UNETR)을 평가하며, 최신의 CNN 기반 모델들과 비교한다. 완전히 새로 학습된 트랜스포머 모델은 평균 DSC 0.736을 달성하여 최고 수준의 CNN 모델들과 경쟁 가능한 성능을 보이며, 의료 영상 분할 분야에서 트랜스포머의 강력한 잠재력을 입증한다.

ABSTRACT

Cancer is one of the leading causes of death worldwide, and head and neck (H&N) cancer is amongst the most prevalent types. Positron emission tomography and computed tomography are used to detect, segment and quantify the tumor region. Clinically, tumor segmentation is extensively time-consuming and prone to error. Machine learning, and deep learning in particular, can assist to automate this process, yielding results as accurate as the results of a clinician. In this paper, we investigate a vision transformer-based method to automatically delineate H&N tumor, and compare its results to leading convolutional neural network (CNN)-based models. We use multi-modal data from CT and PET scans to perform the segmentation task. We show that a solution with a transformer-based model has the potential to achieve comparable results to CNN-based ones. With cross validation, the model achieves a mean dice similarity coefficient (DSC) of 0.736, mean precision of 0.766 and mean recall of 0.766. This is only 0.021 less than the 2020 competition winning model (cross validated in-house) in terms of the DSC score. On the testing set, the model performs similarly, with DSC of 0.736, precision of 0.773, and recall of 0.760, which is only 0.023 lower in DSC than the 2020 competition winning model. This work shows that cancer segmentation via transformer-based models is a promising research area to further explore.

연구 동기 및 목표

• 다중 모달리티 CT 및 PET 영상에서 비전 트랜스포머 기반 모델의 성능을 두경부 종양 자동 분할에 대해 조사한다.
• 분할 정확도 측면에서 트랜스포머 기반 모델과 최신 CNN 기반 모델(예: SE 기반 U-Net, nnU-Net)을 비교한다.
• 이 문제에서 트랜스포머 기반 모델 성능에 대한 데이터 증강의 영향을 평가한다.
• 기존의 응용 분야 외부의 다른 의료 영상 맥락에서 새로 제안된 UNETR 아키텍처의 일반화 능력과 강건성을 평가한다.
• 트랜스포머가 복잡하고 저대비, 잡음이 많은 의료 영상 분할 작업에서 CNN을 따라하거나 능가할 수 있는지 탐색한다.

제안 방법

• 3D ViT 인코더를 사용해 전역적 문맥 특징을 추출하고, U-Net 스타일 디코더를 활용해 분할을 수행하는 비전 트랜스포머 기반 아키텍처인 UNETR를 채택한다.
• 크로스 엔트로피 손실과 DSC 손실을 사용해 HECKTOR 데이터셋의 다중 모달리티 3D CT 및 PET 스캔에서 모델을 엔드 투 엔드로 학습시킨다.
• 일관성 향상을 위해 강도 변동, 회전, 반전, 탄성 변형 등의 광범위한 데이터 증강 기법을 적용한다.
• 다양한 임상 센터와 스캔 프로토콜을 고려해 안정적인 평가를 확보하기 위해 한 센터를 제외한 교차 검증 전략을 사용한다.
• 두 개의 선도적인 CNN 기반 모델인 SE 기반 U-Net(2020년 경쟁 대회 수상자)과 nnU-Net(자동 설정, 최신 U-Net 변종)과 결과를 비교한다.
• 예측 마스크를 진단 마스크와 비교해 정성적 분석을 수행하며, 저대비 종양, 소형 종양, CT 아티팩트와 관련된 실패 케이스를 식별한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머 기반 모델이 두경부 종양 분할에서 최신 CNN 기반 모델과 유사한 성능을 달성할 수 있는가?
• RQ2 이 의료 영상 작업에서 데이터 증강이 트랜스포머 기반 모델의 성능에 어떤 영향을 미치는가?
• RQ3 다중 모달리티 CT/PET 스캔에서 두경부 종양을 분할할 때 트랜스포머 기반 모델의 주요 실패 원인은 무엇인가?
• RQ4 UNETR의 성능은 다양한 임상 센터와 스캔 획득 프로토콜 간에 어떻게 일반화되는가?
• RQ5 다중 모달리티 CT 및 PET 입력이 종양 정확한 국소화 능력에 얼마나 기여하는가?

주요 결과

• 비전 트랜스포머 기반 UNETR 모델은 교차 검증에서 평균 DSC 0.736(±0.043)을 달성했으며, 이는 2020년 대회 수상자인 SE 기반 U-Net보다 0.021 낮은 수준이지만 경쟁 가능한 성능이다.
• 독립적인 HECKTOR 테스트 세트에서 모델은 DSC 0.736, 정밀도 0.773, 재현율 0.760을 기록했으며, DSC 기준으로 2020년 수상자와 0.023 이내로 유사한 성능을 보였다.
• 세 모델(UNETR, SE 기반 U-Net, nnU-Net) 모두 유사한 분할 마스크를 생성했으며, 대부분의 케이스에서 미세한 차이가 관찰되었다.
• 실패 케이스는 주로 PET 스캔에서의 저대비 종양, 소형 종양, CT 아티팩트(예: 치과 임플랜트로 인한 스트로크 아티팩트) 때문이었다.
• 모델들은 CT(解剖학적 구조 제공)와 PET(강도 차별화) 모두에 강하게 의존함을 보이며, 다중 모달리티 입력의 중요성을 시사했다.
• 자기 지도 사전 학습이 없는 상황에서 데이터 증강이 학습 안정성 향상과 일반화 능력 향상에 필수적임을 확인했다.