数字图像中风格迁移Python代码

以下是数字图像中风格迁移的Python代码实现：

import tensorflow as tf
import numpy as np
import PIL.Image
import time
import functools

# 加载图像
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)

# 加载模型
def load_model():
    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
    vgg.trainable = False

    # 获取每个中间层的输出
    content_layers = ['block5_conv2']

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

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

    return vgg, style_layers, content_layers

# 计算Gram矩阵
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)

# 提取特征
def get_feature_representations(model, content_path, style_path):
    content_image = load_img(content_path)
    style_image = load_img(style_path)

    # 提取内容图像的特征
    content_outputs = model(content_image)
    content_features = [content_outputs[layer_name] for layer_name in content_layers]

    # 提取风格图像的特征
    style_outputs = model(style_image)
    style_features = [style_outputs[layer_name] for layer_name in style_layers]

    # 计算风格图像的Gram矩阵
    style_feature_outputs = [gram_matrix(style_feature) for style_feature in style_features]

    # 将内容图像和风格图像的特征合并
    content_dict = {content_name:value for content_name,value in zip(content_layers, content_features)}
    style_dict = {style_name:value for style_name,value in zip(style_layers, style_feature_outputs)}

    return {'content':content_dict, 'style':style_dict}

# 计算内容损失
def get_content_loss(base_content, target):
    return tf.reduce_mean(tf.square(base_content - target))

# 计算风格损失
def get_style_loss(base_style, gram_target):
    height, width, channels = base_style.get_shape().as_list()
    gram_style = gram_matrix(base_style)

    return tf.reduce_mean(tf.square(gram_style - gram_target))

# 计算总损失
def compute_loss(model, loss_weights, init_image, gram_style_features, content_features):
    style_weight, content_weight = loss_weights

    # 提取初始图像的特征
    model_outputs = model(init_image)

    # 将特征分为内容特征和风格特征
    content_output_features = model_outputs[content_layers]
    style_output_features = model_outputs[style_layers]

    # 计算内容损失
    content_loss = tf.add_n([get_content_loss(content_output_features[name], content_features[name]) for name in content_output_features.keys()])
    content_loss *= content_weight / len(content_layers)

    # 计算风格损失
    style_loss = tf.add_n([get_style_loss(style_output_features[name], gram_style_features[name]) for name in style_output_features.keys()])
    style_loss *= style_weight / len(style_layers)

    # 计算总损失
    loss = content_loss + style_loss
    return loss

# 计算梯度
def compute_grads(cfg):
    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

# 进行风格迁移
def run_style_transfer(content_path, style_path, num_iterations=1000, content_weight=1e3, style_weight=1e-2):
    # 加载模型
    model, style_layers, content_layers = load_model()

    # 提取内容图像和风格图像的特征
    feature_representations = get_feature_representations(model, content_path, style_path)
    content_features = feature_representations['content']
    style_features = feature_representations['style']

    # 计算风格图像的Gram矩阵
    gram_style_features = {name:gram_matrix(style_features[name]) for name in style_features.keys()}

    # 初始化图像
    init_image = load_img(content_path)
    init_image = tf.Variable(init_image, dtype=tf.float32)

    # 定义优化器
    opt = tf.optimizers.Adam(learning_rate=5, beta_1=0.99, epsilon=1e-1)

    # 定义损失权重
    loss_weights = (style_weight, content_weight)

    # 进行风格迁移
    start_time = time.time()

    best_loss, best_img = float('inf'), None

    cfg = {
        'model': model,
        'loss_weights': loss_weights,
        'init_image': init_image,
        'gram_style_features': gram_style_features,
        'content_features': content_features
    }

    # 迭代优化
    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, clip_value_min=0.0, clip_value_max=1.0)
        init_image.assign(clipped)

        if loss < best_loss:
            best_loss = loss
            best_img = init_image.numpy()

        if i % 100 == 0:
            print("Iteration: {}".format(i))
            print("Total loss: {:.4e}, "
                  "style loss: {:.4e}, "
                  "content loss: {:.4e}".format(loss, style_score, content_score))

    end_time = time.time()
    print("Total time: {:.1f}".format(end_time - start_time))

    # 显示结果
    best_img = tf.squeeze(best_img, axis=0)
    imshow(best_img, 'Output Image')

# 运行风格迁移
content_path = 'content.jpg'
style_path = 'style.jpg'
run_style_transfer(content_path, style_path, num_iterations=1000, content_weight=1e3, style_weight=1e-2)