[Paper Review] U-Net 논문 리뷰

U-Net: Convolutional Networks for Biomedical Image Segmentation

기존 연구의 문제점

• 이미지에서 하나의 patch를 추출하는데, 각 patch를 모두 다른 network에서 학습을 진행했다.
  • 이는 overlapping하는 부분 때문에 연산이 낭비되는 문제가 있었다.
• 크기가 큰 patch의 경우 max-pooling을 더욱 많이 거치는데,
  • 이 경우 localization accuracy(즉, 각 pixel 별 class 예측의 정확도)가 떨어진다.
• 크기가 작은 patch의 경우 context에 대한 인식률이 떨어지게 된다.

U-net의 Architecture

• Fully Convolutional Network이 발전한 모델이다.
• end-to-end로 train이 가능한 구조이다.
• Unet은 크게 두가지 path로 구성되어있다.
  • 하나의 path는 contracting path이다. U 구조에서 좌측 부분이다.
    • context 정보를 담고 있다.
    • 일반적인 convolutional network이다. 반복적인 3x3 upadded convolution, ReLU, 2x2 stride 2 maxpooling, and doubling the channel
  • 하나의 path는 expansive path이다. U 구조에서 우측 부분이다.
    • precise localization 정보를 담고 있다. 즉, 각 pixel이 가지고 있는 class 정보를 의미한다.
    • 일반적인 upsampling network로 구현된다. channel을 절반으로 줄이는 2x2 deconvolution, contracting path에서 copy하고 crop한 feature map, 3x3 convolution, ReLU
  • 마지막 1x1 convolution은 feature vector를 원하는 class로 대응되도록 하기 위해서 존재한다.
• Contracting path가 이미 큰 feature channel을 통해서 context information을 전달하기에, expansive path가 contracting path와 비대칭적이고, U-형태를 지니게 된다.

Unet의 Architecture. channel은 위에, x,y 사이즈는 아래에 기술. 하얀 상자는 복사된 feature map을 의미한다.

Overlap-Tile Strategy

• 또한 자연스러운 경계선 형성을 위하여, 다음과 같이 Overlap-Tile Strategy를 사용한다.
• 이미지의 가장자리 부분의 정확도를 올리기 위해서, 사진 촬영에 의해서 잘려진 바깥 부분을 mirroring하여 extrapolate해준다.
  • 이 방법은 특히 큰 이미지에 유효하다.
• 본 방법론이 잘 적용되려면, input tile size가 중요하다. 모든 2x2 maxpooling 결과가 even해야한다.

Weighted Loss to Touching Objects

• 아래 그림과 같이 서로 접해있는 경우, 구별하기 어렵다는 문제가 있다.
• 이를 위해서 접촉하는 셀 사이의 배경 lable이 큰 가중치를 얻는 weighted loss를 사용했다.

Training

• large input tile, large batch size 선호
• SGD 사용
  • momentum 0.99

Energy Function

• 마지막 feature map에서 pixel-wise softmax를 계산한 후, cross entropy를 구한다.
• softmax 공식

• $a_k(x)$은 image position x에서 feature channel k에서의 activation을 의미한다.
• $K$는 class의 개수
• $p_k(x)$는 approxiamted maximum-function.
  • $a_k(x)$이 최대인 곳에서 $p_k(x) ≈ 1$ 이고, 나머지 는 $p_k(x) ≈ 0$
• cross entropy는 다음과 같이 계산된다.

• 여기서 cross entropy가 각 image position에서 $p_{l(x)}(x)$의 deviation을 penalize한다고 한다.
• $l(x)$는 position x에서의 true label을 의미한다. 즉, $l : Ω → {1, . . . , K}$
• $w : Ω → R$에서 $w$은 weight map이다. 이들은 training에서 더욱 강조하기 위한 픽셀이 있을 때 사용한다. 이는 다음에 이어서 설명한다.

Weight Map

• weight map은 학습 전 미리 계산하는데, 이는 각각의 ground truth segmentation에 대해서 수행한다.
  • 이는 각 class가 각 픽셀에서 출현하는 정도가 달라서 이에 대해서 보상하기 위해서 수행하는 방법이다.
  • model이 touching하는 cell 사이에 작게 존재하는 separation border를 잘 학습하게 하려는 목적으로 사용된다.
    • separation border는 morphological operation을 이용해서 구해진다. 그렇게 해서 weight map을 구하면 다음과 같다.

• $w_c : Ω → R$은 calss frequency를 균형 잡기 위한 weight map
• $d_1 : Ω → R$은 border에서 가장 가까운 cell까지의 거리이고, $d_2 : Ω → R$은 border에서 두번째로 가까운 cell까지의 거리이다.
• $w_0 = 10, σ ≈ 5$ pixels로 두고 실험했다.
• 초기 weight map은 variance가 1이 되도록 초기화 되는 것이 이상적이다.
• 이 모델에서는 Grassian distribution을 std_dev를 $\sqrt(2/N)$로 두면서 해결했다. 여기서 N은 incoming nodes of NN이다.
  • 만약 3x3 convolution에 64 feature channel이 직전의 layer라면, $N = 9 \cdot 64 = 576$

Data Augmentation

• 의료 데이터 특성상 숫자가 부족하여, data augmentation을 수행하였다.
  • shift, rotation, deformation, gray value variataion, elastic deformation을 수행하였다.
  • data augmentation을 수행하면, 적은 data로도 invariance와 robustness property를 학습시킬 수 있다.
• smooth deformation은 random displacement vector를 coarse 3x3 grid에 사용하여 구현했다.
  • displacement는 Gaussian distribution을 토해서 구했다. std_dev는 10 pixel
  • pixel 당 displacement는 bicubic interpolation을 통해서 구했다.
• Data Augmentation을 통하여 약 20개의 sample만으로도 80퍼센트의 accuracy를 보였다.

Experiments

• 크게 세가지 실험을 했다고 한다.
  • segmentation of neuronal structures in electron microscopic recordings(EM segmentation challenge)
  • cell segmentation task in light microscopic images
    • PhC-U373
    • DIC-HeLa
• segmentation of neuronal structures in electron microscopic recording
  • Evaluation은 map을 10단계로 나눠서 thresholding을 하고, warping error와 rand error와 pixel error를 계산함으로써 수행했다.
  • 그 결과 sliding-window convolutional network보다 더 나은 성능을 보였다.

• PhC-U373

• DIC-HeLa

• cell segmentation task in light microscopic images에서의 결과표