Deep Learning/Object Segmentation

U-Net: Convolutional Networks for Biomedical Image Segmentation

728x90

0. 시작하기에 앞서

본 게시물은 아래 참고링크의 게시글을 매우 적극적으로 활용하였으므로, 원 저작자에게 감사의 뜻을 남깁니다.

1. 개념

U-Net은 오토인코더(autoencoder)와 같이 데이터의 차원을 축소했다가 다시 확장하는 방식의 모델로, Semantic Segmentation을 수행할 수 있다.

그러나 오토인코더는 인코딩 단계에서 차원 축소를 거치면서 이미지 객체에 대한 자세한 위치 정보를 잃게 되고, 디코딩 단계에서도 저차원의 정보만을 이용하기 때문에 위치 정보 손실을 회복하지 못하는 단점이 존재한다.
이러한 단점을 극복하기 위해 U-Net은 고차원 정보도 함께 이용하여 이미지의 특징을 추출함과 동시에 정확한 위치 파악을 하기위해 인코딩 단계의 각 레이어에서 얻은 특징을 디코딩 단계의 각 레이어에 합치는(concatenation) 방법을 사용하며, 인코더 레이어와 디코더 레이어의 직접 연결을 스킵 연결(skip connection)이라고 한다.

2. U-Net의 구조

2.1. 전체 구조

이름과 같이 'U'처럼 좌우가 대칭이 되도록 차원을 축소하는 부분과 차원을 확장하는 부분의 레이어를 배치

2.2. Contracting Path

넓은 범위의 이미지를 보며, context information을 추출하는 부분

2.2.1. ConvBlock

전체 구조 이미지에서 나타나는 파란색 박스
- contracting path 구조 이미지에서 보라색 테두리 안에 있는 파란색 박스와 동일

2.2.2. EncoderBlock

contracting path 구조 이미지에서 보라색 테두리 전체
- ConvBlock의 output은 디코더 부분으로 복사하기 위한 것
- Max Pooling의 output은 인코더의 다음 단계로 보내지는 다운샘플링된 값

2.3. Bridge

1개의 ConvBlock 레이어로 표현 가능

2.4. Expanding Path

context information을 위치정보와 결합하여 각 위치가 어떤 객체에 속하는지를 구분

2.4.1. DecoderBlock

expanding path 구조 이미지에서 회색 테두리 전체
- 노란색 박스는 ConvTranspose 를 통해 이전 layer의 output(feature map)에 대해 해상도를 2배로 늘리고 채널수를 절반으로 줄인 것을 복사한 것
- 녹색 박스는 skip connection을 통해 대칭되는 위치의 인코더 맵의 output을 복사한 것
- 녹색 박스와 노란색 박스를 concatenation한 것을 ConvBlock을 통해 채널 수를 절반으로 줄임
다만 Contracting Path의 Feature map이 Expanding Path의 Feature map보다 큰 이유는 Contracting Path에서 여러 번의 패딩이 없는 3×3 Convolution Layer를 지나면서 Feature map의 크기가 줄어들기 때문
- Contracting Path의 Feature map의 테두리 부분을 자른 후 크기를 동일하게 맞추어 두 feature map을 합쳐 줌

3. 전체 코드

3.1. 클래스를 이용해 구현

# U-Net model
# coded by st.watermelon

import tensorflow as tf
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Conv2DTranspose
from tensorflow.keras.layers import Activation, BatchNormalization, Concatenate


""" 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


""" 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


""" 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 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

3.2. 함수로 구현

# 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