이번에 리뷰할 논문은 2015년 ISBI cell tracking chellenge에 처음 등장하여 2등과 큰 격차로 우승한 U-Net 모델에 대한 논문입니다. 2015년 MICCAI에서 발표되었으며 본 논문에 대한 링크는 여기를 클릭하시면 됩니다.

Authors: Olaf Ronneberger, Philipp Fischer, Thomas Brox
Conference: MICCAI
Paper: U-Net: Convolutional Networks for Biomedical Image Segmentation
Year: 2015

역자의 말

영어를 한글로 변역하면서 변역된 단어가 적절하지 못한 경우 영어를 그대로 적용하였습니다. 부족한 부분이 있을 수 있으니 언제라도 잘못된 정보에 대해서는 comment나 메일을 통해서 지적 부탁드리겠습니다.

wogur379@gmail.com

Introduction

Deep convolutional networks는 여러 이미지 분야에서 뛰어난 성능(state of the art)을 보이고 있다. Convolutional Network는 오래전부터 있었지만 학습데이터의 크기와 네트워크의 사이즈에 제한이 있었다. 그러나 Krizhevsky의 AlexNet이 2012년 수백만개의 파라미터를 가진 8개 레이어의 모델로 백만개의 ImageNet을 학습하여 성공적인 결과를 낸 이후부터 더 크도 더 깊은 네트워크가 생기기 시작했다.

대부분 convolutional networks의 사용은 이미지 분류에 많이 사용된다. 그러나 많은 이미지 분야의 일들은, 특히 biomedical image processing 분야에서는 localization 하기를 원한다. 예를 들면 각각의 픽셀이 클래스로서 분류하는 걸 말한다. Ciresan의 연구: Deep neural networks segment neuronal membranes in electron microscopy images 에서 나온 모델은 각 이미지의 해당하는 부분에 대해서 픽셀별로 클래스를 예측하고도록 학습하였다. 첫 번째로 이 모델은 지역화(localize)가 가능하고, 둘 때로는 학습 데이터는 큰 이미지에 대한 패치(patch) 이미지를 사용하여 더 많은 학습데이터를 만들었다. (여기서 patch image란 이미지의 일부분을 말한다.)

그러나 Ciresan의 모델은 두 가지 단점이 있다. 첫 번째는 모델이 각 픽셀마다 실행이 되어야하기 때문에 패치이미지의 겹쳐진부분도 있어서 느리다는 것이고, 두 번째는 지역적 정확도(localization accuracy)와 질감의 사용(?)(the use of context) 간의 trade-off가 있다는 점이다. 패치이미지가 크다면 pooling을 통해 차원을 낮추게되여 지역적 정보가 많이 손실이 되고, 작은 패치이미지의 경우 일부분 밖에 볼 수 없기때문에 context의 의미를 알기 쉽지가 않게된다. 좋은 지역화(localization)과 the use of context 모두 동시에 이루어져야 한다.

이번 논문에서 소개하는 모델은 fully convolutional network 이다. 우리는 이 모델이 적은 학습데이터만으로도 더 정확한 분할(segmentation)이 가능하게했다. 메인 아이디어는 fig. 1에서 볼 수 있듯이 layer간 연결할 수 있는 부분을 사용한다는 것이다. 여기서 연결되는 부분의 pooling이 었던 곳은 up-convolution으로 바뀌어 있는 부분이다.

차원이 축소되었던 부분을 다시 upsampling을 통해서 복원하여 높은 해상도의 결과물을 얻을 수 있다. 이 연결된 부분을 통해서 모델은 정보를 전달받아 더 정확한 결과물을 얻을 수 있게된다.

Figure 1. U-net architecture에 대한 소개

이 모델에서 중요한 부분 중 하나는 upsampling 부분에서 높은 해상도의 부분을 전달받기위해 더 많은 feature에 대한 채널 수를 가지게 된다는 것이다. 결과적으로 점점 채널을 늘리는 경로와 줄이는 경로는 대칭적인 모양을 이루게되고 이 모양이 u자 모양을 나타낸다하여 U-Net으로 불리게 되었다.

이 모델은 더 이상 fully connected layer (FCN)을 사용하지 않고 전부 convolutional layer를 사용한다. 그 결과 각 패치이미지의 중복된 부분을 사용하여 커다란 이미지도 원활하게 분할할 수 있게되었다. 이미지의 가장자리 부분에 있는 픽셀들을 예측하기위해서 fig.2 처럼 빈 공간대신 입력값의 이미지를 거울처럼 대칭으로하여 추가하였다.

Figure 2. Overlap-tile strategy for seamless segmentation of arbitrary large images. 노란 영역의 segmentation하기 위해 파란 부분이 입력값으로 들어간다. 여기서 가장자리 밖에 있는 이미지의 부분은 입력값을 거울처럼 반대로 확장하여 사용한다.

본 논문에서는 학습 데이터에 대해서 elastic deformation을 사용하여 augmentation하였기 때문에 적은양의 학습 데이터만으로도 충분히 학습 하는 것이 가능했다. 이 방법은 이미지의 전체를 바꿀 필요없이 변화에 대한 불변성(invariance)을 학습할 수 있다. 이 방법은 특히 biomedical segmentation 부분에서 중요하게 사용될 수 있는데, 이 변형(deformation)은 실제로 조직인 변화하는 과정과 비슷하게 작용하기 때문이다.

많은 cell segmenation task의 어려운 점은 각 클래스가 서로 거의 붙어있다는 점이다. 이러한 것을 해결하기 위해서 본 논문에서는 서로 인접한 세포 사이의 배경에 대해서는 큰 loss를 얻을 수 있도록 weigted loss를 사용한다.

Figure 3. HeLa cells on glass recorded with DIC (differential interference contrast) microscopy. (a) raw image. (b) overlay with ground truth segmentation. Different colors indicate different instances of the HeLa cells. (c) generated segmentation mask (white: foreground, black: background). (d) map with a pixel-wise loss weight to force the network to learn the border pixels.

앞서 소개한 방법들은 segmentation의 여러분야에서 사용할 수 있고, 본 논문에서는 Ciresan의 모델보다 더 좋은 결과를 낸 the segmentation of neuronal structures in EM stacks에서의 결과를 보여준다. 거기다 2015년 ISBI cell tracking challenge에서 light microscopy 이미지에 대한 결과도 보여준다.

Network Architecture

모델 구조는 fig.1 에서 확인할 수 있고, 이미지 차원을 축소하는 부분(contracting path)과 확장하는 부분(expansive path)로 구성되어있다. 모든 convolution의 필터는 3x3으로 하였고, ReLU를 사용하였다. Max pooling을 2x2 필터와 stride 2를 사용하였다. 하나씩 downsampling할때마다 feature의 채널 수를 두배로 늘렸다. 확장하는 부분에서는 2x2 필터를 통해 up-convolution하였고, 이미지 차원을 줄이는 부분으로부터 feature map의 연결하여 이어 붙이기 때문에 이후에는 채널 수를 절반으로 줄였고, 두번 3x3 convolution filter를 거쳤다. 마지막 부분은 64개의 feature에 대한 채널 수로부터 각 클래스 수만큼 1x1 convolution을 하였다. 전체 네트워크는 23개의 convolutional layers로 구성되었다.

segmentation map 결과값이 좋게(a seamless tiling) 나오기 위해서는 2x2 max-pooling이 x축과 y축이 모두 짝수 값을 같는 레이어에 적용되도록 입력값의 사이즈를 정하는 것이 중요하다.

Training

학습과정은 stochastic gradient descent (SGD) 방법을 통해서 학습 되었다. padding을 하지않은 convolution을 사용하였기 때문에 출력이미지는 입력이미지보다 더 작게 된다. 또한 GPU를 효율적으로 최대한 사용하기위해 하나의 이미지에 대해서 패치이미지를 크게하여 배치 사이즈를 줄였다. momentum은 0.99로 사용했다.

loss function(본 논문에서는 energy function)은 각 픽셀에 대해서 softmax한 후 cross entropy를 사용하였다. softmax는 아래와 같이 정규화식을 말한다.

\[p_k(x) = \frac{exp(a_k(x))}{\sum_{k^{'}}^K exp(a_{k^{'}}(x))}\]

여기서 a_k(x)는 feature 채널 K에 대한 활성화(activation)을 말한다. K는 class의 수이고, p_k(x)는 x에 대한 확률 값이다. Cross entropy는 각 픽셀에 대해서 1과 p_l(x) (x)의 차이만큼 패널티를 더한다.

본 논문에서는 fig. 3에서 처럼 클래스간 서로 구분하는 테두리가 매우 좁은 경우를 해결하고 방법으로 학습데이터에서 특정 클래스의 빈도 차이만큼 보정하기위해 정답 데이터로부터 미리 weight map을 계산하였다.

여기서는 w_0을 10으로 하고 sigma는 5 픽셀 정도로 하였다.

많은 deep network에서는 가중치의 초기값을 잘 정하는 것이 중요하다. 이상적인 초기값은 각 feature map에 unit variance가 있도록 설정하는 것이다. 본 논문에서는 가우시안 분포로부터 가중치를 할당했다. 여기서 표준편차는 \root{2/N} (문장에 수식을 못넣는 노션…)이고, N은 하나의 뉴런이 갖고있는 노드의 수이다. 예를 들면 이전 레이어에서 3x3 convolution에 64 feature 채널을 갖는다고하면 N은 9 x 64 = 576이다.

Data Augmentation

Data Augmentation은 학습데이터가 얼마 없을때, 모델이 일관성(robust)이 있고 불변성(invariance)를 학습하도록 하기위해서 필요한 작업이다. 현미경 이미지에서 주로 회전하거나 이동시키거나 약간 변화를 주거나 회색값에 변화를 주거나 하였는데 몇 없는 학습데이터를 학습하기위한 augmentation 방법의 핵심은 elastic deformation을 적용하는 것이였다. 본 논문에서는 3x3 랜덤하게 변하는 벡터값을 통해 자연스러운 변화를 만들어내었다. 각 픽셀의 변화값은 bicubic interpolation을 통해서 계산되었다. Drop out 또한 data augmentation의 역할을 하였다.

Experiments

여기서는 3가지 task에 대하여 U-Net의 성능을 설명한다. 첫 번째 task는 전자현미경 데이터의 뉴런 구조의 segmentation이다. 데이터에 대한 예시와 결과는 fig. 2에서 확인할 수 있다. 데이터는 2012 ISBI EM segmentation challege에서 얻었고, 이 현재도 계속 결과를 제출할 수 있다. (현재 leader board 1위는 거의 99%의 성능을 가진다..) 학습데이터는 512x512로 30개의 이미지이고 serial section transmission electron microscopy of th Drosophila first instar larva ventral nerve cord (VNC)의 이미지이다. 각 이미지는 모두 정답이 있고 세포는 하얀색, 세포막은 검정색으로 구분되어져 있다. 평가데이터도 사용은 가능하지만 정답은 알 수 없다. 평가는 주최측에 예측값을 보내면 결과를 받을 수 있다. 10개의 클래스에 대해 구분되어 평가가 완료되면 “warping error”, “Rand error” 그리고 “pixel error”를 얻을 수 있다. (현재는 새로운 평가지표가 생김)

warping error, rand error, pixel error란?

u-net의 결과는 어떠한 전처리나 후처리도 없이 warping error는 0.0003529, rand-error는 0.0382로 Ciresan의 결과(IDSIA)보다 훨씬 좋은 결과를 나타냈다. IDSIA에 후처리를 적용한 방법이 rand-error에서 가장 좋은 결과를 나타냈다.

Table 1. Ranking on th EM segmentation challenge (march 6th, 2015)

다른 task인 light 현미경 이미지에서의 세포 segmentation에 u-net을 적용했다. 2014 그리고 2015년에 ISBI cell tracking challenge 중 하나였고 첫 번째 데이터 셋은 phase constast microscopy로 기록된 a polyacrylimide substrate에 있는 Glioblastoma-astrocytoma U373 cells을 포함한 “PhC-U373” 이라는 데이터 셋(fig. 4의 a,b)이다. 여기에는 부분적으로만 정답이 그려져있는 35개의 이미지가 있고 평가는 평균 IOU (Intersection Over Union)을 통해서 계산된다. U-Net은 92%의 IOU를 얻었고 2등이 83%라는 것을 통해 큰 차이로 성능을 뛰어넘었다는 것을 알 수 있다. 두 번째 데이터 셋은 differential interference contrast (DIC) microscopy에 의해 기록된 a flat glass에 있는 HeLa cells인 “DIC-HeLa” 라는 데이터 셋(fig. 4의 c,b)이다. 20개의 학습데이터와 정답이 있고, 여기서는 77.5%의 평균 IOU를 얻었다. 2등은 46% 밖에 되지않에 매우 큰 차이임을 알 수 있다.

Figure 4. Result on the ISBI cell tracking challenge. (a)는 "PhC-U373"이라는 데이터 셋이다. (b)는 그에 대한 정답데이터이다. (c)는 "DIC-HeLa"라는 데이터 셋이다. (d)는 노란색 테두리가 정답이고 각 segmentation부분은 무작위로 색상을 칠한 것이다.

Table 2. Segmentation results (IOU) on the ISBI cell tracking challenge 2015

Conclusion

U-Net 모델은 다른 biomedical segmentation 방법들보다 훨씬 더 좋은 결과를 보여주었다. 그리고 elastic deformation 덕분에 적은 데이터만으로도 좋은 결과를 얻을 수 있었다.