Initial commits

app.py

@@ -0,0 +1,129 @@
+# Install gradio
+# Import necessary libs
+import tensorflow as tf
+import keras
+import numpy as np
+
+# Load the VGG19 model
+from keras.applications.vgg19 import VGG19
+from keras.applications.vgg19 import preprocess_input, decode_predictions
+
+model = VGG19(include_top=False, weights='imagenet')
+
+# Define style layers and content layers for VGG19
+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)
+
+# Model for extracting style feature
+style_models = [
+    keras.Model(inputs=model.input, outputs=model.get_layer(layer).output) for layer in style_layers
+]
+
+# Model for extracting content feature
+content_model = keras.Model(inputs=model.input, outputs=model.get_layer(content_layers).output, name='Content_model')
+
+# Style models summary
+for m in style_models:
+    print(f'Style model: {m.name}')
+    print(m.summary())
+
+# Content model summary
+print(f'Content model: {content_model.name}')
+content_model.summary()
+
+# Compute content loss
+def content_loss(gen, content_img):
+    return tf.reduce_sum(tf.square(content_img - gen))
+
+# Gram matrix
+def gram_matrix(image):
+    # Reshape image to 2-D array
+    channel = image.shape[-1] # Number of channels
+    a = tf.reshape(image, [-1, channel])
+    # Compute Gram matrix corresponding to the input image
+    dimension = a.shape[0] # Dimension M*N of an MxNxc image
+    return tf.matmul(a, a, transpose_a=True) / tf.cast(dimension, tf.float32)
+
+# Style loss
+def style_loss(gen, style_img):
+    # Compute Gram matrix of generated image and style image
+    gen_gram = gram_matrix(gen)
+    style_gram = gram_matrix(style_img)
+    # Compute style loss
+    return tf.reduce_sum(tf.square(gen_gram-style_gram))/4
+
+# Training parameter
+lr = 10.
+optimizer = keras.optimizers.Adam(lr)
+alpha = 1e-7
+beta = 1e-10
+style_weight = len(style_layers)
+
+# Train step
+@tf.function(reduce_retracing=True)
+def train_step(image, process_content_image, process_style_image):
+    with tf.GradientTape() as tape:
+        # Calculate content loss
+        gen_content_feature = content_model(image)
+        content_feature  = content_model(process_content_image)
+        loss_content = content_loss(gen_content_feature, content_feature)
+        # Calculate style loss
+        loss_style = 0
+        for style_model in style_models:
+            gen_style_feature = style_model(image)
+            style_feature = style_model(process_style_image)
+            loss_style += style_loss(gen_style_feature, style_feature)
+        # Calculate total loss
+        total_loss = alpha*loss_content + beta*loss_style/style_weight
+        # Calculate gradient
+        grad = tape.gradient(total_loss, image)
+        # Apply gradient
+        optimizer.apply_gradients([(grad, image)])
+    # Return loss for visualize
+    return total_loss
+
+# Deprocess image
+def deprocess_image(x):
+    # Util function to convert a tensor into a valid image
+    x = x.reshape((256,256,3))
+    # Remove zero-center by mean pixel
+    x[:, :, 0] += 103.939
+    x[:, :, 1] += 116.779
+    x[:, :, 2] += 123.68
+    # 'BGR'->'RGB'
+    x = x[:, :, ::-1]
+    x = np.clip(x, 0, 255).astype("uint8")
+    return x
+
+def load_image(image, RESO=256):
+  image = tf.image.convert_image_dtype(image, tf.float32) * 255
+  image = tf.image.resize(image,[RESO, RESO],method=tf.image.ResizeMethod.BICUBIC)
+  image = tf.expand_dims(preprocess_input(image), 0)
+  return image
+
+def style_transfer(content_image, style_image, step=100):
+  process_content_image = load_image(content_image)
+  process_style_image = load_image(style_image)
+  generate_image = tf.Variable(process_content_image, trainable=True)
+  for step in range(1, step+1):
+      # Train step
+      loss = train_step(generate_image, process_content_image, process_style_image)
+  return deprocess_image(generate_image[0].numpy())
+
+import gradio as gr
+
+demo = gr.Interface(
+    fn=style_transfer,
+    inputs=[gr.Image(label='Input Image'), gr.Image(label='Style Image'), gr.Slider(10, 200, 50, step=10, label='Step', show_label=True)],
+    outputs=gr.Image(label='Style Transfer Image'),
+)
+
+demo.launch(debug=True, share=True)

requirements.txt

@@ -0,0 +1,2 @@
+tensorflow
+keras