gan图像风格迁移代码

GAN 图像风格迁移是一种利用生成对抗网络（GAN）实现的图像风格变换技术。下面是一个使用 TensorFlow 实现 GAN 图像风格迁移的示例代码：

import tensorflow as tf
import numpy as np
import argparse
import os
import sys
import time
import datetime
import random
import cv2

def build_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='model.ckpt', help='Model checkpoint to load')
    parser.add_argument('--input_img', type=str, help='Input image file path')
    parser.add_argument('--output_img', type=str, help='Output image file path')
    parser.add_argument('--style_img', type=str, help='Style image file path')
    parser.add_argument('--content_weight', type=float, default=1.0, help='Weight of content loss')
    parser.add_argument('--style_weight', type=float, default=5.0, help='Weight of style loss')
    parser.add_argument('--tv_weight', type=float, default=1e-3, help='Weight of total variation loss')
    parser.add_argument('--learning_rate', type=float, default=1e-3, help='Learning rate')
    parser.add_argument('--num_iters', type=int, default=1000, help='Number of iterations')
    parser.add_argument('--save_every', type=int, default=100, help='Save checkpoint every N iterations')
    parser.add_argument('--print_every', type=int, default=10, help='Print loss every N iterations')
    parser.add_argument('--gpu', type=int, default=0, help='GPU device ID')
    return parser

def build_vgg19(input_tensor):
    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet', input_tensor=input_tensor)
    vgg.outputs = [vgg.layers[9].output, vgg.layers[13].output, vgg.layers[17].output, vgg.layers[21].output]
    return vgg

def gram_matrix(x):
    features = tf.keras.backend.batch_flatten(tf.keras.backend.permute_dimensions(x, (2, 0, 1)))
    gram = tf.keras.backend.dot(features, tf.keras.backend.transpose(features))
    return gram

def content_loss(content, generated):
    return tf.reduce_mean(tf.square(content - generated))

def style_loss(style, generated):
    S = gram_matrix(style)
    G = gram_matrix(generated)
    channels = 3
    size = 256 * 256
    return tf.reduce_mean(tf.square(S - G)) / (4. * (channels ** 2) * (size ** 2))

def total_variation_loss(x):
    a = tf.square(x[:, :255, :255, :] - x[:, 1:, :255, :])
    b = tf.square(x[:, :255, :255, :] - x[:, :255, 1:, :])
    return tf.reduce_mean(tf.pow(a + b, 1.25))

def build_model(content, style, generated):
    content_loss_val = content_loss(content, generated)
    style_loss_val = style_loss(style, generated)
    tv_loss_val = total_variation_loss(generated)
    loss = args.content_weight * content_loss_val + args.style_weight * style_loss_val + args.tv_weight * tv_loss_val
    optimizer = tf.train.AdamOptimizer(args.learning_rate).minimize(loss)
    return loss, optimizer

def preprocess_img(img):
    img = cv2.resize(img, (256, 256))
    img = img.astype(np.float32) / 255.0
    img = np.expand_dims(img, axis=0)
    return img

def postprocess_img(img):
    img = np.squeeze(img, axis=0)
    img = np.clip(img * 255.0, 0, 255).astype(np.uint8)
    return img

def main(args):
    # Set GPU device
    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)

    # Create session
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)

    # Build VGG19 model
    input_tensor = tf.placeholder(tf.float32, shape=(None, 256, 256, 3))
    vgg = build_vgg19(input_tensor)

    # Load images
    content_img = cv2.imread(args.input_img)
    style_img = cv2.imread(args.style_img)

    # Preprocess images
    content_img = preprocess_img(content_img)
    style_img = preprocess_img(style_img)

    # Build model
    content_tensor = vgg(input_tensor)[0]
    style_tensors = [vgg(style_img)[i] for i in range(4)]
    generated_tensor = tf.Variable(content_img, dtype=tf.float32)
    loss_op, optimizer_op = build_model(content_tensor, style_tensors, generated_tensor)

    # Load checkpoint
    saver = tf.train.Saver()
    saver.restore(sess, args.model)

    # Train model
    for i in range(args.num_iters):
        _, loss_val = sess.run([optimizer_op, loss_op])
        if i % args.print_every == 0:
            print('[{:04d}/{:04d}] loss = {:.4f}'.format(i+1, args.num_iters, loss_val))
        if i % args.save_every == 0:
            saver.save(sess, args.model)
            generated_img = sess.run(generated_tensor)
            generated_img = postprocess_img(generated_img)
            cv2.imwrite(args.output_img, generated_img)

    # Save final result
    saver.save(sess, args.model)
    generated_img = sess.run(generated_tensor)
    generated_img = postprocess_img(generated_img)
    cv2.imwrite(args.output_img, generated_img)

if __name__ == '__main__':
    parser = build_parser()
    args = parser.parse_args()
    main(args)

使用方法：

1. 下载并安装 TensorFlow。
2. 下载预训练的 VGG19 模型（[vgg19.npy](https://github.com/machrisaa/tensorflow-vgg/blob/master/vgg19.npy)）并将其放在代码所在目录下。
3. 执行以下命令：

python gan_image_style_transfer.py --input_img input.jpg --output_img output.jpg --style_img style.jpg

其中 `input.jpg` 是输入图像，`style.jpg` 是风格图像，`output.jpg` 是输出图像。可以通过调整命令行参数来改变模型训练的参数和行为。