U-net

UNet이란?

image segmentation은 이미지의 모든 픽셀이 어떤 카테고리에 속하는지 분류하는 것을 말한다.

이미지 전체에 대해 단일 카테고리를 예측하는 image classification과는 달리, image segmentation은 픽셀 단위의 분류를 수행하므로 일반적으로 더 어려운 문제로 인식된다.

일반적으로 LeNet, AlexNet, VGG, GoogleLeNet : Image Classification

semantic segmentation은 이미지 내에 있는 객체들을 의미 있는 단위로 분할해내는 것. instance segmentation은 같은 카테고리에 속하는 서로 다른 객체까지 더 분할하여 semantic segmentation 범위를 확장한 것.

Unet 구조는 의료용 이미지에 주로 많이 사용된다.

- U-Net은 'U-Net: Convolutional Networks for Biomedical Image Segmentation' 이라는 논문에서 제안한 구조

- 매우 적은 수의 학습 데이터로도 정확한 image segmentation 성능을 보여주었다. 

 

UNet 구조

1. 기본 구조

- autoencoder와 같은 encoder-decoder 기반 모델

- 인코딩 단계 : 입력 이미지의 특징을 포착할 수 있도록 채널의 수를 늘리면서 차원 축소

- 디코딩 단계 : 저차원으로 인코딩된 정보만 이용하여 채널의 수를 줄이고 차원을 늘려서 고차원의 이미지를 복원

But 인코딩 단계에서 차원 축소를 거치면서 이미지 객체에 대한 자세한 위치 정보를 잃게 되고, 디코딩 단계에서도 저차원의 정보만을 이용하기에 위치 정보 손실을 회복하지 못하게 된다.

 

기본 idea : 저차원 뿐만 아니라 고차원 정보도 이용하여 이미지의 특징을 추출함과 동시에 정확한 위치파악도 가능하게 하자 !

-> encoding 단계의 각 레이어에서 얻은 특징을 decoding 단계의 각 레이어에 합치는 concatenation 방법을 사용

-> encoder 레이어와 decoder 레이어의 직접 연결을 skip connection이라고 한다.

U-Net은 인코더 또는 축소경로 (contracting path)와 디코더 또는 확장경로(expending path)로 구성되며 두 구조는 서로 대칭적이다. encoder와 decoder를 연결하는 부분을 bridge라고 한다.

 

Contracting Path

 

- 그림에서 세로 방향 숫자는 맵(map)의 차원을 표시하고 가로 방향 숫자는 채널 수를 표시

ex) 세로 방향 숫자 256x256 과 가로 방향 숫자 128 은 해당 레이어의 이미지가 256x256x128 임을 의미. 입력 이미지는 512x512x3 이므로 RGB 3개 채널을 갖고 크기가 512x512 인 이미지를 나타냄.

 

ConvBlock

그림에서 파란색 박스가 인코더의 각 단계마다 계속 반복하여 나타나는 것을 볼 수 있는데, 이 박스는 3x3 Convolution, Batch Normalization, ReLU 활성화 함수가 차례로 배치된 것을 나타낸다. 3x3 Convolution을 거치면서 사이즈(해상도)가 좌우상하 각각 1씩 줄어드는데, padding 옵션을 주어 해상도를 줄이지 않고 그대로 유지할 수도 있다. 이 박스 두 개를 한데 묶어서 한 개의 레이어 블록으로 구현하여 사용하면 편리하다. 

""" Conv Block """
class ConvBlock(tf.keras.layers.Layer):
    def __init__(self, n_filters):
        super(ConvBlock, self).__init__()

        self.conv1 = Conv2D(n_filters, 3, padding='same')
        self.conv2 = Conv2D(n_filters, 3, padding='same')

        self.bn1 = BatchNormalization()
        self.bn2 = BatchNormalization()

        self.activation = Activation('relu')

    def call(self, inputs):
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.activation(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = self.activation(x)

        return x

 

EncoderBlock

인코더 그림을 보면 보라색 박스안에 한 개의 ConvBlock이 있고 이 박스가 인코더의 각 단계마다 나타나는 것을 볼 수 있다. 이 박스에서 나오는 출력이 2개인데, 한 개의 출력은 U-Net의 디코더로 복사하기 위한 연결선이며(녹색 화살선), 또 한 개의 출력은 2x2 max pooling 으로 다운 샘플링(down sampling)하여 인코더의 다음 단계로 내보내는 빨간색 화살선이다. 이 박스도 한 개의 레이어 블록으로 구현하여 사용하면 편리하다. 

 

""" Encoder Block """
class EncoderBlock(tf.keras.layers.Layer):
    def __init__(self, n_filters):
        super(EncoderBlock, self).__init__()

        self.conv_blk = ConvBlock(n_filters)
        self.pool = MaxPooling2D((2,2))

    def call(self, inputs):
        x = self.conv_blk(inputs)
        p = self.pool(x)
        return x, p

 

Bridge

브릿지는 두개의 파란색 박스로만 구성되어 있으므로 1개의 ConvBlock 레이어로 표현 가능

 

Expanding Path

DecoderBlock

디코더의 자세한 구조 ↓

- 노란색 박스 : bridge 또는 이전 layer의 출력으로부터 넘어온 feature map을 transposed convolution layer를 통해서 해상도는 2배가 늘어나고 채널수는 절반으로 줄어든 것

- 녹색 박스 : 스킵 연결을 통해서 대칭되는 인코더에 있는 맵의 출력을 복사한 것

두 개의 맵을 서로 합쳐서(concatenation) 저차원 이미지 정보뿐만 아니라 고차원 정보도 이용할 수 있는 것이다.

이를 다시 ConvBlock을 통해서 채널수를 절반으로 줄인다.
디코더의 그림에서도 회색 박스가 반복적으로 나타나므로 한 개의 레이어 블록으로 구현하여 사용하면 편리하다.

이 블록 이름을 DecoderBlock 이라고 할 수 있다.

 

""" Decoder Block """
class DecoderBlock(tf.keras.layers.Layer):
    def __init__(self, n_filters):
        super(DecoderBlock, self).__init__()

        self.up = Conv2DTranspose(n_filters, (2,2), strides=2, padding='same')
        self.conv_blk = ConvBlock(n_filters)

    def call(self, inputs, skip):
        x = self.up(inputs)
        x = Concatenate()([x, skip])
        x = self.conv_blk(x)

        return x

디코더 그림의 맨 상단의 오른쪽 부분은 U-Net의 출력으로서 1x1 컨볼루션으로 특징 맵을 처리하여 입력 이미지의 각 픽셀을 분류하는 세그멘테이션 맵을 생성하는 부분이다. 컨볼루션 필터의 개수는 분류할 카테고리 개수이며 활성 함수로는 카테고리 수가 1개라면 sigmoid 함수를, 여러 개라면 softmax 함수를 사용한다.

EncoderBlock과 DecoderBlock 을 사용하면 U-Net을 다음과 같이 간단히 코드로 구현할 수 있다.

""" U-Net Model """
class UNET(tf.keras.Model):
    def __init__(self, n_classes):
        super(UNET, self).__init__()

        # Encoder
        self.e1 = EncoderBlock(64)
        self.e2 = EncoderBlock(128)
        self.e3 = EncoderBlock(256)
        self.e4 = EncoderBlock(512)

        # Bridge
        self.b = ConvBlock(1024)

        # Decoder
        self.d1 = DecoderBlock(512)
        self.d2 = DecoderBlock(256)
        self.d3 = DecoderBlock(128)
        self.d4 = DecoderBlock(64)

        # Outputs
        if n_classes == 1:
            activation = 'sigmoid'
        else:
            activation = 'softmax'

        self.outputs = Conv2D(n_classes, 1, padding='same', activation=activation)

    def call(self, inputs):
        s1, p1 = self.e1(inputs)
        s2, p2 = self.e2(p1)
        s3, p3 = self.e3(p2)
        s4, p4 = self.e4(p3)

        b = self.b(p4)

        d1 = self.d1(b, s4)
        d2 = self.d2(d1, s3)
        d3 = self.d3(d2, s2)
        d4 = self.d4(d3, s1)

        outputs = self.outputs(d4)

        return outputs

 

 

전체적인 네트워크 구조

전체 코드

- EncoderBlock, DecoderBlock class를 사용하지 않고 함수 구조로 구현

- 100*75*3 의 이미지를 입력으로 받아서 100*75*1 의 Image Segmentation한 이미지를 출력으로 내보내는 모델

- 해상도가 홀수라서 decoding하는 부분에 해상도를 맞추기 위해 약간의 수정

- 그리고 3*3 convolution할 때 padding 옵션을 주어서 convolution후에도 해상도는 줄어들지 않고 입력의 해상도를 그대로 유지하도록

def ConvBlock(n_filter, inputs):
  x = Conv2D(n_filter, 3, padding='same')(inputs)
  x = BatchNormalization()(x)
  x = Activation('relu')(x)
  x = Conv2D(n_filter, 3, padding='same')(x)
  x = BatchNormalization()(x)
  x = Activation('relu')(x)
  return x

def unet_like():
  # Encoder
  inputs = Input(shape=(100, 75, 3))
  c1 = ConvBlock(64, inputs) # 100*75*64
  p1 = MaxPool2D(2)(c1) # 50*37*64

  c2 = ConvBlock(128, p1) # 50*37*128
  p2 = MaxPool2D(2)(c2) # 25*18*128

  c3 = ConvBlock(256, p2) # 25*18*256
  p3 = MaxPool2D(2)(c3) # 12*9*256

  c4 = ConvBlock(512, p3) # 12*9*512
  p4 = MaxPool2D(2)(c4) # 6*4*512

  # Bridge
  b = ConvBlock(1024, p4) # 6*4*1024

  # Decoder
  d1 = Conv2DTranspose(512, (2,2), strides=2, output_padding=(0,1))(b) # 12*9*512 (12*8 -> 12*9)
  d1 = Concatenate()([c4, d1]) # [12*9*512][12*9*512]
  d1 = ConvBlock(512, d1) # 12*9*512

  d2 = Conv2DTranspose(256, (2,2), strides=2, output_padding=(1,0))(d1) # 25*18*256
  d2 = Concatenate()([c3, d2]) # [25*18*256][25*18*256]
  d2 = ConvBlock(256, d2) # 25*18*256

  d3 = Conv2DTranspose(128, (2,2), strides=2, output_padding=(0,1))(d2) # 50*37*128
  d3 = Concatenate()([c2, d3]) # [50*37*128][50*37*128]
  d3 = ConvBlock(128, d3) # 50*37*128

  d4 = Conv2DTranspose(64, (2,2), strides=2, output_padding=(0,1))(d3) # 100*75*64
  d4 = Concatenate()([c1, d4]) # [100*75*64][100*75*64]
  d4 = ConvBlock(64, d4) # 100*75*64

  outputs = Conv2D(1, 1, padding='same', activation='sigmoid')(d4) # 100*75*1

  model = keras.Model(inputs=inputs, outputs=outputs)
  return model