1. Introduction

• FCNN 기반의 접근법

  receptive fields 때문에 long-range dependency 구축에 한계를 가짐

  locality of the receptive fields in convolutional layers still limits their learning capabilities to relatively small regions.

• self-attention module

  • 전체적인 문맥 파악 & long-range dependency 구축 가능해짐
  • ViT(Vision Transformer)에서 vision분야에 transformer 적용

• UNETR(UNEt TRansformer) 제안

  

  • 3D medical image segmentation에 transformer 적용
  • 3D segmentation을 1D sequence-to-sequence 문제로 재정의하여 해결
  • Transformer를 encoder에 적용
  • CNN 기반의 decoder 사용 : transformer는 지역적인 정보를 잘 잡지 못하기 때문(전역적인 정보는 잘 잡는 반면)

  <aside> 🏅 contributions

  • 3D image segmentation을 위한 Transformer 기반의 모델
  • 구조
    1. long-range dependency를 구축하기 위해 3D volume의 transformer encoder를 적용
    2. skip-connected decoder는 다른 해상도의 feature들을 합쳐서 segmentation mask를 도출
  • BTCV dataset과 MSD dataset에서 좋은 성능을 냄 </aside>

3. Methodology

3.1 Architecture

• Encoder-Decoder형태의 **UNET 구조**를 가짐

  • Encoder : 수축 경로, 이미지 특징 추출, transformer으로 구성
  • Deocder : 확장 경로, 마스크 예측,
  • skip-connection을 통해, decoder에 각 단계별 추출된 정보를 전달

• INPUT 준비

  1. 3D volume($x \in R^{H.W.D.C}$) → 1D sequence($x_v \in R^{N.(P^3.C)}$)로 생성
     • (H,W,D)를 (P,P,P)크기의 patch로 쪼개기 ⇒ 패치 개수(N) = $(H.W.D)/P^3$
  2. linear layer를 통해 $C$ → $K$ 차원으로
     • ($x_v \in R^{N.(P^3.C)}$) ⇒ ($x_v \in R^{N.(P^3.K)}$)
  3. 1D learnable positional embedding 수행 - 공간정보 유지
     • 1D learnable positional embedding : $E_{pos}\in R^{N.K}$
     • projected patch embedding : $E \in R^{N.(P^3.C)}$
     • $Z_0 = [X_v^1E;...;X_v^NE] + E_{pos}$
     • segmentation task이기 때문에 [class] token은 더하지 않음

• Encoder

  • transformer block을 stask처럼 쌓음

    

• Skip-connection

  • Encoder에서 나온 여러 해상도의 feature들이 deocder에서 합쳐짐
  • features : $z_i(i \in{3,6,9,12})$ - 크기 : ${H\over P},{H\over P}, {H\over P} ,K$

• Decoder

  • UpSampling 진행하면서, encoder에서 나온 feature들 합치기

• 최종 mask

  • 1x1x1 conv layer & softmax activation을 통해 voxel-wise semantic prediction 생성

3.2 Loss 함수

• voxel단위의 soft dice loss & cross-entropy loss

  

  • J : class 개수
  • I : voxel 개수
  • $Y_{i,j}$ : j번째 class에 대한 i번째 voxel의 예측
  • $P_{i,j}$ : j번째 class에 대한 i번째 voxel의 정답(GT)

4. Experiments

4.1 Datasets

• BTCV(CT)
  • 13개의 장기가 annotation된 복부 CT 이미지 사용
  • (512x512 pixels)*(80~225 slices)
  • 13 class segmentation task with 1-channel input
• MSD(MRI/CT)
  • 뇌종양 segmentation 과제
  • 3 class segmentation task with 4-channel input
• spleen segmentation task