[Segmentation] Convolutional Networks for Biomedical Image Segmentation : U-Net 논문 리뷰

투빅스 15기 오주영

INDEX

Prior Work - FCN(Fully Convolutional Networks)

Contribution

Introduction

Network Architecture

Training

Conclusion

Unet-family

Prior Work - FCN(Fully Convolutional Networks)

• fully convolutional network를 사용함으로서 임의의 크기의 이미지를 입력 받아 그에 상응하는 크기의 segmentation을 수행하고자 함

• FCN은 Semantic Segmentation 모델을 위해 기존 이미지 분류에서 우수한 성능을 보인 AlexNet, VGG, GoogLeNet 등을 목적에 맞게 변형한 모델

  

• Semantic Segmentation의 관점에서 기존의 Fully Connected Layer를 가진 모델의 한계
  1. 이미지의 위치 정보가 사라짐
  2. 입력 이미지의 크기가 고정되어야함

• Architecture

  • Skip Architecture : FCN 논문에서 정확하고 상세한 Segmentation을 얻기 위해 deep, coarse, semantic 정보와 shallow, fine, appearance 정보를 결합하는 Skip Architecture를 정의함

    

    

  • 전체 convolution이 진행된 후 deconv하여 upsampling을 진행한 경우 coarse한 위치 정보만 포함하고 있기 때문에 segmentation을 하기에는 부족함

  • 세부적인 정보를 추가하기 위해 이전 stage의 featuremap을 stride를 다르게 주어 더해주게됨

    

Contribution

• 얕은 층의 fine location 정보와 깊은 층의 global semantic 정보를 결합하는 U 모양의 architecture를 제시
• 서로 근접한 객체 경계를 잘 구분하도록 학습하기 위하여 Weighted Loss를 제시

Introduction

• 기존 연구의 문제점

  1. sliding window 방식을 사용하면 patch별로 네트워크를 별도로 실행해야하고 patch의 중복으로 인한 중복성이 많기 때문에 속도가 느리다.
  2. localization accuracy와 context 사이의 trade-off가 존재 : patch가 크면 context 정보를 더 많이 활용할 수 있지만 localization accuracy는 감소, patch가 작으면 localization accuracy는 증가하지만 context를 거의 볼 수 없음

• FCN을 기반으로 architecture를 수정하고 확장하여 극소수의 영상으로 훈련 가능하고 보다 정확한 segmentation을 산출

• FCN과는 다르게 Upsampling 파트에서도 많은 수의 featuremap을 가지고 있어 네트워크가 context 정보를 더 높은 해상도를 가진 layer에 전파할 수 있음

  • Contracting path와 Expanding path의 대칭구조를 사용

Network Architecture

• 네트워크는 Context 정보를 얻기 위한 Contracting path, 정확한 Localization을 위한 Expanding path로 이루어져 있음

  • Contracting path

    1. 3x3 Conv(unpadded, stride 1) + ReLU + BN
    2. 3x3 Conv(unpadded, stride 1) + ReLU + BN
    3. 2x2 Max-Pooling (stride 2)

    

  • Expanding path

    1. 2x2 Deconv(stride 2)
    2. Contracting path에서 동일한 level의 featuremap을 1.에서 생성된 featuremap과 concat
    3. 3x3 conv(unpadded, stride 1) + ReLU + BN
    4. 3x3 conv(unpadded, stride 1) + ReLU + BN

    

  • Unet에서는 Convolution layer에서 padding을 사용하지 않기 때문에 지속적으로 featuremap의 크기가 줄어들게되므로 Contracting path의 featuremap과 concat하기 전 같은 크기로 cropping해줌

  • 마지막 layer는 1x1 convolution 연산을 해주어 segmentation 수행

Training

• Overlap-tile strategy & Mirroring Extrapolate
• Weighted Loss
• Data Augmentation

• Overlap-tile strategy

  

  • U-Net은 unpadded conv를 사용하여 output image의 크기가 input image의 크기보다 작기 때문에 위 그림처럼 왼쪽의 파란색을 Input으로 넣으면 오른쪽의 노란색이 output이 됨

  • input으로 사용하는 patch가 겹치는 부분이 존재하도록 이미지를 자르고 segmentation을 수행하는 Overlab-tile strategy를 취함

    

  • output image의 크기가 작아지므로 실제로 segmentation을 수행하기 위해 image 경계에 위치한 image를 복사하고 상하 또는 좌우 반전을 통해 Mirroring image를 모델에 input으로 사용

• Weighted Loss

  • biomedical 분야의 segmentation 작업에서 중요한 과제 중 하나는 맞닿아 있는 세포 사이의 경계를 분리하는 것

  • $Loss = \sum_{x\in\Omega}w(x)log(p_{l(x)}(x))$

    • $l$ : the true label of each pixel
    • $w$ : a weight map that we introduced to give some pixels more importance in the training

  • $p_k(x) = exp(a_k(x)) / \sum_{k'=1}^{K} exp(a_{k'}(x))$

    • $p_k(x)$ : the probability that k class is in x pixel
    • K : number of classes
    • $a_k(x)$ : activation in feature channel k at the pixel position x

  • $w(x) = w_c(x) + w_0 · exp(-\frac{(d_1(x) + d_2(x))^2}{2\sigma^2})$

    • $w_c(x)$ : weight map to balance toe class frequencies
    • $d_1(x)$ : distance to the border of the nearest cell
    • $d_2(x)$ : distance to the border of the second nearest cell
    • pixel x와 경계의 거리가 가까우면 큰 값을 가지게 되어 해당 pixel의 weight가 커지게됨

• Data Augementation

  • Biomedical domain의 특성 상 모델을 학습시키기에는 데이터가 충분하지 않음
  • Shift, Rotation 등의 기본적인 Data Augmentation
  • Elastic Distortion : 픽셀별로 이미지가 다른 방향으로 뒤틀리게 됨(세포 조직의 실질적인 형과 유사하다고 함)

  

Conclusion

• Unet architecture는 biomedical segmentation application에서 매우 우수한 성능을 달성

• FCN을 확장시켜 context 정보와 local 정보를 충분히 활용할 수 있는 U모양의 architecture를 제안

• Elastic distortion으로 Data Augmentation을 했기 대문에 적은 dataset에서도 훈련 가능

Unet & Unet-family

• 데이터셋에 맞는 모델의 최적 깊이를 알 수 없음

  • 비용을 들여 최적의 architecture를 찾아내거나 여러 깊이의 모델을 앙상블하는 비효율적인 작업이 필요함
  • 이는 인코더를 공유하지 않고 각각 돌아간다는 점에서 비효율적

• Skip Connection이 동일한 깊이를 가지는 인코더와 디코더만 연결되는 제한적인 구조

  • 같은 크기의 featuremap만 concat할 수 있으므로 encoder와 decoder의 결합이 약함

  Unet++

• 인코더를 공유하는 다양한 깊이의 U-net을 만들어서 deep supervision을 이용해서 함께 학습하고 앙상블하는 형태를 제안

• Skip connection을 동일한 깊이에서의 featuremap들이 모두 결합하도록 하여 유연한 featuremap을 만들어줌

  

• 이후 MultiResUNet, $U^2$Net, 등 많은 변형

Reference

• https://joungheekim.github.io/2020/09/28/paper-review/
• https://kuklife.tistory.com/117?category=872135
• https://kuklife.tistory.com/119
• https://daljoong2.tistory.com/179
• https://github.com/ShawnBIT/UNet-family
• https://arxiv.org/pdf/1411.4038.pdf
• https://arxiv.org/pdf/1505.04597.pdf