Contents

[Paper Review]Convolutional Networks for Biomedical Image Segmentation(UNet)

GoodYoung included in Paper Review DL

2024-05-17 1019 words 5 minutes

Contents

개요

Image Segmentation 문제를 해결하기 위해 제안된 대표적인 딥러닝 모델인 UNet에 대한 논문 리뷰를 해보려고 한다.

Abstract

딥러닝 네트워크에선 성공적인 훈련을 위해선 수많은 annotated training samples들이 필요한데 이것들을 효율적으로 사용할 수 있는 강력한 data augmentation 방법에 대해서 소개한다.

그 다음으론 Unet의 구조에 대해서 간략한 소개가 나온다.

Unet의 구조는 context들을 포착할 수 있는 contracting path와 localization을 정확하게 해주는 expanding path가 서로 대칭적으로 이루어져 있다는 것을 알려준다.

그리고 또한 Unet은 적은 양의 image로 학습을 할 수 있고 ISBI challenge에서 그 전 모델보다 더 좋은 성능을 보여주었다.

이런 Unet에 대해서 좀 더 자세히 Introduction의 파트에서 설명하게 될 것이다.

Introduction

본 논문에서는 지난 2년 간 deep convolutional net이 성공을 이룬 이유는 Krizhevsky의 연구(Alex Net) 때문이라고 설명했다.

보통 convolutional network의 사용은 classification task(하나의 이미지를 single class label) 였지만 biomedical image task는 localization(class label이 픽셀 단위로 분류) 이 필요했다.

또한 biomedical tasks에서는 수많은 양의 dataset들을 구하기 어려운 한계도 있었다.

/images/paper/unet/sliding-window.png — [그림 1] sliding-window

따라서 위 그림 처럼 이전에 sliding-window 방식으로 픽셀의 주변 patch가 한 단위가 되어 pixel predict를 하여 localization을 수행하는 방법이 있었지만 두가지의 단점이 있었다.
- 모든 patch 별로 연산을 하므로 시간이 오래 걸린다.
- 수많은 중복된 패치들이 많다.
- patch 크기가 크면 max-pooling layer가 많아져서 localization accuracy가 떨어지고 반대로 patch 크기가 작아지면 littel context를 하여서 localization accuracy와 use of context 간의 trade off가 있다.

따라서 이러한 단점들을 해결하기 위해 Fully convolutional network의 구조를 변형하여 더욱 elegant한 구조로 모델을 만들었다.
- 적은 데이터셋으로도 work가 가능하고
- 더욱 더 정밀한 segmentation이 가능해짐

/images/paper/unet/unet-main-idea.png — [그림 2] unet main idea

Unet의 main idea는 contracting layer(downsampling) 에서 pooling operator가 upsampling operator로 교체되는 연속적인 layer들을 보완하는 것이다.
- contracting path로 부터 high resolution feature들이 upsampled output과 합쳐진다.
- 그런 다음 연속적인 합성곱 층들이 더 precise한 output을 조합한다.
- output의 해상도를 증가시킨다.

Upsampling part에서 더 많은 feature channel들을 가지고 있게 된다. 이는 context 정보가 high resolution으로 잘 전달이 되게 한 것이다.
- 따라서 contracting path와 expansive path가 U-shaped architecture가 된다.

또한 Unet은 fully connected layer를 사용하지 않고 오직 각각의 convolution의 필요한 부분만 사용을 하였다.

이것은 overlap-tile strategy으로 임의의 큰 이미지의 seamless segmentation을 가능하게 한다.
- 다른 방식은 GPU memory의 한계 때문에 거대한 이미지를 처리할 수 없었다.
- Sliding Window 방식과 유사하지만 1개의 픽셀만 classification을 하는 것이 아닌 패치 영역을 classification을 한다.
- 또한 전체 영역을 다 예측하는 것이 아닌 572의 입력을 받아 388 크기의 아웃풋을 생성한다.
  - fully connected layer를 사용하지 않고 오직 convolution만 사용하여 임의의 이미지를 입력받아도 문제가 없다.
  - 큰 이미지를 한정된 GPU에서 학습하게 된다.

Border region을 예측하기 위해서 missing context는 input image를 mirroring 하여 extrapolated(나타나게 만들 수 없는 부분을 예측)로 예측이 된다.
- 위의 오른쪽 그림을 보면 노란색 박스 영역을 예측 하려면 파란색 박스 영역의 이미지(노란 부분보다 더 큰 크기)가 필요한데 이렇게 누락되는 부분은 mirroring의 전략을 가져간다.
  - 실제로 오른쪽 그림의 왼쪽을 자세히 보면 테두리의 이미지가 원본 이미지랑 대칭임을 볼 수 있게 된다.
- 또한 그 다음 그림은 단계로 넘어가는 과정인데 이전에 예측에 사용했던 patch와 겹치기 때문에 overlap-tile이라고 불린다.
- 또 만약 572의 입력이 들어올 때 388로 예측을 하기 때문에 외각 부분은 예측할 수 없어서 이 때문에 mirroring 전략을 취하는 것이다.

/images/paper/unet/elastic-deformations.png — [그림 3] elastic deformation

Unet은 little training data에서도 가능했다. 모델 학습 시 data augmentation을 image의 invariance(이미지 불변성)를 학습하는 elastic deformation방법을 적용하였다.
- deformation: 변형
- image를 다양하게 변형하여 augmentation을 하는 방법이다.
- Data Augmentation 파트에서 자세히 다루게 된다.

이 방법은 label이 있는 image를 학습하는 방법(annotated image corpus)이 아니어도 이런 deformation을 통해 image의 invariance를 학습하면서 더욱 네트워크가 robust하게 된다.

또한 weighted loss를 사용하여 접촉이 있는 세포들 사이의 background label은 loss function에서 큰 가중치를 얻게 하여 접촉되어 있는 객체를 구분하였다고 나온다.
- Training 파트에서 자세히 다루게 된다.

Network Architecture

이번엔 UNet의 Network 구조에 대해서 알아보도록 하겠다. 기본적으로 UNet은 Contracting path와 Expansive path로 나뉘는 구조를 가진다.

Contracting path는 일반 convolutional network랑 동일하다.

Expansive path는 up-sample을 하고 Contracting path에서 추출한 feature map을 copy and crop을 하여 붙인 후 2번의 convolution을 거친다.
- 이때 crop을 하는 이유는 contracting path와 concatenation을 하려는 feature map과 해상도를 맞춰야 하기 때문이다.

마지막 layer에서는 class의 개수로 mapping하기 위해 1*1 convolution이 사용된다.

over-lap tile을 수행하기 위하여 max pooling에 들어가는 x,y가 모두 짝수가 되도록 input tile size를 생각해줘야 한다고 나와 있다.

둘 다 짝수여야 각 영역이 정확히 나누어 떨어지게 되기 때문이다.

더 자세한 코드에 대한 구현은 여기에서 보면 될 것 같다.

Training

다음으론 UNet을 어떻게 학습을 시켰는지에 대해서 나온다.

본 논문에서는 오버헤드를 최소화 하고 GPU자원을 최대한 활용하기 위해 많은 작은 이미지를 한 번에 처리하는 대신, 큰 크기의 이미지를 하나씩 처리하도록 배치 크기를 줄이는 것이다.
- 고해상도 이미지를 다루는 작업에서는 mini batch 큰 이미지를 하나씩 처리하는 전략을 취한다.
모멘텀을 0.99로 설정하여 이전에 조정한 가중치를 현재 시점의 가중치를 조정하는데 많이 반영하도록 했다.
- 모멘텀 계수가 높을수록 과거의 기울기 정보가 현재의 가중치 업데이트에 더 큰 영향을 미치게 된다.
- 진행하던 속도에 관성이 적용되니까 local minima에 빠지더라도 그 지점을 벗어날 가능성이 있게 된다.

/images/paper/unet/soft-entropy.png — [그림 4] energy function

그리고 energy function을 사용했다. 이는 맨 마지막 layer에 pixel별로 soft-max(p_l(x))를 수행하고 여기에 cross entropy loss를 적용하는 식이다.

또한 weight map을 Cross Entropy에 곱하였는데, 더 중요한 pixel을 더욱 더 강조하기 위해서이다. weight map수식은 아래 그림과 같다.

/images/paper/unet/weight-map.png — [그림 5] weight map

weight map에서 특정 클래스가 가지는 fixel freuency 차이를 보완해주고, 접촉하는 cell간의 small separation border를 학습을 시킨다.

touching cell의 경우에 그 두 클래스 간의 픽셀 수에 불균형이 있을 경우, 모델은 빈도가 많은 클래스에 집중할 가능성이 있게 된다. 따라서 이런 차이를 보상해주는 것이 중요하게 된다.

이런 separation border는 morphological operation을 사용하여 계산이 된다.

Wc는 Class Frequency의 balance를 위한 것이고, d1은 근처 cell의 border 거리, d2는 두번째로 가까운 cell의 border 거리이다.

따라서 모델이 다양해지면서 weight initialization이 굉장히 중요해졌다.
- 그렇지 않으면 일부가 지나치게 활성화되어 다른 부분은 전혀 기여하지 않을 수 있다.
- 본 논문에선 Gaussian Distribution을 사용하였다.

Data Augmentation

Data Augmentation은 few training sample일 때 모델이 특정 invariance을 알게 되고 robust해지기 위해 중요하다.

microscopical image들은 특히 이런 Data augmentation이 필요하게 되는데 UNet은 shift, rotation등을 수행하였다.

또한 random elastic deformation그림 3을 사용한 것이 key concept이다.
- elastic(탄성): 속체에서 어떤 힘이나 시간 흐름으로 인해 변화가 발생하는 경우. 이 힘이 제거된 후 변형이 원래처럼 돌아오게 되면 이 변형을 탄성이라고 합니다. 이렇게 탄성이 있는 경우는 같은 물체라 해도 촬영 방법이나 각도 등에 의해서 다른 결과를 가져올 수 있으므로 이럴 때 사용하면 좋다고 하나 이 외의 경우에도 사용하게 된다.

contracting path의 end 부분에 drop-out layer을 추가하여 더욱 절대적인 data augmentation을 수행하였다.

Experiments

UNet은 왼쪽 그림의 EM segmentation challenge에서 Warping Error에 대해 가장 좋은 성능을 보였다.

또한 오른쪽 그림의 Cell Segmentation Task(ISBI cell tracking challenge)에서의 IOU가 2등과 확연한 차이를 보여주어 우승을 하게 하였다.

Implement

코드에 대한 구현은 여기에 해두었다.

Reference