본문 바로가기
Visual Intelligence/Generative Model

[Generative Model] Neural style transfer

by goatlab 2022. 12. 9.
728x90
반응형
SMALL

데이터 로드

 

import keras
from google.colab import auth # Google File Drive Stream 접근 허용
auth.authenticate_user()

from google.colab import drive
drive.mount('/content/gdrive')

 

사용할 데이터는 구글 드라이브에 저장하고 Colaboratory에서 연동하여 파일을 불러와 사용한다.

 

%cd gdrive/My Drive/Datasets/sample/

 

이미지를 작성하기 전에 colaboratory의 terminal 위치를 이미지 샘플이 있는 곳으로 바꿔준다.

 

from keras.preprocessing import image
from tensorflow.keras.utils import load_img, img_to_array, save_img

# 변환하려는 이미지 경로
target_image_path = '/content/cat.jpeg'

# 스타일 이미지 경로
style_reference_image_path = '/content/style.png'

# 생성된 사진의 차원
width, height = load_img(target_image_path).size
img_height = 400
img_width = int(width * img_height / height)

 

처리할 이미지들의 크기가 다르면 style transfer를 구현하는 것이 더 어려워지게 되기 때문에 같은 높이로 크기를 변경한다.

 

import tensorflow.keras.backend as K
import tensorflow as tf
import numpy as np
from keras.applications import vgg19

tf.compat.v1.disable_eager_execution()

# Utility Function
def preprocess_image(image_path):
    img = load_img(image_path, target_size = (img_height, img_width))
    img = img_to_array(img)
    img = np.expand_dims(img, axis = 0)
    img = vgg19.preprocess_input(img)
    
    return img

def deprocess_image(x):
    # ImageNet 평균 픽셀 값 (cf)vgg19.preprocess_input conversion
    x[:, :, 0] += 103.939
    x[:, :, 1] += 116.779
    x[:, :, 2] += 123.68
    
    #'BGR -> RGB'
    x = x[:, :, ::-1]
    x = np.clip(x, 0, 255).astype('uint8')
    
    return x

 

VGG 컨볼루션넷에 입출력할 이미지의 로드, 전처리, 사후 처리를 위한 유틸리티 함수를 정의한다.

 

Style Transfer 네트워크 구성

 

# Loading pretrained VGG19 & applying to 3 image
target_image = K.constant(preprocess_image(target_image_path))
style_reference_image = K.constant(preprocess_image(style_reference_image_path))
combination_image = K.placeholder((1, img_height, img_width, 3)) # placeholder for generated image

# 입력된 훈련 데이터를 미니 배치로 나누어 계산 그래프에 주입하기 위해 보통 플레이스홀더의 첫 번째 차원 (배치 차원)을 비움
# style transfer은 하나의 이미지만 사용하므로 combination_image placeholder의 배치 차원을 1로 고정
input_tensor = K.concatenate([target_image, style_reference_image, combination_image], axis = 0)

# 타깃 , 스타일 참조, 생성된 이미지를 행으로 쌓아서(axis = 0) minibatch처럼 생성
# input_tensor의 차원은 (3, 400, 300, 3)
model = vgg19.VGG19(input_tensor = input_tensor, weights = 'imagenet', include_top = False)
model.summary()

print('Loading Model Finish')

 

콘텐츠 손실에서 VGG19 컨볼루션 네트워크의 상위 층은 타깃 이미지와 생성된 이미지를 동일하게 바라봐야 한다.

 

# Definition of Loss
def content_loss(base, combination):
    return K.sum(K.square(combination - base))

 

스타일 손실은 유틸리티 함수를 이용해 입력 행렬의 그람 행렬을 계산한다. 이 행렬은 원본 특성 행렬의 상관 관계를 기록한 행렬이다.

 

def gram_matrix(x):
    features = K.batch_flatten(K.permute_dimensions(x, (2, 0, 1)))
    # batch_flatten : Turn a nD tensor into a 2D tensor with same 0th dimension.
    # permute_dimensions : Permutes axes in a tensor
    
    # 특성 맵의 차원을 배치 차원처럼 맨 앞으로 옯겨서 특성맵을 기준으로 벡터를 펼침
    gram = K.dot(features, K.transpose(features))
    
    return gram

def style_loss(style, combination):
    S = gram_matrix(style)
    C = gram_matrix(combination)
    channels = 3
    size = img_height * img_width
    
    return K.sum(K.square(S - C)) / (4. * (channels ** 2) * (size ** 2))

# Total Variation Loss
def total_variation_loss(x):
    a = K.square(x[:, :img_height - 1, :img_width - 1, :] - x[:, 1:,:img_width - 1, :])
    b = K.square(x[:, :img_height - 1, :img_width - 1, :] - x[:, img_height - 1, 1:, :])
    
    return K.sum(K.pow(a + b, 1.25))
    # 두 손실에 하나를 더 추가한다. 생성된 이미지의 픽셀을 사용하여 계산하는 총 변위 손실
    # 이는 생성된 이미지가 공간적인 연속성을 가지도록 도와주며 픽셀의 격자 무늬가 과도하게 나타나는 것을 막아주는 일종의 규제 항

 

마지막으로 생성된 이미지의 픽셀을 사용해 계산하는 총 변위 손실이다. 이 세 손실의 가중치 평균을 최소화해야 한다. 콘텐츠 손실은 block5_conv2층 하나만 사용해서 계산한다. 스타일 손실을 계산하기 위해서는 하위 층과 상위 층에 걸쳐 여러 층을 사용한다. 그리고 마지막 에 총 변위 손실을 추가한다.

 

사용하는 스타일 참조 이미지와 콘텐츠 이미지에 따라 content_weight 계수 (전체 손실에 기여하는 콘텐 츠의 손실의 정도)를 조정하는 것이 좋다. Content_weight가 높으면 생성된 이미지에 타깃 콘텐츠가 더 많이 나타나게 된다.

 

# Definition of Final Loss to Minimize
outputs_dict = dict([(layer.name, layer.output) for layer in model.layers]) # 층 이름과 활성화 텐서 mapping

content_layer = 'block5_conv2' 
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1',
                'block4_conv1',
                'block5_conv1']

total_variation_weight = 1e-4    # 손실 항목의 가중치 평균에 사용할 가중치
style_weight = 1.
content_weight = 0.025

loss = K.variable(0.) # 모든 손실 요소를 더하여 하나의 스칼라 변수로 손실을 정의
layer_features = outputs_dict[content_layer]
target_image_features = layer_features[0, :, :, :]
combination_features = layer_features[2, :, :, :]
loss = loss + content_weight * content_loss(target_image_features, combination_features) # 콘텐츠 손실 합산

for layer_name in style_layers: # 스타일 손실 합산
    layer_features = outputs_dict[layer_name]
    style_reference_features = layer_features[1, :, :, :]
    combination_features = layer_features[2, :, :, :]
    sl = style_loss(style_reference_features, combination_features)
    loss = loss + (style_weight / len(style_layers)) * sl
    
loss = loss + total_variation_weight * total_variation_loss(combination_image) # 총 변위 손실 합산

 

마지막으로 경사 하강법 단계를 설정한다. 게티스의 논문에서 L-BFGS 알고리즘을 사용하여 최적화를 수행하였으므로 이를 사용한다. K-BFPS 알고리즘은 scipy 라이브러리에 구현이 되어 있는데 두 가 지 제약 사항이 있다.

 

1) 손실 함수의 값과 기울기 값을 별개의 함수로 전달해야 한다.

2) 이 함수는 3D 이미지 배열이 아니라 1차원 벡터만 처리할 수 있다.

 

손실 함수의 값과 기울기 값을 따로 계산하는 것을 비효율적이다. 두 계산 사이에 중복되는 계산이 많기 때문이다. 한꺼번에 계산하는 것보다 거의 두 배가량 느리다. 이를 피하기 위해 손실과 기울기 값을 동시에 계산하는 Evaluator란 이름의 클래스를 만들어 이용한다. 처음 호출할 때 손실 값을 반환하면서 다음 호출을 위해 기울기를 캐싱한다.

 

# Gradient Descent (L-BFGS Algorithm - quasi-Newton Method)
grads = K.gradients(loss, combination_image)[0] # 손실에 대한 생성된 이미지의 그래디언트

fetch_loss_and_grads = K.function([combination_image], [loss, grads]) # 현재 손실과 그래디언트 값 추출

class Evaluator(object):
    def __init__(self):
        self.loss_value = None
        self.grads_value = None
        
    def loss(self, x):
        assert self.loss_value is None
        x = x.reshape((1, img_height, img_width, 3))
        outs = fetch_loss_and_grads([x])
        loss_value = outs[0]
        grad_values = outs[1].flatten().astype('float64')
        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()

 

마지막으로 scipy L-BFGS 알고리즘을 사용하여 경사 하강법 단계를 수행하며 각 반복마다 생성된 이미지를 저장한다.

 

from scipy.optimize import fmin_l_bfgs_b
import time

# Style Trasfer Iteration Loop
result_prefix = 'style_transfer_result'
iterations = 21

x = preprocess_image(target_image_path) # 초기값은 타깃 이미지
x = x.flatten() # scipy.optimize.fmin_l_bfgs_b 함수가 벡터만 처리할 수 있으므로 이미지를 벡터로 펼침

for i in range(iterations):
    print('반복 횟수 : ',i)
    
    start_time = time.time()
    x, min_val, info = fmin_l_bfgs_b(evaluator.loss, x, fprime = evaluator.grads, maxfun = 20)
    print('현재 손실 값 : ',min_val)
    
    img = x.copy().reshape((img_height, img_width, 3))
    img = deprocess_image(img)
    fname = result_prefix + '_at_iteration_%d.png' % i
    
    if i%5 ==0:
        save_img(fname, img)
        print('저장 이미지 : ', fname)
        
    end_time = time.time()
    
    print('%d 번째 반복 완료 : %d초' % (i, end_time - start_time))
import matplotlib.pyplot as plt

# 콘텐츠 이미지
plt.imshow(load_img(target_image_path, target_size=(img_height, img_width)))
plt.figure()

# 스타일 이미지
plt.imshow(load_img(style_reference_image_path, target_size=(img_height, img_width)))
plt.figure()

# 생성된 이미지
plt.imshow(img)
plt.show()

728x90
반응형
LIST