[paper review] U-Net: Convolutional Networks for Biomedical Image Segmentation

date

Sep 7, 2023

slug

paper-review-u-net

author

status

Public

1. introduction

image classification

AlexNet, VGG16과 같은 네트워크는 image classification에 사용되는 네트워크이다. 일반적인 CNN 분류 모델은 보통 single classification task에 사용된다.
하지만, biomedical image processing 분야에서는 한 이미지 내에 모든 pixel을 classification을 하는 semantic segmentation task가 중요하게 여겨진다.

Semantic Segmentation

그림 1 : sematic segmentation 로 분류한다.

예) 아이폰의 인물사진 모드 등에서 필요한 작업, →인물 사진 모드는 사람을 강조함. 여기서 semantic segmentation을 학습한 모델로 카메라에 담긴 이미지에서 사람을 파악하고 나머지 객체들은 흐릿하게 처리함.

그림 2 : 사물 인식하는 방법에는 다양한 유형이 존재한다. segmentation은 각 픽셀마다 클래스(class)를 할당하는 작업을 의미한다.

semantic segmentation 정의 : 이미지 내에 있는 물체(object)들을 의미 있는(semantic) 단위로 분할(segmentation)하는 작업을 의미한다.

일반적인 분류(classification)에서는 단일 이미지를 하나의 클래스로 분류했다.
분할(segmentation)에서는 각 픽셀마다 하나의 클래스

해당 논문은 이러한 semantic segmentation을 의료용 biomedical image 분석에 사용한 모델 u-net을 소개한 논문이다. 좌측 그림으로 부터 세포를 구분하여 우측그림과 같이 나타내는 것이 u-net이다.

2. Architecture of U-net : contraction path + expansive path

u-net의 네트워크 모양이 u자 형태를 가지고 있어서 u-net으로 이름이 붙여졌다.

빨간색으로 박스친 부분이 contracting path 부분, 우측 파란색 박스친 부분이 expansive path 이다.

contraction path

일반적인 convolution 연산이 진행된다. 각 contracting path의 step마다 3*3 convolution을 두 차례씩 반복한다. 연산마다 패딩이 없으므로 feature map이 조금씩 줄어든다. (convolution 연산에 활성함수인 ReLU도 포함함)

각 contracting path마다 2*2 max-pooling 연산을 수행해줌으로써 feature map 크기가 절반으로 줄어든다.

각 step에서 down sampling을 해주면서 채널의 수가 2배씩 늘어난다.

정리)

2*2 max pooling 사용 : 해상도(너비와 높이) 2배 감소
conv 연산으로 채널 크기는 2배 증가
일반적인 CNN모델처럼 다음의 형식을 반복

conv 연산들 → ReLU → max pooling

expanding path

각 expanding step마다 2*2 up convolution을 수행한다. 그렇게 해줌으로써 feature map의 크기가 2배씩 늘어나게 된다.

expanding step마다 3*3 convolution을 두 차례씩 반복한다. (활성함수는 ReLU)

이 구간에서도 패딩이 없으므로 feature map이 조금씩 줄어들게 된다. 그리고 up-sampling 마다 채널의 수가 절반씩 줄어들게 된다.

expansive path에서 매칭되는 contracting path 부분의 이미지를 연산한다. 그림에서는 회색 화살표를 나타내고 이것을 concatenationg이라고 한다.

이러한 연산을 하는 이유는 semantic segmantation에서는 위치정보가 중요하고, 그 위치정보의 소실을 방지하기 위해 이전 연산했던 값을 더해준다. 이것을 skip connection이라고도 한다.

정리)

2*2 convolution(up-convolution) 사용 : 해상도(너비와 높이) 2배 증가
conv 연산으로 채널 크기는 2배 감소
수축 경로(contracting path)에서 처리된 feature채널 레벨에서 연결(concatenation map)을 잘라내어(cropping) 가져와 수행

3. Training : objective function

세포(cell)를 명확히 구분하기 위해 작은 분리 경계를 학습한다.

초기 weight 값 = 10

: 픽셀 x의 위치에서 가장 가까운 경계와 거리

: 픽셀 x의 위치에서 두 번째로 가까운 경계와 거리

: 표준편차 5 pixel

해당 식에서 즉, 학습 데이터의 x 픽셀이는 x의 위치에 해당하는 클래스의 빈도수에 의해 결정되는 식이다.

만약, background라면 상대적으로 픽셀수가 많을 테니 해당 값은 매우 작을 것이다. 뒤에 보이는 은 픽셀 x의 위치에서 가장 가까운 경계와 거리를 나타내며, 다음 는 픽셀 x의 위치에서 두 번째로 가까운 경계와의 거리를 나타낸다.

energy function = cross entropy

u-net은 segmentation을 위한 네트워크이므로,픽셀 단위로 softmax를 사용한다.

하늘색 : weight map -픽셀별로 가중치를 부과한다.

노란색 : softMax에서 나온 정답 레이블

연두색 : feature map에 있는 각 픽셀

softmax 함수부터 보면 K는 클래스 최종적으로 cross entropy인 energy function을 구할 수 있다.

따라서, energy function에 보이시는 노란색에서 feature map에 픽셀 단위로 softmax를 수행한다. 초록색 부분을 보면 x는 feature map에 있는 각 픽셀을 의미한다. 그리고 w(X)는 픽셀 별로 가중치를 부과하는 역할로 위에서 계산한 값을 대입한다.

data augmentation

의료 데이터에서는 학습 데이터의 수가 적은 경우가 많으므로, data augmentation이 필요한다.

본 논문에서는 일반적인 data augmention 기술을 사용한다. 다양한 구성의 데이터 셋으로 만들어주니 견고한 성질을 가지는 모델로 학습된다.

현미경 등으로 촬영하는 사진들은 색상이 다양하지 않고 회색빛깔로 이루어져 있고 객체 간 구별도 선명하지 않기 때문에 data augmentation을 이용해 풍부한 데이터셋을 만드는 것이 더욱 필요하다.

세포에서는 deformation 즉, 기형현상이 흔하기 때문에 이러한 random elastic deformation을 사용해 더 좋은 성능을 내었다.

data augmentation method(elastic deformation)

overlap-tile strategy

하지만, 앞에 말했던 학습과정에서의 문제가 있다. input image의 크기와 ouput image의 크기가 다르다. 그렇기 때문에 실제 segmentation되는 부분은 input image의 일부만 된다는 것이다.

이를 해결하기 위해 input data를 넣을 때 논문에서 overlap-tile 기법을 사용하였다.

위 그림에서 bule tile과 yellow tile에서 overlap을 확인할 수 있고, blue area에 기반하여 yellow area의 segmentation을 예측했음을 확인할 수 있다.

그리고 missing data는 mirroring을 통해 보강하였다. (파랑 사각형 안의 빈 부분은 인접한 부분을 거울처럼 복사해 채웠다.)

이 부분을 mirroring 기법을 이용해 missing context부분을 채워주었다.

dataset

학습데이터는 위의 그림처럼 좌측에 보이는 것이 input 이미지, 우측 이미지가 라벨이다.

30개의 이미지 중 24개는 학습데이터 3개는 검증데이터 나머지 3개는 테스트데이터로 사용하였다. input data의 size는 512*512이다.

해당 데이터 셋을 활용하여 현미경으로 찍은 세포를 segmentation하였다.

4. experiments

3가지 기준(warping error, rand error, pixel error)에서 u-net이 rank 1을 차지한 것을 확인할 수 있음.

u-net은 다양한 biomedical segmentation applications에서 아주 좋은 성능을 보여줬다. 저자는 아까 Data Augmentation 에서 소개해 드렸던 elastic deformation이 Data augmentation 덕분에 적은 사이즈의 데이터셋만 요구했고, 논문에서 약 10시간 정도 학습 시간이 소요되었다.

마지막으로 u-net구조가 다양한 task에 쉽게 응용될 수 있을 거라 확신한다고 말하며 논문을 끝냈다.