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

Title: U-Net: Convolutional Networks for Biomedical Image Segmentation
paper: https://arxiv.org/abs/1505.04597

UNet은 제목에서도 알 수 있다시피 본래 biomedical 분야에서 image segmentation을 위해서 만들어진 구조이다.

여기서 image segmentation이란,

(사전 지식)Image segmentation

이미지 내에 있는 각 물체들을 의미 있는 단위로 분할 하는 작업으로, 각 픽셀마다 하나의 클래스로 분류해서(예. 배경=0, 강아지=1, 고양이=2, 오리=3, ...) 한 장의 segmentation map을 만드는 것이다.

 

이 논문의 계기

세포 이미지를 학습하기 위해 사람이 이미지에서 세포 하나하나를 segmentation하려면, 세포 이미지를 픽셀 단위로 분류할 수 있는 고급 작업자가 필요한데, 이러면 데이터 단가가 너무 비싸져서 이를 해결할 수 있는 방안으로, biomedical에서 쓰일 수 있는 data augmentation 기법을 위해 만들어짐.

(그래서 논문 제목도 U-Net: Convolutional Networks for Biomedical Image Segmentation)

1. Abstract

저자는 레이블 정보가 있는 데이터(annotated data)가 적을 때 효율적인 data augmentation과 함께 효율적인 학습 전략을 제안했다. 이 구조는 feature map의 사이즈가 작아졌다가 커지는, 대칭적인 구조인데, 작아지는 부분은 context를 파악할 수 있는 contracting path, 커지는 부분은 localization이 가능하게 하는 expanding path이다.

참고) feature map이란?

- input에서 kernel(또는 filter)를 통과하면서 나오는 output을 feature map이라고 함
저자는 이 네트워크가 적은 이미지로부터 end-to-end로 학습될 수 있다고 하며, 전자 현미경 사진에서 중성자의 구조를 분해하는 ISBI challenge에서 가장 좋은 method(sliding window convolutional network로 되어 있는 모델)의 성능을 넘었다고 함. 이 구조는 이 때 당시 최신 GPU에서 512x512 이미지의 segmentation이 1초도 걸리지 않을만큼 inference 속도도 빠르다고 함.

2. Introduction

 

고해상도 이미지에 대해서 전체 이미지를 한 번에 처리하려면 많은 연산량을 필요로 하기 때문에 고해상도의 의료 이미지를 효율적으로 처리하기 위한 방안으로 이미지를 나누어서 차례대로 처리하고자 한다. 그런데, 기존의 sliding window 탐색 방식은 탐색 구간 별로 겹치는 구간이 있어, 겹치는 구간에 대해서 불필요한 반복학습을 수행하게 된다. 반면 Unet은 학습 영역을 일정 단위(patch)로 쪼개서 한 번 학습한 구역은 다시 학습하지 않는 방식으로 구현이 되었다

2.1. architecture

 UNet의 architecture는 FCN(fully convolutional network)구조를 변형하고 늘려서 만든 구조이다.

참고) FCN

간단하게 짚고 넘어가자면, 이름 그대로 convolutional layer로만 이루어진 network이다.
FCN 구조의 특징은 크게 세가지로 나눌 수 있다.

1. convolutional layer
2. deconvolutional layer(upsampling)
3. skip-connection

2.1.1. Convolutional layer

보통 CNN(convolutional neural network)같은 경우, task와 관계없이 뒤에 dense layer(==fully connected layer)를 쓰는데, 

VGGNet

이러한 경우, 두 가지 문제가 있었다.

1. 이미지의 위치 정보, 즉 segmentation과 같은 task에서 필요한 객체들의 위치 정보가 dense layer에 들어가면서 사라지게 된다.
2. dense-layer의 크기가 고정되어 있기 때문에, 한 번 짜여진 구조에서, input image의 크기가 고정된다.
   1. convolutional network끼리 연결되어 있는 kernel filter의 경우, input과 output의 channel에는 영향을 받으나, filter의 size는 독립적
   2. dense layer와 연결되어 있는 weight matrix의 경우 input과 output의 size에 연관이 되어 있으므로, dense layer가 섞인 구조에서 filter size가 고정되면, 이에 맞춰서 input size도 결정이 되고, 따라서 input image의 크기가 영향을 받는다.

상단의 w와 h는 kernel size이며, input과 output에 독립적이다.

이를 해결하기 위해 뒷단에도 dense layer에서 convolutional layer로 대체한다.
VGGnet을 예로 들면 다음과 같은 그림이 나온다.

그런데, 이 상태에서 바로 출력해버리면, 출력 feature map은 픽셀 단위로 classification까지는 하기 힘들 것이다. 논문에서는 이를 coarse(거칠다)하다고 말한다. 따라서 이러한 coarse map을 원본 이미지의 크기와 가까운 dense map으로 바꿔줄 필요가 있다.

2.1.2. deconvolutional layer(upsampling)

coarse map에서 dense map으로 바꿔주는 method에는 여러가지가 있다고 한다.(FCN 정리 블로그 참조)

• Interpolation
• Deconvolution
• Unpooling
• Shift and stitch

Interpolation은 간단하게 말하면, feature map을 벌려서 그 사이 값들을 중간값으로 메꾸는 방법이다.

자세한 내용은 상단 블로그 참조.
FCNs에서는 Bilinear interpolation과 Deconvolution방법을 함께 사용하여 Coarse Feature map으로부터 Dense prediction을 구했다.
반면, U-net은 학습 가능한 parameter를 활용한 Deconvolution을 통해서 upSampling을 수행하였다. Deconvolution은 쉽게 생각하면, convolutional layer를 거꾸로 연산하는 것이다.

kernel를 곱해서 convolution 연산과는 반대로 커지는 방향으로 output을 구성한다. 이러한 작업을 통해서 coarse map에서 dense map을 구현할 수 있었지만, 원래 feature map이 크기가 작은 상태였기 때문에 dense map은 아직 거칠고, 정보가 부족한 상태이다.

따라서 보다 정교한 segmentation을 위해서 추가 작업이 필요한데, 그것은 바로 resnet에서도 나온, skip-connection이다.

2.1.3. skip-connection

이 논문에서는 추상적이면서, 선이 굵은 의미론적 정보들(semantic)과 더 디테일한 appearance 정보들을 결합한 skip architecture를 정의한다.

시각화 모델에서 입력 이미지로부터 얕은 층에서는 주로 직선, 곡선, 색상 등의 낮은 수준의 특징을 잘 파악하고, 깊은 층에서는 보다 복잡하고 포괄적인 개체 정보를 파악할 수 있다고 한다. (내가 이해한 바로는, 깊은 층에서 이미지에 있는 객체들의 전체적인 형태를 뭉뚱그려서 파악하고, 얕은 층에서 그 이미지의 상세한 정보(위치에 따른 색상, 직선, 곡선 등)을 파악하는 것 같다.)
 따라서 저자는 이러한 직관을 바탕으로 앞에서 구한 Dense map에 얕은 층에서의 정보를 결합하는 방식으로 segmenation의 품질을 개선했다.

FCN - 16s

FCN - 8s

이러한 skip architecture를 통해 아래와 같은 결과를 얻을 수 있었다.

Unet

UNet 논문에서 왜 이렇게 FCN 설명을 길게 했냐면, UNet도 FCN의 구조에서 확장과 변화를 준 구조이기 때문이다. UNet의 구조는 아래와 같다.

FCN과 마찬가지로 input image에서부터 사이즈가 점점 작아지면서 global하고 coarse한 정보를 파악하는 contraction path, 그리고 다시 커지면서 동시에, skip connection을 통해서 detail하고 local 정보들을 파악하는 expansive path로 나뉜다. 이 때 FCN과 다른 점은 다음과 같다.

FCN vs Unet

FCN

UNet

1. FCN은 upsampling을 할 때 bilinear interpolation도 쓰지만, UNet은 학습 가능한 파라미터를 활용한 Deconvolution을 사용한다.
2. FCN은 skip-connection 과정에서 두 feature map을 더하지만(add), UNet은 두 개를 합친다(concat)
3. UNet은 up-sampling을 단계별로 수행한다. 하지만, FCN은 한 번에 진행한다.

overlap-tile strategy

이미지를 UNet에 통과시키면 결과적으로 이미지 사이즈가 작아진다. 노란색 박스 영역을 segmentation 하기 위해서는 사이즈가 더 큰 파란색 영역을 들고와야 하는데,

전체 이미지를 한번에 processing하는 것이 아니고, 일부분을 조금씩 차례대로 처리한다.

 그러면 원래 이미지보다 예측 범위가 작아진다. 따라서 원래 이미지보다 예측 범위가 작아져서 전체 이미지를 처리할 수 없게 된다.

이를 해결하기 위해 이미지 미러링을 활용한 Extrapolation 기법을 사용한다. 즉, 이미지 바깥으로 원본 이미지를 데칼코마니한 것처럼 더 그려서 input으로 넣는다는 것이다. 이 때, 위에 사진을 보면 '파란색' 영역이 겹치다 보니 overlap-tile 방식이라고 표현하는 것 같다. 아래는 extrapolatation(외삽)을 할 때 어떻게 미러링을 하는지 보여주는 사진이다.

3. model architecture

localize를 위해서 고화질의 feature들이(from contracting path) upsampling이 된 output과 결합한다.(회색 화살표) 이러한 연속된 convolution layer는 이 정보를 바탕으로 더 적절한 output을 모으기 위해 학습한다. 그림에서 볼 수 있듯이 contracting path는 이미지의 context 포착. 즉, 주로 큰 특징을 학습하는데 효과적인 구조이며

expanding path는 input image와 같은 해상도의 segmentation map을 그리기 위해 다시 복원하는 작업이다. 

이 때, contracting path의 feature map을 expanding path에 연결하는, skip-connection이 추가된다. 이렇게 함으로써 contracting하기 전의 이미지를 학습에 녹여줌으로써 이미지의 세밀한 localization을 할 수 있게 된다. 여기서 살짝 짚고 넘어가야할 건, expanding path를 거친 feature map과 contracting path에 있는 feature map을 concat해야 하는데, 문제는 이 둘이 크기가 다르다는 것이다. 그 이유는 중간에 convolutional 연산으로 인해 조금식 크기가 작아졌기 때문이다. 따라서 contracting path에 있는 feature map을 잘라서(crop) 붙여줘야 한다.
 그런데, 최종 출력인 segmentation map의 크기는 input image 크기보다 작다. 왜냐하면 convolution 연산에서 padding을 쓰지 않았기 때문에 전체적으로 크기가 조금씩 작아지기 때문이다.

 

4. training

stochastric gradient descent를 최적화하는데 사용한다.
overhead를 최소화하기 위해 batch를 단일 이미지 크기로 줄인다.
높은 momentum을 사용한다.
네트워크의 출력은 픽셀 단위의 softmax로 예측되며, 최종 feature map(channel k)에 대한 픽셀 x의 예측 값은 다음과 같다.

a_k(x)는 x 픽셀의 activation이며, k는 클래스의 수이다.

참고 softmax 함수

따라서 Loss function은 cross-entropy 함수가 사용된다.

P_{l(x)}(x): 정답 레이블에 해당하는 k값을 반환하는 함수

정답 label에 해당하는 함수에 log를 취하기 때문에 확률값을 반환한다.
참고) cross entropy 함수

4.1. Touching cells separation

보통 사진에서 독사진이나 물체끼리 겹쳐지지 않은 사진도 많겠지만, 특히 세포 사진들 같은 경우는 다른 세포끼리 맞붙어있거나 하는 경우가 많다. 이러한 경우 때문에 논문에서는 세포가 인접해있는 경우에 세포의 경계 사이를 구분하도록 손실함수를 재정의했다.

이 때, 서로 맞닿아 있는 cell들의 경계를 명확하게 구분하기 위해 Weight map Loss가 추가된다.

• w_c(x): 픽셀 x에 있는 클래스의 빈도수에 따라서 결정되는 weight map. 
• w_0: 초기값. 상수항
• d_1: 픽셀 x와 가장 가까운 개체(여기서는 cell이라고 표현)와 거리
• d_2: 픽셀 x와 두번째로 가까운 개체와의 거리

따라서 d1과 d2가 클수록(멀수록) w(x)는 작아지고, d1과 d2가 작을수록(가까울수록) w(x)는 커지는 경향을 보여 가까울수록 loss를 높게 적용해서 더 민감하게 반응하도록 한다.

4.1. data augmentation

data augmentation - elastic deformation

출처: https://medium.com/secure-and-private-ai-writing-challenge/data-augmentation-increases-accuracy-of-your-model-but-how-aa1913468722

elastic deformation은 data augmentation 기법 중 하나라고 한다. data augmentation은 dataset을 새로 추가시키지 않고, 기존의 training dataset에 약간의 변형을 주어 부족한 training dataset으로 인한 성능 저하를 막는 방법이다. 아래 사진이 elastic deformation이다. 

보통 cell들은 왼쪽 그림처럼 반듯한 모양으로 있지 않는다. 항상 움직이고, 랜덤한 모습을 하고 있기 때문에 그러한 특성을 반영해서 사진에 사실적인 왜곡을 주면서 부족한 데이터 갯수를 충원하는 방식이 elastic deformation이라고 한다.
 
5. Experiment

참고) IoU(Intersection over Union)

설명 출처: https://ballentain.tistory.com/12

object detection이나 segmentation model의 성능 평가를 하는 과정에서 사용되는 도구로 생각하면 될 것 같다. ground-truth의 영역과 predict 영역이 얼마나 겹쳐있는지 숫자로 보고 싶을 때 쓴다고 생각하면 될 듯. 수식은 아래와 같다.

두 영역의 교집합의 넓이 / 두 영역의 합집합의 넓이

참고)

https://medium.com/@msmapark2/u-net-%EB%85%BC%EB%AC%B8-%EB%A6%AC%EB%B7%B0-u-net-convolutional-networks-for-biomedical-image-segmentation-456d6901b28a
https://dacon.io/en/forum/405849
https://www.youtube.com/watch?v=O_7mR4H9WLk 
https://89douner.tistory.com/297