[논문 리뷰] U-Netmer: U-Net meets Transformer for medical image segmentation

Sheng He, Rina Bao외 2명|arXiv (Cornell University)|2023. 04. 03.

Radiomics and Machine Learning in Medical Imaging인용 수 6

한 줄 요약

U-Netmer는 의료 영상 분할을 위한 U-Net과 Transformer 아키텍처의 새로운 융합을 제안하며, '토큰 평탄화' 및 '스케일 민감성' 문제를 해결하기 위해 국소 패치에 U-Net을 적용하고, 전역적 맥락을 모델링하기 위해 Transformer를 사용한다. 이는 7개의 데이터셋과 모odalities에서 최신 기술 수준의 성능을 달성하며, 다양한 스케일 간 분할 불일치는 지도 없이도 어려운 테스트 이미지를 순위 매기는 신뢰도 있는 신뢰도 점수로 기능한다.

ABSTRACT

The combination of the U-Net based deep learning models and Transformer is a new trend for medical image segmentation. U-Net can extract the detailed local semantic and texture information and Transformer can learn the long-rang dependencies among pixels in the input image. However, directly adapting the Transformer for segmentation has ``token-flatten" problem (flattens the local patches into 1D tokens which losses the interaction among pixels within local patches) and ``scale-sensitivity" problem (uses a fixed scale to split the input image into local patches). Compared to directly combining U-Net and Transformer, we propose a new global-local fashion combination of U-Net and Transformer, named U-Netmer, to solve the two problems. The proposed U-Netmer splits an input image into local patches. The global-context information among local patches is learnt by the self-attention mechanism in Transformer and U-Net segments each local patch instead of flattening into tokens to solve the `token-flatten" problem. The U-Netmer can segment the input image with different patch sizes with the identical structure and the same parameter. Thus, the U-Netmer can be trained with different patch sizes to solve the ``scale-sensitivity" problem. We conduct extensive experiments in 7 public datasets on 7 organs (brain, heart, breast, lung, polyp, pancreas and prostate) and 4 imaging modalities (MRI, CT, ultrasound, and endoscopy) to show that the proposed U-Netmer can be generally applied to improve accuracy of medical image segmentation. These experimental results show that U-Netmer provides state-of-the-art performance compared to baselines and other models. In addition, the discrepancy among the outputs of U-Netmer with different scales is linearly correlated to the segmentation accuracy which can be considered as a confidence score to rank test images by difficulty without ground-truth.

연구 동기 및 목표

• 시각 기반 Transformer에서 1D 토큰화 과정 중 국소 패치 간 상호작용이 손실되는 '토큰 평탄화' 문제를 해결하기 위해.
• 고정된 패치 크기로 인해 분할 정확도가 크게 변동하는 패치 기반 Transformer의 '스케일 민감성' 문제를 극복하기 위해.
• 다양한 패치 스케일을 지원하면서도 공유 파라미터를 사용하는 통합 모델 구조를 개발하기 위해.
• 다양한 스케일에서의 분할 일관성 기반으로 테스트 이미지에 대한 신뢰도 점수를 생성하여 지도 없이도 어려운 케이스를 순위 매길 수 있도록 하기 위해.
• 다양한 의료 영상 모달리티와 기관에서 최신 기술 수준의 분할 성능를 입증하기 위해.

제안 방법

• 입력 영상은 여러 스케일(s = 1, 2, 4, 8)에서 겹치지 않는 국소 패치로 분할되며, s=1은 전체 영상이다.
• 각 국소 패치는 U-Net 기반 구조를 사용하여 독립적으로 분할되며, 패치 내부의 공간적 상호작용을 유지한다.
• 전역 맥락은 모든 국소 패치의 임bedding을 고려하는 Transformer 인코더를 통해 모델링된다.
• 최종 분할 결과는 U-Net의 출력을 전역 주의 특징을 사용하여 융합함으로써 생성된다.
• 신뢰도 점수는 서로 다른 스케일에서의 분할 맵 간의 Dice 계수(예: s=1과 s=2)로 계산되며, 이는 분할 정확도의 대체 지표로 기능한다.
• 모델은 스케일 간 공유 파라미터를 사용하여 엔드 투 엔드로 훈련되며, 재학습 없이 다중 스케일 추론이 가능하다.

실험 결과

연구 질문

• RQ1 U-Net과 Transformer의 전역-국소 융합이 표준 U-Net 및 하이브리드 모델보다 의료 영상 분할 성능에서 뛰어나게 되는가?
• RQ2 다른 스케일 간 분할 맵 간 불일치가 실제 분할 정확도와 상관관계를 가지는가? 이는 지도 없이도 신뢰도 추정이 가능한가?
• RQ3 제안된 모델은 재학습 없이 다양한 패치 스케일에서 일관된 성능을 유지할 수 있는가?
• RQ4 스케일 불일치 기반으로 유도된 신뢰도 점수는 기존 방법보다 이미지 어려움 순위 매기기에서 더 우수한가?
• RQ5 모델은 다양한 기관과 영상 모달리티에 대해 일반화 가능한가?

주요 결과

• U-Netmer는 MRI, CT, 초음파, 내 endoscopy를 포함한 4종의 영상 모달리티와 7개의 기관을 포함한 7개의 공개 데이터셋에서 최신 기술 수준의 성능를 달성한다.
• 모델은 평가된 모든 데이터셋에서 4개의 베이스라인 및 6개의 최신 기술 모델보다 뚜렷한 향상을 보였다.
• 다른 스케일 간 분할 맵 간 불일치(예: s=1과 s=2)는 실제 Dice 정확도와 강한 선형 상관관계를 보이며, 7개 데이터셋 중 6개에서 피어슨 상관계수 0.8 초과를 기록했다.
• 스케일 불일치 기반의 신뢰도 점수는 7개 데이터셋 중 6개에서 ProtoSeg의 평균 분할 능력 점수보다 이미지 어려움 순위 매기기에서 뛰어난 성능를 보였다.
• 낮은 신뢰도 점수를 가진 이미지들은 작은, 대비가 낮거나 흐릿한 목표 영역을 가지며, 높은 신뢰도 점수를 가진 이미지들은 명확하고 크며 잘 정의된 구조를 가진다.
• 모델이 다양한 스케일에서 일관된 예측을 생성할 수 있는 능력 덕분에, 인간이 개입하는 리뷰를 위한 어려운 분할 케이스를 신뢰성 있게 식별할 수 있다.