Thesis/Paper Reading

U-Net: Convolutional Networks for Biomedical Image Segmentation

1bang9 2024. 4. 24. 12:23

이번 리뷰는 세그멘테이션을 하기 위해 많이 쓰이고 있는 U-Net에 대해 리뷰하고자 합니다.

목차는 다음과 같이 구성됩니다.

0. BackGround

1. Coarse Feature map to Dense prediction

2. Architecture Detail

3. Strategic

4. Experiment

 

0.BackGround

UNet은 FCN을 기반으로 하여 개발되었으며, 설명에 앞서 간단하게 FCN의 개념과 U-Net의 기반이 된 이유를 설명할 것입니다.

U-Net은 FCN(Fully convolutional network)을 토대로 확장한 개념이기 때문에 명확한 이해를 돕기 위해 우선적으로 FCN을 이해하는것이 중요합니다.

Image classification 모델들은 기본적으로 내부 구조와 관계없이 모델의 근본적인 목표를 위해 출력층이 Fully-connected layer로 구성되어 있습니다. 이는 각 픽셀에 대한 특징값을 모두 가져오기 위함이자, 각 카테고리에 대한 확률을 출력해야 하므로, 각 클래스에 해당하는 확률 점수를 제공하는 단일 차원의 벡터를 출력합니다.

 

반면 Semantic segmentation model은 본 이미지의 각 픽셀에 대해 클래스를 구분하고 인스턴스 및 배경을 분할하는 것으로 위치 정보가 매우 중요합니다. 하지만 Semantic Segmentation 관점에서는 fc layer가 갖는 한계점이 존재합니다.

FC Layer의 경우 가중치 개수가 고정되어 있기 때문에 바로 앞 레이어의 Feature Map의 크기도 고정되며, 연쇄적으로 각 레이어의 Feature Map 크기와 Input Image 크기 역시 고정됩니다.

 

 

하지만 Flatten으로 공간적 배열을 배제한 Fc layer와 달리, Conv layer는 입력을 Flatten하지 않으며, 각 입력 픽셀 주변의 지역적인 정보를 활용하여 출력을 생성합니다.  첫 번째 fully-connected layer를 (7x7x512) 4096filter conv로 변경하면 가중치의 수가 유지됩니다. 단일 차원의 벡터를 출력 대신에 원본 이미지의 위치 정보를 담은  Feature map으로 출력이 가능합니다.

 

Semantic segmentation의 최종 목적인 픽셀 단위 예측과 비교했을 때, FCN의 출력 Feature map은 너무 coarse합니다224*224 -->> 1/32 -->> (7*7)

 

위 summary처럼 레이어가 5개 있다면, 다운샘플링 요소는 2^5가 됩니다. 이는 출력 특성 맵의 차원이 입력 이미지 차원의 1/32가 됨을 의미합니다.

 

따라서, Coarse map을 원본 이미지 크기에 가까운 Dense map으로 변환해줄 필요가 있습니다

1. Coarse Feature map to Dense prediction

  1. 업샘플링 방법: U-Net은 업샘플링을 위해 **트랜스포즈 합성곱(Transposed Convolution)**을 주로 사용합니다. 이 방법은 FCN에서 사용되는 Bilinear Interpolation이나 Backward Convolution(Deconvolution)과는 다르게, 학습 가능한 파라미터를 이용하여 업샘플링을 수행합니다. 
  2. 스킵 연결(Skip Connections): U-Net의 또 다른 중요한 특징은 스킵 연결을 사용하여 하위 계층에서의 고해상도 피처 맵을 업샘플링 과정의 해당 계층으로 직접 전달한다는 점입니다. 하위 계층의 공간적 정보를 상위 계층의 업샘플링된 출력과 결합함으로써, 특히 객체의 경계를 더 선명하게 만드는 데 도움을 줍니다.

U-Net 아키텍처에서는 모든 부분에 업샘플링을 하는 것은 아닙니다. U-Net의 구조는 두 주요 부분으로 구성되어 있습니다: 수축 경로(Contracting Path)와 확장 경로(Expansive Path).

Contracting Path  ( Encoder )

Contracting path 는 일반적인 합성곱 네트워크와 유사하게, 여러 합성곱 계층과 풀링 계층을 포함하여 특성 맵의 크기를 점차 줄여 갑니다. 이 과정에서 특성 맵은 점점 더 작아지지만, 동시에 특성은 더 추상화되고, 높은 수준의 의미 정보를 담게 됩니다.

Expanding Path 

확장 경로에서는 업샘플링이 주로 이루어집니다.  Contracting path 를 통해 얻은 집약된 특성 맵을 다시 원래의 이미지 크기와 유사한 크기로 확장하여, 픽셀 단위의 정밀한 예측을 가능하게 하는 것입니다. 업샘플링은 트랜스포즈 합성곱(또는 역합성곱)을 통해 이루어지며, 각 업샘플링 단계마다 해당 수준의 Contracting path 에서의 특성 맵과 결합(Concatenate)하여 세밀한 위치 정보를 복원합니다. 여기서 해당 수준이란 Skip Connection 을 의미합니다.

 

Skip Connection

각 업샘플링 단계에서, U-Net은 스킵 연결을 사용하여 Contracting path 에서의 해당 계층의 출력을 Expanding Path의 입력과 결합합니다. 이렇게 함으로써, 업샘플링 과정에서 잃어버린 세부 정보를 회복하고, 정확한 픽셀 위치 정보를 유지할 수 있습니다.

결론적으로, U-Net에서는 확장 경로에서만 주로 업샘플링이 이루어지며, 이는 고해상도의 특성 맵을 생성하여 세밀한 의미론적 분할을 달성하기 위함입니다. Contracting path 에서는 데이터의 크기를 줄이는 다운샘플링이 주로 이루어집니다.

 

 

2. Architecture Detail

  • 3x3 convolutions을 두 차례씩 반복 (패딩 없음)

패딩이 없이,3x3 합성곱 커널을 입력에 적용할 경우, 출력 차원이 소폭 감소하는데, 이는 커널이 입력의 가장자리를 완전히 커버할 수 없기 때문인데, 예로 들어 5X5 의 입력이미지에 3X3 커널을 적용할때,

  1. 합성곱 커널 크기: 3x3 커널을 예로 들면, 이 커널은 중심이 되는 픽셀 주변의 8개 픽셀(총 9개 픽셀)을 고려하여 계산을 수행합니다.
  2. 입력 특성 맵 가장자리 처리: 입력 특성 맵의 가장자리에는 커널을 적용할 충분한 픽셀이 없습니다. 예를 들어, 입력 특성 맵의 크기가 5x5일 경우, 가장자리를 제외한 중앙의 3x3 영역만이 완전한 합성곱 계산에 사용될 수 있습니다.
  3. 출력 특성 맵 크기: 따라서, 합성곱 연산 후의 출력 특성 맵의 크기는 입력보다 각 차원에서 2픽셀씩 작아집니다 (5x5 입력 -> 3x3 출력).
  • 2x2 max-pooling (stride: 2) 

커널이 2 픽셀씩 이동하므로, 입력 특성 맵의 가로와 세로 크기는 각각 반으로 줄어듭니다. 예를 들어, 입력 특성 맵이 8x8 크기라면, 2x2 max-pooling 후의 출력 특성 맵은 4x4 크기가 됩니다.

 

다운샘플링 마다 채널수가 2배가되는데 2배가된 후 빈공간은 가중치로 채워지며,

이 가중치는 Contracting path 에서의 해당 계층의 출력이됩니다.

Contracting path 에서 온 스킵 연결을 통해 전달된 특징 맵은 종종 업샘플링된 특징 맵과 공간 크기가 다를 수 있습니다. 이를 해결하기 위해, 스킵 연결을 통해 온 특징 맵을 크롭하여 업샘플링된 특징 맵과 크기를 맞춥니다. 이 크롭 과정은 테두리를 자르는 것을 의미하며, 주로 특징 맵의 가장자리 부분을 제거합니다.

 

위와 같은 구성으로 총 23-Layers Fully Convolutional Networks 가 이루어집니다.

Convolution 연산에서 패딩을 사용하지 않았기 때문 최종 출력인 Segmentation map의 크기는 Input Image 크기보다 작습니다.

 

3. Strategic

 

Overlap-Title Input

이미지의 크기가 너무 커서 한 번에 처리할 수 없거나 모델의 아키텍처 제한으로 인해 전체 데이터 세트를 동시에 처리할 수 없는 경우 닿아 있는 세그멘테이션 사이의 경계부분을 분리하기 위해 사전에 학습된 가중치 맵을 사용합니다.

는 중심 가중치를, 는 조정 가능한 가중치 상수(여기서는 10으로 주어짐), 는 가우시안 함수의 표준 편차(여기서 대략적으로 5 픽셀로 추정됨)를 나타냅니다. 와 는 주어진 픽셀 에서 두 지점 사이의 거리를 측정하는 함수입니다.

일반적으로 거리에 따라 지수적으로 감소하는 가중치를 생성하여, 중심에서 멀어질수록 가중치가 작아지도록 합니다

특정 픽셀이 객체의 중심에 가까울수록 더 높은 가중치를 받아, 해당 객체에 더 많은 주의를 기울이도록 만들기 위해 사용될 수 있습니다.

식에서 값을 조절함으로써, 가중치가 얼마나 빠르게 감소할지를 제어할 수 있습니다. 가 크면 가중치 감소가 느리고, 가 작으면 빠르게 감소합니다.

 

Training

 

여기서 a(x)는 픽셀 x의 activation, K는 클래스의 수로 출력 값은 픽셀 단위의 위 식처럼 Softmax로 예측됩니다. 

즉,  최종 특징 맵(채널 k)에 대한 픽셀 x의 예측 값이 계산됩니다.

따라서, Loss Function은 Cross-Entropy 함수가 사용된다.

객체들이 서로 닿아있는경우 Weight map Loss가 포함하여, 경계 주변에 더높은 가중치를 포함시킵니다.

결국 저자는 Bio분야에서 동일한 label을 가지지만 서로 붙어있는 경우에도 segmentation을 잘 수행할 수 있는 방법을 찾았습니다.

Experiment

Warping Error : 객체 분할 및 병합이 잘 되었는지 세크멘테이션과 관련된 에러
Warping Error를 기준으로 “EM Segmentation” 데이터에서 U-Net 모델이 가장 좋은 성능을 보이고 있습니다.

 

U-Net 모델은 “PhC-U373” 데이터에서 92% IOU Score를 획득하였으며 2등 모델이 획득한 점수 83% 와 현격한 차이를 보이고 있습니다