Convolutional Networks for Biomedical Image Segmentation (U-Net)

U-Net : Convolutional Networks for Biomedical Image Segmentation

U-Net은 Encoder-Decoder 구조를 가지지만, Biomedical 부분에서 나온 세그멘테이션 논문입니다. 실제 구조가 좋고, 성능이 잘나와서 Biomedical 부분만 아니라 다른 분야에서도 많이 활용되는 네트워크입니다. 인용수 또한 20000이 넘을 정도로 굉장히 높습니다. 네트워크의 구조가 U자 형태로 되어있어서 U-Net이라고 불리는데, 네트워크가 어떻게 구성되었는지 한번 살펴보도록 하겠습니다.

• Paper : https://arxiv.org/abs/1505.04597
• Code : https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/

0. Abstract

U-Net의 논문은 의료 계열에서의 두가지 문제 상황을 해결하기 위해서 나왔습니다.

첫째, 데이터의 부족문제입니다. 딥러닝의 경우 충분히 많은 학습 데이터를 필요로 합니다. 하지만, biomedical의 특성상 data의 수가 많이 부족할 수 밖에 없습니다. 그 이유는 환자의 개인정보등에 의한 민감한 문제들로 인하여 Training data를 구하기가 어렵고, 만일 데이터를 구한다 하더라도 레이블링을 하기 위해서는 일반인이 봐도 모르기에 전문가들이 직접 해야하기 때문입니다.

 

둘째, 같은 클래스가 인접한 셀 구분이 어려운데 단순 세그멘테이션은 이를 수행하기가 어렵습니다. 아래 그림 (a)에서 “A” 와 “B”는 같은 클래스지만, (b)처럼 서로 다른 인스턴스(셀)로 구분을 요구하며, 이를 구분하기 위해 (c) 처럼 인스턴스(셀)의 경계를 테두리로 만들어야 하는 작업 필요합니다.

 

1. Introduction

• Visual Recognition 분야에서 딥러닝 모델 성능의 향상을 보였고 이를 Biomedical 분야에 적용하려는 시도의 연구입니다.
• Biomedical image processing 분야의 경우 각각의 픽셀이 클래스를 가지는 "localization"이 필요합니다.
• 특히, Biomedical의 특성상 학습 데이터의 수가 많이 부족할 수 밖에 없습니다.
  • 전문가가 라벨링을 해야하고 환자의 데이터이기 때문에 라벨된 이미지의 수가 부족합니다.
• 위의 한계를 극복하기 위해서 Ciresan [1]의 경우 Sliding Window 방식을 도입합니다.

• Sliding Window 방식의 경우 아래의 장점이 존재합니다.
  • Localize를 가능하게 합니다.
  • 패치 단위의 학습을 할 경우 데이터의 양이 많아지는 효과가 있습니다.
  • EM segmentation challenge에서 좋은 결과를 가져옵니다.

단, 이때의 Sliding Window는 픽셀을 모두 예측하는게 아니라 단 중앙 하나의 픽셀에 대해서만 Classification을 수행합니다. 그렇기에, 이러한 Sliding Window 방식의 Patch 단위 학습은 2가지의 문제점이 있습니다.

• 각각에 Patch에 대해 학습을 하기에 속도가 느리고 패치마다 많은 영역이 겹치는 부분이 생깁니다.
• Localization 정확도와 context의 사용간에 trade-off가 존재합니다.
  • Patch가 크면 주변정보도 같이 학습이 가능하지만 이미지의 크기를 줄이기위해 많은 Pooling이 필요하고 이는 localization accuracy을 떨어트립니다.
  • Patch가 작으면 Sub sampling에 의한 정보 손실은 작아지지만 작은 context만 보는 문제점이 있습니다.
  • 이는 classifier의 output을 multiple layer의 features로 고려해서 Good Localization과 use of context가 가능해집니다.

이를 해결해주기 위해서, 논문에서는 Sliding Window 방식을 사용하지만 1개의 픽셀만 Classification하는게 아닌 패치 영역을 Classification하는 특징을 보입니다. 그리고, 이때 전체 영역을 다 예측하는게 아닌 padding을 0으로 줘서 572의 입력을 받아 388 크기의 아웃풋을 생성합니다. (즉, 572의 입력이 들어올때 1개의 픽셀이 아닌 388의 픽셀을 예측을 수행합니다.) 이를 Overlap-tile 전략이라고 하고 아래와 같은 모습을 보입니다.

• fully connected layers를 사용하지 않고 오직 convolution만 사용하여 Overlap tile을 통해 임의의 이미지를 입력으로 받아도 문제가 없게합니다.
• 큰 이미지를 한정된 GPU에서 학습하도록 합니다.

보이는 것과 같이 파란색 영역이 들어와서, 노란색 영역이 출력으로 나옵니다. 그리고, 노란색 출력과 초록색 출력 부분이 겹치지 않기에 Overlap tile이라고 불립니다. 단, 이때 U-Net의 경우 572의 입력을 388까지 줄이기에 외각 부분의 예측을 수행시에 손실되는 부분이 생깁니다. 내부의 부분은 문제가 없으나 외각부분의 양쪽으로 92칸의 픽셀 부분은 예측을 할 수가 없습니다. 이를 해결하기위해 논문에서는 외각부분에 대해 예측시에 Mirror padding을 적용합니다. 이미지의 테두리 영역에서 픽셀을 예측하기 위해 입력 이미지를 미러링하여 누락된 부분을 추정한다고 생각하면 됩니다.

2. Network

네트워크의 모습은 다음과 같습니다.

• Unet의 경우 FCN을 확장한 U 모습의 네트워크를 제안합니다.
  • 이는, 적은 학습 데이터를 가지고도 높은 성능을 가집니다.
  • FCN의 주요 공헌이었던 Contracting network를 이용하여 출력의 Resolution을 증가시키는 방법 사용합니다.
• Upsampling시에 채널의 수를 크게함으로서 context information을 높은 resolution layer에 전파합니다.
• FC Layer을 제거하고 Convolution만을 이용하여 Overlap tile 전략에의해 임의의 크기의 이미지가 들어와도 Segmentation이 가능하게 합니다.
• 결론적으로 contracting path는 expanding path와 대칭이고 U-shape의 모양을 가집니다.
• 이를 이용하여 "the segmentation of neuronal structures in EM stacks"와 ISBI cell tracking challenge 2015에서 SOTA를 달성합니다.

2.1 Contracting Path

2.2 Expanding Path

3. Training

Optimizer

• SGD with momentum 0.99
  • 배치를 키우는 것보다 타일의 크기를 키우는걸 선호합니다.
  • 단, 배치가 작아서 학습이 잘 안되는 문제가 발생하기에 이를 해결하려고 모멘텀을 키워 이전의 샘플도 같이 사용하도록합니다.

Energy function

• pixel-wise soft-max를 수행하고 cross-entropy loss 함수를 통해서 학습을 진행합니다.
  • pk(x) : approximated maximum-function (k개의 클래스에 대해서 분류한 확률)
  • a(x) : 픽셀 x의 활성화 함수
  • K : 클래스의 개수

Touching cells separation

• pixel-wise loss weight를 계산하기위한 weight map을 생성합니다.
  • 같은 클래스를 가지는 인접한 셀을 분리하기위해서 해당 경계부분에 가중치를 제공합니다.
  • 가중치의 계산은 아래의 식을 이용해서 계산합니다.
  
  • Wc(x) : 클래스별 빈도에 따른 가중치
  • d1 : 가장 가까운 셀과의 거리
  • d2 : 두번째로 가까운 셀과의 거리
  • 인접 셀들간의 거리가 가까워서 d1 + d2가 작으면 지수함수 부분이 크게되어서 가중치를 크게 받습니다.
    • w0 : 10, 시그마 : 5으로 설정

• Energy function과 weight map을 결합해서 최종적인 로스함수 생성합니다.

Weight initialize

• Gaussian distribution 학습 파라미터 초기화합니다.
  • standard deviation : (2/N) ** 0.5 (N : number of incoming nodes of one neuron)
    • 예) 3x3x64 Convolution의 경우 9x64=576의 N을 가집니다.

3.1 Data Augmentation

• Random Elastic deformations를 통해서 Augmentation을 수행합니다. 이는, 클래스의 불균형을 해결해주고 적은 양으로도 많은 학습을 가능하게 만듭니다.
  • 모델이 invariance와 robustness를 학습할 수 있도록 하는 방법입니다.
  • shift and rotation invariance와 deformations에 대한 robustness, gray value variations을 충족시킵니다.

4. Experiments

4.1 EM segmentation challenge

• 30 images (512x512 pixels) from serial section transmission electron microscopy of the Drosophila first instar larva ventral nerve cord (VNC) - 초파리 유충 1단계의 무척추 중추 신경계
• cells (white) and membranes (black) - 세포와 세포막을 각각 Segmentation 하는 테스크
• 평가 지표는 Warping Error, Rand Error, Pixel Error

• 7개 버전으로 회전시킨 입력데이터를 학습한 u-net에 전처리와 후처리를 통해 1등 달성했습니다.

4.2 ISBI cell tracking challenge

• PhC-U373 : Giloblastoma-astrocytoma U373 Cells on a polyacrylimide substrate recorded by phas contrast microscopy (Fig.4.a, b)
  • 35개의 annotated training images
• DIC-HeLa : HeLa cells on a flat glass recorded by differential interference contrast (DIC) microscopy (Fig.4.c, d)
  • 20개의 annotated training images

• IoU 평가지표로 둘 모두 압도적으로 1등을 달성했습니다.

5. Conclusion

• 3개의 데이터셋에서 압도적인 성능을 달성한 네트워크입니다.
• 학습 데이터셋의 부족을 극복하기위해서 elastic deformations을 사용합니다.
• Titan GPU (6 GB)로 10시간의 학습시간과 수초의 추론시간을 가집니다.

5.1 Advantages

• 바이오 메디컬 이미지가 가지는 여러가지 한계를 극복하기 위해서 다양한 기법들을 도입합니다.
  • 데이터 셋의 부족
    • Overlap-tite strategy
    • Mirroring Extrapolation
    • Elastic Deformations
  • 겹치는 세포들의 분리
    • Weight Loss
• Concatenation과 High Channel을 통해서 High Resolution의 세그멘테이션 출력을 생성합니다.
• 3개의 데이터셋에서 모두 SOTA를 달성했을 정도로 압도적인 성능을 보입니다.

5.2 Disadvantages

• 압도적인 성능에 대한 근거가 많이 부족함. 비록 딥러닝이 블랙박스 모형이기는 하지만 이전의 DeconvNet의 논문처럼 시각적으로 해석하려는 시도가 들어갔으면 더 좋았을 것 같습니다.
• Biomedical으로 데이터셋을 한정해서 실험했는데 VOC2012 등의 대회에도 제출해서 결과를 확인해봤으면 더 좋았을 것 같습니다.
• Zero-padding을 안해서 발생하는 문제를 Mirroring을 통해서 해결하려는 이유를 모르겠습니다.
  • 실제 세포의 이미지를 보면 대칭형태가 아닌데 차라리 타일을 더 크게해서 학습하고 Crop하는 형식이 더 맞지 않을까 생각이듭니다.
  • 아니면 DeconvNet에서 결과들을 앙상블한 것처럼 제로페딩을 수행하고 앙상블하는 방법도 있었을 것 같습니다.

6. Apendix

• http://brainiac2.mit.edu/isbi_challenge/home
• https://towardsdatascience.com/review-u-net-biomedical-image-segmentation-d02bf06ca760
• https://joungheekim.github.io/2020/09/28/paper-review/