Paper: U-Net: Convolutional Networks for Biomedical Image Segmentation

기존의 딥러닝 모델
: 주로 이미지 전체에 1개의 정답을 예측하는 데 사용되었음.
(ex. 개 ? 고양이 ? 사람 ?)
이는 이미지를 전체적으로 하나의 class로 분류하는 데 적합한 방법임.
이것이 Classification.
다만, 딥러닝이 다양한 field에서 사용되고, 발전함에 따라...
의료 영상 분야에서는 단순 class 분류만으로는 부족함.
훨씬 더 정교하고 정확한 예측이 필요함.
(ex. 종양이 '어디에' 있는지, '어떤 크기로' 있는지 등)
이럴 때 쓰이는 것이 Semantic Segmentation.
= 이미지 내 모든 픽셀 하나하나마다 어떤 class에 속하는지 예측함.
(ex. 이 픽셀은 종양이다. 이 픽셀은 배경이다. 등)
Deep convolution networks는 최근 많은 visual recognition task에서 높은 성과를 보여옴.
단, convolutional networks는 사용 가능한 데이터셋의 크기와, 네트워크 구조의 한계가 있었음.
이는, krizhevsky의 연구(AlexNET)를 통해 ImageNet 데이터셋에서 8개 layer와 수백만 개 parameters를 가진 대규모 네트워크를 훈련하여 극복할 수 있었고, 이 이후로 더 크고, 깊은 신경망이 훈련되어 왔음.
일반적으로, Convolutional networks는 classification tasks(분류 작업)에 사용됨.
이는 이미지에 single class label을 출력함.(ex. 개 or 고양이 구분 → '개'라는 하나의 클래스 label을 출력 = 이미지의 전반적인 영역에서 분류 작업을 하는 것. 이미지 내의 개별 픽셀이나 부분 정보에 대해서는 고려 X)
하지만, Biomedical image processing(의료 영상 처리)같은 분야의 visual task에서는 각 픽셀에 클래스 라벨을 할당해야 하는 'localization이 포함된 output'을 필요로 함. 또한, 의료 영상 처리 task의 경우, 수천 장의 학습 이미지를 얻는 것이 현실적으로 어렵다는 문제 존재.
그래서, sliding-window(슬라이딩 윈도우)텍스트 방식의 네트워크를 훈련시켜 문제를 해결하고자 했음.
픽셀 주변 local region(patch)을 input으로 써서 각 픽셀의 클래스 라벨을 예측하고자 함.
→ localization 수행 가능했고, patch 단위의 학습 데이터는 실제 학습 이미지 수보다 더 많았음.
→ 이 network는 ISBI 2012의 EM segmentation 부문에서 우승함.
다만, 이 방법에는 두 가지 문제가 존재함.
1. 속도가 느림
: network가 각 patch에 대해서 따로따로 실행되어야 하기 때문에, 속도가 꽤 느림.
2. trade-off 문제
: localization 정확도와 context 사용 간 trade-off 현상 문제가 존재함.
본 논문에서는, Fully Convolution Network(FCN) 아키텍처를 기반으로, 아주 적은 훈련 이미지를 사용하며, 더 정밀한 segmentation(세분화)를 할 수 있도록 개선했음.
↑ U-net architecture(ex. 32 X 32 pixel)
기존 아키텍처에서는 일반적인 contracting network(수축 network)에 연속된 layer를 추가해서, pooling operators를 upsampling operators로 대체하는 것임.
이를 통해, output의 해상도를 높일 수 있음.
그리고, localization을 하기 위해, contracting path의 고해상도 feature를 upsampling한 output과 결합.
이렇게 되면, 연속적인 convolution layer는 이 정보를 기반으로 더 정밀한 output을 생성할 수 있음.
U-Net, 업샘플링 단계에서의 수정(개선) 사항
업샘플링 단계에서 많은 피처 채널을 추가해서, 네트워크가 context information을 더 높은 해상도 layers로 전달할 수 있음.
그렇게 되면, 업샘플링 단계와 contracting path가 대칭형이 되어서, U자형의 아키텍처를 만들게 됨.
이 네트워크에는 fully connected layers가 없고, 각 합성곱 연산에서 valid한 부분에서만 사용.
즉, segmentation map이 input image에서 모든 context를 활용할 수 있는 픽셀만 포함한다는 것.
→ overlap-tile strategy*를 통해 임의의 큰 이미지를 seamlessly(균일하게) segmentation할 수 있도록 함.

overlap-tile strategy*
: 임의의 크기를 가진 이미지를 끊김 없이 segmentation할 수 있는 방법으로, 큰 이미지를 여러 개 작은 tile로 나누어서, 타일 간 경계를 고려해서 overlap 시키는 방식.
경계 pixel 예측 시 부족한 context(데이터 좀 부족한 영역에서)는 input 이미지를 mirroring(거울 대칭)해서 보완. GPU 메모리 제한된 환경에서도 이미지 크기 상관 없이 높은 해상도/정확도로 이미지 처리 가능.

Biomedical tasks에서는 사용 가능한 학습 데이터 수가 적음. 따라서 elastic deformations(탄성 변형)를 적용해서, data augmentation 진행
→ network가 해당 변형에 대해 invariance(불변성) 학습하게 되는 것
→ Biomedical segmentation에서 흔한 tissue(조직) 변형을 효율적으로 시뮬레이션 할 수 있으니, 효과적임
Cell segmentation tasks에서는 같은 class의 touching objects(접촉하는 객체)를 분리하는 것이 중요.
따라서 weighted loss(가중치 손실 함수) 를 통해, 접촉하는 세포(객체) 사이를 구분하는 배경 labels에 높은 weight를 부여하여 해결함.
결과적으로 해당 방법이 다양한 Biomedical Segmentation 문제를 해결하는데 활용됨.
성능도 높고, 대회에서도 훌륭한 성과로 우승을 하기도 함.

↑ U-net architecture(ex. 32 X 32 pixel)
네트워크 아키텍처를 살펴보자면,
좌측 contracting path(수축 경로) & 우측 expansive path(확장 경로)로 구성.
contracting path(수축 경로-특징 추출)
expansive path(확장 경로-해상도 복원)
final layer
1 X 1 convolution을 사용해서 각 64차원 특징 벡터를 필요한 클래스 개수로 mapping.
총 network는 23개의 convolutional layers.
tiling할 때, output segmentation map을 매끄럽게 처리하려면, input tile size를 잘 선택해야 함.
모든 max-pooling 연산이 x와 y size가 짝수인 layer에 적용되도록 input tile size를 설정해야한다는 것. (U-Net에서 2 X 2 max-pooling이 기본적으로 쓰이는 편이라)

에너지 함수는 최종 feature map에 대해 pixel 단위로 soft-max 계산하고, 이걸 cross entropy loss function이랑 결합한 것을 뜻함.
pixel 위치 x ∈ Ω에서 feature 채널 k의 a_k(x)를 기반으로 확률 p_k(x)를 계산함.

각 pixel에서 p_l(x)(X)가 실제 label l(x)와 얼마나 일치하는지를 평가함.


↑ weight map 수식Data Augmentation는 적은 수의 학습 sample로도 network가 invariance(불변성) & robustness(강인성)을 학습하라 수 있도록 하기 위해 필수적임.
microscopical images에서는 shift(이동) & rotation(회전)에 대해 불변성을 가져야하고, deformation(변형) & gray value(gray scale 값) 변화에 대해서 강인성을 가져야 함.
random elastic deformations(랜덤 탄성 변형)는 미세한 영상 segmentation 작업에서 network 성능 극대화를 위한 핵심 기법임. 이미지가 적어도 network가 다양한 변형에 대해 강인하도록 학습시킴.(ex. 세포 조직이 왜곡/변형된 이미지도 처리를 정확히 잘할 수 있는 것. 세포에서는 조직 왜곡/변형/기형 현상이 흔해서 이런 작업을 해서 성능을 높이도록 함)

Drop-out layer는 수축 경로에서 data augmentation의 역할을 하며 network의 성능을 향상시킴.
본 실험에서는 U-Net을 3가지의 segmentation tasks에 적용하였음.
전자 현미경 기록에서 neuronal structures(신경 구조)를 세분화하는 것.
↑ Ranking on the EM segmentation challenge
결과적으로, input data를 7가지 회전 버전으로 평균화했더니,
추가 전/후처리 없이 Warping Error: 0.0003529/Rand Error: 0.0382라는 성능을 기록함.
Rand Error가 다른 알고리즘에 비해 성능이 좋지 못하지만, 다른 알고리즘들은 U-Net과 달리 후처리 기법을 사용했음.
light microscopic images(광학 현미경 이미지)에서 cell segmentation을 작업을 수행하는 데에도 사용됨.

결과적으로,
PhC-U373
평균 IOU(Intersection Over Union, 교집합) 값: 92% 달성.
DIC-HeLa
평균 IOU(Intersection Over Union, 교집합) 값: 77.5% 달성.
두 데이터 셋 모두에서 U-Net이 타 알고리즘 성능에 비해 월등히 뛰어남을 확인할 수 있음.
U-Net은 다양한 biomedical segmentation 작업에서 우수한 성능을 보임.
특히 elastic deformations(탄성 변형)를 data augmentation 단계에서 활용하여, annotated image가 적은 상황에서도 우수한 성능을 낼 수 있었음.