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

목차
0. Abstract
1. Introduction
    1.1 논문에서 지적한 문제점
    1.2 논문에서 제안한 방법
2. Network Architecture
3. Training
    3.1 Data Augmentation
4. Experiments
5. Conclusion
6. U-Net Code

 

0. Abstract

 본 논문에서는 한정된 수의 데이터를 더 효율적으로 이용하기 위한 data augmentation에 의존한 네트워크와 학습 방식을 소개한다. 논문에서 제시하는 아키텍처는 context 정보를 얻을 수 있는 contracting path와 localization을 위한 symmetric expanding path로 이루어져 있다. 이러한 아키텍처를 가진 end-to-end 모델은 적은 개수의 이미지를 가지고도 기존 방식들의 성능을 뛰어넘었다.

Image Segmentation이란?

 Image Classification의 경우 이미지 자체에 대한 정답 레이블을 얻어내는 방식이었다면 Image Segmentation은 픽셀 단위로 관심 객체를 추출하는 방법으로 일반적으로 더 어려운 문제라고 할 수 있다.

Image Segmentation의 종류
1) Semantic Segmentation

 위 사진과 같이 각 픽셀을 특정 클래스에 대응시키는 방식이 semantic segmentation이다. 이 방식은 같은 분류의 두 물체가 있더라도 이를 서로 구분해서 인식하지 않는다는 특징이 있다.

2) Instance Segmentation
 Instance segmentation 방식은 semantic segmentation과 다르게 클래스가 아닌 instance를 기반으로 픽셀을 범주로 분류한다. 그렇게 때문에 서로 다른 객체에 대해서는 서로 다른 label을 가진다.

 

1. Introduction

1.1 논문에서 지적한 문제점

 지난 2년 동안 깊은 CNN 네트워크들은 많은 최신 visual recognition task를 뛰어넘었다. 그러나 이용 가능한 학습 데이터의 크기나 네트워크의 크기로 인해서 한계가 있었다. 특히 biomedical image processing의 경우에는 출력에 localization이 포함되어 있어야 하며 많은 양의 학습 데이터를 확보하기 힘들다는 한계가 존재한다.

 이를 해결하기 위해서 Ciresan이 시도한 방식은 각 픽셀에 대한 입력으로 픽셀 주변의 local region 넣어줌으로써 class label을 예측할 수 있도록 sliding window로 학습시키는 방법이다. 위 그림이 sliding window를 그림으로 표현한 것인데, 이 방식을 이용하면 localization이 가능해지고 학습 데이터를 패치 단위로 쪼개기 때문에 데이터의 양이 학습 이미지의 양보다 훨씬 많아진다.

 그러나 이 방식에는 두 가지 단점이 존재한다. 첫 번째는 각 패치에 대해서 따로 실행시켜야 하기 때문에 속도가 느리고 overlapping되는 패치들로 인해서 불필요한 중복이 생긴다. 또한, localization 정확성과 use of context 사이의 trade off가 존재한다.

 

1.2 논문에서 제안한 방법

 본 논문에서는 fully convolutional network를 이용해서 적은 학습 이미지를 이용해서도 더 높은 정확도를 낼 수 있는 네트워크를 설계하였다. 여기서의 핵심은 contracting network에서 pooling operator를 upsampling operator로 대체해서 사용하는 것이다. 이는 출력 이미지의 해상도를 증가시킬 수 있다. 또한, localization을 위해서 contracting path에서의 높은 해상도의 feature를 upsampled output과 결합하였다. 또 다른 중요한 점은 upsampling 부분에서 많은 양의 feature channel을 가지기 때문에 context 정보를 높은 해상도로 전달시키는 것이 가능하다. 이로써 expansive path가 contracting path와 대칭이 되어서 U자 모양의 네트워크를 형성하게 된다.

 

1) Overlap-Tile Strategy & Mirroring Extrapolate

 U-Net에서는 위 그림과 같이 이미지를 패치로 나눠서 입력으로 사용하였다. 이때 입력 이미지로 파란색 영역에 해당하는 부분을 넣게 되면 노란색 범위가 semented 이미지로 나오게 된다. 이때 출력에서 겹치는 범위가 없도록 입력의 범위를 조절하기 위한 방법이 overlap-tile strategy이다. 이러한 점을 고려해야 하는 이유는 모델의 구조 상 convolution 연산에 padding을 넣지 않았기 때문에 출력 이미지가 입력 이미지에 비해서 크기가 줄어들기 때문이다.

 이때 가장자리 부분의 missing data는 위 그림과 같은 mirroring extrapolation을 이용해서 채우게 된다. 즉, 빈 공간을 입력 이미지를 거울로 반사시킨 형태로 채우는 방식이다. 

 

2) Data Augmentation

 매우 적은 학습 데이터를 가지고 있기 때문에 elastic deformation이라는 data augmentation 방식을 적용시킴으로써 이 문제를 해결하였다. 이는 특히 biomedical segmentation에서 중요한데 그 이유는 tissue에서 deformation이 가장 흔한 variation이고, 현실적인 deformation을 효율적으로 실험 가능하기 때문이다.

 

2. Network Architecture

 본 논문에서 제안한 모델의 아키텍처는 위 사진과 같다. 크게 왼쪽의 contracting path와 오른쪽의 expansive path, 그리고 둘 사이의 bottleneck 부분으로 나누어 볼 수 있다.

 

1) Contracting Path

 Contracting path는 일반적인 CNN 모델 형태를 따른다. 두 번의 3 * 3 unpadded convolution과 downsampling을 위한 stride가 2인 2 * 2 maxpooling 연산을 거친다. 각 downsampling 과정은 feature map의 채널 수를 2배로 증가시킨다. 이 구간은 점점 넓은 범위의 이미지 픽셀을 보면서 context 정보를 추출하는 단계이다.

1) 3 * 3 conv + BatchNorm (no padding, stride 1) + ReLU
2) 3 * 3 conv + BatchNorm (no padding, stride 1) + ReLU
3) 2 * 2 max pooling (stride 2) → 해상도(너비와 높이) 2배 감소

 

2) Bottleneck

 Bottleneck 구간은 contracting path에서 expanding path로 전환되는 구간이다. 여기에서도 두 번의 3 * 3 unpadded convolution을 거친 뒤에 dropout이 일어난다.

1) 3 * 3 conv + BatchNorm (no padding, stride 1) + ReLU
2) 3 * 3 conv + BatchNorm (no padding, stride 1) + ReLU
3) Dropout

 

3) Expanding Path

 Expanding path의 모든 단계는 2 * 2 up-convolution을 통해서 feature channel의 개수를 반으로 줄인다. 또한, contracting path의 cropped된 feature map과 concatenation을 거친다. 이후에는 두 번의 3 * 3 unpadded convolution을 거치게 된다. 이때 output size가 input size보다 작은 이유는 convolution 과정에서 padding을 사용하지 않았기 때문이다. expanding path는 정밀한 지역화를 가능하게 하는 단계이다.

1) upsampling : 2 * 2 up-conv (stride 2) → 해상도(너비와 높이) 2배 증가
2) contracting path의 cropped된 feature map과 concatenate
3) 3 * 3 conv + BatchNorm (no padding, stride 1) + ReLU
4) 3 * 3 conv + BatchNorm (no padding, stride 1) + ReLU
5) 마지막 단계에서만 1 * 1 conv 추가적으로 적용

 

3. Training

 입력 이미지와 그에 대응되는 segmentation maps를 SGD를 이용해서 학습시킨다. Unpadded convolution으로 인해서 출력 이미지의 크기가 입력 이미지에 비해서 작아진다. 이때 실행 시간을 줄이고 GPU 메모리의 사용 효율을 최대화하기 위해서 input tile의 해상도를 키우고, batch를 단일 이미지로 축소한다. 또한, momentum의 크기를 0.99로 설정하여서 이전의 학습 데이터들이 현재의 최적화 과정에 많이 반영될 수 있도록 한다.

 

1) Energy Function

 cross entropy loss function에 결합된 최종 feature map에 pixel-wise softmax를 적용시킨 형태이다. 여기서 l(x)는 위치 x에 대한 ground truth label이다.

 위 식은 energy function에서 p_k(x)를 구하기 위한 식으로 a_k(x)는 픽셀의 위치 x에서 k 채널의 activation을 나타낸다. 이때 a_k(x)가 최댓값을 가지는 k에서 p_k(x)는 1에 근접한 값을 가지고, 그 이외의 값들에서는 0에 근접한 값을 가진다. 

 

2) Weight Map

 Energy function에서 w(x)에 해당하는 부분으로 학습 데이터별로 class frequency가 다른 것을 조절해주는 weight map이다. 이는 touching cells 사이의 경계를 학습시키기 위해서 넣어주는 부분이다. 만약 픽셀이 경계선 픽셀일 경우에 더 큰 손실값을 갖게 하였다. 식에서 d_1과 d_2는 각각 가장 가깝거나 두 번째로 가까운 cell 경계까지의 거리를 나타낸다.

 

3.1 Data Augmentation

 의료 데이터를 사용한다는 특성상 학습에 많은 데이터를 이용할 수 없었기 때문에 data augmentation을 이용하였다. Microscopical 이미지들의 경우, 특히 적은 양의 이미지를 이용하는 상황에서 random elastic deformation을 이용하는 것이 학습의 핵심이었다. 이를 위해서 3 * 3 grid에 랜덤한 변위 벡터들을 이용해서 smooth deformations를 생성하였다. 해당 변위들은 10 pixel의 표준편차를 가지는 gaussian distribution으로부터 sample되었다. per-pixel 변위들은 bicubic interpolation을 이용해서 구하였다.

 

4. Experiments

1) EM Segmentation challenge

 Electron microscopic 기록에서 뉴런의 구조를 segmentation하는 task로 EM segmentation challenge에서 제공된 데이터셋을 이용하였다. 해당 데이터셋은 512 * 512 크기의 30개의 training 데이터를 제공한다. 각 이미지는 그에 대응되는 cell과 membrane을 구분하는 ground truth segmentation map이 주어진다.

 위 표에서도 확인할 수 있듯이 u-net은 특별한 전처리나 후처리 없이도 warping error가 0.0003529이라는 작은 값이 나왔다. 이는 sliding-window를 이용한 convolutional network보다는 유의미하게 좋은 성능을 보인다.

Warping Error란?

 Segmentation metric 중 하나로 topological error를 의미한다. 위 사진에서 A와 B의 pixel error는 동일하지만 B의 warping error가 더 작다. 즉, warping error는 객체 분할 및 병합이 잘 되었는지와 관련된 error이다.

 

 

2) ISBI cell tracking challenge

 Light microscopic 이미지를 cell segmentation하는 task이다. 첫 번째로 “PhC-U373”이라는 데이터를 이용하였는데 이는 위상차 현미경으로 기록한 35개의 이미지에 대한 데이터로 92% IOU Score를 획득하였다. 두 번째는 “DIC-HeLa” 데이터셋으로 HeLa 세포를 현미경을 통해 기록하고 20개의 이미지를 데이터로 제공한 것이다. 이 데이터셋으로 77.5% IOU Score를 획득하였다.

 

 

 

 

IOU란?

IOU는 Intersection over Union의 약자로 객체 인식 모델의 정확도를 평가하는 지표이다. 위 사진과 같이 ground truth에 해당하는 bounding box와 prediction에 해당하는 bounding box에 대한 교차 영역의 겹치는 정도를 이용해서 값을 구한다. 이때 겹치는 영역이 넓을수록 IOU 값이 거치고 잘 예측한 것으로 평가할 수 있다.

 

 

5. Conclusion

 U-net 아키텍처는 여러 biomedical segmentation 분야의 적용에서 좋은 성능을 보여주었다. 특히 elastic deformation 방식을 이용한 data augmentation을 통해서 매우 적은 이미지 데이터를 가지고도 학습을 시킬 수 있었다.

 

6. U-Net Code

class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()

        # 논문의 파란색 화살표 부분 
        def CBR2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
            layers = []
            layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels, 
                                 kernel_size=kernel_size, stride=stride, padding=padding,
                                 bias=bias)]
            layers += [nn.BatchNorm2d(num_features=out_channels)]
            layers += [nn.ReLU()]

            cbr = nn.Sequential(*layers)

            return cbr
        
        self.enc1_1 = CBR2d(in_channels=1, out_channels=64)
        self.enc1_2 = CBR2d(in_channels=64, out_channels=64)
        self.pool1 = nn.MaxPool2d(kernel_size=2)

        self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
        self.enc2_2 = CBR2d(in_channels=128, out_channels=128)
        self.pool2 = nn.MaxPool2d(kernel_size=2)

        self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
        self.enc3_2 = CBR2d(in_channels=256, out_channels=256)
        self.pool3 = nn.MaxPool2d(kernel_size=2)

        self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
        self.enc4_2 = CBR2d(in_channels=512, out_channels=512)
        self.pool4 = nn.MaxPool2d(kernel_size=2)

        self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)

        # Expanding Path
        self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)
        self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=2, stride=2, padding=0, bias=True)

        self.dec4_2 = CBR2d(in_channels=2 * 512, out_channels=512)
        self.dec4_1 = CBR2d(in_channels=512, out_channels=256)
        self.unpool3 = nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=2, stride=2, padding=0, bias=True)

        self.dec3_2 = CBR2d(in_channels=2 * 256, out_channels=256)
        self.dec3_1 = CBR2d(in_channels=256, out_channels=128)
        self.unpool2 = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=2, stride=2, padding=0, bias=True)
        
        self.dec2_2 = CBR2d(in_channels=2 * 128, out_channels=128)
        self.dec2_1 = CBR2d(in_channels=128, out_channels=64)
        self.unpool1 = nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=2, stride=2, padding=0, bias=True)

        self.dec1_2 = CBR2d(in_channels=2 * 64, out_channels=64)
        self.dec1_1 = CBR2d(in_channels=64, out_channels=64)

        self.fc = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1, stride=1, padding=0, bias=True)

    def forward(self, x):
        # Contracting Path
        enc1_1 = self.enc1_1(x)
        enc1_2 = self.enc1_2(enc1_1)
        pool1 = self.pool1(enc1_2)

        enc2_1 = self.enc2_1(pool1)
        enc2_2 = self.enc1_2(enc2_1)
        pool2 = self.pool2(enc2_2)

        enc3_1 = self.enc3_1(pool2)
        enc3_2 = self.enc3_2(enc3_1)
        pool3 = self.pool3(enc3_2)

        enc4_1 = self.enc4_1(pool3)
        enc4_2 = self.enc4_2(enc4_1)
        pool4 = self.pool4(enc4_2)

        # Bottleneck
        enc5_1 = self.enc5_1(pool4)
        dec5_1 = self.dec5_1(enc5_1)

        # Expanding Path
        unpool4 = self.unpool4(dec5_1)
        cat4 = torch.cat((unpool4, enc4_2), dim =1)
        dec4_2 = self.dec4_2(cat4)
        dec4_1 = self.dec4_1(dec4_2)

        unpool3 = self.unpool3(dec4_1)
        cat3 = torch.cat((unpool3, enc3_2), dim =1)
        dec3_2 = self.dec3_2(cat3)
        dec3_1 = self.dec3_1(dec3_2)

        unpool2 = self.unpool2(dec3_1)
        cat2 = torch.cat((unpool2, enc2_2), dim =1)
        dec2_2 = self.dec2_2(cat2)
        dec2_1 = self.dec2_1(dec2_2)

        unpool1 = self.unpool1(dec2_1)
        cat1 = torch.cat((unpool1, enc1_2), dim =1)
        dec1_2 = self.dec1_2(cat1)
        dec1_1 = self.dec1_1(dec1_2)

        x = self.fc(dec1_1)

        return x