U-Net: Convolutional Networks for Biomedical Image Segmentation

https://arxiv.org/abs/1505.04597

U-Net: Convolutional Networks for Biomedical Image Segmentation

 

오랜만에 예전 논문으로...

 

요약: 딥러닝 네트워크의 성공적인 훈련에는 수천 개의 주석이 달린 훈련 샘플이 필요하다는 데에 널리 동의합니다. 이 논문에서는 주석이 달린 샘플을 더 효율적으로 활용하기 위해 데이터 증강을 강하게 사용하는 네트워크와 훈련 전략을 제시합니다. 이 아키텍처는 문맥을 포착하기 위한 수축 경로와 정확한 위치 지정을 가능하게 하는 대칭적인 확장 경로로 구성됩니다. 우리는 이러한 네트워크가 매우 적은 이미지로 끝에서 끝까지 훈련될 수 있으며, 전자 현미경 스택에서 신경 구조의 세분화에 대한 ISBI 챌린지에서 이전의 최고 방법(슬라이딩 윈도우 컨볼루션 네트워크)을 능가한다는 것을 보여줍니다. 동일한 네트워크를 전도성 광학 현미경 이미지(위상차 및 DIC)로 훈련시켜 2015년 ISBI 세포 추적 챌린지에서 이 부문을 크게 앞질렀습니다. 또한 이 네트워크는 빠릅니다. 512x512 이미지의 세분화는 최근 GPU에서 1초 이내에 수행됩니다. 전체 구현(카페 기반) 및 훈련된 네트워크는 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net에서 사용할 수 있습니다.

 

소개

지난 2년 동안, 딥 컨볼루션 네트워크는 여러 시각 인식 작업에서 최신 기술을 능가해왔습니다(e.g. [7,3]). 컨볼루션 네트워크는 이미 오래전부터 존재해왔지만 [8], 사용 가능한 훈련 세트의 크기와 고려된 네트워크의 크기 때문에 성공이 제한적이었습니다. Krizhevsky 등 [7]의 돌파구는 800만 개의 훈련 이미지를 포함한 ImageNet 데이터셋에서 8개의 층과 수백만 개의 매개변수를 가진 대형 네트워크를 감독 학습한 덕분이었습니다. 그 이후로, 더 크고 깊은 네트워크가 훈련되었습니다 [12].

 

컨볼루션 네트워크의 일반적인 사용은 이미지에 대한 단일 클래스 레이블을 출력하는 분류 작업에 있습니다. 그러나 많은 시각 작업, 특히 생의학 이미지 처리에서 원하는 출력은 위치 지정(localization)을 포함해야 하며, 즉 각 픽셀에 클래스 레이블이 할당되어야 합니다. 게다가, 생의학 작업에서 수천 개의 훈련 이미지는 대부분 도달할 수 없는 수준입니다. 따라서, Ciresan 등 [1]은 슬라이딩 윈도우 설정에서 네트워크를 훈련하여 해당 픽셀 주변의 지역(패치)을 입력으로 제공함으로써 각 픽셀의 클래스 레이블을 예측하도록 했습니다. 첫째, 이 네트워크는 위치 지정을 할 수 있습니다. 둘째, 패치 형태의 훈련 데이터는 훈련 이미지의 수보다 훨씬 많습니다. 결과적으로 이 네트워크는 ISBI 2012의 EM 세분화 챌린지에서 큰 차이로 우승했습니다.

 

그림 1. U-net 아키텍처 (가장 낮은 해상도에서 32x32 픽셀 예시). 각 파란색 상자는 다중 채널 특징 맵을 나타냅니다. 채널 수는 상자 위에 표시되어 있습니다. x-y 크기는 상자의 왼쪽 하단 모서리에 제공됩니다. 흰색 상자는 복사된 특징 맵을 나타냅니다. 화살표는 다양한 작업을 나타냅니다.

 

Ciresan 등 [1]의 전략에는 두 가지 단점이 명백히 존재합니다. 첫째, 각 패치에 대해 네트워크를 개별적으로 실행해야 하므로 속도가 매우 느리고, 겹치는 패치로 인한 중복이 많습니다. 둘째, 위치 지정 정확도와 문맥 사용 간에 절충이 필요합니다. 더 큰 패치는 위치 지정 정확도를 감소시키는 더 많은 최대 풀링 계층을 필요로 하며, 작은 패치는 네트워크가 문맥을 거의 볼 수 없게 합니다. 최근 접근법 [11,4]은 여러 계층의 특징을 고려한 분류기 출력을 제안하여 좋은 위치 지정과 문맥 사용을 동시에 가능하게 했습니다.

 

이 논문에서는 더 우아한 아키텍처인 '완전 컨볼루션 네트워크' [9]를 기반으로 합니다. 우리는 이 아키텍처를 수정하고 확장하여 매우 적은 훈련 이미지로도 작동하고 더 정밀한 세분화를 제공하도록 합니다. 그림 1을 참조하십시오. [9]의 주요 아이디어는 일반적인 수축 네트워크에 연속적인 계층을 추가하여, 풀링 연산자를 업샘플링 연산자로 대체하는 것입니다. 따라서 이러한 계층은 출력의 해상도를 증가시킵니다. 위치 지정을 위해, 수축 경로의 고해상도 특징을 업샘플링된 출력과 결합합니다. 그 후 연속적인 컨볼루션 계층이 이 정보를 기반으로 더 정밀한 출력을 조립하는 법을 배울 수 있습니다.

 

그림 2. 임의의 큰 이미지의 원활한 세분화를 위한 오버랩 타일 전략 (여기서는 EM 스택에서 신경 구조의 세분화). 노란색 영역의 세분화 예측에는 파란색 영역 내의 이미지 데이터가 입력으로 필요합니다. 누락된 입력 데이터는 반사되어 보간됩니다.

 

우리 아키텍처의 중요한 수정 사항 중 하나는 업샘플링 부분에서 많은 수의 특징 채널을 가지고 있다는 것입니다. 이를 통해 네트워크는 문맥 정보를 고해상도 계층으로 전달할 수 있습니다. 그 결과, 확장 경로는 수축 경로와 대칭적으로 U자형 아키텍처를 형성하게 됩니다. 네트워크는 완전 연결된 계층이 없고, 각 컨볼루션의 유효 부분만 사용하므로, 세분화 맵은 입력 이미지에서 전체 문맥을 사용할 수 있는 픽셀만 포함합니다. 이 전략은 오버랩 타일 전략을 통해 임의의 큰 이미지를 원활하게 세분화할 수 있게 합니다(그림 2 참조). 이미지의 경계 영역 픽셀을 예측하기 위해, 부족한 문맥은 입력 이미지를 반사하여 보간됩니다. 이러한 타일링 전략은 네트워크를 큰 이미지에 적용하는 데 중요하며, 그렇지 않으면 해상도가 GPU 메모리에 의해 제한될 것입니다.

 

우리의 작업에는 사용할 수 있는 훈련 데이터가 매우 적기 때문에, 사용할 수 있는 훈련 이미지에 탄성 변형을 적용하여 과도한 데이터 증강을 사용합니다. 이를 통해 네트워크는 주석이 달린 이미지 데이터셋에서 이러한 변형을 보지 않고도 이러한 변형에 대한 불변성을 학습할 수 있습니다. 이는 생의학 세분화에서 특히 중요한데, 조직에서 가장 일반적인 변형이 변형이기 때문이며, 현실적인 변형을 효율적으로 시뮬레이션할 수 있습니다. 불변성을 학습하기 위한 데이터 증강의 가치는 비지도 특징 학습의 범위에서 Dosovitskiy 등 [2]에 의해 입증되었습니다.

 

많은 세포 세분화 작업에서 또 다른 도전 과제는 동일한 클래스의 접촉하는 객체를 분리하는 것입니다(그림 3 참조). 이를 위해, 우리는 접촉하는 세포 사이의 분리 배경 레이블이 손실 함수에서 큰 가중치를 받는 가중치 손실을 사용할 것을 제안합니다.

 

결과적으로, 이 네트워크는 다양한 생의학 세분화 문제에 적용할 수 있습니다. 이 논문에서는 ISBI 2012에서 시작된 진행 중인 대회인 EM 스택에서 신경 구조의 세분화 결과를 보여주며, 여기서 우리는 Ciresan 등 [1]의 네트워크를 능가했습니다. 또한, 2015년 ISBI 세포 추적 챌린지의 광학 현미경 이미지에서 세포 세분화 결과를 보여줍니다. 여기서 우리는 두 개의 가장 어려운 2D 전송 광 데이터셋에서 큰 차이로 우승했습니다.

 

2. 네트워크 아키텍처

네트워크 아키텍처는 그림 1에 나와 있습니다. 이 아키텍처는 수축 경로(왼쪽)와 확장 경로(오른쪽)로 구성되어 있습니다. 수축 경로는 전형적인 컨볼루션 네트워크의 아키텍처를 따릅니다. 두 개의 3x3 컨볼루션(패딩되지 않은 컨볼루션)을 반복적으로 적용하며, 각각은 Rectified Linear Unit (ReLU)과 스트라이드 2의 2x2 맥스 풀링 연산으로 다운샘플링이 뒤따릅니다. 각 다운샘플링 단계에서 특징 채널의 수를 두 배로 늘립니다. 확장 경로의 각 단계는 특징 맵의 업샘플링, 2x2 컨볼루션("업-컨볼루션")을 통해 특징 채널의 수를 절반으로 줄이는 것, 수축 경로에서 대응하는 크기로 잘린 특징 맵과의 연결, 그리고 두 개의 3x3 컨볼루션과 각 ReLU로 구성됩니다. 모든 컨볼루션에서 경계 픽셀의 손실 때문에 크롭이 필요합니다. 마지막 층에서는 1x1 컨볼루션을 사용하여 각 64-컴포넌트 특징 벡터를 원하는 클래스 수로 매핑합니다. 총 23개의 컨볼루션 층이 있습니다.

출력 세분화 맵의 원활한 타일링을 허용하기 위해(그림 2 참조), 모든 2x2 맥스 풀링 연산이 x-와 y-크기가 짝수인 층에 적용되도록 입력 타일 크기를 선택하는 것이 중요합니다.

3. 훈련

입력 이미지와 해당 세분화 맵을 사용하여 Caffe [6]의 확률적 경사 하강법 구현으로 네트워크를 훈련시킵니다. 패딩되지 않은 컨볼루션 때문에 출력 이미지는 일정한 경계 너비만큼 입력보다 작습니다. 오버헤드를 최소화하고 GPU 메모리를 최대한 활용하기 위해 큰 배치 크기보다 큰 입력 타일을 선호하여 배치를 단일 이미지로 줄입니다. 따라서 높은 모멘텀(0.99)을 사용하여 이전에 본 많은 훈련 샘플이 현재 최적화 단계에서 업데이트를 결정하도록 합니다.

 

에너지 함수는 최종 특징 맵에 대한 픽셀 단위 소프트맥스와 교차 엔트로피 손실 함수를 결합하여 계산됩니다. 소프트맥스는 다음과 같이 정의됩니다:

 

 

이를 통해 네트워크는 각 픽셀의 클래스 확률 분포를 예측하고, 실제 클래스와의 차이를 최소화하도록 훈련됩니다. 교차 엔트로피 손실 함수는 다음과 같이 정의됩니다:

 

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

y_k(x) -> w_k(x)

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

 

 

그림 3. 유리 위에 기록된 DIC(차분 간섭 대비) 현미경으로 촬영한 HeLa 세포. (a) 원본 이미지. (b) 실제 세분화와 오버레이된 이미지. 다른 색상은 HeLa 세포의 다른 인스턴스를 나타냅니다. (c) 생성된 세분화 마스크(흰색: 전경, 검은색: 배경). (d) 네트워크가 경계 픽셀을 학습하도록 강제하기 위한 픽셀 단위 손실 가중치 맵.

 

각 실제 세분화에 대해 가중치 맵을 미리 계산하여 훈련 데이터 세트에서 특정 클래스의 픽셀 빈도의 차이를 보상하고, 서로 접촉하는 세포 사이에 도입한 작은 분리 경계를 네트워크가 학습하도록 강제합니다(그림 3c 및 d 참조).

분리 경계는 형태학적 연산을 사용하여 계산됩니다. 가중치 맵은 다음과 같이 계산됩니다:

 

 

 

 

3.1 데이터 증강

데이터 증강은 소수의 훈련 샘플만 사용할 수 있을 때 네트워크에 원하는 불변성과 강인성 속성을 가르치는 데 필수적입니다. 현미경 이미지의 경우, 주로 이동 및 회전 불변성, 변형 및 그레이 값 변동에 대한 강인성이 필요합니다. 특히, 훈련 샘플의 랜덤 탄성 변형은 매우 적은 주석이 달린 이미지로 세분화 네트워크를 훈련하는 핵심 개념인 것으로 보입니다. 우리는 거친 3x3 격자에서 랜덤 변위 벡터를 사용하여 매끄러운 변형을 생성합니다. 변위는 표준 편차가 10픽셀인 가우시안 분포에서 샘플링됩니다. 픽셀 단위 변위는 바이큐빅 보간법을 사용하여 계산됩니다. 수축 경로 끝의 드롭아웃 계층은 추가적인 암묵적 데이터 증강을 수행합니다.

 

표 1. EM 세분화 챌린지 [14] 순위(2015년 3월 6일 기준), 워핑 오류로 정렬

 

그림 4. ISBI 세포 추적 챌린지 결과. (a) "PhC-U373" 데이터 세트의 입력 이미지 일부. (b) 수동으로 주석된 실제 정답(노란색 경계)과 세분화 결과(청록색 마스크). (c) "DIC-HeLa" 데이터 세트의 입력 이미지. (d) 수동으로 주석된 실제 정답(노란색 경계)과 세분화 결과(랜덤 색상 마스크).

 

표 2. ISBI 세포 추적 챌린지 2015의 세분화 결과(IOU)

 

4. 실험

우리는 세 가지 다른 세분화 작업에 U-Net의 적용을 시연합니다. 첫 번째 작업은 전자 현미경 기록에서 신경 구조의 세분화입니다. 데이터 세트의 예와 우리가 얻은 세분화 결과는 그림 2에 표시되어 있습니다. 전체 결과는 보충 자료로 제공합니다. 이 데이터 세트는 ISBI 2012에서 시작된 EM 세분화 챌린지 [14]에서 제공되며, 여전히 새로운 기여를 받고 있습니다. 훈련 데이터는 초파리의 제1령 유충 복부 신경 코드(VNC)의 연속 절단 전자 현미경 이미지(512x512 픽셀) 30개로 구성됩니다. 각 이미지는 세포(흰색)와 막(검은색)에 대한 완전히 주석이 달린 실제 세분화 맵과 함께 제공됩니다. 테스트 세트는 공개적으로 이용 가능하지만, 그 세분화 맵은 비밀로 유지됩니다. 예측된 막 확률 맵을 주최자에게 보내 평가를 받을 수 있습니다. 평가는 10개의 다른 수준에서 맵을 임계값으로 설정하고 "워핑 오류", "Rand 오류" 및 "픽셀 오류"를 계산하여 수행됩니다 [14].

 

U-Net은 입력 데이터의 7개의 회전 버전 평균을 사용하여 추가 전처리나 후처리 없이 워핑 오류 0.0003529(새로운 최고 점수, 표 1 참조)와 Rand 오류 0.0382를 달성합니다. 이는 Ciresan 등 [1]의 슬라이딩 윈도우 컨볼루션 네트워크 결과보다 상당히 더 좋습니다. Ciresan 등 [1]의 최고 제출물은 워핑 오류 0.000420과 Rand 오류 0.0504를 기록했습니다. Rand 오류 측면에서 이 데이터 세트에서 더 나은 성능을 보인 알고리즘은 Ciresan 등 [1]의 확률 맵에 매우 특화된 후처리 방법을 사용합니다.

 

우리는 또한 U-Net을 광학 현미경 이미지의 세포 세분화 작업에 적용했습니다. 이 세분화 작업은 ISBI 세포 추적 챌린지 2014 및 2015 [10,13]의 일부입니다. 첫 번째 데이터 세트 "PhC-U373"은 위상차 현미경으로 기록된 폴리아크릴아미드 기질 위의 성상세포종 U373 세포를 포함합니다(그림 4a, b 및 보충 자료 참조). 여기에는 부분적으로 주석이 달린 훈련 이미지 35장이 포함됩니다. 여기서 우리는 92%의 평균 IOU("intersection over union")를 달성했으며, 이는 두 번째로 좋은 알고리즘의 83%보다 상당히 높은 수치입니다(표 2 참조). 두 번째 데이터 세트 "DIC-HeLa"는 평평한 유리 위에 기록된 차등 간섭 대비(DIC) 현미경으로 기록된 HeLa 세포를 포함합니다(그림 3, 그림 4c, d 및 보충 자료 참조). 여기에는 부분적으로 주석이 달린 훈련 이미지 20장이 포함됩니다. 여기서 우리는 77.5%의 평균 IOU를 달성했으며, 이는 두 번째로 좋은 알고리즘의 46%보다 상당히 높은 수치입니다.

 

 

5. 결론

U-Net 아키텍처는 매우 다양한 생의학 세분화 응용 분야에서 매우 좋은 성능을 발휘합니다. 탄성 변형을 사용한 데이터 증강 덕분에 주석이 달린 이미지가 매우 적게 필요하며, NVidia Titan GPU(6GB)에서 10시간의 매우 합리적인 훈련 시간만 필요합니다. 우리는 전체 Caffe[6] 기반 구현과 훈련된 네트워크를 제공합니다. U-Net 아키텍처는 더 많은 작업에 쉽게 적용될 수 있을 것이라고 확신합니다.

감사의 글

이 연구는 독일 연방 및 주 정부의 우수성 이니셔티브(EXC 294)와 BMBF(Fkz 0316185B)의 지원을 받았습니다.

 

1505.04597v1.pdf

1.57MB

 

 

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

결국 window에 따른 영향을 받으니 전체적으로 보는 것이 아니라 지역적으로 본다. 그러니 segmentation과 같이 전체 window를 볼때는 다른 방식이 좋을것. 기억상 SAM2가 FPN을 사용한다.

-----------------------------------------------------------------------------------------------------------------------------------------------------------------