圖像風格轉移

介紹

什麼是圖像風格遷移，我想圖片比文字更具表現力。上圖中“output image”就是圖像風格遷移後得到的結果。那麼它是如何實現的呢？首先讓我們看下CNN每層學習到了什麼。

如圖所示，CNN網絡最開始會學習到圖像的“紋理”，“邊緣“等信息，隨着層數的加深將會學習到更加豐富的信息。其實，在圖像風格轉移中我們就是使用卷積的前幾層作爲圖像的”風格“。至於”content image“方法一樣，只不過我們使用較高的層作爲輸出。
顯而易見，我們需要有一個強大的CNN網絡用來提取特徵，爲此，我們利用遷移學習使用VGG19模型。有關遷移學習，VGG16模型介紹，可以查看通過遷移學習實現OCT圖像識別這篇文章。

將”style image“,"content image","init image(要生成的目標圖像)"輸入VGG19網絡,提取特徵構建模型。分別計算”content loss“，”style loss“並與係數相乘，然後將兩個損失相加得到總損失。得到總損失後就可以計算對”init image“的梯度，然後使用梯度下降更新。

項目的細節要求，將會在對應代碼裏介紹。這裏極力推薦使用”google colab“，當然，前提是”科學上網“。

數據處理

加載圖片：

import os

img_dir='/tmp/nst'

if not os.path.exists(img_dir):
    os.makedirs(img_dir)

import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['figure.figsize']=(10,10)
mpl.rcParams['axes.grid']=False
import numpy as np
from PIL import Image
import time
import functools
import tensorflow as tf
import tensorflow.contrib.eager as tfe
from tensorflow.python.keras.preprocessing import image as kp_image
from tensorflow.python.keras import models
from tensorflow.python.keras import losses
from tensorflow.python.keras import layers
from tensorflow.python.keras import backend as K

# 開啓eager模式，開啓後不能關閉
tf.enable_eager_execution()

# content圖片路徑
content_path='/tmp/nst/Green_Sea_Turtle_grazing_seagrass.jpg'
# style圖片路徑
style_path='/tmp/nst/The_Great_Wave_off_Kanagawa.jpg'

def load_img(path_to_img):
    max_dim=512
    img=Image.open(path_to_img)

    # img.size:
    # return:width,height
    long=max(img.size)

    # 縮放比
    scale=max_dim/long

    # img.size[0]:width img.size[1]:height
    # round:返回四捨五入的值
    # Image.ANTIALIAS:抗鋸齒
    img=img.resize((round(img.size[0]*scale),round(img.size[1]*scale)),Image.ANTIALIAS)

    img=kp_image.img_to_array(img)

    # expand dim:batch_size
    # axis:對於2維來說，0:列，1：行，對於大於2維來說：維度從外向里加，如5維度：0，1，2，3，4
    img=np.expand_dims(img,axis=0)

    return img

顯示照片：

def imgshow(img,title=None):
    # load_img fn:增加了batch_size 維度
    # 這裏顯示照片不需要此維度
    out=tf.squeeze(img,axis=0)

    out=out.astype('uint8')

    if title is not None:
        plt.title(title)
    plt.imshow(out)

    # 顯示content圖像和style圖像

    plt.figure(figsize=(10,10))
    content_img=load_img(content_path).astype('uint8')
    style_img=load_img(style_path).astype('uint8')

    plt.subplot(1,2,1)
    imgshow(content_img,'content_img')

    plt.subplot(1,2,2)
    imgshow(style_img)

    plt.show()

將圖片轉爲適合VGG19的輸入格式：

def load_and_process_img(img_path):
    img=load_img(img_path)
    # vgg提供的預處理，主要完成（1）去均值（2）RGB轉BGR（3）維度調換三個任務。
    img=tf.keras.applications.vgg19.preprocess_input(img)

    return img

將圖片由BGR轉到RGB並將像素值限制到[0，255]：

def deprocess_img(processed_img):
    x=processed_img.copy()

    if len(x.shape) == 4:
        x=np.squeeze(x,0)
    assert len(x.shape) == 3
    # 如果是RGB轉BGR，此處改爲”-=“
    x[:, :, 0] += 103.939
    x[:, :, 1] += 116.779
    x[:, :, 2] += 123.68
    # 'BGR'->'RGB'
    x = x[:, :, ::-1]

    x=np.clip(x,a_min=0,a_max=255).astype('uint8')

    return x

創建模型

指定使用VGG19模型中的哪些層作爲”content image“特徵層，”style image“特徵層，並以此來構建新模型。

# content層
content_layers=['block5_conv2']

# style層
style_layers=[
    'block1_conv1',
    'block2_conv1',
    'block3_conv1',
    'block4_conv1',
    'block5_conv1'
]

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

# 創建模型

# 使用vgg19中間層作爲模型輸出
def get_model():
    vgg=tf.keras.applications.vgg19.VGG19(
        # 不使用最後全連接層
        include_top=False,
        # 使用imagenet數據集
        weights='imagenet'
    )
    # 因爲vgg19我們僅是用來提取特徵
    vgg.trainable=False

    # 獲取對應層輸出
    style_outputs=[ vgg.get_layer(name).output for name in style_layers]
    content_outputs=[ vgg.get_layer(name).output for name in content_layers]

    model_outputs=style_outputs+content_outputs

    return models.Model(vgg.input,model_outputs)

損失函數

content loss：

模型的”content loss“就是輸入圖像X和原始圖像P之間的歐氏距離，損失函數如下圖所示：

style loss：

我們將l層第i個feature map和第j個feature map的內積，表示模型提取的”風格特徵“，然後依然使用歐氏距離來計算損失。
一層損失計算：

我們的”style loss“，一般具有多層，所以總”style loss“需要累加：

總損失：
模型總損失就是”content loss“與“style loss”相加：

# content 損失
def get_content_loss(base_content,target):

    # 歐式距離計算損失
    return tf.reduce_mean(tf.square(base_content - target))

# style 損失
# 使用gram矩陣來表示風格特徵
def gram_matrix(input_tensor):
    # (batch_size,height,width,channel)
    channels=int(input_tensor.shape[-1])
    a=tf.reshape(input_tensor,shape=[-1,channels])
    n=tf.shape(a)[0]
    gram=tf.matmul(a=a,b=a,transpose_a=True)

    return gram/tf.cast(n,tf.float32)
    
def get_style_loss(base_style,gram_target):
    gram_style=gram_matrix(base_style)
    
    # 歐氏距離計算損失
    return tf.reduce_mean(tf.square(gram_style - gram_target))

計算損失函數，自然需要獲取模型輸出，下面獲取“content output”和“style output”：

def get_feature_representtations(model,content_path,style_path):

    # 將content img，style img 轉爲適合VGG19的輸入
    content_img=load_and_process_img(content_path)
    style_img=load_and_process_img(style_path)

    # 創建content，style模型
    content_outputs=model(style_img)
    style_outputs=model(content_img)

    # model output feature 注意此處取值區間
    # model output == content out + style out
    content_features=[ content_layer[0] for content_layer in content_outputs[num_style_layers:]]
    style_features=[ style_layer[0] for style_layer in style_outputs[:num_style_layers]]

    return content_features,style_features

梯度計算

def compute_loss(model,loss_weight,init_image,gram_style_features,content_features):
    # “style image”損失函數係數，“content image”損失函數係數
    # 此係數的作用是讓“output image”內容更像誰一些，比如：
    # content image係數更大，那麼“output image”內容與“content image”相似度更高
    style_weight,content_weight=loss_weight
    
    # 將“init image”輸入VGG19模型，得到“init_image_output features”
    model_outputs=model(init_image)
    
    # 根據上面的設置獲取對應區間層的“feature output”
    style_output_features=model_outputs[:num_style_layers]
    content_output_features=model_outputs[num_style_layers:]
    
    # “style image” loss
    style_score=0
    
    # “content image” loss
    content_score=0
    
    # 先計算每層損失，並設定每層的損失權重相同（當然，可以設置每層權重不同值）
    
    # 設定每層損失權重相同
    weight_per_style_layer=1.0/float(num_style_layers)
    
    # 累加每層損失
    for target_style,comb_style in zip(gram_style_features,style_output_features):
        style_score+=weight_per_style_layer*get_style_loss(comb_style[0],target_style)
        
    # 與“style_score”損失同理
    weight_per_content_layer=1.0/float(num_content_layers)
    for target_content,comb_content in zip(content_features,content_output_features):
        content_score+=weight_per_content_layer*get_content_loss(comb_content[0],target_content)
  
  # 損失函數*對應係數   
  style_score *= style_weight
  content_score *= content_weight

  # 相加得到總損失
  loss = style_score + content_score 
  return loss, style_score, content_score

def compute_grads(cfg):
    # eager模式下，先記錄
    with tf.GradientTape() as tape:
        # 參數輸入形式是字典
        all_loss=compute_loss(**cfg)
    total_loss=all_loss[0]
    return tape.gradient(total_loss,cfg['init_image']),all_loss

模型訓練

import IPython.display

def run_style_transfer(content_path, 
                       style_path,
                       num_iterations=1000,
                       content_weight=1e3, 
                       style_weight=1e-2): 
  
  # 此處我們的模型主要是用來提取特徵，做損失函數
  model = get_model() 
  for layer in model.layers:
    layer.trainable = False
  
  #  獲取模型“style feature”和“content feature”，注意此函數的取值區間
  style_features, content_features = get_feature_representations(model, content_path, style_path)
  
  # 將“style feature”轉爲可用於計算損失的gram矩陣形式
  gram_style_features = [gram_matrix(style_feature) for style_feature in style_features]
  
  # 目標圖像設置，此處使用的是“content image”
  # 此圖像的初始化對結果影響不大
  init_image = load_and_process_img(content_path)
  
  # eager模式下變量使用“tfe.Variable”
  init_image = tfe.Variable(init_image, dtype=tf.float32)
  
  # 優化器設置
  # beta1：一階矩估計的指數衰減率
  opt = tf.train.AdamOptimizer(learning_rate=5, beta1=0.99, epsilon=1e-1)
  
  # 初始化模型結果
  # float('inf') 正無窮 float('-inf')負無窮
  best_loss, best_img = float('inf'), None
  
  # 損失函數參數配置
  loss_weights = (style_weight, content_weight)
  cfg = {
      'model': model,
      'loss_weights': loss_weights,
      'init_image': init_image,
      'gram_style_features': gram_style_features,
      'content_features': content_features
  }
    
  # 設置訓練結果
  num_rows = 2
  num_cols = 5
  display_interval = num_iterations/(num_rows*num_cols)
  start_time = time.time()
  global_start = time.time()
  
  norm_means = np.array([103.939, 116.779, 123.68])
  min_vals = -norm_means
  max_vals = 255 - norm_means   
  
  imgs = []
  for i in range(num_iterations):
    
    # 梯度計算及參數更新
    grads, all_loss = compute_grads(cfg)
    loss, style_score, content_score = all_loss
    
    opt.apply_gradients([(grads, init_image)])
    clipped = tf.clip_by_value(init_image, min_vals, max_vals)
    init_image.assign(clipped)
    end_time = time.time() 
    
    if loss < best_loss:
      # 損失更新
      best_loss = loss
      # 轉爲RGB顯示
      best_img = deprocess_img(init_image.numpy())

    if i % display_interval== 0:
      start_time = time.time()
      
      # 顯示訓練過程
      plot_img = init_image.numpy()
      
      # 轉爲RGB顯示
      plot_img = deprocess_img(plot_img)
      
      imgs.append(plot_img)
      IPython.display.clear_output(wait=True)
      IPython.display.display_png(Image.fromarray(plot_img))
      print('Iteration: {}'.format(i))        
      print('Total loss: {:.4e}, ' 
            'style loss: {:.4e}, '
            'content loss: {:.4e}, '
            'time: {:.4f}s'.format(loss, style_score, content_score, time.time() - start_time))
  print('Total time: {:.4f}s'.format(time.time() - global_start))
  IPython.display.clear_output(wait=True)
  plt.figure(figsize=(14,4))
  for i,img in enumerate(imgs):
      plt.subplot(num_rows,num_cols,i+1)
      plt.imshow(img)
      plt.xticks([])
      plt.yticks([])
      
  return best_img, best_loss 

訓練結果展示：

best, best_loss = run_style_transfer(content_path, 
                                     style_path, num_iterations=1000)
Image.fromarray(best)

總結

我們利用遷移學習使用VGG19模型提取“style feature”和“content feature”，都使用歐氏距離計算損失函數。其中，使用gram矩陣計算“style loss”。
最近開始使用”google colab“訓練模型，感覺不錯，推薦給大家。

本文代碼部分來自Raymond Yuan，在次表示感謝。