import gradio as gr import tensorflow as tf import numpy as np from PIL import Image from tensorflow.keras.applications.vgg19 import VGG19, preprocess_input from tensorflow.keras.preprocessing.image import img_to_array from tensorflow.keras.models import Model # Load VGG19 model model = VGG19(include_top=False, weights='imagenet') model.trainable = False # Define content and style layers content_layer = 'block5_conv2' content_model = Model(inputs=model.input, outputs=model.get_layer(content_layer).output) style_layers = ['block1_conv1', 'block3_conv1', 'block5_conv1'] style_models = [Model(inputs=model.input, outputs=model.get_layer( layer).output) for layer in style_layers] weight_of_layer = 1. / len(style_models) def process_image(img): # Convert to array and preprocess img = img_to_array(img) img = preprocess_input(img) img = np.expand_dims(img, axis=0) return img def deprocess(img): # Perform the inverse of the preprocessing step img = img.copy() # Create a copy to avoid modifying the original img[:, :, 0] += 103.939 img[:, :, 1] += 116.779 img[:, :, 2] += 123.68 # Convert BGR to RGB img = img[:, :, ::-1] img = np.clip(img, 0, 255).astype('uint8') return img # Gram matrix def gram_matrix(A): channels = int(A.shape[-1]) a = tf.reshape(A, [-1, channels]) n = tf.shape(a)[0] gram = tf.matmul(a, a, transpose_a=True) return gram / tf.cast(n, tf.float32) # Content loss def content_loss(content, generated): a_C = content_model(content) a_G = content_model(generated) loss = tf.reduce_mean(tf.square(a_C - a_G)) return loss # Style loss def style_cost(style, generated): J_style = 0 for style_model in style_models: a_S = style_model(style) a_G = style_model(generated) GS = gram_matrix(a_S) GG = gram_matrix(a_G) current_cost = tf.reduce_mean(tf.square(GS - GG)) J_style += current_cost * weight_of_layer return J_style # Total Loss Function def compute_total_loss(content, style, generated, alpha=10, beta=1000): J_content = content_loss(content, generated) J_style = style_cost(style, generated) return alpha * J_content + beta * J_style def ensure_pil_image(img): if isinstance(img, np.ndarray): return Image.fromarray(img.astype('uint8')) return img def neural_style_transfer(content_img, style_img, iterations=50, alpha=10, beta=1000): try: # Ensure we have PIL images content_img_pil = ensure_pil_image(content_img) style_img_pil = ensure_pil_image(style_img) # Resize images to a manageable size content_img_pil = content_img_pil.resize((300, 300), Image.LANCZOS) style_img_pil = style_img_pil.resize((300, 300), Image.LANCZOS) # Process images content = process_image(content_img_pil) style = process_image(style_img_pil) # Initialize with content image generated = tf.Variable(content, dtype=tf.float32) # Optimizer opt = tf.keras.optimizers.Adam(learning_rate=0.7) progress_images = [] for i in range(iterations): with tf.GradientTape() as tape: total_loss = compute_total_loss( content, style, generated, alpha, beta) # Get gradients and apply grads = tape.gradient(total_loss, generated) opt.apply_gradients([(grads, generated)]) if i % 10 == 0 or i == iterations - 1: # Save progress image current_img = generated.numpy() img_squeezed = np.squeeze(current_img, axis=0) img_deprocessed = deprocess(img_squeezed) progress_images.append(Image.fromarray(img_deprocessed)) print(f"Iteration {i}, Loss: {total_loss.numpy()}") # Get final image final_img = generated.numpy() final_img = np.squeeze(final_img, axis=0) final_img = deprocess(final_img) return Image.fromarray(final_img), progress_images except Exception as e: print(f"Error in neural_style_transfer: {e}") # Return a default error image error_img = Image.new('RGB', (300, 300), color='red') return error_img, [] def style_transfer_interface(content_img, style_img, iterations=50, content_weight=10, style_weight=1000): # Check if images are provided if content_img is None or style_img is None: return None # Perform style transfer result_img, _ = neural_style_transfer( content_img, style_img, iterations=iterations, alpha=content_weight, beta=style_weight ) return result_img # Example images content_path = "content/images/content" style_path = "content/styles/style" example_content_1 = f"{content_path}1.jpg" example_content_2 = f"{content_path}2.jpg" example_content_3 = f"{content_path}3.jpg" example_style_1 = f"{style_path}1.jpg" example_style_2 = f"{style_path}2.jpg" example_style_3 = f"{style_path}3.jpg" examples = [ [example_content_1, example_style_1, 10, 5, 1000], [example_content_2, example_style_2, 20, 10, 1500], [example_content_3, example_style_3, 50, 15, 2000], ] with gr.Blocks(title="Neural Style Transfer") as app: gr.Markdown("# Neural Style Transfer App") gr.Markdown( "Upload a content image and a style image to generate a stylized result") with gr.Row(): with gr.Column(): content_input = gr.Image(label="Content Image", type="pil") style_input = gr.Image(label="Style Image", type="pil") with gr.Row(): iterations_slider = gr.Slider( minimum=10, maximum=100, value=50, step=10, label="Iterations" ) with gr.Row(): content_weight_slider = gr.Slider( minimum=1, maximum=20, value=10, step=1, label="Content Weight" ) style_weight_slider = gr.Slider( minimum=500, maximum=2000, value=1000, step=100, label="Style Weight" ) submit_btn = gr.Button("Generate Stylized Image") with gr.Column(): output_image = gr.Image(label="Stylized Result") gr.Examples( examples=examples, inputs=[content_input, style_input, iterations_slider, content_weight_slider, style_weight_slider], outputs=output_image, fn=style_transfer_interface, cache_examples=False, ) submit_btn.click( fn=style_transfer_interface, inputs=[content_input, style_input, iterations_slider, content_weight_slider, style_weight_slider], outputs=output_image ) # Launch the app if __name__ == "__main__": app.launch(share=True, debug=True)