U-Net: Convolutional Networks for Biomedical Image Segmentation

본 논문은 Image Segmentaion을 위한 논문이다. 기존의 FCN(Fully Convolutional Network)를 발전시킨 모델이다. 

Abstract

• Data augmentation을 사용
• Symmeteric path(constracting path, expanding path)을 통하여 정확한 localization 가능

Introduction

U-Net은 FCN을 계승한 모델인데 주요 차이점은 아래와 같다. 

• Upsampling: U shape의 대칭형의 구조를 갖고 있음 (propagate context information to higher resolution layers)
• Overlap-tile strategy: tile(patch)단위로 image를 모델에 넣어서 segmentation 가능 
  * 경계면의 pixel을 예측하기 위해서 경계부의 이미지를 mirroring 시킴 
  

  

  
  

  

• Data Augmentation: 기존 훈련 이미지에 deformation(변형)을 가하여 데이터 생성 
  -> learn invariance to such deformation (변형에 강건한 모델 생성)
• Weighted loss: 같은 클래스들끼리의 인접한 부분을 분리

Network Architecture

U-Net

Contracting path

• 3x3 convolutions 두 차례씩 반복 (no padding)
• ReLU Activations
• 2x2 max pooling (stride:2)
• Down sampling마다 channel 수 2배씩 늘어남 

Expanding path

• 2x2 up convolution
• 3x3 convolutions 두 차례씩 반복 (no padding)
• Up sampling마다 channel 수 반으로 줄어듦
• ReLU Activations
• Up Conv된 feature map은 Contracting path의 테두리가 cropped된 feature map과 concatenation 함 
• 1x1 convolution in last layer

-> 최종 출력인 Segmentation map의 크기는 Input image의 크기보다 작게 출력된다. 

-> tile(patch) 단위로 넣을 때 끊기지 않는 segmenation map을 얻기 위해서는 input tile size가 2x2 max pooling에 들어가는 input의 x,y의 크기가 짝수가 되도록 만들어야 한다. 

 

Training

• Segmentation map이 산출되면 그것을 기준으로 network가 훈련을 한다.
• 큰 배치 사이즈보다는 큰 input tiles를 선호 (single image를 batch로 사용) 
• high momentum(0.99): 이전에 보았던 훈련 샘플들이 현재 최적화에 가장 많은 영향을 미침 

Loss function

• Softmax function 
  final feature map에 대한 pixel-wise softmax를 한 후에 cross entropy를 적용 
  
  ※     $p_{k}(\mathbf{x})=\exp \left(a_{k}(\mathbf{x})\right) /\left(\sum_{k^{\prime}=1}^{K} \exp \left(a_{k^{\prime}}(\mathbf{x})\right)\right)$
  $a_{k}(x)$: channel k에서의 특정 pixel position x의 activation 값 
  $K$: number of classes
  
  
• Cross Entropy
  

  

  cross entropy

  
  $\ell: \Omega \rightarrow\{1, \ldots, K\}$
  $w$: weight map. 특정 픽셀에 더 많은 가중치 부여
  
  
• Weight map
  각각의 ground truth segmentation에 대한 weight map은 미리 계산해 놓는다. touching cell 간의 경계선을 학습하기 위해서 훈련 데이터 셋의 특정 클래스에 대한 different frequencey of pixels($w_{c}(x)$)을 보상하는 방식으로 weight map을 계산한다. (이미지 내에서 클래스마다 차지하는 픽셀 개수가 상이할 것이다) 
  또한, 가까운 cell의 경계와 가까우면 가중치를 높게 부여한다. 
  
  

  

  weight map

  $w_{c}$: weight map to balance class frequencies
  $d_{1}$: distance to the border of the nearest cell
  $d_{2}$: distnace to the border of the second nearest cell 
  (여기서 cell은 하나의 클래스 객체를 지칭. 의학에서의 cell)

Data Augmentation

매우 적은 수의 샘플을 가지고 있을 때 network에 invariance와 robustness를 부여하기 위한 목적으로 사용된다. 본 논문에서 사용한 데이터처럼 아주 미세한 이미지의 경우 shift, rotation, deformation, gray value variation를 적용한다. 특히, random deformation이 훈련 데이터 수가 적은 경우 가장 중요하다. 

 

Conclusion

• U-Net은 다양한 biomedical segmentation task에서 좋은 성능을 보였다. 
• Image augmentation 방법으로 인해 매우 적은 양의 이미지로도 학습이 가능해졌다.