📚[논문 보러가기](https://arxiv.org/abs/1505.04597) # Abstract - **U-Net**은 적은 양의 학습 데이터로도 높은 정밀도의 세분화를 수행할 수 있는 **CNN 구조**이다. - U-Net의 목표는 **Classfication**이 아니라 **Segmentation**이다. 즉, 각 픽셀이 어떤 의미(클래스)에 속하는가를 학습하여 **Mask 이미지**를 결과물로 반환한다. (Ex, 세포 세분화, 뇌 MRI 영역 구분, 도로 차선 탐지 etc..) - 기존의 딥러닝 모델은 수천~수만 장의 라벨링된 이미지를 요구하지만, 의학 영상 같은 분야에서는 많은 데이터 확보가 거의 불가능하기 때문에 U-Net은 **데이터 증강**을 강하게 활용해 **제한된 데이터를 효과적으로 학습**하도록 설계되었다. - U-Net은 컨텍스트(context)를 포착하는 **Contracting path(수축 경로)와** 픽셀 단위의 정확한 위치를 복원하는 **Expanding Path(확장 경로)로** 구성되어 있다. - 추가로 512 x 512 이미지 한 장을 GPU에서 1초 이내에 분할 할 수 있다. - U-Net은 ISBI 2015 Cell Tracking Challenge에서 1등을 달성했다. # 1. Introduction - 기존 Deep CNN의 경우 이미지 Classfication 분야에 큰 성장을 이루었으나, 이는 **수십만~수백만 개의 라벨된 이미지**와 **대규모 네트워크 구조** 때문이었다. - 하지만 이러한 CNN은 의료 영상처럼 **데이터가 적은 분야**에는 부적합하다. - 일반 CNN의 목표는 Classfication인 반해 U-Net의 목표는 **Segmentation(세분화)이다.** - 기존에는 Segmentation을 수행하기 위해 **sliding window** 방식을 사용했다. - **sliding window** : 각 픽셀 주변의 작은 패치(patch)를 잘라 CNN에 입력하여 그 패치의 중심 픽셀이 어떤 클래스인지 예측하는 방법 - 하지만 **sliding window** 방식에는 **두 가지 큰 문제점**이 있었다. 1. **Ineffciency** : 모든 픽셀마다 네트워크를 따로 돌려야 하므로 중복 계산이 많음 2. **Localization ↔ Context** Trade-off : 큰 Patch를 쓰면 주변 Context는 잘 보지만 경계가 흐릿해지고 작은 패치를 쓰면 경계는 잘 잡지만 전체 맥락을 보지 않는다. 즉, **정확한 위치**와 **넓은 문맥 정보**를 동시에 얻기 어렵다. - **U-Net**은 **Fully Convolutional Network(FCN) 개념**을 기반으로 설계 되었다. - 기존 CNN은 FC Layer을 사용하기 때문에 feature map을 **모두 flatten** 해버리는 과정에서 (x, y) 위치 정보가 사라지는 문제가 생긴다. - FC Layer를 제거하고 모든 층을 Convolution과 Upsampling으로 대체했다. ![[U-Net architecture.png]] - U-Net architecture은 이름 그대로 대칭 구조(U-형)이다. - 왼쪽 - **Contracting path** : downsampling으로 **context 추출** - 오른쪽 - **Expanding path** : upsampling으로 **localization 복원** - 두 경로를 **skip connection**으로 연결하여 저수준 featrue을 다시 합쳐 **경계 세부 정보 손실**을 방지했다. - 기존 **FNC**는 upsampling하면서 채널 수를 줄였기 때문에 context가 약해졌지만, U-Net은 **채널을 풍부하게 유지**해서 **문맥(context) + 세부(localizatoin)을** 함께 전달할 수 있었다. - **Elastic deformation augmentation(탄성 변형 증강)을** 적극 사용하여 적은 수의 라벨링 데이터로도 학습 가능하게 설계했다. ![[Cell Boundary Image.png]] - 하지만 서로 맞닿은 세포들은 같은 클래스라서 구분하기 어렵다. 그래서 **경계(background border)에 더 큰 가중치를 줘서** 네트워크가 **'여기가 세포 경계다'를** 더 강하게 학습하도록 유도했다. - (d) 값이 후에 나올 네트워크가 경계 픽셀을 학습하도록 **픽셀 별 손실 가중치를 매핑**한 이미지이다. # 2. Network Architecture - U-Net은 크게 **Contracting Path**와 **Expanding Path**로 구성되어 있다. ## Contracting Path - Contracting Path은 context을 추출하는 경로로 일반적인 CNN 구조로 구성되어 있다, - 두 개의 3x3 convolution을 거치고 **padding 없이** 수행되어 feature map의 크기가 2씩 줄어든다. - 그 후 **max pooling(stride=2)로 다운샘플링**을 거치고 정보 손실을 방지하기 위해 feature map의 크기를 2배 증가하였다. ## Expanding Path - Expanding Path은 **upsmapling(해상도 복원)을** 수행하는 부분이다. - **up-conv (2×2 transposed conv)을** 통해 feature map의 해상도를 2배로 복원하고 후의 skip connection과 합치기 위해 **feature map의 크기를 절반**으로 줄인다. - 그 후 contracting path의 대응 단계 feature map을 복사해서 붙이는 **copy & crop (skip connection)** 단계를 진행한다. - 다만 합치는 과정에서 padding 없이 convolution을 진행해 두 feature map 크기가 다르기 때문에 **이를 맞추기 위해 crop으로 크기를 조정**한다. - 합쳐진 feature map을 3x3 convolution 연산을 진행하는데 이 과정에서 **왼쪽의 고해상도 + 세밀한 위치 정보**와 **오른쪽 경로의 저해상도이지만 의미적인 context**인 두 정보를 섞어 정제(fusion)하기 위해 **feature map 크기를 절반**으로 줄인다. (+ 연산 효율 증가) ## 1x1 conv - 마지막 단계에서는 1x1 convolution을 적용하여 **feature map 수를 클래스 수(K)로** 줄인다. (Ex, 배경 - 0, 세포 - 1 → 출력 채널 2개) - 이후 softmax를 적용하여 픽셀 별 클래스 확률을 구한다. ## Even size - 총 4번의 2x2 max pooling이 수행되기 때문에 feature map의 입력 이미지의 크기가 가로, 세로 둘 다 2⁴의 배수가 되어야 **U-Net의 다운샘플링/업샘플링 대칭이 정확히 유지**된다. - 따라서 논문에서도 **572x572 이미지**를 사용한다. # 3. Training ## 학습 세팅 - 최적화 알고리즘은 **SGD**를 사용했다. - 일반적으로는 batch size를 키우려고 하는데, U-Net은 **batch size를 1로 두고 입력 이미지 크기를 크게 잡는다.** - segmentation은 출력도 이미지 형태이기 때문에 **큰 타일 하나가 정보가 많아**, GPU 메모리 부족으로 여러 장을 넣기 어렵다. - batch size가 1일 경우 gradient를 계산할 때 딱 한 장의 이미지에만 의존하기 때문에 **한 장의 데이터가 가진 특이한 패턴이나 노이즈가 그대로 gradient에 반영**된다. - 따라서 논문에서는 **momentum을 0.99**로 줘 step들의 gradient를 누적해서 현재 step의 변화량을 **평균처럼 부드럽게** 만들었다. ## 픽셀 별 softmax + cross entropy ### Softmax $p_x(x) = \frac{exp(a_k(x))}{\sum_{i=1}^{K}exp(a_i(x))}$ - a_k​(x) : 마지막 1×1 conv 출력 (logit) → 픽셀 x 위치에서 클래스 k 에 대한 점수 - K : 클래스 수 (예: 배경+세포 → K=2) - x : 2D 좌표 - 해당 softmax는 **각 픽셀마다 독립적**으로 적용된다. - 최종적으로 각 픽셀마다 **클래스 확률 분포**를 갖게 된다. ### Cross entropy $E = -\sum_{x\in \Omega }^{}w(x)log(p_{l(x)}(x)) $ - ℓ(x): 픽셀 x의 **정답 클래스 인덱스** - 예: 세포=1, 배경=0 이런 식 - p_ℓ(x)(x): 정답 클래스에 대한 softmax 확률 - **w(x)**: 픽셀 별 가중치 - 픽셀마다 **예측 확률**과 **정답 클래스**의 cross-entorpy를 계산하고 이를 가중치 w(x)로 조정해 전체 loss를 합산한다. ## 가중치 w(x) - 의료 영상에서는 배경 픽셀이 대부분이고 세포 경계 영역은 매우 적다. - 이러한 클래스 불균형을 해소하기 위해 U-Net은 **각 픽셀 마다 Loss 기여도를 다르게 주는** 가중치 맵 w(x)를 사용한다. $w(x) = w_{c}(x) + w_0 * exp(-\frac{(d_1(x)+d_2(x))^2}{2\sigma ^2})$ - w_c(x) : 클래스 비율에 따라 정해지는 기본 가중치 (소수 클래스에 더 큰 값을 준다.) - d₁(x),d₂​(x): 픽셀 x에서 가장 가까운 두 세포 경계까지의 거리, 경계 근처 픽셀일 수록 d₁(x) + d₂​(x)이 작아져 exp 항이 커진다. - w₀, σ: 하이퍼 파라미터 (논문 예시로 w₀=10,σ=5) ## He initialization - ReLU 같은 활성화 함수를 쓸 때는 입력 값의 분산이 layer를 거쳐가도 거의 일정하게 유지되도록 초기 weight를 다음과 같이 설정한다. $w∼N(0, \sigma ^2), \sigma =\sqrt{\frac{2}{N}}$ - N : 한 뉴런(또는 한 필터)에 들어오는 **입력 수** - 예를 들어 3x3 conv에 입력 채널이 64라면, N = 3 x 3 x 64 = 576 - σ = √(2/576) ≈ 0.059. - 즉, **평균 0, 표준편차 0.059인 정규분포에서 weight**를 뽑아 초기화 ## Data Augmentation - 의료 데이터는 데이터 자체가 별로 없기 때문에 데이터 증강이 매우 중요하다. - 기존의 데이터를 **shift, rotation, flip, gray value variation (명암 변화)** 뿐만 아니라 **Elastic deformation (탄성 변형)을** 사용했다. - **Elastic deformation (탄성 변형)** : 세포/조직은 실제로 모양이 휘고 찌그러졌기 때문에 **random displacement field**를 만들어서 이미지를 휘게 했다. - coarse 3x3 gird를 통해 이미지를 3x3 격자로 나눈 뒤, 격자점마다 임의로 displacement vector를 만들어 이미지 공간을 변형하고 **bicubic interpolation(2D엣거 흔히 쓰는 부드러운 보간법)을** 써서 이동 벡터를 부드럽게 연결한다. - 이렇게 변경된 데이터의 라벨(mask) 또한 같은 변형 필드로 휘게 만들어줘야 한다. # 4. Experiments ## ISBI 2012 EM Segmentation Challenge ### 데이터 셋 - ISBI EM Segmentation Challenge - 크기: 512×512 grayscale 이미지 - Train Image : 30장, Test Image : 30장 - 라벨 : 뉴런(1) / 배경(0) ### 데이터 증강 - **elastic deformation** + shift, rotation, flip, gray value variation 등 여러 증강을 사용함 - 데이터를 10⁴배 수준으로 증강 ### 학습 설정 - Loss : weighted cross-entropy - Optimizer: SGD (momentum=0.99) - 타일 크기: 572×572 입력, 388×388 출력 - Batch size: 1 ### 결과 ![[ISBI 2012 EM Segmentation Challenge Result.png]] - Rand Error : segmentation이 실제 boundary와 얼마나 일치하는가 - 낮을 수록 좋음 - Warwick F-score : 분할된 영역의 정밀도·재현율 - 높을 수록 좋음 - Pixel Error : 단순한 픽셀 단위 오차율 - 낮을 수록 좋음 - 이 3가지 지표에서 모든 모델보다 **성능이 월등히 높았다.** ## 세포 분할 (PhC-U373 & DIC-HeLa) - 현광 현미경 이미지에서 세포를 찾는 문제이다. (각 세포가 어디 있는지, 서로 어떻게 붙어 있는지) ### 데이터 셋 - **PhC-U373**: 인간 뇌교세포(glioblastoma) - **DIC-HeLa**: 자궁경부암 세포(HeLa) - 각각 라벨이 매우 적다. (20장 이하 수준) ### 학습 방법 - 두 모델 다 동일한 네트워크 구조와, 동일한 weight map 계산식을 이용했다. - w(x)의 d₁(x),d₂​(x)이 경계 부분의 Loss를 강조하여 경계를 더 잘 파악했다. ### 결과 ![[PhC-U373 & DIC-HeLa Result.png]] - 기존의 **CellProfiler, Ilastik, Random Forest 기반 방법들보다 월등히 높다.** - 특히 **붙은 세포 간 경계 식별 정확도**가 매우 높았다. - 사람 눈으로 봐도 경계가 뚜렷하게 복원됨을 알 수 있다. ![[ISBI, DIC-HeLa image segmentation result.png]] # 5. Conclusion - **Elastic Deformation Augmentation(탄성 변형 증강)** 덕분에 U-Net은 수십 장 수준의 라벨 데이터로도 **전자현미경(EM)** 영상과 광학현미경 (Light Microscopy) 데이터 셋에 매우 높은 정확도를 달성했다. - NVidia Titan GPU (6 GB 메모리)에서 **약 10시간**이면 충분히 학습 가능했고 추론 속도 역시 512x512 이미지 한 장을 **1초 내**로 분할 할 수 있었다.