[Segmentation] U-Net

키워드

1. U자형 구조를 도입하여 작은 데이터셋으로도 높은 정확도의 세그멘테이션을 가능하게 함.

2. "U-Net: Convolutional Networks for Biomedical Image Segmentation"

3. U-Net, copy and crop, mirroring, random elastic augmentation

U-Net

  Segmentation 분야는 다른 computer vision 분야와 다르게 데이터셋의 크기가 크지 않다. 이러한 문제는 segmentation 기술이 여러 응용 분야에서 활용되기 어렵게 한다. 이에 적은 수의 샘플을 보다 효율적으로 사용하기 위해 새로운 U-Net을 고안하였으며, 강력한 데이터 augmentation과 훈련 전략을 함께 활용하였다. 이를 통해 극소수의 이미지로부터 end-to-end 학습이 가능하며, 512x512 이미지가 1초 이하로 소요되는 뛰어난 속도를 보인다.

  소위 "fully convolutional network"라고 불리는 보다 아키텍처를 기반으로 구축되었으며, 매우 적은 수의 훈련 이미지와 정확한 분할에 초점을 맞추어 수정하고 확장하였다. 아키텍처는 context를 캡처하기 위한 축소(contracting) 경로와 정확한 localization을 가능하게 하는 대칭적인 확장(expansive) 경로로 구성되어있다. 이런 대칭 구조때문에 U-Net이라 부른다.

  또한 사용 가능한 훈련 데이터가 거의 없으므로, 사용 가능한 훈련 이미지에 탄성(elastic) 변형을 적용하여 수많은 sample을 생성하는 augmentation 기법을 사용하였다. 특히 실제 조직의 변형을 효율적으로 시뮬레이션할 수 있기 때문에 생체 의학 분야에서 유의미한 효과를 보였다. 마지막으로 segmentation의 또 다른 어려움은 붙어있는 동일한 클래스의 서로 다른 개체를 분리하는 것이다. 이를 위해 배경 레이블에 가중치를 부여하는 손실을 사용하였다.

Architecture

  기본적으로 전체 네트워크는 수축 경로(왼쪽)와 확장 경로(오른쪽)로 구성된다. 수축 경로는 일반적인 convolutional 네트워크의 구성을 하고 있다. 미세하게 다른 부분은 3x3 conv에 padding이 사용되지 않아 conv 레이어를 통과할때마다 크기가 줄어든다는 것이다. 확장 경로는 feature map의 업샘플링을 수행하는 2x2 conv 레이어("up convolution"), 축소 경로의 일부 영역에서 그대로 가져온 feature map이 연결되는 특별한 구조가 있다. 이후 마지막 레이어에서 1x1 conv를 통해 클래스 수 만큼 feature를 구한 뒤, softmax로 segmentation map을 얻는다.

  특별한 부분은 FC레이어가 없어 최종 conv 레이어까지 receptive field에 해당하는 부분만 관찰한다는 것이다. 이를 통해 receptive field가 중첩되며 임의의 큰 이미지를 원활하게 분할할 수 있는 효과를 얻을 수 있다. 그리고 최대한 많은 context 정보를 더 큰 해상도의 레이어에 전파할 수 있도록 upsampling 영역에 많은 수의 채널이 있다는 점도 특징이다. Upsampling은 또한 단순 pooling이 아닌 conv 연산으로 수행되어 네트워크가 연속적인 레이어로 이루어질 수 있게 하였으며, localization를 위해 수축 경로의 고해상도 feature가 업샘플링된 데이터와 결합하게 된다. 이를 통해 연속적인 conv 레이어가 다양한 정보를 기반으로 보다 정확한 출력을 조합하는 방법을 학습하게 된다.

728x90

  최종 출력보다 더 큰 이미지가 필요하다. 이를 구현하기 위해 이미지의 경계 영역에 있는 픽셀을 미러링하여 입력 이미지의 누락된 context를 추정합니다.이 타일링 전략은 네트워크를 설계에 중요한 역할을 하며, 이를 통해  패딩되지 않은 conv 레이어를 활용하여 일정하게 크기를 줄여나갈 수 있게 도와준다.

Training

  손실 함수는 softmax를 취한 최종 feature map에 cross-entropy loss를 활용한다. ak(x)는 feature의 출력 값이며,  w(x)는 픽셀 별 가중치로 이에 대한 수식은 아래와 같다. 

  이는 훈련 세트의 클래스 불균형을 위한 가중치와 붙어있는 객체의 경계를 네트워크가 학습하도록 강제하기 위한 분할 가중치로 구성되어 있다. 경계에 대한 항은 형태학적(morphological) 연산을 사용하여 계산된다. d1은 가장 가까운 셀의 경계까지의 거리이며, d2는 두 번째로 가까운 셀의 경계까지의 거리이다. w0와 σ는 하이퍼파라미터로, 실험에서는 각각 10과 5가 사용되었다.

  추가적으로 GPU 메모리를 최대한 활용하기 위해 큰 배치 크기보다 단일 배치를 사용한다.

Data Augmentation

  사용할 수 있는 훈련 샘플이 거의 없는 경우 네트워크의 robustness를 위해 augmentation이 필수적이다. 미세한 이미지에서는 shift 및 rotation에 대한 invariant 뿐만 아니라 deformations 과 gray value 변화에 대한 robustness가 필요하다. 무작위 탄성(Random elastic) 변형은 Segmentation 네트워크를 훈련하는 핵심으로 볼 수 있다.

  이는 거친(coarse) 3x3 그리드에 무작위 displacement vector를 사용하여 smooth한 변형을 생성한다. Displacement는 10픽셀 표준편차를 갖는 가우스 분포에서 샘플링된다. 이후 보간(interpolation)을 사용하여 픽셀당 displacement를 계산한다. 또한 이외에도 dropout을 통해  augmentation 효과를 추가한다.

결과

  광학 현미경 이미지의 세포 분할 작업에 U-Net을 적용하였다. 첫 번째 “PhC-U373”는 훈련 이미지가 35개만 포함되어 있다. 이러한 상황에서 92%의 mean IoU를 달성했다. 두 번째 “DIC-HeLa”는 훈련 이미지가 더 적은 20개가 포함되어 있다. 여기서 77.5%의 mean IoU로 준수한 성능을 보였다.

Reference

논문 링크 : https://arxiv.org/abs/1505.04597