[논문 리뷰] U-Net: Convolutional Networks for BiomedicalImage Segmentation

 

 

UNet 에 대한 논문을 리뷰해 보겠습니다. 이것 역시 깊이 들어가기보다는

1. 무엇을 하는 녀석인가
2. 구조가 어떻게 되어 있는가
3. 어떻게 로스를 계산하는가

에 중점을 두도록 하겠습니다.

무엇을 하는 녀석인가

이 논문에서는 세포들이 있는 이미지가 입력으로 주어지면 서로 다른 세포들을 분류해 주는 인공지능입니다. 하지만 다양한 곳에 쓰일 수 있습니다.

구조가 어떻게 되어 있는가

CNN을 이용해서 데이터를 수축시키고 여기에 up-conv를 적용해 다시 팽창시킵니다. 그림을 보면 U자 형태를 띄고 있습니다.

 

첫 번째

572*572*3이미지가 주어집니다. *3은 색을 나타냅니다. 3*3 CNN을 2번 적용하되 padding을 주지 않아 크기가 2씩 줄어듭니다. 채널은 64개를 사용합니다. 이후 max pool을 해서 284*284*64로 만듭니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

두 번째

284*284*64에 3*3 CNN을 2번 적용해 280*280*128로 만듭니다. 마찬가지로 padding을 주지 않고 채널은 128로 만듭니다. 이후 max pool을 해서 140*140으로 만듭니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

세 번째

140*140*128에 3*3 CNN을 2번 적용해 136*136*256로 만듭니다. 마찬가지로 padding을 주지 않고 채널은 256로 만듭니다. 이후 max pool을 해서 68*68로 만듭니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

네 번째

68*68*256에 3*3 CNN을 2번 적용해 64*64*512로 만듭니다. 마찬가지로 padding을 주지 않고 채널은 512로 만듭니다. 이후 max pool을 해서 32*32로 만듭니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

다섯 번째

32*32*512에 3*3 CNN을 2번 적용해 28*28*1024로 만듭니다. 마찬가지로 padding을 주지 않고 채널은 1024로 만듭니다. 이후 up sampling을 통해 56*56*512로 만듭니다. 이때, 채널도 줄어든 것을 알 수 있습니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

여섯 번째

56*56*512에 네 번째에서 만든 64*64*512 이미지를 붙입니다. 이때 이미지의 크기가 맞지 않으니 중앙에서 cropping을 진행합니다. 그럼 56*56*1024 이미지가 나옵니다. 3*3 CNN을 2번 적용해 52*52*512로 만듭니다. 이후 up sampling을 통해 104*104*256으로 만듭니다. 이때, 채널도 줄어든 것을 알 수 있습니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

일곱 번째

104*104*256에 네 번째에서 만든 136*136* 256 이미지를 붙입니다. 이때 이미지의 크기가 맞지 않으니 중앙에서 cropping을 진행합니다. 그럼 104*104*512이미지가 나옵니다. 3*3 CNN을 2번 적용해 100*100*256로 만듭니다. 이후 up sampling을 통해 200*200*128로 만듭니다. 이때, 채널도 줄어든 것을 알 수 있습니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

여덟 번째

200*200*128 에 네 번째에서 만든 280*280*128 이미지를 붙입니다. 이때 이미지의 크기가 맞지 않으니 중앙에서 cropping을 진행합니다. 그럼 200*200*256이미지가 나옵니다. 3*3 CNN을 2번 적용해 196*196*128로 만듭니다. 이후 up sampling을 통해 392*392*64로 만듭니다. 이때, 채널도 줄어든 것을 알 수 있습니다. 각 CNN마다 ReLU를 적용해 줍니다.

 

아홉 번째

392*392*64 에 네 번째에서 만든 568*568*64  이미지를 붙입니다. 이때 이미지의 크기가 맞지 않으니 중앙에서 cropping을 진행합니다. 그럼 392*392*128이미지가 나옵니다. 3*3 CNN을 2번 적용해 388*388*64로 만듭니다. 마지막으로 up sampling을 통해 388*388*2로 만듭니다. 이미지를 보면 초록색 화살표가 있음을 알 수 있습니다.

 

어떻게 로스를 계산하는가

각 픽셀 별로 cross-entrophy를 계산합니다. 각 픽셀마다 배경인지, 1번째 세포인지, 2번째 세포인지, ... k번째 세포인지 라벨링을 합니다. 모델의 출력 부분에서 각 픽셀 별로 softmax를 취한 다음 라벨링한 값이랑 cross-entrophy를 구합니다. 이때, 실제 세포는 세포 간 간격이 좁을수록 큰 가중치를 둬서 경계를 더 잘 예측하게 해줍니다. 구체적인 수식은을 적어보겠습니다. $p_k(x)=\exp(a_k(x))/(\sum^K_{k^\prime=1}\exp(a_{k^\prime}(x)))$ 로 softmax를 적용합니다. 여기서 $x$는 픽셀이라고 생각하면 되고 $k$는 클래스 수입니다. $a_k(x)$는 픽셀 $x$에서 클래스 $k$가 활성화된 정도입니다. 여기에서 $E=\sum_{x\in \Omega}w(x)\log(p_{l(x)}(x))$를 이용해 cross-entrophy를 계산합니다. 여기서 $w(x)$가 가중치입니다. 가중치는 $w(x)=w_c(x)+w_0\cdot\exp \left(-\frac{(d_1(x)+d_2(x))^2}{2\sigma^2} \right)$를 이용해 계산됩니다. 여기서 $w_c$는 클래스 빈도의 균형을 맞추기 위한 가중치 맵이고 $d_1$은 가장 가까운 셀의 경계까지의 거리, $d_2$는 두 번째로 가까운 셀의 경계까지의 거리입니다.논문에서는 $w_0=10$, $\sigma \approx5$를 사용합니다.