style-transfer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For this implementation I used keras layers over tensorflow in eager mode, theoretical aspects can be [found here](https://cs.stanford.edu/people/jcjohns/papers/eccv16/JohnsonECCV16.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import IPython.display as display\n",
    "import functools\n",
    "mpl.rcParams['axes.grid'] = False\n",
    "mpl.rcParams['figure.figsize'] = (7,7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Defining helper functions for images. Maximum lenth is 512. Each dimension is scaled down to number divisible by spatial downgradings in our network which guarantees same output and input shapes. Network can learn on images of arbitrary sizes, still it is best to provide training data of same(proportional) sizes for faster learning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_img(path_to_img, scale=True, div=4):\n",
    "    img = tf.io.read_file(path_to_img)\n",
    "    img = tf.image.decode_jpeg(img, channels=3)\n",
    "    img = tf.image.convert_image_dtype(img, tf.float32)\n",
    "    shape = tf.cast(tf.shape(img)[:-1], tf.float32)\n",
    "    long_dim = max(shape.numpy())\n",
    "    \n",
    "    if scale:\n",
    "        max_dim = 512\n",
    "    else:\n",
    "        max_dim = 1000\n",
    "\n",
    "    scale = max_dim / long_dim\n",
    "    new_shape = tf.cast(tf.round((shape * scale)/div) * div, tf.int32)\n",
    "    img = tf.image.resize(img, new_shape)\n",
    "    img = np.expand_dims(img, axis=0)\n",
    "    return img\n",
    "\n",
    "def imshow(image, title=None, scale=False):\n",
    "    if len(image.shape) > 3:\n",
    "        image = tf.squeeze(image, axis=0)\n",
    "\n",
    "    show = image\n",
    "    plt.imshow(tf.cast(show, tf.uint8))\n",
    "    if title:\n",
    "        plt.title(title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adding functions to preprocess and deprocess images for our loss model(VGG19 in this case). Basically deprocessing consists of shifting channels means and rearranging the order. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocessed_img(image):\n",
    "    image = image.copy()\n",
    "    image = tf.keras.applications.vgg19.preprocess_input(image)\n",
    "    return image\n",
    "\n",
    "def deprocess_img(image, shifts=[103.939, 116.779, 123.68]):\n",
    "    img = image.copy()\n",
    "    if len(img.shape)==4:\n",
    "        img = np.squeeze(img, 0)\n",
    "        \n",
    "    assert len(img.shape)==3\n",
    "    \n",
    "    for i in tf.range(3):\n",
    "        img[:,:,i] += shifts[i]\n",
    "        \n",
    "    #reverse channels\n",
    "    img = img[:, :, ::-1]\n",
    "    img = np.clip(img, 0, 255).astype('uint8')\n",
    "    return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loss model initializer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vgg19_layers(layers):\n",
    "    vgg19 = tf.keras.applications.VGG19(include_top=False, weights='imagenet')    \n",
    "    outputs = [vgg19.get_layer(layer).output for layer in layers]\n",
    "    \n",
    "    return keras.Model([vgg19.input], outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Names of VGG19 layers which will be used to compute content and style losses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "content_layers = ['block5_conv2'] \n",
    "style_layers = ['block1_conv2',\n",
    "                'block2_conv2',\n",
    "                'block3_conv2', \n",
    "                'block4_conv2', \n",
    "                'block5_conv1']\n",
    "\n",
    "content_layers_len = len(content_layers)\n",
    "style_layers_len = len(style_layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculating gram matrix which will be used for computing style loss (check tf.linalg.einsum for more info)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gram_matrix(image):\n",
    "    gram = tf.linalg.einsum('bijc,bijd->bcd', image, image)\n",
    "    shape = image.get_shape()\n",
    "    i, j = shape[1], shape[2]\n",
    "    \n",
    "    return gram / tf.cast(i * j, tf.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Transform model consist of:\n",
    "  * input convolution layer\n",
    "  * 2 downsampling convolutions\n",
    "  * 5 residual blocks\n",
    "  * 2 upsampling convolutions\n",
    "  * output convolution layer\n",
    "\n",
    "#### In some cases like this downsampling can have positive effect because we can have much larger number of filters without performance decrease."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "DefaultConv2D = partial(keras.layers.Conv2D, kernel_size=3, strides=1, padding='SAME', kernel_initializer='lecun_normal')\n",
    "\n",
    "class ResidualUnit(keras.layers.Layer):\n",
    "    def __init__(self, filters, strides=1, activation='elu', **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.activation = keras.activations.get(activation)\n",
    "        self.main_layers = [\n",
    "            DefaultConv2D(filters, strides=strides),\n",
    "            keras.layers.BatchNormalization(),\n",
    "            self.activation,\n",
    "            DefaultConv2D(filters),\n",
    "            keras.layers.BatchNormalization(),\n",
    "        ]\n",
    "    \n",
    "    \n",
    "    def call(self, inputs):\n",
    "        y = inputs\n",
    "        for layer in self.main_layers:\n",
    "            y = layer(y)\n",
    "            \n",
    "        return self.activation(y + inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StyleContentModel(keras.Model):\n",
    "    def __init__(self, style_layers, content_layers):\n",
    "        super(StyleContentModel, self).__init__()\n",
    "        self.vgg19.trainable = False\n",
    "        self.vgg19 = vgg19_layers(style_layers + content_layers)\n",
    "        self.activation = keras.activations.get('elu')\n",
    "        self.style_layers = style_layers\n",
    "        self.content_layers = content_layers\n",
    "        self.num_style_layers = len(style_layers)\n",
    "        \n",
    "        self.transform_layers = [ # <-- transform part\n",
    "            DefaultConv2D(32, kernel_size=9),\n",
    "            keras.layers.BatchNormalization(),\n",
    "            self.activation,\n",
    "            DefaultConv2D(64, strides=2, kernel_size=2),\n",
    "            keras.layers.BatchNormalization(),\n",
    "            self.activation,\n",
    "            DefaultConv2D(128, strides=2, kernel_size=2),\n",
    "            keras.layers.BatchNormalization(),\n",
    "            self.activation,\n",
    "            #Lecun with ELU may improve normalization \n",
    "            keras.layers.Conv2DTranspose(64, 2, strides=2, kernel_initializer='lecun_normal'), \n",
    "            keras.layers.BatchNormalization(),\n",
    "            self.activation,\n",
    "            keras.layers.Conv2DTranspose(32, 2, strides=2, kernel_initializer='lecun_normal'),\n",
    "            keras.layers.BatchNormalization(),\n",
    "            self.activation,\n",
    "            DefaultConv2D(3, kernel_size=9),\n",
    "            keras.activations.get('tanh')\n",
    "        ]\n",
    "        for i in range(5):\n",
    "            self.transform_layers.insert(9, ResidualUnit(128, 1))\n",
    "        \n",
    "        \n",
    "    def call(self, image, transform=True):\n",
    "        img2 = image\n",
    "        outputs = image\n",
    "        if transform:\n",
    "            for layer in self.transform_layers:\n",
    "                img2 = layer(img2)\n",
    "\n",
    "            img2 *= 255. #increasing back to *255 for vgg19\n",
    "            outputs = img2\n",
    "                \n",
    "        outputs = self.vgg19(outputs) # <-- loss part\n",
    "        style_outputs, content_outputs = (outputs[:self.num_style_layers],\n",
    "                                          outputs[self.num_style_layers:])\n",
    "        \n",
    "        style_outputs = [gram_matrix(style) for style in style_outputs]\n",
    "        \n",
    "        styles = {name:style for name, style in zip(self.style_layers, style_outputs)}\n",
    "        contents = {name:content for name, content in zip(self.content_layers, content_outputs)}\n",
    "        \n",
    "        return {'styles':styles, 'contents':contents, 'image':img2}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### tf.Dataset may not work eagerly in some tf versions, if not, replace it by custom python generator. Channel division by 255 improve transform network precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_paths(images_path, num):\n",
    "    filepaths = []\n",
    "    for i in range(num):\n",
    "        filepaths.append(images_path + '/target' + str(i+1) + '.jpg')\n",
    "    \n",
    "    return filepaths\n",
    "\n",
    "def create_train_dataset(filepaths, repeat=None, shuffle_buffer_size=20, n_parse_threads=5):\n",
    "    dataset = tf.data.Dataset.list_files(filepaths).repeat(repeat)\n",
    "    dataset = dataset.map(lambda x: preprocessed_img(load_img(x))/255., num_parallel_calls=n_parse_threads)\n",
    "    dataset = dataset.map(lambda x: (x, sc_model(x, transform=False)['contents']), num_parallel_calls=n_parse_threads)\n",
    "    dataset = dataset.shuffle(shuffle_buffer_size)\n",
    "    return dataset.prefetch(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def style_loss(a, b, s_weight=1e-5):\n",
    "    n = len(a)\n",
    "    loss = tf.add_n([tf.reduce_mean(tf.square(a[name] - b[name])) for name in a.keys()]) / n\n",
    "    loss *= s_weight\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def content_loss(a, b, c_weight=1e1):\n",
    "    n = len(a)\n",
    "    loss = tf.add_n([tf.reduce_mean(tf.square(a[name] - b[name])) for name in a.keys()]) / n\n",
    "    loss *= c_weight\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Defining artifact loss algorithm which basically calculates differences between nearest image pixels. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def artifact_filter(image):\n",
    "    x_pass = image[:,:,1:,:] - image[:,:,:-1,:]\n",
    "    y_pass = image[:,1:,:,:] - image[:,:-1,:,:]\n",
    "    return x_pass, y_pass\n",
    "\n",
    "def artifacts_loss(image, weight):\n",
    "    image = tf.cast(image, tf.float32)\n",
    "    x_p, y_p = artifact_filter(image)\n",
    "    loss = weight * (tf.reduce_mean(tf.square(x_p)) + tf.reduce_mean(tf.square(y_p)))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style_path = '/Users/user/Documents/style-transfer/images/styles/style1.jpg'\n",
    "images_path = '/Users/user/Documents/style-transfer/images/targets'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc_model = StyleContentModel(style_layers, content_layers)\n",
    "\n",
    "style_image = preprocessed_img(load_img(style_path, scale=False))\n",
    "\n",
    "style_learn = sc_model(style_image, transform=False)['styles']\n",
    "train_dataset = create_train_dataset(load_paths(images_path, 10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Found Nadam working best for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = keras.optimizers.Nadam(0.00001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function()\n",
    "def train_step(image, content_learn):\n",
    "    with tf.GradientTape() as tape:\n",
    "        outputs = sc_model(image)\n",
    "        style_l = style_loss(outputs['styles'], style_learn)\n",
    "        content_l = content_loss(outputs['contents'], content_learn)\n",
    "        artifact_l = artifacts_loss(outputs['image'], 2)\n",
    "        loss = style_l + content_l + artifact_l\n",
    "    \n",
    "    img = outputs['image']    \n",
    "    vars = sc_model.trainable_variables\n",
    "    gradients = tape.gradient(loss, vars)\n",
    "    opt = optimizer.apply_gradients([(g, v) for g, v in zip(gradients, vars)])\n",
    "    return style_l, content_l, artifact_l, img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 100000\n",
    "steps_per_epoch = 40\n",
    "\n",
    "step = 0\n",
    "for n in range(epochs):\n",
    "    for c_image, c_learn in train_dataset.take(steps_per_epoch):\n",
    "        step += 1\n",
    "        s_l, c_l, a_l, img = train_step(c_image, c_learn)\n",
    "    \n",
    "    tmp = deprocess_img(img.numpy())\n",
    "    imshow(tmp)    \n",
    "    display.clear_output(wait=True)\n",
    "    plt.title(\"Train step: {} Loss_s: {} Loss_c: {} Loss_a: {}\".format(step, s_l, c_l, a_l))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}