1. U-Net이란?
보통의 CNN에 의해 실행되는 이미지의 클래스 분류(이미지 인식)에서는 Convolution층이 물체의 국소적인 특징을 추출하는 역할을 담당하고, Pooling층이 물체의 전반적인 위치 정보를 애매하게 하는 (위치의 어긋남을 허용) 역할을 담당한다. 그러므로 더욱 층이 깊어질수록 추출되는 특징은 보다 더 국소적이 되고, 그 특징의 전반적인 위치 정보가 보다 더 애매하게 되어버린다. 즉, Pooling층으로 인해 위치의 어긋남과 크기의 차이에 영향을 받지 않고 견고한 패턴 인식이 가능하다는 것이다.
한편, 영역 추출에서는 '물체의 국소적인 특징과 전반적인 위치정보'의 둘 다를 원본 이미지에서 특정해야한다. 즉, Pooling층에서 애매하게 된 국소적 특징의 위치정보를 원본 이미지에서 pixel단위로 정확히 복원해야할 필요성이 있다.
그러므로 '물체의 국소적인 특징과 전반적인 위치 정보' 양쪽 모두 종합하여 학습시키기 위해 개발된 것이 U-Net이다. 아래의 그림과 같이 U자 모양으로 네트워크를 구성하여 U-Net이라는 이름이 붙었다.
U-Net의 구성
아래로 향할수록 Convolution + Pooling에 의해 깊은 층일수록 특징이 국소적이고 위치정보가 애매하며, 얕은 층일수록 특징의 전반적인 위치 정보가 정확하다. 한편, 위로 향할수록 Convolution+Up sampling에 의해 특징을 보존한 채로 이미지가 점점 크게 복원하는 것이 가능하다. 이러한 두 경로로 인해, 이미지 사이즈가 같은 것을 깊은 층에서부터 단계적으로 통합함으로써 국소적 특징을 보존하면서 전체적으로 위치 정보를 복원할 수 있는 것이다.
참고로 아래로 향하는 경로를 Contracting Path, 위로 향하는 경로를 Expanding Path라고 부른다.
파란색 박스 : 이미지, 특징 map
흰색 박스 : 복사된 특징map
박스 위의 숫자 : 채널 수
박스 좌하단의 숫자 : 수직-수평 사이즈
파란색 화살표 : kernel 사이즈 3x3, padding 0의 Convolution, ReLU
회색 화살표 : 특징 map의 복사를 crop
빨간색 화살표 : kernel 사이즈 2x2의 max-pooling
초록색 화살표 : kernel 사이즈 2x2, stride2의 Deconvolution
청록색 화살표 : kernel 사이즈 1x1의 Convolution
U-Net은 일반적인 CNN과 다른 세 가지 특징이 있다. 그것은 바로 'Upsampling(Upサンプリング)', 'Merge（マージ）', '전체 결합층(FCN)이 존재하지 않는것'이다.
1) Upsampling
위로 향하는 경로에 있는 Upsampling이라고 불리는 작업은 차원을 증가시키는(이미지를 크게 복원하는) 의미가 있다. 차원을 감소시키는(이미지를 작게 압축시키는) Pooling의 반대되는 효과를 가지고 있으므로 Uppooling이라고도 불린다.
2) Merge
문자 그대로 정보를 통합한다. U-Net의 아래로 향하는 경로는 층이 깊어질수록 특징이 국소적인 위치 정보가 애매하게 되고, 얕은 층일수록 특정은 전반적인 위치 정보가 정확하게 된다. 위로 향해는 경로는 특징을 보존한 채로 이미지의 크기를 크게 복원하는 것이 가능하므로 양쪽의 경로로 인해 이미지 사이즈가 같은 것을 깊은층부터 단계적으로 merge하면서 전체적인 위치 정보를 복원할 수 있다.
Merge에는 다양한 종류가 있지만, U-Net에서는 '채널로 추가'한다는 의미로 Merge를 한다. 예를 들어, 깊은 층에서의 Merger는 Keras를 사용하여 아래와 같이 표현한다. 5단계째의 Convolution 후(conv5)의 출력과 6단계째의 Convolution후 출력(conv6) 을 Upsampling(2x2)한 것을 concat형으로 (채널로써)통합한다는 의미이다.
up7 = merge([UpSampling2D(size=(2, 2))(conv6), conv5], mode='concat', concat_axis=3)3) FCN가 존재하지 않는다는 것
일반적인 CNN과 달리 FCN이 없다. 클래스 분류에서는 FCN가 필요하지만, 영역추출에서는 이미지가 출력되기 때문에, FCN가 필요하지 않다. 출력층에서는 Sigmoid함수를 사용하므로 0~1사이의 값이 되므로 임계치를 두고 흑백의 이미지를 출력한다.
2. U-Net 논문에 대한 이해
기존의 CNN은 단순 Classification에 주로 사용되었다면, 앞서 말했듯 U-Net은 Classification + Localization에 사용된다. U-Net은 아래와 같은 기존의 Segmentation network들의 문제점을 해결한다.
1) 속도 개선
속도의 향상이 가능한 이유는 overlap 의 비율이 적기 때문이다. sliding window 방식을 사용하면 내가 이미 사용한 patch(;이미지 인식단위) 구역을 다음 sliding window에서 다시 검증한다. 이 과정은 이미 검증이 끝난 부분을 다시 검증하는것이기 때문에 똑같은 일을 반복하는 것이라 할 수 있다. 이에 반해 U-net은 검증이 끝난 곳을 다시 검증하는 Sliding window방식보단 이미 검증이 끝난 부분은 아예 건너뛰고 다음 patch부분부터 검증을 하기 때문에 속도면에서 우위를 가질 수 있다.
2) Trade off의 늪
Patch size가 커진다면 더 넓은 범위의 이미지를 한번에 인식해 context 인식에는 탁월한 효과가 있지만 Localization에서 패널티를 가지게 된다. 즉 너무 넓은 범위를 한번에 인식하다보니 localization에서 약한 모습을 보이게 된다. 한편 Patch size가 작아진다면 반대 효과를 가지게 된다.
U-Net은 여러 layer의 output을 동시에 검증하면서 localization과 context 인식 두가지 토끼를 다 잡을 수 있다.
또한 논문을 살펴보면, input image 사이즈가 572x572 인 반면 output image 사이즈는 388x388이다. 이는 contracting path에서 padding이 없었기 때문에 점점 이미지 외곽 부분이 없어진 결과이다. 그렇다보니 이미지가 단순 작아진것이 아니라 외곽 부분이 잘려나간것과 같게 되었다. 이를 해결하기 위해 zeropadding이 아닌 mirroring이라는 것을 도입하였다. mirror padding을 진행할때 손실되는 path를 살리기 위해서 contracting path의 데이터를 적당한 크기로 crop(잘라냄)한 후 concat하는 방식으로 이미지를 유지한다. 위 구성도(2) 이미지의 회색선이 바로 그러한 부분이다.
2. Keras를 이용한 U-Net 모델
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation, Input, Concatenate, MaxPool2D, Conv2DTranspose, Add
from tensorflow.keras.models import Model
def create_block(input, chs):
x = input
for i in range(2):
x = Conv2D(chs, 3, padding="same")(x) # オリジナルはpaddingなしだがサイズの調整が面倒なのでPaddingを入れる
x = BatchNormalization()(x)
x = Activation("relu")(x)
return x
def create_unet(use_skip_connections, grayscale_inputs=False):
if grayscale_inputs:
input = Input((96,96,1))
else:
input = Input((96,96,3))
# Encoder
block1 = create_block(input, 64)
x = MaxPool2D(2)(block1)
block2 = create_block(x, 128)
x = MaxPool2D(2)(block2)
block3 = create_block(x, 256)
x = MaxPool2D(2)(block3)
block4 = create_block(x, 512)
# Middle
x = MaxPool2D(2)(block4)
x = create_block(x, 1024)
# Decoder
x = Conv2DTranspose(512, kernel_size=2, strides=2)(x) # TPUだとUpsamplingやK.resize_imageが使えない
if use_skip_connections: x = Concatenate()([block4, x])
x = create_block(x, 512)
x = Conv2DTranspose(256, kernel_size=2, strides=2)(x)
if use_skip_connections: x = Concatenate()([block3, x])
x = create_block(x, 256)
x = Conv2DTranspose(128, kernel_size=2, strides=2)(x)
if use_skip_connections: x = Concatenate()([block2, x])
x = create_block(x, 128)
x = Conv2DTranspose(64, kernel_size=2, strides=2)(x)
if use_skip_connections: x = Concatenate()([block1, x])
x = create_block(x, 64)
# output
x = Conv2D(3, 1)(x)
x = Activation("sigmoid")(x)
return Model(input, x)- Contracting Path의 Add가 같은 사이즈여야 하므로, padding = same을 하여 입력되는 사이즈가 같도록 조정
- ConvTranspose2D는 격자모양이 되기 쉬으므로, Upsampling 2D -> Conv의 쪽이 품질이 좋은 것 같지만, TPU로 Upsampling2D를 사용하려고 하면 오류가 발생하기 때문에 ConvTranspose2D로 하였다. kernel_size = strides로 한다.
파라미터 수는 아래와 같다. 이것은 컬러화 + Contracting path의 예이다.
Total params: 31,054,275
Trainable params: 31,042,499
Non-trainable params: 11,7763. U-Net을 이용한 Image Segmentation (Tensorflow 예제)
여기서 사용할 dataset은 Oxford-lllT Pet Dataset이다. 이 dataset은 이미지, 해당 레이블 및 픽셀 단위 마스크로 구성된다. 각 픽셀은 세 가지 카테고리 중 하나가 부여된다.
Class 1 : 애완동물에 속하는 픽셀
Class 2 : 애완동물 테두리(경계) 픽셀
Class 3 : 없음 또는 위 / 주변 픽셀
pip install -q git+https://github.com/tensorflow/examples.gitimport tensorflow as tffrom tensorflow_examples.models.pix2pix import pix2pix
import tensorflow_datasets as tfds
tfds.disable_progress_bar()
from IPython.display import clear_output
import matplotlib.pyplot as plt
1) Oxford-lllT Pet Dataset다운로드
dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)아래의 코드는 이미지을 뒤집기를 하며, [0,1]로 정규화한다. 또한 위에서 언급한 것과 같이 the segmentation mask는 해당되는 Class에 대해 {1,2,3}으로 표시된다. 편의상 the segmentation mask라벨을 {0, 1, 2}의 레이블로 만든다.
def normalize(input_image, input_mask):
input_image = tf.cast(input_image, tf.float32) / 255.0
input_mask -= 1
return input_image, input_mask@tf.function
def load_image_train(datapoint):
input_image = tf.image.resize(datapoint['image'], (128, 128))
input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))
if tf.random.uniform(()) > 0.5:
input_image = tf.image.flip_left_right(input_image)
input_mask = tf.image.flip_left_right(input_mask)
input_image, input_mask = normalize(input_image, input_mask)
return input_image, input_maskdef load_image_test(datapoint):
input_image = tf.image.resize(datapoint['image'], (128, 128))
input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))
input_image, input_mask = normalize(input_image, input_mask)
return input_image, input_maskDataset는 이미 요구되는 테스트 및 훈련의 분할이 되어 있으므로 동일한 분할을 계속 사용한다.
TRAIN_LENGTH = info.splits['train'].num_examples
BATCH_SIZE = 64
BUFFER_SIZE = 1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZEtrain = dataset['train'].map(load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE)
test = dataset['test'].map(load_image_test)train_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()
train_dataset = train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
test_dataset = test.batch(BATCH_SIZE)Dataset에서부터 이미지 예제와 그 이미지에 해당되는 마스크를 출력해보자.
def display(display_list):
plt.figure(figsize=(15, 15))
title = ['Input Image', 'True Mask', 'Predicted Mask']
for i in range(len(display_list)):
plt.subplot(1, len(display_list), i+1)
plt.title(title[i])
plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
plt.axis('off')
plt.show()for image, mask in train.take(1):
sample_image, sample_mask = image, mask
display([sample_image, sample_mask])
2) 모델 정의
모델은 수정된 U-Net을 사용하여 만들 것이다. 앞서 설명했듯, U-Net은 encoder(downsampler)과 decoder(upsampler)로 구성되어 있다. 강력한 기능을 배우고 학습 가능한 parameter의 수를 감소시키기 위해 사전 훈련된 모델을 encoder로 사용할 수 있다. 그러므로, 이 작업을 위한 encoder는 사전 훈련 된 MobileNetV2 모델이며 중간 출력이 사용된다. 그리고 decoder는 Pix2Pix 예제에서 이미 구현한 upsample block이다.
세 개의 채널을 출력하는 이유는 각 픽셀당 세 개의 라벨을 사용할 수 있기 때문이다. 이것을 각 픽셀이 세 개의 클래스로 분류되는 multi-classification으로 생각하라.
OUTPUT_CHANNELS = 3앞서 말했듯, encoder은 tf.keras.applications에 있는 사전 훈련된 MobileNetV2 모델이 될 것이다. encoder은 모델의 중간층에서 출력된 특정한 아웃풋으로 구성되어 있다. encoder은 훈련과정 중에서는 훈련되지 않는다는 것을 주의하라.
base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)
# Use the activations of these layers
layer_names = [
'block_1_expand_relu', # 64x64
'block_3_expand_relu', # 32x32
'block_6_expand_relu', # 16x16
'block_13_expand_relu', # 8x8
'block_16_project', # 4x4
]
layers = [base_model.get_layer(name).output for name in layer_names]
# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=layers)
down_stack.trainable = Falsedecoder/upsampler은 단순히 TensorFlow 예제에서 구현된 일련의 upsample 블록들이다.
up_stack = [
pix2pix.upsample(512, 3), # 4x4 -> 8x8
pix2pix.upsample(256, 3), # 8x8 -> 16x16
pix2pix.upsample(128, 3), # 16x16 -> 32x32
pix2pix.upsample(64, 3), # 32x32 -> 64x64
]def unet_model(output_channels):
inputs = tf.keras.layers.Input(shape=[128, 128, 3])
x = inputs
# Downsampling through the model
skips = down_stack(x)
x = skips[-1]
skips = reversed(skips[:-1])
# Upsampling and establishing the skip connections
for up, skip in zip(up_stack, skips):
x = up(x)
concat = tf.keras.layers.Concatenate()
x = concat([x, skip])
# This is the last layer of the model
last = tf.keras.layers.Conv2DTranspose(
output_channels, 3, strides=2,
padding='same') #64x64 -> 128x128
x = last(x)
return tf.keras.Model(inputs=inputs, outputs=x)
3) 모델 훈련
이제 모델을 compile하고 훈련할 일만 남았다. loss는 losses.SparsecCategoricalCrossentropy(from_logits=True)를 사용할 것이다. 그 이유는 이 loss 함수를 사용하는 이유는 multi-class prediction과 같이 네트워크는 각 픽셀에 라벨을 할당하려고 하기 때문이다. the true segmentation mask에서는 각 픽셀은 {0, 1, 2} 중의 값을 가진다. 여기서 네트워크는 세개의 채널을 출력한다. 기본적으로, 각 채널은 클래스 예측을 배우려고하고,losses,SparseCategoricalCrossentropy(from_logits=True)는 이러한 시나리오에 추천되는 loss이다. 그러한 네트워크의 출력을 사용함으로써, 가장 높은 값을 가진 채널의 픽셀에 라벨을 할당한다. 이것이 create_mask 함수가 하는 일이다.
model = unet_model(OUTPUT_CHANNELS)
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])모델의 구조를 살펴 보자.
tf.keras.utils.plot_model(model, show_shapes=True)
훈련 전에 예측한 것을 확인해보기 위해 모델을 시험해보자.
def create_mask(pred_mask):
pred_mask = tf.argmax(pred_mask, axis=-1)
pred_mask = pred_mask[..., tf.newaxis]
return pred_mask[0]def show_predictions(dataset=None, num=1):
if dataset:
for image, mask in dataset.take(num):
pred_mask = model.predict(image)
display([image[0], mask[0], create_mask(pred_mask)])
else:
display([sample_image, sample_mask,
create_mask(model.predict(sample_image[tf.newaxis, ...]))])show_predictions()
훈련하는 동안 모델이 어떻게 향상되는지 살펴보자. 이 작업을 수행하기 위한 callback 함수는 아래에 정의해두었다.
class DisplayCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
clear_output(wait=True)
show_predictions()
print ('\nSample Prediction after epoch {}\n'.format(epoch+1))EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS
model_history = model.fit(train_dataset, epochs=EPOCHS,
steps_per_epoch=STEPS_PER_EPOCH,
validation_steps=VALIDATION_STEPS,
validation_data=test_dataset,
callbacks=[DisplayCallback()])
loss = model_history.history['loss']
val_loss = model_history.history['val_loss']
epochs = range(EPOCHS)
plt.figure()
plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0, 1])
plt.legend()
plt.show()
4) 예측 만들기
몇 개의 예측을 만들어보자. 시간을 아끼기 위해 epochs의 숫자는 작게 조정하였다. 그러나 높은 정확도의 결과를 확인할 수 있을 것이다.
show_predictions(test_dataset, 3)
참고자료
https://qiita.com/koshian2/items/603106c228ac6b7d8356
https://blog.negativemind.com/2019/03/15/semantic-segmentation-by-u-net/