[Review] U-Net: Convolutional Networks for Biomedical Image Segmentation

Abstract

딥 네트워크의 성공적인 훈련은 수많은 어노테이션된 훈련 샘플이 필요하다는 것에 많은 사람들이 동의한다. 이 논문에서는, 네트워크와 사용가능한 어노테이션된 샘플을 좀 더 효과적으로 사용하기 위해 데이터 어그멘테이션을 강력하게 사용하는 훈련 전략을 보여준다. 해당 아키텍처는 컨텍스트를 캡처하기 위한 수축 가능한 path와 정확한 localization이 가능하도록 하는 대칭적 확장 path로 구성된다. 본 논문에서는 매우 적은 수의 이미지로 end-to-end 학습이 가능한 네트워크와 전자 현미경 스택에서 신경 구조의 세그멘테이션에 대한 ISBI 챌린지의 최고의 방법(sliding-window cnn)보다 성능이 뛰어나다는 것을 보여준다.

1. Introduction

생의학 이미지 처리에서 목표하는 출력은 localization이 포함되어 있어야 한다. 즉, class label이 각 픽셀에 할당되어야 한다. 게다가 생의학 작업에서는 일반적으로 수천 개의 훈련 이미지를 처리할 수 없다. 따라서 논문[1]은 슬라이딩 윈도우 setup에서 네트워크를 훈련시켜 해당 픽셀 주변의 로컬 영역(patch)을 입력으로 제공하여 각 픽셀의 클래스 레이블을 예측했다. 첫째, 이 네트워크는 로컬라이제이션이 가능하다. 둘째, 패치 측면에서 훈련 데이터는 훈련 이미지의 수보다 훨씬 더 많다. 그래서 이 네트워크는 ISBI 2012의 EM 세그 멘테이션 챌린지에서 우승하였다.

그렇지만 논문[1]의 전략에는 두 가지 단 점이 있다.

• 첫 쨰, 패치마다 네트워크를 별도로 실행해야 하기 때문에 속도가 상당히 느리고 패치가 겹쳐서 중복되는 부분이 많다는 점.
• 둘 째, 로컬라이제이션 정확도와 컨텍스트 사용 간에 상충되는 부분이 있다. 패치가 클수록 더 많은 max-pooling 레이어가 필요하므로 로컬라이제이션 정확도가 낮아지는 반면, 패치가 작으면 네트워크가 컨텍스트를 거의 볼 수 없다.

최근의 접근 방식 [11, 4]에서는 여러 계층의 특징을 고려한 classifier output을 제안했다. 우수한 로컬라이제이션과 컨텍스트 사용이 동시에 가능하다.

이 논문은 fully-convolutional-network[9]라고 불리는 좀 더 괜찮은 아키텍처를 기반으로 한다. 저자들은 이 아키텍처를 수정하고 확장하여 매우 적은 수의 학습 이미지로 작동하고 더 정확한 세그멘테이션을 생성하도록 한다(그림 1 참조). 논문[9]의 주요 아이디어는 pooling 연산이 upsampling 연산으로 대체되는 연속적인 레이어를 통해 일반적인 수축 네트워크를 보완하는 것이다. 따라서 이러한 레이어는 출력의 해상도를 높인다.(그러니까 원래 CNN은 계속 conv, pooling 반복하다가 마지막에 flatten을 만드는 형태인데, 논문[9]는 pooling을 통해 작아지는 것이 upsampling연산으로 대체한다는 것 같다.) localize를 위해 contracting path(이미지의 해상도가 점점 작아짐)의 높은 해상도의 피처가 upsampled된 출력과 결합된다. 그런 다음 후속 컨볼루션 레이어는 이 정보를 기반으로 더 정확한 출력을 조합하는 방법을 학습할 수 있다. 즉, Upsample된 출력과, 이전에 있던 같은 계층의 정보를 결합한 뒤 다시 conv연산을 2번 수행하여 더 정확한 출력을 조합하는 방법을 학습한다.

동빈나 U-net 리뷰

• 각 픽셀마다 원-핫 인코딩(one-hot encoding) 형식을 따른다.
• 즉, Plants/Grass 클래스의 픽셀에 1이 있다면 동일 위치의 다른 클래스의 픽셀은 모두 0이어야 한다.

semantic segmentation 작업을 위한 데이터 세트 생성 비용은 매우 비쌈. 그러나 일반적인 CNN 분류 모델의 형식을 크게 바꾸지 않고 학습할 수 있다.

U-net

일반적인 분류 태스크에서는 해상도가 지속적으로 줄어들어 마지막엔 클래스의 개수와 같은 dimension을 갖도록 한다. 그러나 semeantic segmentation에서는 입력 이미지와 같은 해상도를 가지는 출력 이미지가 나와야 하기에 조금 다른 구조를 사용한다.

• 세그멘테이션에서는 up-sampling을 통해 너비와 높이가 증가하도록 만들어야 한다.
• up-sampling 목적으로는 Transposed 컨볼루션을 사용해야 하는데, 대표적인 예시로 3x3 해상도에서 5x5 해상도로 증가시킬 수 있다.

• U-net의 특징은 수축 경로에서 사용했던 피처맵을 그대로 전달해서 확장 경로에서도 사용될 수 있도록 한다.
• 하얀색은 앞쪽에서 사용하였던 피처맵을 그대로 붙여서 사용한다는 의미이다.
• 이렇게 되면 앞쪽 레이어에서 추출했던 특징들을 재사용할 수 있기에 보다 좋은 출력 결과를 얻을 수 있다.
• 그리고 Unet은 세그멘테이션을 위해 사용되는 네트워크이므로 별도의 fc layer가 필요하지 않다.
• 그래서 전체 네트워크는 Fully Convolutional Network 아키텍처로 구성된다.

수축 경로

• 앞부분은 일반적인 분류 모델과 동일하다.
• 그래서 실제로 U-net기반의 세그멘테이션 모델을 만들 때는 사전에 잘 학습되어있는 분류 모델을 그냥 인코더 형태로 사용하는 경우도 많다.
• 그래서 앞부분을 인코더, 뒷부분을 디코더로 보았을 때 인코더는 이미 학습이 잘 되있는 분류 모델을 가져와서 사용할 수도 있다.

확장 경로

U-Net의 overlap-tile 전략

• u-net의 구조상 출력 이미지의 해상도가 입력 이미지보다 작기에 일부러 입력을 크게 넣는다.
• 노란색 바운딩 박스만큼 세그멘테이션 결과가 필요하다면 더 큰 파란색 영역을 패치로 넣어서 결과를 뽑은 뒤에 출력 결과 패치로서 사용한다.
• 그래서 타일A와 타일 B를 사용한다고 하면 일부 겹치는 부분이 생길 수 있다.
• 노란색 바운딩 박스의 왼쪽위는 이미지가 존재하지 않는 부분이니 이미지가 존재하는 부분을 사용해 미러링을 통해 이미지 패치를 만들어 준 뒤 네트워크의 입력으로 넣는 것을 확인가능하다. 일종의 extrapolation

objective function

• 각각의 픽셀마다 확률값을 예측할 수 있도록 모든 픽셀에 대해 확률값을 구하는 소프트맥스를 사용한다.
• 실제 학습을 위해서 일반적으로 사용되는 cross-entropy loss를 사용한다.
• 밑의 식은 해당 정답 레이블에 대한 확률값에 로그를 씌워 그 확률값이 증가할 수 있는 형태로 학습을 진행한다.
• 일반적인 cross-entropy와 유사하지만 한 가지 다른 점은 앞쪽에 추가적인 가중치 함수가 사용된다는 것이다.
• 각각의 픽셀마다 가중치를 부여해서 더 학습이 잘 수행되거나 덜 수행되도록 만드는 것이다.

• 세포같은 경우에는 서로 붙어있는 경우가 있어, 작은 분리 경계를 위해 가중치 함수를 사용하는 것이다.
• \( \)
• \(W_c\): 클래스마다 등장하는 빈도가 다를 수 있으므로 그것을 조율해주는 역할
• \(d_1\): 세포와 가장 가까운 경계까지의 거리
• \(d_2\): 두 번째로 가장 가까운 세포의 경계까지의 거리

식을 보면 지수함수의 입력값에 -가 붙어있는 것을 확인할 수 있다.

• 이를 통해 거리값이 작아지면 가중치가 커지게 되고, 거리값이 커지면 가중치가 작아진다.

u-net의 데이터 증진

• 세그멘테이션에서는 입력 데이터와 출력 데이터가 둘 다 이미지 형태이므로, 실제 구현상 이런 어그멘테이션을 사용할 때는 입력이미지 그리고 출력이미지가 되는 마스크 이미지에 대해서도 같은 transformation을 적용하는 것이 일반적이다.