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

논문 - U-Net_Convolutional Networks for Biomedical Image Segmentation, laf Ronneberger, Philipp Fischer, and Thomas Brox

(arxiv.org/abs/1505.04597)

1. Motivations/Problems

DCNN은 classification 분야에 많이 사용되고 좋은 성능을 보여준다. 하지만 DCNN을 Biomedical semantic segmentation에 DCNN을 적용하는데 문제가 있다. 

 

1. Biomedical image 처리는 네트워크의 결과로 픽셀 레벨의 localization을 요구한다.

즉 각 픽셀당 class label이 할당되어야 한다. 

2. biomedical 분야에서 충분한 훈련 데이터를 얻기 어렵다.

 

이를 해결한 연구로 sliding window 기법을 사용한 IDSIA 네트워크가 있다. 해당 기법은 픽셀 주위의 local region(패치)를 입력으로 하여 각 픽셀의 label을 예측한다. 

이를 통해 localization이 가능하고, 패치를 통해 학습하기 때문에 훈련 이미지가 적다는 문제를 해결했다.

 

하지만 이 네트워크는 두 가지 단점이 존재한다. 

1. 훈련 속도가 느리다. 

슬라이딩 윈도우 방식은 네트워크가 각 패치 별로 동작하면서, 이전 패치에서 검증한 부분을 다음 패치에서 다시 검증하는 중복이 많다. 이런 연산 낭비로 training 속도가 매우 느리다. 

2. localization accuracy와 context(semantic information) 사이 trade off 관계이다.

패치가 클수록 max-pooling layer가 많이 필요하게 되고, localization 성능이 감소한다. (max-pooling layer는 localizatioin accuracy를 감소시킨다.) 

반대로 패치가 작아지면 localization 성능은 향상하지만, 네트워크는 매우 작은 context만 보게 되어 context 파악 성능은 떨어진다. 

*context: 이웃한 픽셀들 간의 관계. 이미지의 일부를 보고 이미지의 문맥을 파악하는 것 (참고: mylifemystudy.tistory.com/87)

 

2. Key ideas

본 논문에서는 biomedical segmentation을 위한 네트워크로 full-convolutional network 기반 U-net을 제안한다. 

 

U-net의 장점

1. 빠른 속도로 훈련이 가능하다.

U-net은 이미 검증이 끝난 부분은 건너뛰고, 다음 patch부터 새로 검증하기 때문에 중복이 적어져 속도가 빠르다. 

2. up-sampling 연산으로 contracting path의 출력 해상도를 보완한다.

localization 정확도를 위해, Expansive paht의 upsampled output과 contracting path의 high resolution feature을 을 결합한다. 그리고 이 결합된 정보를 기반으로 연속 컨볼루션 계층에서 더 정밀한 ouput을 출력한다.

U-net에서 주목할 점은 up-sampling 에서도 많은 수의 feautre channel을 가지고 있어 네트워크가 context 정보를 더 높은 해상도 계층으로 전파할 수 있게 해 준다는 것이다. 이를 통해 localization과 context 사이 trade off를 해결했다. 

 

3. Design

U-net은 contracting path와 expansive path로 구성된 대칭 형태의 u자형 구조이다.

 

[출처] https://tech.socarcorp.kr/data/2020/02/13/car-damage-segmentation-model.html#index6

Contracting path

입력 이미지의 context(의미 정보) 파악을 위한 네트워크

• 전형적인 convolution network 구조

아래와 같은 과정으로 Down-sampling을 반복해서 feature map을 생성한다. 

• 3x3 convolution을 두 차례씩 반복(패딩 없음), 활성화 함수 ReLU
• 2x2 max-pooling

down-sampling 마다 feature 채널의 수가 2배로 늘어난다.

 

Expansive path

feature map을 up-sampling하고, 각 expansive path 단계에 대응되는 contracting 단계에서의 feature map과 결합해 더 정확한 localization을 수행하는 네트워크

• 전형적인 convolution network + up-sampling

다음과 같은 과정으로 feature map을 up-sampling하고, 각 expansive path 단계에 대응되는 contracting 단계에서의 feature map과 결합해 contextual 정보와 localization 정보를 결합한다. 

이때 중앙부분을 crop 한 contracting path의 feature map과 concat 하는데, 이는 컨볼루션에서 가장자리 픽셀의 손실이 발생하는 것을 보정하기 위해서이다. 

 

• 2x2 convolution (Up-conv)
• 3x3 convolution을 두 차례씩 반복(패딩 없음), 활성화 함수 ReLU

(각 conv block에 up-conv을 추가했는데, up-conv를 통해 feature 채널의 수를 절반씩 줄인다. )

 

최종 layer에서 desired number of class에 매핑하기 위해서 1x1 convolution을 수행한다. 

 

정리하면 U-net은 총 23개의 layer로 구성된다. unpadded convolution을 사용했기 때문에, 최종 출력 segmentation map의 크기는 input image보다 작다. 

 

4. Training

U-net은 효과적인 학습을 위해 다양한 전략을 사용했다. 

overlap-tile strategy

큰 이미지의 segmentation을 위해 ovelap-tile 전략을 사용했다.

 

[출처] 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

이미지를 타일(patch)로 나눠서 입력으로 사용한다. 파란 박스의 입력으로 노란 박스의 segmentation 결과를 얻는다. 다음 타일의 segmentation을 위한 입력에 이전 타일의 일부분이 포함되어 overlap-tile 전략이라고 한다.

[출처] 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

이미지의 경계 부분에 위치한 픽셀의 경우, 누락된 픽셀은 Mirroring extrapolation을 통해 추정한다. 

 

Touching cell separation

cell segmentation에서 주요한 과제 중 하나는 동일한 class의 인접한 개체들을 분리하는 것이다. 생물학에서 세포들이 인접해 있는 경우가 많기 때문에 이는 중요한 문제이다. 

 

U-net은 개체 사이 작은 경계를 포작하기 위해, ground-truth에 대한 weight map을 미리 만들고 네트워크 학습에 활용했다.

separation border는 mopoligical operation을 사용해 계산한다. 다음은 weight map을 계산하는 식이다. 

• d1(x): 픽셀 x의 위치에서 가장 가까운 경계와 거리
• d2(x): 픽셀 x의 위치에서 두번째로 가까운 경계와 거리

w(x)는 픽셀 x와 경계의 거리가 가까울수록 큰 값을 가지게 되어 해당 픽셀의 loss 비중이 커진다. 즉, 학습 시 경계에 해당하는 픽셀을 잘 학습하게 된다. 

Excessive data augmentatioin

Biomedical 분야는 대부분의 경우 훈련 데이터가 충분하지 않다. 따라서 U-net은 적은 데이터로 충분한 학습을 하기 위해 data augmentation을 적용했다.

• shift & rotation
• elastic deformation

shift와 rotation에 대한 invariance를 학습하기 위해서, coarse 3x3 grid에 random displayment vector를 사용해 smooth deformation을 생성했다.

 

실제 세포의 경우 deformation은 흔한 현상이기 때문에, data augmentation을 사용해서 더 좋은 성능을 얻을 수 있었다.   

 

Training

네트워크의 출력값은 픽셀 단위의 softmax에 의해 계산된다. 

 

최종 feature map에 대한 픽셀 x의 예측값은 다음 식으로 계산된다. 

• ak(x): 픽셀 위치 x에서 feature 채널 k의 활성화를 나타냄
• K: 클래스 수
• pk(x): approximate maximum function, 측 최대 활성화 ak(x)를 가진 k의 경우, pk(x)=1이고, 다른 모든 k의 경우 pk(x)=0이다.  

loss function으로 cross-entropy 함수와 touching cell을 분리하기 위한 weight map loss가 사용된다.

• l: 각 픽셀의 true label
• w: training에서 일부 픽셀에 more importance를 주기 위한 가중치 맵

5. Performance

다음은 U-net을 2개의 데이터셋으로 검증한 결과이다.

1) Segmentation of neuronal structures in electron microscopic recordings

U-net(input의 rotation 버전 7개의 평균)은 다른 전, 후처리 없이 warping error 0.0003529, rand error 0.0382로 슬라이딩 윈도우 컨볼루션 네트워크(IDSIA) 보다(IDSIA) 좋은 성능을 보인다.

*warping error: 위상 불일치(topological disagreements)(개체 분할 및 병합)에 페널티를 주는 세크멘테이션 메트릭 *rand error: 두 세크멘테이션 간의 유사성 측정  (출처: medium.com/@codecompose/cumedvision1-2e45b94e7b79)

2) Cell segmentation task in light microscopic image

‘PhC-U373’ 데이터 세트에서 U-net은 평균 IOU(intersection over union) 92%를 달성한다.

‘DIC-HeLa’ 데이터 세트에서는 평균 IOU 77.5%를 달성했다.

 

6. Conclusion

Biomedical segmentation을 위한 full-convolutional network 기반 U-net을 제시한다.

 

U-net 특징

• full-Caffed 기반 네트워크를 제공한다. 
• data augmentation을 통해 적은 양의 training 데이터로도 효율적으로 학습한다. 
• context와 localixation의 trade off 문제를 해결했다. 
• 다양한 biomedical segmentation application에 사용될 수 있다. 

[참고]

논문 - U-Net_Convolutional Networks for Biomedical Image Segmentation, laf Ronneberger, Philipp Fischer, and Thomas Brox

(arxiv.org/abs/1505.04597)

joungheekim.github.io/2020/09/28/paper-review/

mylifemystudy.tistory.com/87

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