U-Net: Convolutional Networks for Biomedical Image Segmentation

U-net은 Biomedical 분야에서 Image segmentation을 목적으로
제안된 Fully-Convolutinal Network(FCN)기반 모델이다.

그리고 왜 이름이 U-net이냐면...

진짜 생긴게 U처럼 생김 ㅋㅋ

 
그리고 이 내용을 이해하기 위해선 기초 지식이 필요하다. >> https://l2024-9999.tistory.com/41

Fully Convolutional Networks for Semantic Segmentation(FCN)

이 내용은 필자가 나름대로 이해해서 정리해둔 글이니 베이스가 없다면 한번 보는걸 추천한다. 
 

U-net의 전반적인 형태

Fully-Convolutional Network(FCN)을 기반을 만들어진 모델이라서
다운샘플링 > Feature map 저장 > 업샘플링(Feature map 적용) 순으로 작동하며 큰 틀은 다르지 않다.
 
하지만 뭐가 다르냐?
FCN은 Down-sampling을 진행하는 각 레이어를 지나면서 일부층의 Feature map을 저장한 뒤 최종적인 이미지에 적용시키지만
U-net은 Down-sampling을 진행하는 각 레이어를 지나면서 모든 층의 Feature map을 저장 한 뒤 각 Up-sampling 단계마다
해당하는 층의 Feature map을 적용시킨다.
 

정리하자면 FCN은 다 하고 나서 때려박고 / U-net은 하나하고 박고 하나하고 박는 방식으로 비교를 할 수 있다.

 
기본적인 FCN과 U-net의 차이점을 간단하게 설명했으니 각설하여 U-net에 대해서 자세히 알아보자.

 
먼저 U-net의 네트워크 구성은 Contracting Path / Expanding Path로 나눌 수 있다.

 

Contracting Path (Down-sampling)

 
각 Contracting Path마다 3x3 Convolution 레이어를 2번씩 거쳐서 이미지의 Feature map을 Down-sampling 한다.
사진을 보면 층이 깊어질수록 레이어가 변형되는 걸 볼 수 있는데, 이는 각 레이어의 마지막에 Max-pooling 연산을 하면서
이미지의 Feature map의 사이즈가 줄어들었기 때문에 다음 층의 Convolution 레이어가 이전에 거쳐온 Feature map의 형태에 맞게 조정된다.

 
 
Expanding Path

각 Expanding Path에서는 각 Convolution 레이어를 2번씩 거치서 이미지의 Feature map을 Up-sampling한다.
사진을 보면 이전 사진과 동일하게 변형되는 걸 볼 수 있는데 이는 Contracting Path에서 해당 레이어와 대응되는 Feature map을 결합한 뒤, Up-Convolution을 진행한다. 이 과정을 원본 크기까지 반복.
 
하지만 우리가 주목해야 할 점은 최종 출력인 segmentation map의 크기는 처음 들어온 원본보다는 작다는 것이다.
왜냐? Conv 연산에서 Padding을 안 썼기 때문이다.
 

Overlap-Tite input

FCN 구조의 특성으로 이미지 크기에 제약이 없다는 점이 있는데, 해당 연구진은 큰 이미지에 대해서 이미지를 다 사용하지 않고
Overlap-tite를 사용하였다.

Segmentation 적용 방식

이미지를 tile로 분리한뒤 입력으로 사용한다. 파란색이 각 타일의 전체 범위이고, 노란색이 Segmentation으로 사용한다.

다음 tile에 대한 Segmentation을 만들기 위해선 이전 tile의 일부가 포함되는 Overlap-tite방식을 사용한다.
 
여기까지 이해했다면 각 tile들의 경계들은 버리는건가? 싶은 생각이 들 것이다.
해당 연구진들은 경계 부분의 Segmentation을 위해서 임의의 Padding 대신 경계부분의 미러링을 활용하는 Extrapolation 기법을 사용했다.
 
쉽게말해서 경계 부분의 픽셀을 인접한 타일의 값을 미러링해서 채워넣는 기법이다라고 이해하면 되겠다.

중심 Segmentation의 모서리 부분이 경계부분에 미러링 된 모습이다.

 
Touching cells seperation

Touching cells Segmentation 작업은 Segmentation에서 중요한 단계로

이미지 b, c, d와 같이 각 세포 사이의 경계를 포착 할 수 있어야 한다.

U-net에서는 인접한 세포들이 중첩되어 보이는 것을 해결하기 위해 정확한 경계선을 필요로 하였다.
이를 위해 학습 데이터의 각 픽셀마다 클래스 분포가 다른 점을 인지하여 이에 맞게 가중치를 학습에 반영하였다.
 

Training

1. 입력 이미지와 해당 Segmentation을 확률적 경사하강법(SGD)을 사용하여 학습을 진행함.
2. Padding이 없는 Convolution 때문에 원본이미지보단 사이즈가 작다. 최적화를 위해 Overlap-tile를 진행하고 Batch-size를 1로 줄임.
3. 모멘텀을 높게 설정해서 이전 샘플로 학습을 진행하며 업데이트에 큰 영향을 주도록 설계
4. Softmax를 최종적인 Feature map에 적용하고, Cross Entropy loss function과 결합하여 함수를 계산.
5. 실제 이미지와 예측 이미지에 따른 확률간의 차이, 즉 loss 값을 계산한다.
6. 클래스의 편향을 해소하기위해 각 Segmentation에 대한 weight를 미리 계산하여 훈련중 특정 픽셀에 반영
 
 

정리

U-net은 의료 이미지 분할 및 세밀한 Segmentation 작업을 위해 설계된 딥러닝 모델로, FCN모델의 구조를 기반으로 설계되었다.
1. Overlap-tile input : U-net은 이미지 크기에는 제약이 없지만 최적화를 위해서 Overlap-tile방식을 사용했다.
이미지를 타일화한 뒤, 각 타일의 경계 부분은 segmentation을 미러링한 부분으로 채워 넣어서 segmentation의 정확도를 높였고,
이전 타일의 일부분이 포함되어 있기에 경계 부분의 정보 손실을 최소화 한다.
2. Touching Cells Separation : Segmentation 과정에서는 인접한 객체 사이의 경계를 정확히 포착하는 것이 굉장히 중요하다.
이를 보완하기 위해 U-net에서는 각 픽셀마다 클래스 분포의 편향을 인지하고, 이에 맞는 가중치를 학습에 반영하였다.
3. Contracting Path - Expanding Path 구조 : Down-sampling 과정에서는 이미지를 점진적으로 줄여나가며 각 Pooling층에서 Conv를 거친 Feature map을 저장하고, Up-sampling과정에서는 해당 Feature map을 활용하여 점진적으로 이미지를 원본 사이즈로 복원하여
최종적인 Segmentation을 출력한다. 

최종적인 Architecture를 간단하고 중심을 관통하게 설명하자면
고해상도 사진을 5 / 저해상도 사진을 1이라고 할 때
고해상도 사진(원본) 5짜리가 들어왔을때 n차 공정(conv)을 거치고
5의 특징을 저장한 뒤 4로 축소, 이 과정을 1까지 반복.
1에서 2로 확대한 뒤 저장 해둔 2의 툭징을 적용 후 n차 공정을 거쳐서
5까지 반복한다. 그리하여 최종적으론 원본 사진이 들어와서 모델을 거쳐 해당 사진의 클래스의 경계, 무엇인지 탐지하여 결과값을 보여준다.
 
참고 : 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

U-Net 논문 리뷰 — U-Net: Convolutional Networks for Biomedical Image Segmentation

출처 : https://link.springer.com/chapter/10.1007/978-3-319-24574-4_28