[Keras] Style Transfer 코드 공부

https://medium.com/data-science-group-iitr/artistic-style-transfer-with-convolutional-neural-network-7ce2476039fd

https://harishnarayanan.org/writing/artistic-style-transfer

위 두 글에 나온 코드를 복붙해가면서 스타일 변환기를 구현해봤습니다.

아직 이론적으로나 코드에서나 이해가 안되는 부분이 있어서 좀 더 공부하고 직접 처음부터 짜봐야겠습니다.

코드:

from __future__ import print_function
import numpy as np
from PIL import Image
import time
 
from keras.applications.vgg16 import VGG16
 
from keras import backend
from keras.models import Model
from scipy.optimize import fmin_l_bfgs_b
 
def get_file_from_drive(file_id):
  """
  구글 드라이브에서 file_id에 해당하는 파일을 가져와 읽습니다.
  Colab에서 돌리기 위해 필요합니다.
  """
  
  # Install the PyDrive wrapper & import libraries.
  # This only needs to be done once per notebook.
  from pydrive.auth import GoogleAuth
  from pydrive.drive import GoogleDrive
  from google.colab import auth
  from oauth2client.client import GoogleCredentials
 
  # Authenticate and create the PyDrive client.
  # This only needs to be done once per notebook.
  auth.authenticate_user()
  gauth = GoogleAuth()
  gauth.credentials = GoogleCredentials.get_application_default()
  drive = GoogleDrive(gauth)
 
  # Download a file based on its file ID.
  #
  # A file ID looks like: laggVyWshwcyP6kEI-y_W3P8D26sz
  downloaded = drive.CreateFile({'id': file_id})
  return downloaded
 
# 구글 드라이브에서 파일을 다운로드 합니다.
get_file_from_drive("1hSALjZqlkzcScSKglb5Zc1YlVgX4TPcT").GetContentFile("content.png")
get_file_from_drive("1ihv3YLEf_vLa5Dn2Yd5opXlsb9Vx_YvU").GetContentFile("style.png")
 
height = 512
width = 512
 
# 이미지를 불러오고 사이즈를 바꿉니다.
content_image = Image.open('content.png')
content_image = content_image.resize((width, height))
 
style_image = Image.open('style.png')
style_image = style_image.resize((width, height))
 
# 이미지의 픽셀 데이터를 얻어오고 차원을 하나 추가합니다.
# 만약 채널이 4개면 마지막 채널을 제거합니다.
content_array = np.asarray(content_image, dtype='float32')
if (content_array.shape[2] == 4):
  content_array = content_array[:, :, :-1]
content_array = np.expand_dims(content_array, axis=0)
 
style_array = np.asarray(style_image, dtype='float32')
if (style_array.shape[2] == 4):
  style_array = style_array[:, :, :-1]
style_array = np.expand_dims(style_array, axis=0)
 
print(content_array.shape)
print(style_array.shape)
 
# 이미지 데이터를 전처리 합니다.
def subtract_mean_rgb(img_arr):
  img_arr[:, :, :, 0] -= 103.939
  img_arr[:, :, :, 1] -= 116.779
  img_arr[:, :, :, 2] -= 123.68
  return img_arr[:, :, :, ::-1]
 
content_array = subtract_mean_rgb(content_array)
style_array = subtract_mean_rgb(style_array)
 
# 이미지 데이터를 케라스 변수로 만듭니다.
content_image = backend.variable(content_array)
style_image = backend.variable(style_array)
combination_image = backend.placeholder((1, height, width, 3))
 
# 신경망에 실제로 넣을 데이터를 만듭니다.
input_tensor = backend.concatenate([content_image, style_image, combination_image], axis=0)
print(input_tensor.shape)
 
# 모델을 불러옵니다.
model = VGG16(input_tensor=input_tensor, weights='imagenet', include_top=False)
 
# 스타일 변환 매개변수 선언.
content_weight = 0.05
style_weight = 5.0
total_variation_weight = 1.0
 
# 모델의 여러 레이어에서 값을 사용할 것이므로 미리 모든 레이어를 가져옴.
layers = dict([(layer.name, layer.output) for layer in model.layers])
 
# 전체 손실을 담을 변수 선언.
loss = backend.variable(0.0)
 
# Content 손실 함수(MSE)
def content_loss(content, combination):
  return backend.sum(backend.square(content - combination))
 
# Content 손실을 계산합니다.
# block2_conv2는 VGG의 2-2번째 합성곱 + Relu 계층입니다.
# 어떤 계층을 사용할지는 하이퍼 파리미터입니다.
layer_features = layers['block2_conv2']
content_image_features = layer_features[0, :, :, :]
combination_features = layer_features[2, :, :, :]
loss = loss + content_weight * content_loss(content_image_features, combination_features)
 
# 이미지에서 스타일에만 집중하도록 그람 행렬을 사용합니다.
def gram_matrix(x):
  features = backend.batch_flatten(backend.permute_dimensions(x, (2, 0, 1)))
  gram = backend.dot(features, backend.transpose(features))
  return gram
 
# Style 손실 함수
def style_loss(style, combination):
  S = gram_matrix(style)
  C = gram_matrix(combination)
  channels = 3
  size = height * width
  st = backend.sum(backend.square(S - C)) / (4. * (channels ** 2) * (size ** 2))
  return st
 
# 스타일 손실을 계산할 계층
feature_layers = ['block1_conv2', 'block2_conv2',
                  'block3_conv3', 'block4_conv3',
                  'block5_conv3']
 
# 스타일 손실을 계산합니다.
for layer_name in feature_layers:
  layer_features = layers[layer_name]
  style_features = layer_features[1, :, :, :]
  combination_features = layer_features[2, :, :, :]
  sl = style_loss(style_features, combination_features)
  loss = loss + (style_weight / len(feature_layers)) * sl
 
# 결과 이미지의 노이즈를 조절하기 위한 손실 함수
def total_variation_loss(x):
  a = backend.square(x[:, :height-1, :width-1, :] - x[:, 1:, :width-1, :])
  b = backend.square(x[:, :height-1, :width-1, :] - x[:, :height-1, 1:, :])
  return backend.sum(backend.pow(a + b, 1.25))
loss = loss + total_variation_weight * total_variation_loss(combination_image)
 
# 학습을 위해 경사도를 계산합니다.
grads = backend.gradients(loss, combination_image)
 
# scipy.optimize는 Loss와 Grads를 따로 요구하는데
# 따로 계산하는건 비효율적이므로 아래 코드가 사용됨.
outputs = [loss]
if isinstance(grads, (list, tuple)):
  outputs += grads
else:
  outputs.append(grads)
f_outputs = backend.function([combination_image], outputs)
 
def eval_loss_and_grads(x):
  x = x.reshape((1, height, width, 3))
  outs = f_outputs([x])
  loss_value = outs[0]
  grad_values = outs[1].flatten().astype('float64')
  return loss_value, grad_values
 
class Evaluator(object):
  def __init__(self):
    self.loss_value = None
    self.grads_values = None
    
  def loss(self, x):
    assert self.loss_value is None
    loss_value, grad_values = eval_loss_and_grads(x)
    self.loss_value = loss_value
    self.grad_values = grad_values
    return self.loss_value
 
  def grads(self, x):
    assert self.loss_value is not None
    grad_values = np.copy(self.grad_values)
    self.loss_value = None
    self.grad_values = None
    return grad_values
  
evaluator = Evaluator()
 
# 노이즈 이미지 데이터 생성.
x = np.random.uniform(0, 255, (1, height, width, 3)) - 128.0
 
# 학습 진행.
iterations = 10
for i in range(iterations):
  print('Start of iteration', i)
  start_time = time.time()
  x, min_val, info = fmin_l_bfgs_b(evaluator.loss, x.flatten(),
                                   fprime=evaluator.grads, maxfun=20)
  print(min_val)
  end_time = time.time()
  print('Iteration %d completed in %ds' % (i, end_time - start_time))
  
# 결과 데이터를 이미지로 바꾸기 위한 처리.
x = x.reshape((height, width, 3))
x = x[:, :, ::-1]
x[:, :, 0] += 103.939
x[:, :, 1] += 116.779
x[:, :, 2] += 123.68
x = np.clip(x, 0, 255).astype('uint8')
 
result_img = Image.fromarray(x)
result_img.save("result.bmp")

결과:

(1, 512, 512, 3)

(1, 512, 512, 3)

(3, 512, 512, 3)

Start of iteration 0

66426397000.0

Iteration 0 completed in 19s

Start of iteration 1

39155966000.0

Iteration 1 completed in 18s

Start of iteration 2

33633920000.0

Iteration 2 completed in 18s

Start of iteration 3

31667130000.0

Iteration 3 completed in 18s

Start of iteration 4

30758005000.0

Iteration 4 completed in 18s

Start of iteration 5

30270147000.0

Iteration 5 completed in 18s

Start of iteration 6

29993350000.0

Iteration 6 completed in 18s

Start of iteration 7

29820692000.0

Iteration 7 completed in 18s

Start of iteration 8

29712857000.0

Iteration 8 completed in 18s

Start of iteration 9

29645433000.0

Iteration 9 completed in 18s

+

=