[Review] U-Net : Convolutional Networks for Biomedical Image Segmentation

U-Net: Convolutional Networks for Biomedical Image Segmentation

 

 

 

2015년 5월에 나온 Biomedical 분야에서 이미지 분할(Image Segmentation)을 목적으로 제안된 End-to-End 방식의 Fully-Convolutional Network 기반 모델인 U-net에 대해 정리해보려 합니다.

 

◼ 논문 링크 : U-Net : Convolutional Networks for Biomedical Image Segmentation

Abstract

• data argumentation을 이용해 few labeled data를 가지고도 적절한 학습이 가능하다.
• contracting path 부분을 통해 context 정보를 얻고, contracting path 부분과 symmetric한 expanding path를 통해 다시 원본 크기에 맞도록 upsampling한 후 segmentation을 진행하게 된다. (FCN모델과 비슷함)
• FCN 구조를 응용한 구조를 사용함으로서 segmentation map을 출력한다.
• FCN모델과 다른점은 contracting path 와 expanding path 를 이어주는 연산의 유무이다.
• ISBI cell tracking challenge 2015에 참가하여 우승하였고, challenge 데이터 셋을 사용하였다.

 

Introduction

Deep convolutional networks는 visual recognition tasks에서 SOTA(State of the Arts)를 달성하였다.

• 기존 CNN은 single classification task를 수행
• 하지만 Biomedical image 처리에서는 localization, 즉 각 픽셀에 클래스 라벨을 할당해야 한다.
• Biomedical image의 수가 충분하지 않다는 문제점이 있다.
• 데이터 문제를 해결하기 위해 patch기반의 학습을 이용해 데이터의 수를 늘리고자 했다.(sliding-window setup)

sliding-window setup을 이용하면,

1.  localization을 할 수 있다.
2.  patch 데이터를 이용하여 데이터의 수를 늘릴 수 있다.
3.  patch별로 네트워크를 따로 운영해야 하기 때문에 속도가 상당히 느리고, patch가 겹치기 때문에 중복이 많다. 
4.  trade-off between localization accuracy & context

 

논문의 U-Net은 fully convolutional network를 기반으로 적은 수의 데이터로부터 학습이 가능하고, 정확한 segmentation을 가능하게 한다. U-Net은 이미지의 전반적인 context 정보를 얻기 위한 네트워크와 정확한 지역화(Localization)를 위한 네트워크가 대칭 형태로 구성되어 있다.

 

 

Fig1. 을 보면 expanding path(upsampling)를 보면 기존의 FCN모델보다 conv filter channel이 사용되었음을 알 수 있다. 이를 통해, propagate context information to higher resolution layers가 가능하다.

 

Over-tile Strategy

Fig2

Fig3

 

논문에서는 input image 사이즈가 572x572, output image 사이즈는 388x388이다.
이는 contracting path에서 padding이 없었기 때문에 점점 이미지 외곽 부분이 없어진 결과이다.
이를 해결하기 위해 이 논문에선 mirroring방식을 사용하였다. 

맨 바깥쪽 없어지는 부분이 안쪽 이미지와 형태는 같고, 거울에 반사된 형태를 가지고 있다. 이처럼 사라지는 부분(missing context)은 zero-padding이 아닌 mirror padding의 형태로 채워준다.

 

또한 U-Net은 Data augmentation도 사용했습니다. elastic deformation augmentation을 적용해 학습 데이터의 수를 늘렸다.

 

Biomedical Image segmentation의 경우 인접한 세포들을 분리하는 segmentation을 적용하는데 어려움이 있었다.

U-Net은 인접한 세포들 사이의 경계가 되는 background를 segemetation할 수 있도록 가중 손실을 loss function으로 사용하였다.

 

 

 

Network Architecture

U-Net은 Fig1에 나와 있듯이 contracting path와 expanding path(expansive path)로 이루어진 U자형 모델입니다.

 

contracting path

 

• 3x3 convolutions을 두 차례씩 반복 (패딩 없음)
• 활성화 함수 : ReLU
• 2x2 max-pooling (stride: 2)
• Down-sampling 마다 채널의 수(kernel)를 2배로 늘림

위의 과정을 반복해서 적용하며, 이렇게 특성 추출을 하며 점점 feature map의 사이즈가 줄어드는 구조를 가진다.

 

 

expansive path

 

Expanding Path는 Contracting Path와 반대의 연산으로 특징맵을 확장한다.

 

• 2x2 convolution (“up-convolution”)
• 3x3 convolutions을 두 차례씩 반복 (패딩 없음)
• Up-Conv를 통한 Up-sampling 마다 채널의 수를 반으로 줄임
• 활성화 함수 : ReLU
• Up-Conv 된 feature map은 Contracting path의 테두리가 Cropped된 feature map과 concatenation 함 (copy & crop)
• 마지막 레이어에 1x1 convolution 연산 → 2개의 클래스 구분

concatenation 과정은 contracting path의 high resolution features를 upsampled output과 합쳐주는 것으로, localization을 위해 필요하다. crop을 하는 이유는 contracting path에 있는 feature map과 expansive path에 있는 feature map의 해상도가 같지 않기 때문이다.

맨 마지막 단계에서 3x3 CNN을 2번 거치고 upsampling 대신 1x1 CNN을 사용해 channel의 개수를 2개로 줄인다. 최종적으로 판단할 클래스 개수가 2개(foreground, background)이기 때문이다.

 

위와 같이 foreground, background 2가지로 전체 이미지에 있는 픽셀들을 분류한다.

U-Net은 23개의 CNN 레이어들로 구성되어 있고, segmentation map에 seamless tiling을 하기 위해서 모든 2x2 max-pooling이 적용되기 위해 feature map size가 x 및 y축 크기가 짝수가 되도록 입력 타일(patch) 크기를 선택하는 것이 중요하다.

 

 

Training

optimizer

• stochastic gradient descent (SGD)
• Momentum : 0.99

Framework

• Caffe

Batch Size

• Single Image

모멘텀을 0.99로 설정해 이전에 조정한 가중치를 현재 시점의 가중치를 조정하도록 설정하였다.

 

 

Loss Function

Energy Function은 마지막에 얻은 feature map에 픽셀 단위로 softmax를 수행하고 여기에 cross entropy loss function을 적용한 것이다.

Energy Function

여기서 x는 feature map에 있는 각 픽셀을 말하며, w(x)는 weight map이라고 픽셀 별로 가중치를 부과하는 역할을 한다. 각 픽셀에서 계산한 걸 다 더하여 산출한다.

 

softmax 함수는 아래의 공식으로 표현이 된다.

 

softmax 공식

 

• K : class의 수
• k : k번째 class
• $a_{k}(x)$ : 픽셀 x 가 class k일 값 (픽셀별 모델의 output)
                      $x\in \Omega$에서 k 채널의 activation($\Omega \subset \mathbb{Z}$)
• $p_{k}(x)$ :  픽셀 x 가 class k일 확률

 

 

weight map

•  $w_{0}$ : Weight hyper-parameter, 논문에서 10으로 설정
• $\sigma$ : Weight hyper-parameter, 논문에서 5로 설정
• $d_{1}(x)$ : 픽셀 x의 위치로부터 가장 가까운 경계와 거리
• $d_{2}(x)$ : 픽셀 x의 위치로부터 두 번째로 가까운 경계와 거리

 

weight map 수식은 해당 픽셀의 중요도를 결정해주는 역할을 한다.

같은 class들을 구분해줄 수 있는 경계에 해당하는 픽셀을 더 잘 구분할 수 있도록 픽셀이 경계에 있으면 작아지고 경계에서 멀어지면 커져서 경계 부근의 픽셀들이 큰 가중치를 얻도록 설계되었다.

또한 weight map 수식의 오른쪽 exp 부분은 gaussian distribution과 유사함을 보인다.

 

논문에서는 경계(border)의 weight map을 계산하기 위해 이미지 처리기법인 morphological operation을 적용했다.

 

 

Weight Initialization 가중치 초기화

Deep network에서는 weight의 초기화가 중요하다.

노드의 개수를 N이라고 하면 $\sqrt{\frac{2}{N}}$ 의 표준편차를 갖는 Gaussian Distribution을 이용하여 모델의 파라미터를 초기화하고 학습한다.

 

 

Data Augmentation

Data Augmentation는 학습용 데이터셋이 적을 때 사용하기 유용한 방법이다.

데이터의 증강으로 invariance하고 rubustness한 성질을 가지는 모델로 학습이 가능해진다.

 

데이터 증강 방법으로 Rotation(회전), Shift(이동), Gray Value, Elastic deformations 등이 있다.

 

본 논문에서 shift, rotation 외에 random-elastic deformation이라는 기법을 사용해 Data Augmentation을 구현하기도 했다.

elastic deformation 기법은 elastic(유연)하게 이미지의 경계를 변형시켜 학습하여 흔들림, 왜곡된 이미지도 잘 구분할 수 있게 해 준다. 

 

elastic deformation

 

 

Data Augmentation을 구현 방법

1. coarse 3 x 3 grid에 random displacement vectors을 이용해 smooth deformation을 수행한다.
   displacement는 10개 픽셀이 가지는 값들의 표준편차를 따르는 가우시안 분포에서 임의로 뽑은 값으로 수행한다.
2. bicubic interpolation을 이용해 픽셀 단위로 displacement를 계산한다.
3. contracting path의 맨 끝에 있는 Drop out layer가 더욱 implicit한 Data augmentation을 수행한다.

 

Experiments

U-net의 성능을 테스트하기 위해 3가지의 segmentation tasks에 적용하였다.

 

1. The segmentation of neuronal structures in electron microscopic recordings (EM segmentation challenge at ISBI 2012)

• Dataset : 전자 현미경으로 찍은 512 x 512 해상도의 이미지 30장
• 데이터는 이미지와 함께 객체와 배경이 구분된(0 or 1) Ground Truth Segmentation Map을 포함하고 있다.
• 평가 지표 : warping error, Rand error, pixel error
  
  • Warping Error : Segmentation metric 중 하나로 topological Error를 의미한다.
  • Rand Error : 참고자료
  • Pixel Error : Prediction과 Ground truth가 불일치하는 pixel의 비율을 의미한다.

 

U-Net 모델이 가장 좋은 성능을 보이고 있으며, Sliding Window 방식을 이용한 IDSIA보다 성능이 좋았다.

자세한 수치는 위의 표를 보고 확인하면 된다.

 

 

2. cell segmentation task in light microscopic images (ISBI cell tracking challenge 2014 and 2015)

 

광학 현미경(light microscopic)에서 얻은 이미지로도 cell segmentation task를 수행했다.

세포 분류 대회인 ISBI cell tracking challeng 2014, 2015에 좋은 결과를 얻었고 결과는 위의 Table 2에서 확인할 수 있다.

 

a, b는 PhC-U373, c, d는 DIC-HeLa 데이터 셋에서 Segmentation 결과이다. 노란선이 ground truth이다.(b, d)

• PhC-U373 Data : 위상차 현미경으로 기록한 35개의 이미지
• DIC-HeLa Data : HeLa 세포를 현미경을 통해 기록한 20개의 이미지

 

이 task에서는 IoU로 성능을 측정했다. IoU는 세포라고 구분한 영역과 실제 세포의 영역이 겹치는 정도로 성능을 측정한 지표이다.

두 데이터에서 모두 좋은 결과를 보였고, 두 번째 모델과의 성능의 차이도 나타내고 있다.

• PhC-U373 :  U-Net 92%  vs  second 83%
• DIC-HeLa :  U-Net 77%  vs  second 46%

 

Conclusion

• U-Net은 다양한 biomedical segmentation applications에서 좋은 성능을 보여줬다.
• Elastic deformation 기법을 이용해 아주 적은 training sample로 좋은 성능을 보여줬다.
• NVidia Titan GPU(6 GB)에서 10 시간 정도의 짧은 시간으로 학습하였다.
• U-Net의 구조가 다양한 task에 쉽게 적용될 수 있을 것이라 확신하며 마무리

 

 

 

 

<참고 자료>

Minguinho_zeze.log

Time Traveler

goe87088.log

Code Journey

https://www.theteams.kr/teams/6500/post/73665