Tensorflow 2.0 Neural Style Transfer 튜토리얼

https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/tutorials/generative/style_transfer.ipynb#scrollTo=eqxUicSPUOP6

Google Colaboratory

Tensor flow 2.0 alpha 버전 Neural Sytle Transfer Notebook 리뷰입니다. 

Neural Style Transfer란? content 이미지가 style reference 이미지로 그려진 것과 같이 만들어주는 알고리즘입니다. 아래 거북이 사진(content)과 스타일 사진(style reference)을 통해 새로운 이미지를 만든 것을 알 수 있습니다.

content

Style Reference

Result

https://www.tensorflow.org/install/gpu(설치 가이드 / GPU 지원)

GPU support  |  TensorFlow

Setup(설치하기)

모듈 임포트하기

from __future__ import absolute_import, division, print_function, unicode_literals

!pip install tensorflow-gpu==2.0.0-alpha0
import tensorflow as tf

import IPython.display as display

import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['figure.figsize'] = (12,12)
mpl.rcParams['axes.grid'] = False

import numpy as np
import time
import functools

설치 이후에도 tensorflow version1을 쓰시려면 

tf.compat.v1

을 통해 기존 함수에 접근하시면 됩니다.

이미지 다운로드하기

content_path = tf.keras.utils.get_file('turtle.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/Green_Sea_Turtle_grazing_seagrass.jpg')
style_path = tf.keras.utils.get_file('kandinsky.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/Vassily_Kandinsky%2C_1913_-_Composition_7.jpg')

거북이와 화가 이미지를 각각 다운로드합니다!

Visualize the input(입력 이미지 시각화 하기)

이미지를 불러오는 함수를 정의하고 최대 512 pixel로 정의합니다.

def load_img(path_to_img):
  max_dim = 512
  img = tf.io.read_file(path_to_img)
  img = tf.image.decode_image(img, channels=3)
  img = tf.image.convert_image_dtype(img, tf.float32)

  shape = tf.cast(tf.shape(img)[:-1], tf.float32)
  long_dim = max(shape)
  scale = max_dim / long_dim

  new_shape = tf.cast(shape * scale, tf.int32)

  img = tf.image.resize(img, new_shape)
  img = img[tf.newaxis, :]
  return img

이미지를 보여주는 함수 역시 정의합니다. 

def imshow(image, title=None):
  if len(image.shape) > 3:
    image = tf.squeeze(image, axis=0)

  plt.imshow(image)
  if title:
    plt.title(title)

아래 코드를 실행시키면 시각화할 수 있습니다.

content_image = load_img(content_path)
style_image = load_img(style_path)

plt.subplot(1, 2, 1)
imshow(content_image, 'Content Image')

plt.subplot(1, 2, 2)
imshow(style_image, 'Style Image')

Define content and Style representations

이미 학습된 vgg 모델에서 인풋 레이어와 가까운 레이어는 낮은 단계의 특징, 에지나 질감을 나타내고 멀어지면 멀어질수록 추상적인 특징을 나타냅니다. 이러한 성질을 통해 Content 이미지와 Style 이미지에서 원하는 특징만 추출할 수 있습니다.

vgg19 모델 load 후 테스트하기

x = tf.keras.applications.vgg19.preprocess_input(content_image*255)
x = tf.image.resize(x, (224, 224))
vgg = tf.keras.applications.VGG19(include_top=True, weights='imagenet')
r = vgg(x)

#tf.keras.applications.vgg19.decode_predictions(r.numpy())
labels_path = tf.keras.utils.get_file(
    'ImageNetLabels.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())

print(imagenet_labels[np.argsort(r)[0,::-1][:5]+1])

아래는 분류한 클래스를 뜻합니다.

 

vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')

print()
for layer in vgg.layers:
  print(layer.name)

include_top=False로 주면 분류하는 맨 끝단 레이어만 제외하고 불러올 수 있습니다.

레이어를 추출해서 원하는 특징을 가져올 수 있습니다.

# Content layer where will pull our feature maps
content_layers = ['block5_conv2'] 

# Style layer we are interested in
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1', 
                'block4_conv1', 
                'block5_conv1']

num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

block5-conv2는 인풋 레이어에서 멀기 때문에 content에 가깝고, 스타일 레이어는 다양하게 추출한 것을 볼 수 있습니다.

Build model(모델 만들기)

아래 함수를 통해 레이어 이름으로 레이어의 아웃풋을 만들 수 있습니다.

def vgg_layers(layer_names):
  """ Creates a vgg model that returns a list of intermediate output values."""
  # Load our model. We load pretrained VGG, trained on imagenet data
  vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
  vgg.trainable = False
  
  outputs = [vgg.get_layer(name).output for name in layer_names]

  model = tf.keras.Model([vgg.input], outputs)
  return model

아웃풋 추출한 결과입니다.

style_extractor = vgg_layers(style_layers)
style_outputs = style_extractor(style_image*255)

#Look at the statistics of each layer's output
for name, output in zip(style_layers, style_outputs):
  print(name)
  print("  shape: ", output.numpy().shape)
  print("  min: ", output.numpy().min())
  print("  max: ", output.numpy().max())
  print("  mean: ", output.numpy().mean())
  print()

 

Calculate Style(스타일 계산하기)

Gram matrix를 통해 스타일을 계산할 수 있습니다.

다른 특징 맵들 간의 평균과 상관관계를 통해 구할 수 있습니다. 아래는 구현 결과입니다.

def gram_matrix(input_tensor):
  result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
  input_shape = tf.shape(input_tensor)
  num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)
  return result/(num_locations)

Extract style and content(스타일과 내용 추출)

class StyleContentModel(tf.keras.models.Model):
  def __init__(self, style_layers, content_layers):
    super(StyleContentModel, self).__init__()
    self.vgg =  vgg_layers(style_layers + content_layers)
    self.style_layers = style_layers
    self.content_layers = content_layers
    self.num_style_layers = len(style_layers)
    self.vgg.trainable = False

  def call(self, inputs):
    "Expects float input in [0,1]"
    inputs = inputs*255.0
    preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)
    outputs = self.vgg(preprocessed_input)
    style_outputs, content_outputs = (outputs[:self.num_style_layers], 
                                      outputs[self.num_style_layers:])

    style_outputs = [gram_matrix(style_output)
                     for style_output in style_outputs]

    content_dict = {content_name:value 
                    for content_name, value 
                    in zip(self.content_layers, content_outputs)}

    style_dict = {style_name:value
                  for style_name, value
                  in zip(self.style_layers, style_outputs)}
    
    return {'content':content_dict, 'style':style_dict}

extractor = StyleContentModel(style_layers, content_layers)

results = extractor(tf.constant(content_image))

style_results = results['style']

print('Styles:')
for name, output in sorted(results['style'].items()):
  print("  ", name)
  print("    shape: ", output.numpy().shape)
  print("    min: ", output.numpy().min())
  print("    max: ", output.numpy().max())
  print("    mean: ", output.numpy().mean())
  print()

print("Contents:")
for name, output in sorted(results['content'].items()):
  print("  ", name)
  print("    shape: ", output.numpy().shape)
  print("    min: ", output.numpy().min())
  print("    max: ", output.numpy().max())
  print("    mean: ", output.numpy().mean())

Run gradient descent(경사 하강법)

extractor로 style_target 레이어와 content_target 레이어를 뽑아냅니다.

style_targets = extractor(style_image)['style']
content_targets = extractor(content_image)['content']

content 이미지의 스타일을 변경하는 것이기 때문에 최적화할 변수로 content_image를 넣습니다.

image = tf.Variable(content_image)

clip 함수로 픽셀 값을 0~1로 맞춰줍니다.

def clip_0_1(image):
  return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

Adam으로 옵티마이저를 선언해줍니다.

opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)

스타일과 콘텐츠의 가중치를 할당해줍니다.

style_weight=1e-2
content_weight=1e4

tf.GradientTape으로 이미지를 업데이트해줍니다.

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    outputs = extractor(image)
    loss = style_content_loss(outputs)

  grad = tape.gradient(loss, image)
  opt.apply_gradients([(grad, image)])
  image.assign(clip_0_1(image))

step 함수로 테스트를 조금 해봅니다.

train_step(image)
train_step(image)
train_step(image)
plt.imshow(image.read_value()[0])

조금 더 오래 학습시켜 봅시다.

import time
start = time.time()

epochs = 10
steps_per_epoch = 100

step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(image)
    print(".", end='')
  display.clear_output(wait=True)
  imshow(image.read_value())
  plt.title("Train step: {}".format(step))
  plt.show()

end = time.time()
print("Total time: {:.1f}".format(end-start))

Total Variation loss

위처럼 구현하면 고주파수가 많이 생긴다는 단점이 있습니다. (ringing 현상)

그래서 total variation loss는 이러한 이미지의 현상을 줄여주는 효과가 있습니다.

def high_pass_x_y(image):
  x_var = image[:,:,1:,:] - image[:,:,:-1,:]
  y_var = image[:,1:,:,:] - image[:,:-1,:,:]

  return x_var, y_var

x_deltas, y_deltas = high_pass_x_y(content_image)

plt.figure(figsize=(14,10))
plt.subplot(2,2,1)
imshow(clip_0_1(2*y_deltas+0.5), "Horizontal Deltas: Original")

plt.subplot(2,2,2)
imshow(clip_0_1(2*x_deltas+0.5), "Vertical Deltas: Original")

x_deltas, y_deltas = high_pass_x_y(image)

plt.subplot(2,2,3)
imshow(clip_0_1(2*y_deltas+0.5), "Horizontal Deltas: Styled")

plt.subplot(2,2,4)
imshow(clip_0_1(2*x_deltas+0.5), "Vertical Deltas: Styled")

이는 edge-detector와 비슷합니다.

plt.figure(figsize=(14,10))

sobel = tf.image.sobel_edges(content_image)
plt.subplot(1,2,1)
imshow(clip_0_1(sobel[...,0]/4+0.5), "Horizontal Sobel-edges")
plt.subplot(1,2,2)
imshow(clip_0_1(sobel[...,1]/4+0.5), "Vertical Sobel-edges")

따라서 high_pass_x_y로 가로, 세로의 변화율을 구해서 제곱한 것을 optimize 해주면 됩니다.

def total_variation_loss(image):
  x_deltas, y_deltas = high_pass_x_y(image)
  return tf.reduce_mean(x_deltas**2) + tf.reduce_mean(y_deltas**2)

Re-run the optimization

가중치를 할당합니다.

total_variation_weight=1e8

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    outputs = extractor(image)
    loss = style_content_loss(outputs)
    loss += total_variation_weight*total_variation_loss(image)

  grad = tape.gradient(loss, image)
  opt.apply_gradients([(grad, image)])
  image.assign(clip_0_1(image))

import time
start = time.time()

epochs = 10
steps = 100

step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(image)
    print(".", end='')
  display.clear_output(wait=True)
  imshow(image.read_value())
  plt.title("Train step: {}".format(step))
  plt.show()

end = time.time()
print("Total time: {:.1f}".format(end-start))image = tf.Variable(content_image)