U-Net: Convolutional Networks for Biomedical Image Segmentation
본격적인 논문 리뷰에 앞서 U-net이란 무엇이고, 어느 분야에 주로 사용되며, 어떤 아이디어를 통해 나오게 되었는지 간략하게 설명해보겠다.
U-net은 이름에서 U가 들어간 것과 같이 모델의 구조 및 아이디어 자체가 U자 형태를 띈다.
이러한 U-net이 나오게 된 계기로는 Instance Segmentation을 진행하기 위해서이다.
앞전에 다룬 YOLO에서는 주로 어떤 object의 종류를 분류하거나 존재와 위치를 detection하는 Object Detection을 위한 모델링이 진행되었지만,
U-net에서는 다르다. U-net에서는 Instance Segmentation을 통해 각 물체들을 의미 있는 단위로 분할하는 작업을 실시한다.
How?) 각 픽셀마다 각 클래스로 분류(배경=0, 사람=1, ...)를 통해 한 장의 segmentation map을 만든다.
즉, 이미지의 pixel 수준에서 각 영역이 어떤 의미를 갖는지 분리하는 방법이다.
Segmentation?
Segmentation 모델은 다양한 분야에서 활용된다. 이미지 내의 영역 분리가 필요한 분야에서 많이 사용되며, 주로 병원, 자율주행 카메라 등에서 사용된다.
이렇게 이미지 내의 영역을 분리하는 접근 방식은 크게 두 가지가 존재한다.
- Semantic Segmentation
- Instance Segmentation
1. Semantic Segmentation은 한 이미지 내에 어떤 객체도 존재하지 않는다.(No Object) 오직 class의 위치를 인식(localization) 하고 판별(classification) 하는 것이다.
2. Instance Segmentation은 비슷하게 생겼지만, Multi Objects가 존재한다. 또한 같은 객체여도 종류가 다르면 따로 구분해준다. 이것은 각 객체의 위치를 정확히 식별하는 객체인식(Object detection) 의 접근법을 보여준다.
아래의 그림을 보면 이해가 쉬울 것이다.
또 다른 예시를 통해 이해해보자!
- Classification – 해당 사진이 특정 object를 나타내는지 판별
- Localization – 해당 사진 중 우리가 집중하고자 하는 부분(위치)를 골라냄
- Object Detection – 사진에 보이는 해당 object들을 각각 골라냄
- Semantic Segmentation – 의미 대로 Image 분할 (Pixel 분류)
- Instance Segmentation – object대로 사진 분할 (Pixel 분류)
왜 U-net을 만들게 되었을까?
논문 제목에서와 같이 for Biomedical Image Segmentation 때문이다.
세포같이 작은 단위를 구별하기에는 Object Detection이 아닌 위에서 언급한 Segmentation이 효과적일 것이다. 하지만 이러한 작은 이미지를 픽셀별로 분류한다면 이는 사람이 하기에 꽤나 어렵고, 비용도 많이 들 것이다.
이를 해결하기 위해 Biomedical에 쓰일 수 있는 딥러닝 모델 U-net이 탄생하게 된 것이다.
1. Abstract
딥러닝 네트워크를 성공적으로 훈련시키기 위해서는 수많은 training set이 존재해야한다. 본 논문에서는 이러한 샘플들을 더욱 효과적으로 학습시키기 위해 data augmentation을 제시한다.
모델링에서는 특징 추출을 위한 "contracting path" 와 정확한 localization을 위한 "expanding path" 로 이루어져있다.
이러한 네트워크로 인해 적은 수의 이미지로도 end-to-end 학습이 가능하며, 기존의 슬라이딩 윈도우 convolutional network보다 뛰어난 성능을 보인다.
2015년 ISBI cell tracking 대회에서 뛰어난 성능과 빠른 속도로 U-net의 우수함을 보였다.
2. Introduction
convolutional network의 일반적인 사용은 이미지에 대한 출력인 single class label의 도출이다. 하지만 Biomedical과 같은 세밀한 작업이 필요한 task에서는 이미지의 class 도출뿐 아니라 localization, 즉 각 픽셀에 클래스 label이 할당되어야 한다.
이를 해결하기 위해 U-net이 나왔으며, 논문에서 언급한 방식은 pixel과 pixel 주변의 영역을 받아 픽셀에 담긴 정보가 어떠한 객체를 나타내는지 판단하는 방식이다. 이렇게 되면 학습 데이터가 이미지 단위가 아닌 이미지 속 일부(patch)가 한 단위가 되기 때문에 훨씬 풍부한 데이터셋을 구현할 수 있다.
-
patch의 단점
하지만 이런 경우 (단점1) patch별로 연산을 하게 되어 속도가 매우 느리며 서로 중복된 영역을 가지는 patch들이 많아지기에 중복된 예측 결과가 나오기도 한다. 또한 (단점2) patch가 클 수록 더 많은 max-pooling 작업이 필요하기에 localization 정확도가 낮아지고, 반대로 patch가 너무 작으면 특징을 제대로 추출할 수 없는 trade-off가 발생하기도 한다. -
단점을 개선한 U-net
위에서 말한 단점들을 개선하여 U-net이 탄생하였다.
U-net은 fully convolutional network라는 구조로 기존의 방식보다 elegant한 구조로 만들어졌다.
U-net은 적은 데이터 셋에 학습시켜도 좋은 성능이 나오는 모델이며, 구조를 살표보면 특징을 추출하여 점차 해상도가 줄어드는 feature map을 up-sampling하고 이전에 얻은 feature map과 이어주는 과정을 거친다.
- 이 네트워크는 localizatoin이 가능하다.
- patch 측면에서 학습 데이터는 학습 이미지 수 보다 훨씬 더 많다.
- U-net의 몇 가지 특징들
1) upsampling과정에서 channel의 수가 더 많기 때문에 higher resolution layer에 context information을 전파할 수 있다.
: 개인적인 생각으로는 결국 왼쪽과 오른쪽을 concat할 텐데 오른쪽의 채널수가 더 많기에 정보전달을 잘 할수 있다고 해석했다.
2) contracting path(점차 해당도가 줄어드는 경로) 와 expansive path(점차 해상도가 증가하는 경로) 의 layer 개수가 거의 대칭이 되기 때문에 U자 형태를 가지고 있다.
3) Fully connected layer를 사용하지 않았기 때문에 patch에서 얻은 정보만을 가지고 해당 patch에서 classification을 수행한다.
: 이렇게 FC layer를 없애고 마지막에 Fully convolution layer로 마무리 짓는 이유가 있다. 이 부분은 "2. model architecture" 부분에서 이야기하겠다.
4) Data augmentation을 사용했다.
: 이러한 Data augmentation은 Biomedical segmentation에서 중요한데, deformation이 조직(근육조직 등)의 공통적 variation으로 가장 많이 쓰이는 방식이고 사실적 deformation은 효과적으로 시뮬레이션 될 수 있기 때문이라고 한다.
5) 세포의 경우 세포가 여러개 붙어있기에 구분하기 힘들다는 문제가 있었다. 이를 U-net은 거리를 이용한 가중치를 통해 만든 loss function으로 극복했다.
6) U-net은 Semantic Segmentation에서 다른 모델과 비교했을 때 좋은 성능을 보였다.
: Experiments에서 자세히 다루겠다.
2. Network Architecture
Network Architecture는 위에서 언급한 그림을 다시 가져오겠다.
이는 수축 경로(왼쪽:contracting path)와 확장 경로(오른쪽:expansive path)로 구성된다.
Contracting path
(출처 : https://www.jeremyjordan.me/convnet-architectures/)
수축 경로는 두 개의 3x3 convolution(패딩되지 않은 conv)을 반복적으로 적용한 다음 각각 활성화 함수(Relu-rectified linear unit)와 다운샘플링을 위한 stride=2의 2x2 max-pooling이 진행된다.(2x2 maxpooling을 위해 input tile size를 x와y가 짝수가 되도록 한다) 추가로 각 다운샘플링 단계마다 피쳐 채널 수를 두 배로 늘린다.
Expansive path
확장경로의 모든 단계는 feature map의 업샘플링과 특징 채널 수를 절반으로 줄이는 2x2 convolution(up convolution)으로 구성된다.
(up sampling은 주로 Interpolation(확장 후 중간값으로 매꿈), Deconvolution(파라미터를 이용해 확장-겹치는 영역은 더한다))
이에 더하여 localization을 더욱 잘 하기 위해, 위치정보가 잘 들어가 있는 수축경로의 feature map을 crop하여 확장경로에 붙여준다.
여기서 큰 궁금증이 생겼다!!
Q) !
위치(경계에 대한) 정보를 학습하고자 앞전의 feature map을 crop하여 concat을 시켰는데, 왜 굳이 crop한 것을 붙였을까?
crop하는 것이 아닌 conv layer를 거쳐서 뒷단의 feature map크기와 맞게 특징추출을 한 후에 이를 concat 시키면 더 괜찮지 않을까?=> 이에 대한 개인적인 생각으로는 외곽의 정보가 크게 필요가 없기에,
최대한 위치 정보를 잘 가지고 있는 앞전의 feature map을 잘라서 사용한 것이 아닌가 싶다.
마지막 레이어에서는 흔히 사용되는 Fully Connected layer가 아닌, U-net은 Fully Convolution layer로 convolutional layer로만 이루어진 network로 구성되어 있다. 맨 마지막 단계에 1x1 convolution을 사용하여 각 64개 컴포넌트 특징 벡터를 원하는 수(논문에서는 2개: foreground, background)의 클래스로 매핑한다.
이와 같이 픽셀별 클래스 분류로 전체 이미지를 2개 클래스로 구분한다.
Q. 그렇다면 U-net은 왜 Fully Convolution layer로 dense layer를 사용하지 않을까?
Segmentation task는 세분화된 픽셀별로 클래스를 정하기에 위치정보가 굉장히 중요하다. 하지만 이러한 정보가 dense layer에 들어가면서 사라지게 된다.
그렇기에 마지막에 Dense layer가 아닌 conv layer를 거치게 됨으로 이를 해결한다.
3. Training
U-net은 Training에서 재밌는 점을 확인해 볼 수 있다.
여기서 GPU를 최대한 활용하기 위해 이미지 파일을 여러개의 batch로 나누었다.
학습용 데이터 셋의 단위가 흔히 사용되는 이미지가 아닌 이미지의 '일부'인 것이다. 그렇기에 필요한 이미지 개수는 줄이면서 데이터 셋 내 입력 파일의 개수는 늘렸다.
Loss function
- a_k(x) : 픽셀 x가 class k일 값(픽셀 별 모델의 Output)
- p_k(x) : 픽셀 x가 class k일 확률(0~1)
- l(x) : 픽셀 x의 실제 label
- w0 : 논문의 Weight hyper-parameter, 논문에서 10으로 설정
- σ : 논문의 Weight hyper-parameter, 논문에서 5로 설정
- d1(x) : 픽셀 x의 위치로부터 가장 가까운 경계와 거리
- d2(x) : 픽셀 x의 위치로부터 두 번째로 가까운 경계와 거리
U-net에서 loss function은 cross entropy를 사용한다. 하지만 그림에서와 같이 세포간의 경계를 명확하게 하기 위한 추가적인 가중치 w(weight map loss)가 포함된다.(거리가 경계와 가까우면 loss값이 커진다.)
ex)
: w를 통해 세포 사이의 간격이 촘촘한 경우, 이를 잘 구분할 수 있게 만들어 준다.
정리하면 loss = weight map(w) X log(픽셀에서 얻은 클래스별 예측값을 softmax한 것) 이라고 볼 수 있다.
Data Augmentation
Data augmentatoin은 주로 데이터 셋이 많지 않을 때 사용된다.
현미경 등으로 촬영하는 이미지들은 색깔이 다양하지 않고, 회색빛으로 이루어져있기에 객체 간에 구별이 선명하지 않다. 따라서 Data augmenatation을 이용하여 풍부한 데이터 셋을 만든다.
=> 이를 통해 invariance하고 rubustness한 성질을 갖는 모델로 학습된다.
U-net은 Data augmentation을 위해 shift, rotation random-elastic deformation을 수행한다.
특히 논문에서 "random-elastic deformation(탄성 변형)" 을 작은 데이터 셋을 가지고 segmentation network를 학습시킬 때 key concept으로 여긴다고 하였기에 주목할 필요가 있다.
- Elastic Deformation 사용을 추천하는 경우 :
연속체에서 어떠한 힘이나 시간의 흐름으로 인해 변화가 발생하는 경우.
이 힘이 제거된 후 변형이 원래처럼 돌아오게 되면 이 변형을 '탄성'이라고 한다. 이렇게 탄성이 있는 경우는 같은 물체라고 해도 촬영 방법이나 각도 등에 의해서 다른 결과를 가져올 수 있으므로 이럴 때 사용하면 좋다고한다.
- Elastic deformation은 주로 다음과 같은 단계로 수행된다.
1) 이미지 그리드 생성: 원본 이미지와 동일한 크기의 그리드를 생성한다. 이 그리드는 각 픽셀의 위치를 나타낸다.
2) 그리드 변형: 생성된 그리드의 각 점을 미세하게 이동시킨다. 이동 벡터는 주로 정규 분포를 따르는 작은 값으로 무작위로 생성된다. 이 작은 이동은 이미지의 픽셀을 약간 찌그러뜨리거나 늘리는 효과를 준다.
3) 이미지 재샘플링: 이동된 그리드에 따라 이미지를 재샘플링한다. 재샘플링은 그리드 내 각 픽셀의 값을 사용하여 새로운 이미지를 생성하는 과정이다.
이렇게 생성된 이미지는 원본 이미지와 형태가 유사하지만 약간의 찌그러짐이나 늘어남을 가지게 된다.
이러한 변형된 이미지들을 훈련 데이터셋에 추가함으로써 모델은 이러한 변형에 대해 더 강인한 특성을 학습할 수 있고, 이는 모델의 일반화 성능을 향상시키며 데이터의 다양성을 높여 새로운 입력에 대한 예측 능력을 향상시킬 수 있다.
4. Experiments
Table1. EM segmentation challenge
전자현미경으로 관찰되는 뉴런 구조에서 cell segmentation task를 수행
데이터셋은 전자 현미경으로 찍은 512 x 512 해상도의 이미지 30장으로 이루어져 있고, 이미지의 각 부분에 세포는 흰색, 세포막은 검은색으로 칠해져있다.
- Pixel Error : prediction과 ground truth에 pixel마다 할당된 class들끼리 불일치하는 pixel들의 전체 pixel 기준 비율
- Warping Error & Rand Error : digital topology error 관련
U-net을 포함한 10개의 모델을 가지고 성능을 평가했고, 결과는 위와 같았다.
Table2. Segmentation results (IOU) on the ISBI cell tracking challenge 2015.
광학 현미경(light microscopic)에서 얻은 이미지로도 cell segmentation task를 수행
a와 c가 입력 이미지이며, b와 d가 ground truth segmentation map이다.
이렇게 a와 c같은 이미지를 통해 세포를 얼마나 잘 구별할 수 있는지 시험해보는 것이다.
성능은 IOU로 측정했고, 결과는 아래와 같이 굉장히 우수한 성능을 나타냈다.
5. Conclusion
U-net은 U자형 아이디어 모델링을 통해 biomedical segmentation 분야에서 굉장한 성능을 보였다. 논문에서는 Data augmentation에서 사용한 Elastic Deformation 기법 덕분에 적은 사이즈의 데이터셋만 요구했고, 이를 통한 합리적인 학습시간을 가졌다고 했다.
또한 이를 통해 U-net의 구조가 다양한 task에서 쉽게 응용될 수 있을 것이라 확신하며 논문을 마쳤다.
