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

임형빈·2024년 1월 29일

CV paper_review

PaperReview

목록 보기

5/12

Biomedical Task에서 Sell Segementation 에 관한 연구를 진행한 U-Net에 관한 리뷰이다.

U-Net에 관한 리뷰에 앞서, 기존 연구인 Fully Convolutional Networks for Semantic Segmentation에 관해 설명하고자 한다.

Fully Convolutional Networks for Semantic Segmentation

모델 핵심 구조

DownSampling

UpSampling

Skip Connection

DownSampling:
AlexNet, Vgg, GoogleNet 등의 합성곱 신경망 기반 모델은 이미지 분류를 위해 신경망 마지막에 FCLayer를 활용했다.
FcLayer는 고정된 크기의 input이 필요하다.
FcLayer를 통과하는 순간 이미지의 위치 정보가 손실된다.
따라서 FcLayer를 포함한 모델은 세분화 테스크에 적합하지 않다.

세분화 테스크에 합성곱 신경망 기반 모델을 활용하기 위해 FcLayer를 Convolution layer로 대체.

대체 결과, 고양이가 있을 법한 영역을 나타내주는 heatmap이 만들어지는 것을 확인.

1 x 1 convolution을 사용하여 위치 정보를 보존하면서 차원을 축소한다.
위치 정보가 보존된 상태에서 Filter들의 조합에 따른 정보 분석을 하기 위함이라고 보면 되겠다.

(너비와 높이 방향으로는 정보 축약을 하지 않고 채널 축으로 정보를 축악한다.)

DownSampling을 진행할 때, convolution과 pooling을 진행하면서 input 이미지보다 작아진 Feature map을 생성한다.
이미지 정보를 압축하는 역할을 수행하며, Feature를 추출했다고 볼 수 있다.

Upsampling은 세분화를 위해 작아진 featuremap을 다시 input 크기로 맞춰주는 과정이다.

Skip Connection:
DownSampling을 통해 나온 feature map은 크기가 줄어 하나의 값이 원래 이미지의 큰 영역을 나타내게 된다.
DownSampling결과에서 바로 Upsampling을 수행해 큰 영역에 해당하는 숫자를 생성하게 된다면 많은 정보가 손실되고 부정확하게 된다.
따라서 이전 Layer의 값을 다른 합성곱 연산을 수행하지 않은 채 마지막에 합쳐서 Upsampling에 활용한다는 idea.

Abstract.

모델 구조는 context를 캡쳐하는 수축 경로(contracting path)와 정확한 localization을 가능하게 하는 확장 경로 (symmetric expanding path)로 이루어진 구조.

강력한 데이터 증강 기법을 사용했다고 한다.
강력한 데이터 증강기법을 사용했기 때문에 적은 이미지 데이터셋에서 학습이 가능하다.
(Unet은 biomedical 테스크에서 세분화를 위한 모델, 의료 데이터는 데이터수가 적다는 한계점이 있다.)

Introduction

Biomedical image processing task의 경우, 이미지 안에 픽셀 단위에 각 클래스의 레이블을 할당해야 한다.
하지만 기존의 합성곱 신경망은 이미지 출력이 단일 클래스 레이블인 분류 작업에 사용되었다.
또한 기존의 합성곱 신경망은 대규모 데이터셋이 필요했으며, 의료 태스크에서는 그만큼 많은 훈련 데이터셋을 얻는 것에 한계가 있다.

Ciresan, D.C., Gambardella, L.M., Giusti, A., Schmidhuber, J.: Deep neural net- works segment neuronal membranes in electron microscopy images. In: NIPS. pp. 2852–2860 (2012) 앞선 연구에서, 이미지 내에서 픽셀 주변의 지역 (패치)를 추출하여 이를 인풋으로 사용한 sliding window 방식의 모델을 제안했다.

슬라이딩 윈도우 방식으로 위치 정보를 추출할 수 있었으며, 패치 데이터셋으로 적은 데이터셋의 한계를 극복했다는 의의를 가진다.

하지만 패치마다 신경망을 각각 실행해서 연산량이 상당히 많아 느렸다는 단점이 있고, 패치를 추출할 때, 겹치는 영역이 많이 발생했다.
또 패치 크기에 따라 위치 정확도와 context 간의 trade off가 발생한다.

더 큰 패치는 많은 max pooling layer가 필요해 위치 정확도에서 정보 손실이 발생하고,
작은 패치는 한정된 context만 볼 수 있기 때문이다. 이후 다른 연구들에서 단점을 보완해 나갔다.

본 저자들은, 이를 기반해 fully convolutional network를 제안했다.
구조는 위 사진과 같다.

구조 특징

이미지를 참고하여 의미(context) 정보를 추출하는 수축 경로 (contracting path)

의미 정보를 픽셀 위치 정보와 결합하여 각 픽셀마다 어떤 클래스 레이블에 속하는지 구분하는 확장 경로 (expanding path)

수축 경로에서 확장 경로로 전환되는 bottle neck 구간

주요 Method

Overlap-tile strategy 메소드를 사용했다. 테두리 부분의 소실된 부분은, mirroring extrapolate를 통해 붙여 넣었다.

의료 분야의 sell segmentation에서, 동일한 클래스의 객체들을 분리하는 것이 어렵다는 단점이 있다. 저자들은 경계가 가까운 픽셀들에 높은 가중치를 부여하는 weight loss를 사용해 이를 개선했다.

적은 데이터셋의 한계를 극복하기 위해 excessive하게 데이터 증강을 사용했다.

Network Architecture

Contracting Path (수축 경로):

주변 픽셀들을 참조하는 범위를 넓혀가며 이미지의 의미 정보를 추출한다.
DownSampling 시 채널의 수를 2배씩 증가 시킨다.

Expanding Path (확장 경로):

수축경로에서 동일한 level의 피쳐맵을 추출해, 크기를 맞추기 위해 cropping한 후 이전 layer에서 생성된 피쳐맵과 연결시킨다. (Concatenation)
합성곱을 거치며 피쳐맵의 크기가 줄어들기 때문에 동일한 수축 경로와 확장경로 Level에서 피쳐맵의 크기가 다르기 때문에 cropping 과정이 필요하다.

crop된 동일 Level의 수축 경로 피쳐맵과, 확장경로의 피쳐맵을 skip connection 시켜 의미 정보와 위치 경보를 결합 시킨다.
Upsampling 부분에서 많은 특징 채널을 가지고 있어, 의미(context) 정보를 더 높은 해상도의 레이어로 전파할 수 있다고 한다.

Overlap-tile strategy

U-Net의 합성곱 연산에서 패딩을 사용하지 않기 때문에 모델의 output 크기는 input의 크기보다 작다는 특징이 있다.
따라서 원하는 크기의 Segmentation Output을 생성하기 위해, 겹치는 부분이 존재하도록 이미지를 잘라 input으로 투입해야 한다.
노란색 사각형의 밑변 + 초록색 사각형의 밑변 크기의 output을 생성하기 위해, 파란색 사각형과 빨간색 사각형의 겹치는 부분이 존재한다.

Mirroring extrapolate

원 사진보다, 모델에 투입해야 할 이미지의 크기가 커야 한다.
따라서 추가적으로 필요한 테투리 부분의 픽셀을 보충하기 위해, 경계 부분의 이미지를 좌우 반전하여 mirror 이미지를 생성해 원본 이미지에 붙여 넣는다.
세포 이미지의 특성을 반영했다고 보면 되겠다.

Weight Loss

Cell Segmentation Task에서는, 동일한 클래스의 객체들을 분리하기 어렵다는 단점이 있다.
위 사진의 d처럼, cell의 경계에 속한 픽셀 부분에 집중하기 위해, 경계에 속하면 큰 가중치를 부여했다.
더 큰 Loss를 만들어 경계에 있는 픽셀을 더 잘 분류하도록 학습하겠다는 아이디어라고 보면 되겠다.

Total loss는 아래 식과 같다. 일종의 cross entropy loss라고 보면 되겠다.
~~본 논문에서는 Energy Function으로 정의했지만, 같은 의미로 해석했다.~~

where l : Ω → {1,...,K} is the true label of each pixel

P 값은 추론 값이라고 보면 된다.
$a_k(x)$ 는 픽셀 위치 $x$ 에서 Feature Channel $K$ 의 활성화 정도를 나타낸다.
이때 $K$ 는 클래스 수이다.

$w_c(x)$ 는 클래스 Frequencies를 밸런싱 하는 Weight Map이다.

저자들은 위 식에서 $w_0$ 를 10, 분산을 약 5로 설정했다.

$d_1(x)$ 는 가장 가까운 세포까지의 거리, $d_2(x)$ 는 두번째로 가까운 세포까지의 거리를 의미한다.
즉 경계에 있을 수록 $(d_1(x) + d_2(x)$ 가 작을수록) 큰 가중치 $w(x)$ 가 만들어진다.

Data Augmentation

현미경 이미지의 경우, 이동 및 회전 불변성과 변형 및 gray value 변동에 대해 invariance함과 robustness가 필요하다.
이를 위해 rotation, shift, random elastic deformation 등을 사용했다.
의료 테스크에 맞는 데이터 증강 기법을 사용했다고 보면 되겠다.