import os import numpy as np import tensorflow as tf from tqdm import tqdm import gradio as gr import typing from huggingface_hub import HfApi, Repository import tempfile # 定义模型和辅助函数 print("Importing necessary libraries and defining functions...") 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} CONTENT_LOSS_FACTOR = 1 STYLE_LOSS_FACTOR = 100 WIDTH = 450 HEIGHT = 300 EPOCHS = 20 STEPS_PER_EPOCH = 100 LEARNING_RATE = 0.03 image_mean = tf.constant([0.485, 0.456, 0.406]) image_std = tf.constant([0.299, 0.224, 0.225]) def normalization(x): return (x - image_mean) / image_std def load_images(image_path, width=WIDTH, height=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. 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. 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 get_vgg19_model(layers): 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=CONTENT_LAYERS, style_layers=STYLE_LAYERS): 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 _compute_content_loss(noise_features, target_features): content_loss = tf.reduce_sum(tf.square(noise_features - target_features)) x = 2. * WIDTH * HEIGHT * 3 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. * (WIDTH * HEIGHT) ** 2 * 3 ** 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 optimizer = tf.keras.optimizers.Adam(LEARNING_RATE) model = NeuralStyleTransferModel() def neural_style_transfer(content_image_path, style_image_path): content_image = load_images(content_image_path) style_image = load_images(style_image_path) target_content_features = model([content_image, ])['content'] target_style_features = model([style_image, ])['style'] noise_image = tf.Variable((content_image + np.random.uniform(-0.2, 0.2, (1, HEIGHT, WIDTH, 3))) / 2) @tf.function def train_one_step(): 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 for epoch in range(EPOCHS): for step in range(STEPS_PER_EPOCH): _loss = train_one_step() output_image_path = tempfile.mktemp(suffix='.jpg') save_image(noise_image, output_image_path) return output_image_path def transfer_style(content_image, style_image): content_image_path = tempfile.mktemp(suffix='.jpg') style_image_path = tempfile.mktemp(suffix='.jpg') content_image.save(content_image_path) style_image.save(style_image_path) output_image_path = neural_style_transfer(content_image_path, style_image_path) return output_image_path # 创建Gradio界面 iface = gr.Interface( fn=transfer_style, inputs=[ gr.inputs.Image(type="pil", label="Content Image"), gr.inputs.Image(type="pil", label="Style Image") ], outputs=gr.outputs.Image(type="file", label="Styled Image"), title="Neural Style Transfer", description="Upload a content image and a style image to perform neural style transfer." ) # 运行Gradio应用 iface.launch()