train.py

# -*- coding: utf-8 -*-
import os
import numpy as np
from tqdm import tqdm
import tensorflow as tf
import cv2
import argparse
import typing
import h5py

# 解析命令行参数
def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument("--content_img_path", type=str, default="./images/1.jpg", help="原图路径")
    parser.add_argument("--style_img_path", type=str, default="./images/style.jpg", help="风格图片路径")
    parser.add_argument("--output_path", type=str, default="./output/1", help="生成图片保存路径")
    parser.add_argument("--epochs", type=int, default=20, help="总训练轮数")
    parser.add_argument("--step_per_epoch", type=int, default=100, help="每轮训练次数")
    parser.add_argument("--learning_rate", type=float, default=0.01, help="学习率")
    parser.add_argument("--content_loss_factor", type=float, default=1.0, help="内容损失总加权系数")
    parser.add_argument("--style_loss_factor", type=float, default=100.0, help="风格损失总加权系数")
    parser.add_argument("--img_size", type=int, default=0, help="图片尺寸，0代表不设置使用默认尺寸(450*300)，输入1代表使用图片尺寸，其他输入代表使用自定义尺寸")
    parser.add_argument("--img_width", type=int, default=450, help="自定义图片宽度")
    parser.add_argument("--img_height", type=int, default=300, help="自定义图片高度")
    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    return opt

def load_images(image_path, width, height):
    """
    加载并处理图片，返回一个张量
    """
    x = tf.io.read_file(image_path)
    x = tf.image.decode_jpeg(x, channels=3)
    x = tf.image.resize(x, [height, width])
    x = x / 255.0
    x = normalization(x)
    x = tf.reshape(x, [1, height, width, 3])
    return x

def load_images_from_list(image_array, width, height):
    """
    从numpy数组加载并处理图片，返回一个张量
    """
    x = tf.convert_to_tensor(image_array, dtype=tf.float32)
    x = tf.image.resize(x, [height, width])
    x = x / 255.0
    x = normalization(x)
    x = tf.reshape(x, [1, height, width, 3])
    return x

def save_image(image, filename):
    """
    保存图片
    """
    x = tf.reshape(image, image.shape[1:])
    x = x * image_std + image_mean
    x = x * 255.0
    x = tf.cast(x, tf.int32)
    x = tf.clip_by_value(x, 0, 255)
    x = tf.cast(x, tf.uint8)
    x = tf.image.encode_jpeg(x)
    tf.io.write_file(filename, x)

def save_image_for_gradio(image):
    """
    将图片保存为numpy数组
    """
    x = tf.reshape(image, image.shape[1:])
    x = x * image_std + image_mean
    x = x * 255.0
    x = tf.cast(x, tf.int32)
    x = tf.clip_by_value(x, 0, 255)
    x = tf.cast(x, tf.uint8)
    numpy_array = x.numpy()  # 将TensorFlow张量转换为numpy数组
    return numpy_array

def get_vgg19_model(layers):
    """
    创建并初始化vgg19模型
    """
    vgg = tf.keras.applications.VGG19(include_top=False, weights="imagenet")
    outputs = [vgg.get_layer(layer).output for layer in layers]
    model = tf.keras.Model(vgg.input, outputs)
    model.trainable = False
    return model

class NeuralStyleTransferModel(tf.keras.Model):
    def __init__(self, content_layers: typing.Dict[str, float], style_layers: typing.Dict[str, float]):
        super(NeuralStyleTransferModel, self).__init__()
        self.content_layers = content_layers
        self.style_layers = style_layers
        layers = list(self.content_layers.keys()) + list(self.style_layers.keys())
        self.outputs_index_map = dict(zip(layers, range(len(layers))))
        self.vgg = get_vgg19_model(layers)

    def call(self, inputs, training=None, mask=None):
        outputs = self.vgg(inputs)
        content_outputs = []
        for layer, factor in self.content_layers.items():
            content_outputs.append((outputs[self.outputs_index_map[layer]][0], factor))
        style_outputs = []
        for layer, factor in self.style_layers.items():
            style_outputs.append((outputs[self.outputs_index_map[layer]][0], factor))
        return {"content": content_outputs, "style": style_outputs}

def normalization(x):
    """
    对输入图片进行归一化处理，返回归一化后的值
    """
    return (x - image_mean) / image_std

def _compute_content_loss(noise_features, target_features):
    """
    计算指定层上两个特征之间的内容损失
    """
    content_loss = tf.reduce_sum(tf.square(noise_features - target_features))
    x = 2.0 * M * N
    content_loss = content_loss / x
    return content_loss

def compute_content_loss(noise_content_features, target_content_features):
    """
    计算并返回当前图片的内容损失
    """
    content_losses = []
    for (noise_feature, factor), (target_feature, _) in zip(noise_content_features, target_content_features):
        layer_content_loss = _compute_content_loss(noise_feature, target_feature)
        content_losses.append(layer_content_loss * factor)
    return tf.reduce_sum(content_losses)

def gram_matrix(feature):
    """
    计算给定特征的格拉姆矩阵
    """
    x = tf.transpose(feature, perm=[2, 0, 1])
    x = tf.reshape(x, (x.shape[0], -1))
    return x @ tf.transpose(x)

def _compute_style_loss(noise_feature, target_feature):
    """
    计算指定层上两个特征之间的风格损失
    """
    noise_gram_matrix = gram_matrix(noise_feature)
    style_gram_matrix = gram_matrix(target_feature)
    style_loss = tf.reduce_sum(tf.square(noise_gram_matrix - style_gram_matrix))
    x = 4.0 * (M**2) * (N**2)
    return style_loss / x

def compute_style_loss(noise_style_features, target_style_features):
    """
    计算并返回图片的风格损失
    """
    style_losses = []
    for (noise_feature, factor), (target_feature, _) in zip(noise_style_features, target_style_features):
        layer_style_loss = _compute_style_loss(noise_feature, target_feature)
        style_losses.append(layer_style_loss * factor)
    return tf.reduce_sum(style_losses)

def total_loss(noise_features, target_content_features, target_style_features):
    """
    计算总损失
    """
    content_loss = compute_content_loss(noise_features["content"], target_content_features)
    style_loss = compute_style_loss(noise_features["style"], target_style_features)
    return content_loss * CONTENT_LOSS_FACTOR + style_loss * STYLE_LOSS_FACTOR

@tf.function
def train_one_step(model, noise_image, optimizer, target_content_features, target_style_features):
    """
    一次迭代过程
    """
    with tf.GradientTape() as tape:
        noise_outputs = model(noise_image)
        loss = total_loss(noise_outputs, target_content_features, target_style_features)
    grad = tape.gradient(loss, noise_image)
    optimizer.apply_gradients([(grad, noise_image)])
    return loss

def main(content_img, style_img, epochs, step_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height):
    global CONTENT_LOSS_FACTOR, STYLE_LOSS_FACTOR, CONTENT_IMAGE_PATH, STYLE_IMAGE_PATH, OUTPUT_DIR, EPOCHS, LEARNING_RATE, STEPS_PER_EPOCH, M, N, image_mean, image_std, IMG_WIDTH, IMG_HEIGHT

    # with tf.device('/cuda:0'):
    CONTENT_LOSS_FACTOR = content_loss_factor
    STYLE_LOSS_FACTOR = style_loss_factor
    CONTENT_IMAGE_PATH = content_img
    STYLE_IMAGE_PATH = style_img
    EPOCHS = epochs
    LEARNING_RATE = learning_rate
    STEPS_PER_EPOCH = step_per_epoch

    # 内容特征层及损失加权系数
    CONTENT_LAYERS = {"block4_conv2": 0.5, "block5_conv2": 0.5}
    # 风格特征层及损失加权系数
    STYLE_LAYERS = {
        "block1_conv1": 0.2,
        "block2_conv1": 0.2,
        "block3_conv1": 0.2,
        "block4_conv1": 0.2,
        "block5_conv1": 0.2,
    }

    if img_size == "default size":
        IMG_WIDTH = 450
        IMG_HEIGHT = 300
    else:
        IMG_WIDTH = img_width
        IMG_HEIGHT = img_height

    print("IMG_WIDTH:", IMG_WIDTH)
    print("IMG_HEIGHT:", IMG_HEIGHT)

    # 我们准备使用经典网络在imagenet数据集上的预训练权重，所以归一化时也要使用imagenet的平均值和标准差
    image_mean = tf.constant([0.485, 0.456, 0.406])
    image_std = tf.constant([0.299, 0.224, 0.225])

    model = NeuralStyleTransferModel(CONTENT_LAYERS, STYLE_LAYERS)

    content_image = load_images_from_list(CONTENT_IMAGE_PATH, IMG_WIDTH, IMG_HEIGHT)
    style_image = load_images_from_list(STYLE_IMAGE_PATH, IMG_WIDTH, IMG_HEIGHT)

    target_content_features = model(content_image)["content"]
    target_style_features = model(style_image)["style"]

    M = IMG_WIDTH * IMG_HEIGHT
    N = 3

    optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)

    noise_image = tf.Variable((content_image[0] + np.random.uniform(-0.2, 0.2, (1, IMG_HEIGHT, IMG_WIDTH, 3))) / 2)

    for epoch in range(EPOCHS):
        with tqdm(total=STEPS_PER_EPOCH, desc="Epoch {}/{}".format(epoch + 1, EPOCHS)) as pbar:
            for step in range(STEPS_PER_EPOCH):
                _loss = train_one_step(model, noise_image, optimizer, target_content_features, target_style_features)
                pbar.set_postfix({"loss": "%.4f" % float(_loss)})
                pbar.update(1)
    
    return save_image_for_gradio(noise_image)

if __name__ == "__main__":
    print(tf.config.list_physical_devices('GPU'))
    opt = parse_opt()
    main(opt.content_img_path, opt.style_img_path, opt.epochs, opt.step_per_epoch, opt.learning_rate, opt.content_loss_factor, opt.style_loss_factor, opt.img_size, opt.img_width, opt.img_height)