{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n        break\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-08-22T09:57:30.204767Z","iopub.execute_input":"2021-08-22T09:57:30.205167Z","iopub.status.idle":"2021-08-22T09:57:38.969821Z","shell.execute_reply.started":"2021-08-22T09:57:30.205084Z","shell.execute_reply":"2021-08-22T09:57:38.968742Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Neural Style Transfer with AdaIn for a Single Content/Style Image Pair","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow.keras import layers\nfrom tensorflow.keras.models import Model, Sequential\nfrom tensorflow.keras.optimizers import Adam\nimport numpy as np\nfrom tensorflow.keras.applications import VGG19\nfrom tensorflow.keras.applications.vgg19 import preprocess_input\nfrom tensorflow.keras.preprocessing.image import load_img, img_to_array, array_to_img\nimport matplotlib.pyplot as plt\nimport cv2\nimport gzip\nimport zipfile\nimport pandas as pd\nimport skimage.io as sio\nfrom PIL import Image\nfrom io import StringIO, BytesIO","metadata":{"execution":{"iopub.status.busy":"2021-08-22T09:58:37.669165Z","iopub.execute_input":"2021-08-22T09:58:37.669529Z","iopub.status.idle":"2021-08-22T09:58:43.840746Z","shell.execute_reply.started":"2021-08-22T09:58:37.669500Z","shell.execute_reply":"2021-08-22T09:58:43.839754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMG_SIZE = (256, 256)","metadata":{"execution":{"iopub.status.busy":"2021-08-22T09:58:43.844371Z","iopub.execute_input":"2021-08-22T09:58:43.844644Z","iopub.status.idle":"2021-08-22T09:58:43.848501Z","shell.execute_reply.started":"2021-08-22T09:58:43.844617Z","shell.execute_reply":"2021-08-22T09:58:43.847499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vgg19 = VGG19(include_top=False, input_shape=(IMG_SIZE[0], IMG_SIZE[1], 3))","metadata":{"execution":{"iopub.status.busy":"2021-08-22T09:58:43.850559Z","iopub.execute_input":"2021-08-22T09:58:43.850855Z","iopub.status.idle":"2021-08-22T09:58:48.423525Z","shell.execute_reply.started":"2021-08-22T09:58:43.850822Z","shell.execute_reply":"2021-08-22T09:58:48.422766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AdaInLayer(layers.Layer):\n  def __init__(self):\n    super(AdaInLayer, self).__init__()\n    #self.batch_size = batch_size\n  \n  def call(self, inputs):\n    cmaps = inputs[0]\n    smaps = inputs[1]\n\n    sh = tf.shape(cmaps)\n    #cmap_reshape = \n    mean_c = tf.math.reduce_mean(cmaps, axis=[1, 2])\n    mean_c = tf.expand_dims(mean_c, axis=1)#tf.reshape(mean_c, (sh[0], 1, 1, sh[3]))\n    mean_c = tf.expand_dims(mean_c, axis=2)\n    std_c = tf.math.reduce_std(cmaps, axis=[1, 2]) + 0.000001\n    std_c = tf.expand_dims(std_c, axis=1)\n    std_c = tf.expand_dims(std_c, axis=2)\n\n    mean_s = tf.math.reduce_mean(smaps, axis=[1, 2])\n    mean_s = tf.expand_dims(mean_s, axis=1) #tf.reshape(mean_s, (sh[0], 1, 1, sh[3]))\n    mean_s = tf.expand_dims(mean_s, axis=2)\n    std_s = tf.math.reduce_std(smaps, axis=[1, 2])\n    std_s = tf.expand_dims(std_s, axis=1)\n    std_s = tf.expand_dims(std_s, axis=2)\n\n    norm_c = tf.divide(tf.math.subtract(cmaps, mean_c), std_c)#tf.linalg.normalize(cmaps, axis=[1, 2])\n\n    out = tf.multiply(norm_c, std_s) + mean_s\n\n    return out\n  ","metadata":{"execution":{"iopub.status.busy":"2021-08-22T09:58:48.425460Z","iopub.execute_input":"2021-08-22T09:58:48.425733Z","iopub.status.idle":"2021-08-22T09:58:48.435980Z","shell.execute_reply.started":"2021-08-22T09:58:48.425706Z","shell.execute_reply":"2021-08-22T09:58:48.435072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def deprocess_image(x):\n    # Util function to convert a tensor into a valid image\n    x = x.reshape((IMG_SIZE[0], IMG_SIZE[1], 3))\n    # Remove zero-center by mean pixel\n    x[:, :, 0] += 103.939\n    x[:, :, 1] += 116.779\n    x[:, :, 2] += 123.68\n    # 'BGR'->'RGB'\n    x = x[:, :, ::-1]\n    x = np.clip(x, 0, 255).astype(\"uint8\")\n    return x","metadata":{"execution":{"iopub.status.busy":"2021-08-22T09:58:48.437265Z","iopub.execute_input":"2021-08-22T09:58:48.437597Z","iopub.status.idle":"2021-08-22T09:58:48.452112Z","shell.execute_reply.started":"2021-08-22T09:58:48.437568Z","shell.execute_reply":"2021-08-22T09:58:48.451179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Optimized for Single Image Stylization\n\nclass AdaInStyleTransfer_S():\n  def __init__(self, decoder_inp_size,\n               decoder_conv_filters,\n               decoder_num_conv_layers,\n               style_loss_weight,\n               style_weight,\n               total_variation_weight,\n               img_size,\n               batch_size,\n               content_feat_shape,\n               style_feat_shape,\n               single_image=True):\n    self.img_size = img_size\n    self.decoder_inp_size = decoder_inp_size\n    self.decoder_conv_filters = decoder_conv_filters\n    self.decoder_num_conv_layers = decoder_num_conv_layers\n    self.style_loss_weight = style_loss_weight\n    self.style_weight = style_weight\n    self.total_variation_weight = total_variation_weight\n    self.content_feat_shape = content_feat_shape\n    self.style_feat_shape = style_feat_shape\n    self.optimizer = Adam(learning_rate=0.0001)\n    self._create_base_model()\n    self._create_model_for_style_loss()\n    self._create_encoder()\n    self._create_decoder_v2()\n    self.create_core()\n    self.op_images = []\n    self.batch_size = batch_size\n    \n    pass\n\n\n  def _create_base_model(self):\n    self.base_model = VGG19(include_top=False)\n    for layer in self.base_model.layers:\n      layer.trainable = False\n\n  def _create_model_for_style_loss(self):\n    input_img = layers.Input(shape=self.img_size, name='style_latent_input')\n    #x = preprocess_input(input_img)\n\n    style_latent_base = Model(inputs=self.base_model.input, outputs=[self.base_model.get_layer('block1_conv1').output,\n                                                                     self.base_model.get_layer('block2_conv1').output,\n                                                                     self.base_model.get_layer('block3_conv1').output,\n                                                                     self.base_model.get_layer('block4_conv1').output,\n                                                                     ])\n    \n    \n    style_latent = style_latent_base(input_img)\n    self.style_loss_model = Model(inputs=input_img, outputs=style_latent)\n    for layer in self.style_loss_model.layers:\n      layer.trainable = False\n    \n    \n  def _create_encoder(self):\n    input_img = layers.Input(shape=self.img_size, name='encoder_input')\n    \n    #x = preprocess_input(input_img)\n    \n    model = Model(inputs=self.base_model.input, outputs=self.base_model.get_layer('block4_conv1').output)\n    fmaps = model(input_img)\n    \n    self.encoder = Model(inputs=input_img, outputs=fmaps)\n    for layer in self.encoder.layers:\n      layer.trainable = False\n\n\n  def _create_decoder(self):\n    input_tensor = layers.Input(shape=self.decoder_inp_size, name='decoder_input')\n    x = input_tensor\n    paddings = tf.constant([[0,0], [1, 1], [1, 1], [0, 0]])\n    for i in range(len(self.decoder_conv_filters)-1):\n      if (i!=len(self.decoder_conv_filters)-1):\n        x = layers.UpSampling2D(interpolation='nearest')(x)\n      for j in range(self.decoder_num_conv_layers[i]):\n        filters = self.decoder_conv_filters[i]\n        if ((j==self.decoder_num_conv_layers[i]-1) and (i!=len(self.decoder_conv_filters)-1)):\n          filters = self.decoder_conv_filters[i+1]\n        x = layers.Conv2D(filters = filters,\n                          kernel_size = (3, 3),\n                          padding='valid',\n                          )(x)\n        x = tf.pad(x, paddings, mode='REFLECT')\n        \n        if i!=len(self.decoder_conv_filters)-2:\n          x = layers.ReLU()(x)\n          #x = layers.Activation('tanh')(x)\n    #x = layers.Activation('tanh')(x)\n    deprocess_output = self._deprocess_decoder_output(x)\n    self.decoder = Model(inputs = input_tensor, outputs=deprocess_output)\n\n  def _create_decoder_v2(self):\n    input_tensor = layers.Input(shape=self.decoder_inp_size, name='decoder_input')\n    x = input_tensor\n    paddings = tf.constant([[0,0], [1, 1], [1, 1], [0, 0]])\n    for i in range(len(self.decoder_conv_filters)-1):\n      if (i!=len(self.decoder_conv_filters)-1) and (self.decoder_num_conv_layers[i]!=1):\n        x = layers.UpSampling2D(interpolation='nearest')(x)\n      for j in range(self.decoder_num_conv_layers[i]):\n        filters = self.decoder_conv_filters[i]\n        if ((j==self.decoder_num_conv_layers[i]-1) and (i!=len(self.decoder_conv_filters)-1)):\n          filters = self.decoder_conv_filters[i+1]\n        x = layers.Conv2D(filters = filters,\n                          kernel_size = (3, 3),\n                          padding='valid',\n                          )(x)\n        x = tf.pad(x, paddings, mode='REFLECT')\n        \n        if i!=len(self.decoder_conv_filters)-2:\n          x = layers.ReLU()(x)\n          #x = layers.Activation('tanh')(x)\n    #x = layers.Activation('tanh')(x)\n    #x = layers.UpSampling2D(interpolation='nearest')(x)\n    #deprocess_output = self._deprocess_decoder_output(x)\n    self.decoder = Model(inputs = input_tensor, outputs=x)\n  \n  def create_core(self):\n    c_input = layers.Input(shape=(self.content_feat_shape), dtype=tf.float32, name='content_image')\n    s_input = layers.Input(shape=(self.style_feat_shape), dtype=tf.float32, name='style_image')\n\n    # t = AdaInLayer()([c_input, s_input])\n    # weighted_t = (1-self.style_weight)*c_input + self.style_weight*t\n    # out_img = self.decoder(weighted_t)\n    # out_t = self.encoder(out_img)\n\n    # content_output = tf.concat([tf.expand_dims(weighted_t, axis=0), tf.expand_dims(out_t, axis=0)], axis=0)\n    # self.nst_model = Model(inputs=[c_input, s_input], outputs=[content_output, out_img])\n\n    t = AdaInLayer()([c_input, s_input])\n    weighted_t = (1-self.style_weight)*c_input + self.style_weight*t\n    out_img = self.decoder(weighted_t)\n    out_t = self.encoder(out_img)\n\n    content_output = tf.concat([tf.expand_dims(weighted_t, axis=0), tf.expand_dims(out_t, axis=0)], axis=0)\n\n    self.nst_model = Model(inputs=[c_input, s_input], outputs=[content_output, out_img])\n\n\n  def calculate_loss(self, model, X, y, training=True):\n    outputs = model(X)\n    y1 = outputs[0]\n    y2 = outputs[1]\n\n    content_loss = self._custom_content_loss(y1[0], y1[1])\n    style_loss = self.style_loss(y, y2)\n\n    total_loss = content_loss + self.style_loss_weight*style_loss #+ self.total_variation_weight*tf.image.total_variation(y2)\n\n    return total_loss\n  \n  @tf.function\n  def _get_grads(self, model, X, y):\n    with tf.GradientTape() as tape:\n      loss = self.calculate_loss(model, X, y, True)\n    return loss, tape.gradient(loss, model.trainable_variables)\n\n  def custom_training_for_single_image(self, content_image, style_image, epochs=100):\n    #content_image = content_image[:, :, :, ::-1]#tf.reverse(content_image, axis=[-1])\n    #style_image = style_image[:, :, :, ::-1]#tf.reverse(style_image, axis=[-1])\n\n    content_feats = self.encoder(preprocess_input(content_image))\n    style_feats = self.encoder(preprocess_input(style_image))\n    style_feats_for_loss = self.style_loss_model(preprocess_input(style_image))\n    for i in range(1, epochs+1):\n      loss_value, grads = self._get_grads(self.nst_model, [content_feats, style_feats], style_feats_for_loss)\n      self.optimizer.apply_gradients(zip(grads, self.nst_model.trainable_variables))\n      if(((i+1)%20)==0):\n        print(\"Epoch: {}, Total Loss: {:.3f}\".format(i, loss_value.numpy()[0]))\n        ops = self.nst_model([content_feats, style_feats])\n        self.op_images.append(ops)\n        \n\n  def _custom_content_loss(self, y, y_pred):\n    sh = tf.shape(y)\n    content_fmaps = tf.reshape(y, (sh[0], sh[1]*sh[2]*sh[3]))\n    encoded_output = tf.reshape(y_pred, (sh[0], sh[1]*sh[2]*sh[3]))\n    return tf.keras.losses.MSE(content_fmaps, encoded_output)\n\n\n  \n   \n  def style_loss(self, y_true, y_pred):\n    encoded_styles = y_true # self.style_loss_model(y_true)\n    output_style_encoded = self.style_loss_model(y_pred)\n    loss = 0\n    for i in range(len(encoded_styles)):\n      #sh = tf.shape(encoded_styles[i])\n      #tf.print(sh)\n      #sh2 = tf.shape(output_style_encoded[i])\n      mu_s = tf.math.reduce_mean(encoded_styles[i], axis=[1, 2])#tf.math.reduce_mean(tf.reshape(encoded_styles[i], (sh[0], sh[3], sh[1]*sh[2])), axis=2)\n      mu_o = tf.math.reduce_mean(output_style_encoded[i], axis=[1, 2])#tf.math.reduce_mean(tf.reshape(output_style_encoded[i], (sh[0], sh[3], sh[1]*sh[2])), axis=2)\n      mu_diff = tf.keras.losses.MSE(mu_s, mu_o)#tf.reduce_sum(tf.sqrt(tf.square(tf.subtract(mu_o, mu_s))), axis=1)#\n\n      std_s = tf.math.reduce_std(encoded_styles[i], axis=[1, 2])\n      std_o = tf.math.reduce_mean(output_style_encoded[i], axis=[1, 2])#tf.math.reduce_std(output_style_encoded[i], axis=[1, 2])\n      std_diff = tf.keras.losses.MSE(std_s, std_o)#tf.reduce_sum(tf.sqrt(tf.square(tf.subtract(mu_o, mu_s))), axis=1)#\n\n      loss = loss + mu_diff + std_diff\n    return loss\n\n\n  def content_loss(self, y_true, y_pred):\n    adain_c = y_pred[0]\n    out_img = y_pred[1]\n    content_fmaps = adain_c#self.encoder(y_true)\n    encoded_output = out_img#elf.encoder(out_img)\n    sh = tf.shape(content_fmaps)\n    content_fmaps = tf.reshape(content_fmaps, (sh[0], sh[1]*sh[2]*sh[3]))\n    encoded_output = tf.reshape(encoded_output, (sh[0], sh[1]*sh[2]*sh[3]))\n    return tf.keras.losses.MSE(content_fmaps, encoded_output)#tf.reduce_sum(tf.sqrt(tf.square(tf.subtract(encoded_output, content_fmaps))))#\n\n  def variation_loss(self, y_true, y_pred):\n    return tf.reduce_mean(tf.image.total_variation(y_pred))\n\n    \n","metadata":{"execution":{"iopub.status.busy":"2021-08-22T10:04:54.266835Z","iopub.execute_input":"2021-08-22T10:04:54.267199Z","iopub.status.idle":"2021-08-22T10:04:54.485987Z","shell.execute_reply.started":"2021-08-22T10:04:54.267169Z","shell.execute_reply":"2021-08-22T10:04:54.484820Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"content_fnames = [\"photo_jpg/00dcf0f1e3.jpg\", \"photo_jpg/047da870f6.jpg\"]\nstyle_fnames = [\"monet_jpg/6043aadea0.jpg\", \"monet_jpg/3d13fe022e.jpg\"]\npath_prefix = \"/kaggle/input/gan-getting-started/\"\nx_cont = np.zeros(shape=(len(content_fnames), IMG_SIZE[0], IMG_SIZE[1], 3), dtype=np.uint8)\nx_styl = np.zeros(shape=(len(style_fnames), IMG_SIZE[0], IMG_SIZE[1], 3), dtype=np.uint8)\nk=0\nfor i,j in zip(content_fnames, style_fnames):\n  cimg = load_img(path_prefix+i, target_size=IMG_SIZE)\n  simg = load_img(path_prefix+j, target_size=IMG_SIZE)\n  x_cont[k, :, :, :] = img_to_array(cimg)\n  x_styl[k, :, :, :] = img_to_array(simg)\n\n  k += 1","metadata":{"execution":{"iopub.status.busy":"2021-08-22T10:04:55.211768Z","iopub.execute_input":"2021-08-22T10:04:55.212138Z","iopub.status.idle":"2021-08-22T10:04:55.240311Z","shell.execute_reply.started":"2021-08-22T10:04:55.212108Z","shell.execute_reply":"2021-08-22T10:04:55.239246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(x_cont[0].astype(np.uint8))\nplt.title(\"Content Image\");\nplt.figure()\nplt.imshow(x_styl[0].astype(np.uint8))\nplt.title(\"Style Image\");","metadata":{"execution":{"iopub.status.busy":"2021-08-22T10:04:57.392692Z","iopub.execute_input":"2021-08-22T10:04:57.393063Z","iopub.status.idle":"2021-08-22T10:04:57.798837Z","shell.execute_reply.started":"2021-08-22T10:04:57.393029Z","shell.execute_reply":"2021-08-22T10:04:57.797749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"newObj_s = AdaInStyleTransfer_S(decoder_inp_size = (32, 32, 512),\n                            decoder_conv_filters = [256, 256, 128, 64, 3],\n                            decoder_num_conv_layers = [1, 4, 2, 2, 1],\n                            style_loss_weight = 2,\n                            style_weight = 0.8,\n                            total_variation_weight=0.00001,\n                            img_size = (IMG_SIZE[0], IMG_SIZE[1], 3),\n                            batch_size=1,\n                            content_feat_shape=(32, 32, 512),\n                            style_feat_shape=(32, 32, 512))","metadata":{"execution":{"iopub.status.busy":"2021-08-22T10:04:57.800598Z","iopub.execute_input":"2021-08-22T10:04:57.800978Z","iopub.status.idle":"2021-08-22T10:04:58.825476Z","shell.execute_reply.started":"2021-08-22T10:04:57.800939Z","shell.execute_reply":"2021-08-22T10:04:58.824553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"newObj_s.custom_training_for_single_image(x_cont[0:1], x_styl[0:1], 1000) ","metadata":{"execution":{"iopub.status.busy":"2021-08-22T10:17:42.145682Z","iopub.execute_input":"2021-08-22T10:17:42.146070Z","iopub.status.idle":"2021-08-22T10:55:26.717110Z","shell.execute_reply.started":"2021-08-22T10:17:42.146034Z","shell.execute_reply":"2021-08-22T10:55:26.716026Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(newObj_s.op_images))\nplt.imshow(deprocess_image(newObj_s.op_images[54][1].numpy()[0]))","metadata":{"execution":{"iopub.status.busy":"2021-08-22T10:56:11.977248Z","iopub.execute_input":"2021-08-22T10:56:11.977882Z","iopub.status.idle":"2021-08-22T10:56:12.197220Z","shell.execute_reply.started":"2021-08-22T10:56:11.977841Z","shell.execute_reply":"2021-08-22T10:56:12.196124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}