FaceSwap_GAN_v2.1_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_Wnn1c4w_6ph"
   },
   "source": [
    "<a id='1'></a>\n",
    "# 1. Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 35,
     "output_extras": [
      {
       "item_id": 1
      }
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 8857,
     "status": "ok",
     "timestamp": 1520927478188,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -480
    },
    "id": "nchj1aCO_6pq",
    "outputId": "beb33124-6983-482d-cc6a-bfc53f75e54f"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "from keras.models import Sequential, Model\n",
    "from keras.layers import *\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras.activations import relu\n",
    "from keras.initializers import RandomNormal\n",
    "from keras.applications import *\n",
    "import keras.backend as K\n",
    "from tensorflow.contrib.distributions import Beta\n",
    "import tensorflow as tf\n",
    "from keras.optimizers import Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "8DbMfY5U_6qX"
   },
   "outputs": [],
   "source": [
    "from image_augmentation import random_transform\n",
    "from image_augmentation import random_warp\n",
    "from utils import get_image_paths, load_images, stack_images\n",
    "from pixel_shuffler import PixelShuffler\n",
    "from instance_normalization import InstanceNormalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "VaDCPEmn_6qq"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import cv2\n",
    "import glob\n",
    "from random import randint, shuffle\n",
    "from IPython.display import clear_output\n",
    "from IPython.display import display\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "e5B5wMqtXbiF"
   },
   "outputs": [],
   "source": [
    "from umeyama import umeyama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 249,
     "output_extras": [
      {
       "item_id": 1
      }
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1523,
     "status": "ok",
     "timestamp": 1520927484089,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -480
    },
    "id": "GCxklH1C_6q9",
    "outputId": "bb9c1f27-62ae-4539-aa5a-981d7c222456"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[name: \"/device:CPU:0\"\n",
       " device_type: \"CPU\"\n",
       " memory_limit: 268435456\n",
       " locality {\n",
       " }\n",
       " incarnation: 10970346582023774686, name: \"/device:GPU:0\"\n",
       " device_type: \"GPU\"\n",
       " memory_limit: 11297803469\n",
       " locality {\n",
       "   bus_id: 1\n",
       " }\n",
       " incarnation: 4030019729566350385\n",
       " physical_device_desc: \"device: 0, name: Tesla K80, pci bus id: 0000:00:04.0, compute capability: 3.7\"]"
      ]
     },
     "execution_count": 5,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "\n",
    "device_lib.list_local_devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oO5lqbfJ_6ta"
   },
   "source": [
    "<a id='3'></a>\n",
    "# 3. Import VGGFace\n",
    "(Skip this part if you don't want to apply perceptual loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MnpodW5t_6th"
   },
   "source": [
    "If you got error ```_obtain_input_shape(...)``` error, this is because your keras version is older than vggface requirement. \n",
    "\n",
    "Modify ```_obtain_input_shape(...)``` in ```keras_vggface/models.py``` will solve the problem. The following is what worked for me:\n",
    "\n",
    "```python\n",
    "input_shape = _obtain_input_shape(input_shape,\n",
    "                                  default_size=224,\n",
    "                                  min_size=197,\n",
    "                                  data_format=K.image_data_format(),\n",
    "                                  include_top=include_top)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "mLRGBnk2_6tk"
   },
   "outputs": [],
   "source": [
    "from keras_vggface.vggface import VGGFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_network_PL = 'vgg16' # 'resnet50' or 'vgg16'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 53,
     "output_extras": [
      {
       "item_id": 4
      }
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 2095,
     "status": "ok",
     "timestamp": 1520927487465,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -480
    },
    "id": "izA4jM-w_6tw",
    "outputId": "4c6a1170-61ba-4b86-9170-fe0895ebc0a9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://github.com/rcmalli/keras-vggface/releases/download/v2.0/rcmalli_vggface_tf_notop_vgg16.h5\n",
      "58916864/58909280 [==============================] - 1s 0us/step\n"
     ]
    }
   ],
   "source": [
    "vggface = VGGFace(include_top=False, model=base_network_PL, input_shape=(224, 224, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "s67eB2mu_6u9"
   },
   "source": [
    "<a id='4'></a>\n",
    "# 4. Training Phase Configuration\n",
    "\n",
    "Default training data directories: `./faceA/` and `./faceB/`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "OTtVZOxY_6vA"
   },
   "outputs": [],
   "source": [
    "K.set_learning_phase(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "NMxZMuPO_6vL"
   },
   "outputs": [],
   "source": [
    "channel_axis=-1\n",
    "channel_first = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "F60PN39B_6vd"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Base model:\n",
    "1. \"XGAN\": standard autoencoder as generator + adversarial losses + encoding losses \n",
    "2. \"VAE-GAN\": variational autoencoder as generator + adversarial losses\n",
    "3. \"GAN\": standard autoencoder as generator + adversarial losses\n",
    "\"\"\"\n",
    "base_model = \"GAN\"\n",
    "\n",
    "\"\"\"\n",
    "Decoder type:\n",
    "1. Shared decoder\n",
    "2. Separate decoders\n",
    "\"\"\"\n",
    "decoder_type = 1 if base_model == \"VAE-GAN\" else 2\n",
    "\n",
    "IMAGE_SHAPE = (64, 64, 3)\n",
    "nc_in = 3 # number of input channels of generators\n",
    "nc_D_inp = 6 # number of input channels of discriminators\n",
    "nz = 2048 if decoder_type == 1 else 1024 # nz=2048 according to VAE-GAN paper https://arxiv.org/abs/1512.09300\n",
    "\"\"\"\n",
    "Notes for latent dimension:\n",
    "1. BicycleGAN use nz=8 for noise\n",
    "2. From UNIT implementation: https://github.com/mingyuliutw/UNIT/blob/master/src/trainers/cocogan_nets_da.py\n",
    "      self.g_vae = GaussianVAE2D(ch * 8, ch * 8, kernel_size=1, stride=1),\n",
    "      where ch=32 by default. (notice that its 2D latent space, i.e., using Conv2D instead of Dense)\n",
    "\"\"\"\n",
    "\n",
    "use_perceptual_loss = False\n",
    "use_lsgan = True\n",
    "use_mixup = True # mixup paper: https://arxiv.org/abs/1710.09412\n",
    "mixup_alpha = 0.2\n",
    "\n",
    "batchSize = 8\n",
    "lrD = 1e-4 # Discriminator learning rate\n",
    "lrG = 1e-4 # Generator learning rate\n",
    "\n",
    "# Path of training images\n",
    "img_dirA = './faceA/*.*'\n",
    "img_dirB = './faceB/*.*'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vObbstXt_6vl"
   },
   "source": [
    "<a id='5'></a>\n",
    "# 5. GAN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "ozFotFix_6vn"
   },
   "outputs": [],
   "source": [
    "conv_init = RandomNormal(0, 0.02)\n",
    "\n",
    "def conv_block(input_tensor, f, k=3, strides=2, dilation_rate=1, use_inst_norm=True):\n",
    "    x = input_tensor\n",
    "    x = Conv2D(f, kernel_size=k, strides=strides, dilation_rate=dilation_rate, kernel_initializer=conv_init, use_bias=True, padding=\"same\")(x)\n",
    "    if use_inst_norm:\n",
    "        x = InstanceNormalization()(x)\n",
    "    x = LeakyReLU(alpha=0.1)(x)\n",
    "    return x\n",
    "\n",
    "def conv_block_d(input_tensor, f, use_inst_norm=True):\n",
    "    x = input_tensor\n",
    "    x = Conv2D(f, kernel_size=4, strides=2, kernel_initializer=conv_init, use_bias=True, padding=\"same\")(x)\n",
    "    if use_inst_norm:\n",
    "        x = InstanceNormalization()(x)\n",
    "    x = LeakyReLU(alpha=0.2)(x)\n",
    "    return x\n",
    "\n",
    "def res_block(input_tensor, f):\n",
    "    x = input_tensor\n",
    "    x = Conv2D(f, kernel_size=5, kernel_initializer=conv_init, use_bias=True, padding=\"same\")(x)\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = LeakyReLU(alpha=0.2)(x)\n",
    "    x = Conv2D(f, kernel_size=5, kernel_initializer=conv_init, use_bias=True, padding=\"same\")(x)\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = add([x, input_tensor])\n",
    "    return x\n",
    "\n",
    "def upscale_ps(filters, dilation_rate=1, use_norm=True):\n",
    "    def block(x):\n",
    "        x = Conv2D(filters*4, kernel_size=3, dilation_rate=dilation_rate, use_bias=True, kernel_initializer=RandomNormal(0, 0.02), padding='same')(x)\n",
    "        x = InstanceNormalization()(x)\n",
    "        x = LeakyReLU(0.1)(x)\n",
    "        x = PixelShuffler()(x)\n",
    "        return x\n",
    "    return block\n",
    "\n",
    "def Discriminator(nc_in, input_size=64):\n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    #x = GaussianNoise(0.05)(inp)\n",
    "    x = conv_block_d(inp, 64, False)\n",
    "    x = conv_block_d(x, 128)\n",
    "    x = conv_block_d(x, 256)\n",
    "    out = Conv2D(1, kernel_size=4, kernel_initializer=conv_init, use_bias=False, padding=\"same\")(x)   \n",
    "    return Model(inputs=[inp], outputs=out)\n",
    "\n",
    "def Discriminator2(nc_in, input_size=64):\n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    x = Conv2D(32, kernel_size=5, kernel_initializer=conv_init, use_bias=False, padding=\"same\")(inp)\n",
    "    x = LeakyReLU(0.1)(x)\n",
    "    #x = GaussianNoise(0.05)(inp)\n",
    "    x = conv_block_d(x, 64, False)\n",
    "    x = conv_block_d(x, 128)\n",
    "    out0 = x\n",
    "    x = conv_block_d(x, 256)\n",
    "    out1 = x\n",
    "    x = conv_block_d(x, 512)\n",
    "    out2 = x\n",
    "    out3 = Conv2D(1, kernel_size=3, kernel_initializer=conv_init, use_bias=False, padding=\"same\")(x)   \n",
    "    return Model(inputs=[inp], outputs=[out0, out1, out2, out3]) \n",
    "\n",
    "def Discriminator_code():\n",
    "    inp = Input(shape=(1024, ))\n",
    "    x = Dense(256)(inp)\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = Activation('relu')(x)\n",
    "    x = Dense(128)(x)\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = Activation('relu')(x)\n",
    "    x = Dense(32)(x)\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = Activation('relu')(x)\n",
    "    out = Dense(1, activation='sigmoid')(x)   \n",
    "    return Model(inputs=[inp], outputs=out)\n",
    "\n",
    "def Encoder(nc_in=3, input_size=64):\n",
    "    def l2_norm(x):\n",
    "        epsilon = 1e-12\n",
    "        x_norm = K.sqrt(K.sum(K.square(x)))\n",
    "        return x / (x_norm + epsilon)  \n",
    "  \n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    x = Conv2D(64, kernel_size=5, kernel_initializer=conv_init, use_bias=False, padding=\"same\")(inp)\n",
    "    x = LeakyReLU(0.1)(x)\n",
    "    x = conv_block(x,128)\n",
    "    x = conv_block(x,256)\n",
    "    x = conv_block(x,512)\n",
    "    x = conv_block(x,1024)\n",
    "    x = Dense(1024)(Flatten()(x))\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = LeakyReLU(0.1)(x)\n",
    "    x = Dense(1024)(x)\n",
    "    x = Lambda(l2_norm)(x)\n",
    "    code = x\n",
    "    x = Dense(4*4*1024)(x)\n",
    "    x = Reshape((4, 4, 1024))(x)\n",
    "    out = upscale_ps(512, dilation_rate=1)(x)\n",
    "    return Model(inputs=inp, outputs=[out, code])\n",
    "\n",
    "def Encoder2(nc_in=3, input_size=64):\n",
    "    def l2_norm(x):\n",
    "        epsilon = 1e-12\n",
    "        x_norm = K.sqrt(K.sum(K.square(x)))\n",
    "        return x / (x_norm + epsilon)  \n",
    "  \n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    x = Conv2D(64, kernel_size=5, kernel_initializer=conv_init, use_bias=False, padding=\"same\")(inp)\n",
    "    x = LeakyReLU(0.1)(x)\n",
    "    x = conv_block(x,128)\n",
    "    x = conv_block(x,256)\n",
    "    x = conv_block(x,512)\n",
    "    x = conv_block(x,1024)\n",
    "    #x = Flatten()(x)\n",
    "    out = x\n",
    "    return Model(inputs=inp, outputs=out)\n",
    "\n",
    "def Encoder_vae(nc_in=3, input_size=64):  \n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    x = Conv2D(64, kernel_size=5, kernel_initializer=conv_init, use_bias=False, padding=\"same\")(inp)\n",
    "    x = LeakyReLU(0.1)(x)\n",
    "    x = conv_block(x,128)\n",
    "    x = conv_block(x,256)\n",
    "    x = conv_block(x,512)\n",
    "    x = conv_block(x,1024)\n",
    "    out = x\n",
    "    return Model(inputs=inp, outputs=out)\n",
    "\n",
    "def Decoder_ps(nc_in=512, input_size=8):\n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    code = Input(shape=(1024,))\n",
    "    \n",
    "    x_code = Dense(4*4*512)(code)\n",
    "    x_code = Reshape((4, 4, 512))(x_code)\n",
    "    x_code = upscale_ps(256)(x_code)\n",
    "    x_code = upscale_ps(128)(x_code)\n",
    "    \n",
    "    x = inp\n",
    "    x = upscale_ps(256)(x)\n",
    "    x = concatenate([x, x_code])\n",
    "    x = upscale_ps(128)(x)\n",
    "    x = upscale_ps(64)(x)\n",
    "    x = res_block(x, 64)\n",
    "    x = res_block(x, 64)\n",
    "    alpha = Conv2D(1, kernel_size=5, padding='same', activation=\"sigmoid\")(x)\n",
    "    bgr = Conv2D(3, kernel_size=5, padding='same', activation=\"tanh\")(x)\n",
    "    out = concatenate([alpha, bgr])\n",
    "    return Model([inp, code], [out, code])  \n",
    "\n",
    "def Decoder_code(nc_in=1024, input_size=4):        \n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    \n",
    "    x = inp\n",
    "    x = Flatten()(x)\n",
    "    x_code = Dense(1024)(x)\n",
    "    x_code = InstanceNormalization()(x_code)\n",
    "    x_code = LeakyReLU(0.1)(x_code)\n",
    "    x_code = Dense(1024)(x_code)\n",
    "    x_code = Dense(4*4*512)(x_code)\n",
    "    x_code = Reshape((4, 4, 512))(x_code)\n",
    "    out = x_code\n",
    "    return Model(inp, out)   \n",
    "  \n",
    "def Decoder_common(nc_in=1024, input_size=4):\n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    \n",
    "    x = inp\n",
    "    x = upscale_ps(512)(x)\n",
    "    x = upscale_ps(256)(x)\n",
    "    x = upscale_ps(128)(x)\n",
    "    x = upscale_ps(64)(x)\n",
    "    x = res_block(x, 64)\n",
    "    x = res_block(x, 64)\n",
    "    alpha = Conv2D(1, kernel_size=5, padding='same', activation=\"sigmoid\")(x)\n",
    "    bgr = Conv2D(3, kernel_size=5, padding='same', activation=\"tanh\")(x)\n",
    "    out = concatenate([alpha, bgr])\n",
    "    return Model(inp, out) \n",
    "\n",
    "def Decoder_code_vae(nc_in=1024, input_size=4, z_dim=32):   \n",
    "    \"\"\"VAE implementation reference: https://github.com/EmilienDupont/vae-concrete \"\"\"\n",
    "    def l2_norm(x):\n",
    "        epsilon = 1e-12\n",
    "        x_norm = K.sqrt(K.sum(K.square(x)))\n",
    "        return x / (x_norm + epsilon)  \n",
    "    def sampling_normal(args):\n",
    "        z_mean, z_log_var = args\n",
    "        epsilon = K.random_normal(shape=(z_dim,), mean=0., stddev=1.)\n",
    "        return z_mean + K.exp(z_log_var / 2) * epsilon\n",
    "      \n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    \n",
    "    x = inp\n",
    "    x = Flatten()(x)\n",
    "    x_code = Dense(1024)(x)\n",
    "    x_code = InstanceNormalization()(x_code)\n",
    "    x_code = LeakyReLU(0.1)(x_code)\n",
    "    x_code = Dense(1024)(x_code)    \n",
    "    x_code = InstanceNormalization()(x_code)\n",
    "    x_code = LeakyReLU(0.1)(x_code)    \n",
    "    z_mean = Dense(z_dim)(x_code) # latent dim = 32\n",
    "    z_log_var = Dense(z_dim)(x_code)    \n",
    "    encoding = Lambda(sampling_normal)([z_mean, z_log_var])    \n",
    "    out = encoding\n",
    "    return Model(inp, [out, z_mean, z_log_var])  \n",
    "\n",
    "def Decoder_code_soft_gate(nc_in=1024, input_size=4, nc_out=512):   \n",
    "    def l2_norm(x):\n",
    "        epsilon = 1e-12\n",
    "        x_norm = K.sqrt(K.sum(K.square(x)))\n",
    "        return x / (x_norm + epsilon)  \n",
    "      \n",
    "    inp = Input(shape=(input_size, input_size, nc_in))\n",
    "    \n",
    "    x = inp\n",
    "    x = Flatten()(x)\n",
    "    x_code = Dense(1024)(x)\n",
    "    x_code = InstanceNormalization()(x_code)\n",
    "    x_code = LeakyReLU(0.1)(x_code)\n",
    "    x_code = Dense(1024)(x_code)    \n",
    "    soft_gate = InstanceNormalization()(x_code)\n",
    "    soft_gate = LeakyReLU(0.1)(soft_gate)\n",
    "    soft_gate = Dense(1024)(soft_gate)  \n",
    "    soft_gate = InstanceNormalization()(soft_gate)\n",
    "    soft_gate = LeakyReLU(0.1)(soft_gate)\n",
    "    soft_gate = Dense(nc_out, activation='sigmoid')(soft_gate)    \n",
    "    x_code = Lambda(l2_norm)(x_code)\n",
    "    x_code = Dense(4*4*nc_out)(x_code)\n",
    "    x_code = Reshape((4, 4, nc_out))(x_code)    \n",
    "    out = multiply([soft_gate, x_code])    \n",
    "    return Model(inp, out)  \n",
    "  \n",
    "def Decoder_common_vae(nc_in=64, input_size=4):\n",
    "    inp = Input(shape=(nc_in, ))\n",
    "    \n",
    "    x = inp    \n",
    "    x = Dense(4*4*1024)(x)\n",
    "    x = InstanceNormalization()(x)\n",
    "    x = LeakyReLU(0.1)(x)\n",
    "    x = Reshape((4, 4, 1024))(x)\n",
    "    x = upscale_ps(512)(x)\n",
    "    x = upscale_ps(256)(x)\n",
    "    x = upscale_ps(128)(x)\n",
    "    x = upscale_ps(64)(x)\n",
    "    x = res_block(x, 64)\n",
    "    x = res_block(x, 64)\n",
    "    alpha = Conv2D(1, kernel_size=5, padding='same', activation=\"sigmoid\")(x)\n",
    "    bgr = Conv2D(3, kernel_size=5, padding='same', activation=\"tanh\")(x)\n",
    "    out = concatenate([alpha, bgr])\n",
    "    return Model(inp, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "GgOUJTFkCMMP"
   },
   "outputs": [],
   "source": [
    "def get_generator_vae(x, y, encoder, decoder_code1, decoder_code2, decoder_common):\n",
    "    enc_x = encoder(x)\n",
    "    enc_y = encoder(y)\n",
    "    latent_z1, z1_mean, z1_log_var = decoder_code1(enc_x)\n",
    "    latent_z2, z2_mean, z2_log_var = decoder_code2(enc_y)\n",
    "    concat_latent_zs = concatenate([latent_z1, latent_z2])\n",
    "    out = decoder_common(concat_latent_zs)\n",
    "    netG = Model([x, y], [out, z1_mean, z1_log_var, z2_mean, z2_log_var])\n",
    "    return netG  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "9iKbuaj8MjmT"
   },
   "outputs": [],
   "source": [
    "if base_model == \"XGAN\":    \n",
    "    encoder = Encoder()\n",
    "    decoder_A = Decoder_ps()\n",
    "    decoder_B = Decoder_ps()\n",
    "    x = Input(shape=IMAGE_SHAPE)\n",
    "    netGA = Model(x, decoder_A(encoder(x)))\n",
    "    netGB = Model(x, decoder_B(encoder(x)))\n",
    "    \n",
    "elif base_model == \"VAE-GAN\":\n",
    "    encoder = Encoder_vae()\n",
    "    decoder_A = Decoder_code_vae(z_dim=nz//2)\n",
    "    decoder_B = Decoder_code_vae(z_dim=nz//2)\n",
    "    # Since VAE has a more representative latent space, \n",
    "    # we introduce bottlneck to AnB decoder in hope of it will encode the most important facial attributes.\n",
    "    # And we can perhaps manipulate these attributes in the future.\n",
    "    # Result: Still under experiment (with beta=0.01).\n",
    "    decoder_AnB = Decoder_code_vae(z_dim=nz//2) # z_dim=nz//64?\n",
    "    \n",
    "    x = Input(shape=IMAGE_SHAPE)\n",
    "    y = Input(shape=IMAGE_SHAPE)\n",
    "    \n",
    "    if decoder_type == 1:\n",
    "        decoder_common = Decoder_common_vae(nc_in=(nz//2 + nz//2))\n",
    "        netGA = get_generator_vae(x, y, encoder, decoder_AnB, decoder_A, decoder_common)\n",
    "        netGB = get_generator_vae(x, y, encoder, decoder_AnB, decoder_B, decoder_common)\n",
    "    elif decoder_type == 2:\n",
    "        decoder_A2 = Decoder_common_vae(nc_in=nz)\n",
    "        decoder_B2 = Decoder_common_vae(nc_in=nz)\n",
    "        netGA = get_generator_vae(x, y, encoder, decoder_AnB, decoder_A, decoder_A2)\n",
    "        netGB = get_generator_vae(x, y, encoder, decoder_AnB, decoder_B, decoder_B2)\n",
    "        \n",
    "elif base_model == \"GAN\":\n",
    "    encoder = Encoder2()\n",
    "    decoder_A = Decoder_code()\n",
    "    decoder_B = Decoder_code()\n",
    "    decoder_AnB = Decoder_code()\n",
    "\n",
    "    x = Input(shape=IMAGE_SHAPE)\n",
    "    y = Input(shape=IMAGE_SHAPE)\n",
    "\n",
    "    if decoder_type == 1:\n",
    "        decoder_common = Decoder_common()\n",
    "        netGA = Model([x, y], decoder_common(concatenate([decoder_AnB(encoder(x)), decoder_A(encoder(y))])))\n",
    "        netGB = Model([x, y], decoder_common(concatenate([decoder_AnB(encoder(x)), decoder_B(encoder(y))])))\n",
    "    elif decoder_type == 2:\n",
    "        decoder_A2 = Decoder_common()\n",
    "        decoder_B2 = Decoder_common()\n",
    "        netGA = Model([x, y], decoder_A2(concatenate([decoder_AnB(encoder(x)), decoder_A(encoder(y))])))\n",
    "        netGB = Model([x, y], decoder_B2(concatenate([decoder_AnB(encoder(x)), decoder_B(encoder(y))])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "FNDEAQqM_6wT"
   },
   "outputs": [],
   "source": [
    "netDA = Discriminator(nc_D_inp)\n",
    "netDB = Discriminator(nc_D_inp)\n",
    "netDA2 = Discriminator2(nc_D_inp//2)\n",
    "netDB2 = Discriminator2(nc_D_inp//2)\n",
    "\n",
    "if base_model == \"XGAN\":\n",
    "    netD_code = Discriminator_code()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "V5gCF2CP_6wv"
   },
   "source": [
    "<a id='6'></a>\n",
    "# 6. Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 34,
     "output_extras": [
      {}
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 5938,
     "status": "ok",
     "timestamp": 1519052064453,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -540
    },
    "id": "ODmintRS_6wx",
    "outputId": "41d05bf3-5a3e-4ebc-970e-1855c31cc258"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model loaded.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    encoder.load_weights(\"models/encoder.h5\")\n",
    "    decoder_A.load_weights(\"models/decoder_A.h5\")\n",
    "    decoder_B.load_weights(\"models/decoder_B.h5\")\n",
    "    if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "        decoder_AnB.load_weights(\"models/decoder_AnB.h5\")\n",
    "        if decoder_type == 1:\n",
    "            decoder_common.load_weights(\"models/decoder_common.h5\")\n",
    "        elif decoder_type == 2:\n",
    "            decoder_A2.load_weights(\"models/decoder_A2.h5\")\n",
    "            decoder_B2.load_weights(\"models/decoder_B2.h5\")\n",
    "    netDA.load_weights(\"models/netDA.h5\") \n",
    "    netDB.load_weights(\"models/netDB.h5\") \n",
    "    netDA2.load_weights(\"models/netDA2.h5\") \n",
    "    netDB2.load_weights(\"models/netDB2.h5\") \n",
    "    if base_model == \"XGAN\":\n",
    "        netD_code.load_weights(\"models/netD_code.h5\")\n",
    "    print (\"model loaded.\")\n",
    "except:\n",
    "    print (\"Failed loading. Something goes wrong with weights files.\")\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4sRxlmly_6xL"
   },
   "source": [
    "<a id='7'></a>\n",
    "# 7. Inputs/Outputs Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Nxm9RkoG_6xR"
   },
   "outputs": [],
   "source": [
    "if base_model == \"XGAN\":\n",
    "    def cycle_variables(netG):\n",
    "        distorted_input = netG.inputs[0]\n",
    "        fake_output = netG.outputs[0]\n",
    "        code = netG.outputs[1]\n",
    "        \n",
    "        alpha = Lambda(lambda x: x[:,:,:, :1])(fake_output)\n",
    "        bgr = Lambda(lambda x: x[:,:,:, 1:])(fake_output)\n",
    "        masked_fake_output = alpha * bgr + (1-alpha) * distorted_input \n",
    "\n",
    "        fn_generate = K.function([distorted_input], [masked_fake_output])\n",
    "        fn_mask = K.function([distorted_input], [concatenate([alpha, alpha, alpha])])\n",
    "        fn_abgr = K.function([distorted_input], [concatenate([alpha, bgr])])\n",
    "        fn_bgr = K.function([distorted_input], [bgr])\n",
    "        return distorted_input, fake_output, code, alpha, fn_generate, fn_mask, fn_abgr, fn_bgr\n",
    "        \n",
    "elif base_model == \"VAE-GAN\":\n",
    "    def cycle_variables(netG, decoder_common=None):      \n",
    "        distorted_input_x = netG.inputs[0]\n",
    "        distorted_input_y = netG.inputs[1]\n",
    "        fake_output = netG.outputs[0]\n",
    "        z1_mean, z1_log_var = netG.outputs[1], netG.outputs[2]\n",
    "        z2_mean, z2_log_var = netG.outputs[3], netG.outputs[4]\n",
    "        zs = [z1_mean, z1_log_var, z2_mean, z2_log_var]\n",
    "        \n",
    "        concat_latent_z_means = concatenate([z1_mean, z2_mean]) # latent code w/o var. noise\n",
    "        fake_output_inference = decoder_common(concat_latent_z_means)\n",
    "        \n",
    "        alpha = Lambda(lambda x: x[:,:,:, :1])(fake_output)\n",
    "        alpha_inference = Lambda(lambda x: x[:,:,:, :1])(fake_output_inference)\n",
    "        bgr_inference = Lambda(lambda x: x[:,:,:, 1:])(fake_output_inference)        \n",
    "        masked_fake_output = alpha_inference * bgr_inference + (1-alpha_inference) * distorted_input_x \n",
    "    \n",
    "        fn_generate = K.function([distorted_input_x, distorted_input_y], [masked_fake_output])\n",
    "        fn_mask = K.function([distorted_input_x, distorted_input_y], [concatenate([alpha_inference, alpha_inference, alpha_inference])])\n",
    "        fn_bgr = K.function([distorted_input_x, distorted_input_y], [bgr_inference])\n",
    "        return distorted_input_x, distorted_input_y, fake_output, alpha, zs, fn_generate, fn_mask, fn_bgr\n",
    "        \n",
    "elif base_model == \"GAN\":\n",
    "    def cycle_variables(netG):\n",
    "        distorted_input_x = netG.inputs[0]\n",
    "        distorted_input_y = netG.inputs[1]\n",
    "        fake_output = netG.outputs[0]\n",
    "        \n",
    "        alpha = Lambda(lambda x: x[:,:,:, :1])(fake_output)\n",
    "        bgr = Lambda(lambda x: x[:,:,:, 1:])(fake_output)\n",
    "        masked_fake_output = alpha * bgr + (1-alpha) * distorted_input_x \n",
    "\n",
    "        fn_generate = K.function([distorted_input_x, distorted_input_y], [masked_fake_output])\n",
    "        fn_mask = K.function([distorted_input_x, distorted_input_y], [concatenate([alpha, alpha, alpha])])\n",
    "        fn_abgr = K.function([distorted_input_x, distorted_input_y], [concatenate([alpha, bgr])])\n",
    "        fn_bgr = K.function([distorted_input_x, distorted_input_y], [bgr])\n",
    "        return distorted_input_x, distorted_input_y, fake_output, alpha, fn_generate, fn_mask, fn_abgr, fn_bgr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "zR6S1RLyPRp5"
   },
   "outputs": [],
   "source": [
    "if base_model == \"XGAN\":\n",
    "    distorted_A, fake_A, code_A, mask_A, path_A, path_mask_A, path_abgr_A, path_bgr_A = cycle_variables(netGA)\n",
    "    distorted_B, fake_B, code_B, mask_B, path_B, path_mask_B, path_abgr_B, path_bgr_B = cycle_variables(netGB)\n",
    "    real_A = Input(shape=IMAGE_SHAPE)\n",
    "    real_B = Input(shape=IMAGE_SHAPE)\n",
    "    \n",
    "elif base_model == \"VAE-GAN\":\n",
    "    if decoder_type == 1:\n",
    "        distorted_x_A, distorted_y_A, fake_A, mask_A, zs_A, path_A, path_mask_A, path_bgr_A = cycle_variables(netGA, decoder_common)\n",
    "        distorted_x_B, distorted_y_B, fake_B, mask_B, zs_B, path_B, path_mask_B, path_bgr_B = cycle_variables(netGB, decoder_common)\n",
    "    elif decoder_type == 2:\n",
    "        distorted_x_A, distorted_y_A, fake_A, mask_A, zs_A, path_A, path_mask_A, path_bgr_A = cycle_variables(netGA, decoder_A2)\n",
    "        distorted_x_B, distorted_y_B, fake_B, mask_B, zs_B, path_B, path_mask_B, path_bgr_B = cycle_variables(netGB, decoder_B2)\n",
    "    real_A = Input(shape=IMAGE_SHAPE)\n",
    "    real_B = Input(shape=IMAGE_SHAPE)\n",
    "    \n",
    "elif base_model == \"GAN\":\n",
    "    distorted_x_A, distorted_y_A, fake_A, mask_A, path_A, path_mask_A, path_abgr_A, path_bgr_A = cycle_variables(netGA)\n",
    "    distorted_x_B, distorted_y_B, fake_B, mask_B, path_B, path_mask_B, path_abgr_B, path_bgr_B = cycle_variables(netGB)\n",
    "    real_A = Input(shape=IMAGE_SHAPE)\n",
    "    real_B = Input(shape=IMAGE_SHAPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nTTUTdhN_6xu"
   },
   "source": [
    "<a id='8'></a>\n",
    "# 8. Loss Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "q2AK9kf0CMN0"
   },
   "outputs": [],
   "source": [
    "# Loss weighting for generators\n",
    "w_D1 = .5\n",
    "w_D2 = .5\n",
    "w_adv_PL = 1.\n",
    "w_recon_loss = 3.\n",
    "w_kl_loss = 0.01 # The hyperparameter beta of beta-VAE https://openreview.net/forum?id=Sy2fzU9gl\n",
    "w_pl_vgg = (0.01, 0.003, 0.003, 0.001, 0.01)\n",
    "w_mask_fo = 0.1\n",
    "\"\"\"\n",
    "Notes for beta (w_kl_loss):\n",
    "beta-VAE suggests using beta > 1 (e.g., 5, 20, 250) to disentangle latent features.\n",
    "However, the KL loss in such case will be much larger than other losses.\n",
    "Also we don't have much interest in disentangled features (at least for now).\n",
    "Thus beta is set to 0.01 (same with BicycleGAN default value).\n",
    "\"\"\"\n",
    "\n",
    "# Loss weighting for discriminators\n",
    "m_adv_PL = 0.65 # heuristically chosing a value between 0.5 and 1: 0.5 + (1-0.5)/3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "IaicdUsC_6xy"
   },
   "outputs": [],
   "source": [
    "if use_lsgan:\n",
    "    loss_fn = lambda output, target : K.mean(K.abs(K.square(output-target)))\n",
    "else:\n",
    "    loss_fn = lambda output, target : -K.mean(K.log(output+1e-12)*target+K.log(1-output+1e-12)*(1-target))\n",
    "    \n",
    "loss_fn_bce = lambda output, target : -K.mean(K.log(output+1e-12)*target+K.log(1-output+1e-12)*(1-target))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "NfuEKGnB2cnU"
   },
   "outputs": [],
   "source": [
    "# https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.cosine.html\n",
    "def cos_distance(x1, x2):\n",
    "    x1 = K.l2_normalize(x1, axis=-1)\n",
    "    x2 = K.l2_normalize(x2, axis=-1)\n",
    "    return K.mean(1 - K.sum((x1 * x2), axis=-1))\n",
    "\n",
    "def first_order(x, axis=1):\n",
    "    img_nrows = x.shape[1]\n",
    "    img_ncols = x.shape[2]\n",
    "    if axis == 1:\n",
    "        return K.abs(x[:, :img_nrows - 1, :img_ncols - 1, :] - x[:, 1:, :img_ncols - 1, :])\n",
    "    elif axis == 2:\n",
    "        return K.abs(x[:, :img_nrows - 1, :img_ncols - 1, :] - x[:, :img_nrows - 1, 1:, :])\n",
    "    else:\n",
    "        return None   \n",
    "\n",
    "# https://github.com/EmilienDupont/vae-concrete/blob/master/util.py\n",
    "def kl_normal(z_mean, z_log_var):\n",
    "    kl_per_example = .5 * (K.sum(K.square(z_mean) + K.exp(z_log_var) - 1 - z_log_var, axis=1))\n",
    "    return K.mean(kl_per_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 35,
     "output_extras": [
      {
       "item_id": 1
      }
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 572,
     "status": "ok",
     "timestamp": 1520931153866,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -480
    },
    "id": "ccs3GB4T_6yI",
    "outputId": "990bceca-7a89-4531-8278-d788f19bdd54"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Not using perceptual loss.\n"
     ]
    }
   ],
   "source": [
    "# ========== Define Perceptual Loss Model==========\n",
    "if use_perceptual_loss:\n",
    "    print (\"Using perceptual loss.\")\n",
    "    vggface.trainable = False\n",
    "    if base_network_PL == \"resnet50\":\n",
    "        # ResNet50 feats\n",
    "        out_size55 = vggface.layers[36].output\n",
    "        out_size28 = vggface.layers[78].output\n",
    "        out_size7 = vggface.layers[-2].output\n",
    "        vggface_feat = Model(vggface.input, [out_size55, out_size28, out_size7])\n",
    "    elif base_network_PL == \"vgg16\":\n",
    "        # VGG16 feats\n",
    "        out_size224 = vggface.layers[1].output\n",
    "        out_size112 = vggface.layers[4].output\n",
    "        out_size56 = vggface.layers[7].output\n",
    "        out_size28 = vggface.layers[11].output\n",
    "        out_size14 = vggface.layers[15].output\n",
    "        vggface_feat = Model(vggface.input, [out_size224, out_size112, out_size56, out_size28, out_size14])\n",
    "    vggface_feat.trainable = False\n",
    "else:\n",
    "    print (\"Not using perceptual loss.\")\n",
    "    vggface_feat = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "U05RmAdU_6x9"
   },
   "outputs": [],
   "source": [
    "def define_loss(netD, netD2, netG, real, fake_abgr, distorted, zs=None, vggface_feat=None, domain=\"A\", netD_code=None, real_code=None):\n",
    "    alpha = Lambda(lambda x: x[:,:,:, :1])(fake_abgr)\n",
    "    fake_bgr = Lambda(lambda x: x[:,:,:, 1:])(fake_abgr)\n",
    "    fake = alpha * fake_bgr + (1-alpha) * distorted\n",
    "    \n",
    "    if use_mixup:\n",
    "        dist = Beta(mixup_alpha, mixup_alpha)\n",
    "        lam = dist.sample()\n",
    "        \n",
    "        mixup = lam * concatenate([real, distorted]) + (1 - lam) * concatenate([fake, distorted])\n",
    "        out_D1_mixup = netD(mixup)\n",
    "        # GAN loss1\n",
    "        loss_D = loss_fn(out_D1_mixup, lam * K.ones_like(out_D1_mixup)) \n",
    "        loss_G = w_D1 * loss_fn(out_D1_mixup, (1 - lam) * K.ones_like(out_D1_mixup))\n",
    "        \n",
    "        lam2 = dist.sample()\n",
    "        mixup2 = lam2 * real + (1 - lam2) * fake_bgr\n",
    "        out3_D2_mixup = netD2(mixup2)[3]  \n",
    "        # GAN loss2\n",
    "        loss_D2 = loss_fn(out3_D2_mixup, lam2 * K.ones_like(out3_D2_mixup))\n",
    "        loss_G += w_D2 * loss_fn(out3_D2_mixup, (1 - lam) * K.ones_like(out3_D2_mixup))    \n",
    "        # Perceptual adversarial loss\n",
    "        out0_D2_real, out1_D2_real, out2_D2_real, _ = netD2(real)\n",
    "        out0_D2_fake, out1_D2_fake, out2_D2_fake, _ = netD2(fake_bgr)\n",
    "        adversarial_perceptual_loss_D2 = K.mean(K.abs(out0_D2_fake - out0_D2_real)) + K.mean(K.abs(out1_D2_fake - out1_D2_real)) + K.mean(K.abs(out2_D2_fake - out2_D2_real))\n",
    "        loss_G += w_adv_PL * adversarial_perceptual_loss_D2\n",
    "        loss_D2 += K.maximum(K.zeros_like(adversarial_perceptual_loss_D2), m_adv_PL - adversarial_perceptual_loss_D2)\n",
    "        \n",
    "        if base_model == \"XGAN\":\n",
    "            rec_code = netG([fake_bgr])[1]\n",
    "            output_real_code = netD_code([real_code])\n",
    "            # Target of domain A = 1, domain B = 0\n",
    "            if domain == \"A\":\n",
    "                loss_D_code = loss_fn_bce(output_real_code, K.ones_like(output_real_code))\n",
    "                loss_G += .1 * loss_fn(output_real_code, K.zeros_like(output_real_code))\n",
    "            elif domain == \"B\":\n",
    "                loss_D_code = loss_fn_bce(output_real_code, K.zeros_like(output_real_code))\n",
    "                loss_G += .1 * loss_fn(output_real_code, K.ones_like(output_real_code))\n",
    "            loss_G += 1. * cos_distance(rec_code, real_code)\n",
    "    \n",
    "    else:\n",
    "        print (\"use_mixup=False is not supported!\") \n",
    "        \n",
    "    # Reconstruction loss. L1 distance between reconstructed image and ground truth image. \n",
    "    loss_G += w_recon_loss * K.mean(K.abs(fake_bgr - real)) \n",
    "    \n",
    "    if base_model == \"VAE-GAN\":\n",
    "        z1_mean, z1_log_var, z2_mean, z2_log_var = zs\n",
    "        # VAE loss2: KL divergence between N(0,1) and N(z_mean, exp(z_log_var)).    \n",
    "        loss_G += w_kl_loss * kl_normal(z1_mean, z1_log_var) / 2 # Dividing by 2 because AnB KL loss will be updated twice within one iter.\n",
    "        loss_G += w_kl_loss * kl_normal(z2_mean, z2_log_var)\n",
    "    \n",
    "    # Perceptual Loss\n",
    "    if not vggface_feat is None:\n",
    "        def preprocess_vggface(x):\n",
    "            if base_network_PL == \"resnet50\":\n",
    "                x = (x + 1)/2 * 255 # channel order: BGR\n",
    "                x -= [91.4953, 103.8827, 131.0912]\n",
    "            elif base_network_PL == \"vgg16\":\n",
    "                x = (x + 1)/2 * 255 # channel order: BGR\n",
    "                x -= [93.5940, 104.7624, 129.1863]\n",
    "            return x\n",
    "        real_sz224 = tf.image.resize_images(real, [224, 224])\n",
    "        real_sz224 = Lambda(preprocess_vggface)(real_sz224)\n",
    "        fake_sz224 = tf.image.resize_images(fake, [224, 224]) \n",
    "        fake_sz224 = Lambda(preprocess_vggface)(fake_sz224)\n",
    "        \n",
    "        # ResNet50 PL\n",
    "        #pl_params = (0.03, 0.2, 0.3)\n",
    "        #real_feat55, real_feat28, real_feat7 = vggface_feat(real_sz224)\n",
    "        #fake_feat55, fake_feat28, fake_feat7  = vggface_feat(fake_sz224)    \n",
    "        #loss_G += pl_params[0] * K.mean(K.square(fake_feat7 - real_feat7))\n",
    "        #loss_G += pl_params[1] * K.mean(K.abs(fake_feat28 - real_feat28))\n",
    "        #loss_G += pl_params[2] * K.mean(K.abs(fake_feat55 - real_feat55))\n",
    "        \n",
    "        # VGG16 PL\n",
    "        pl_params = w_pl_vgg\n",
    "        real_feat224, real_feat112, real_feat56, real_feat28, real_feat14 = vggface_feat(real_sz224)\n",
    "        fake_feat224, fake_feat112, fake_feat56, fake_feat28, fake_feat14  = vggface_feat(fake_sz224)    \n",
    "        loss_G += pl_params[0] * K.mean(K.abs(fake_feat14 - real_feat14))\n",
    "        loss_G += pl_params[1] * K.mean(K.abs(fake_feat28 - real_feat28))\n",
    "        loss_G += pl_params[2] * K.mean(K.abs(fake_feat56 - real_feat56))\n",
    "        loss_G += pl_params[3] * K.mean(K.abs(fake_feat112 - real_feat112)) \n",
    "        loss_G += pl_params[4] * K.mean(K.abs(fake_feat224 - real_feat224))   \n",
    "    if base_model ==\"GAN\" or base_model ==\"VAE-GAN\":\n",
    "        return loss_D, loss_D2, loss_G\n",
    "    elif base_model == \"XGAN\":\n",
    "        return loss_D, loss_D2, loss_G, loss_D_code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "9CG48U0M_6yU"
   },
   "outputs": [],
   "source": [
    "if base_model == \"XGAN\":\n",
    "    loss_DA, loss_DA2, loss_GA, loss_DA_code = define_loss(netDA, netDA2, netGA, real_A, fake_A, distorted_A, \n",
    "                                                           zs=None, vggface_feat=vggface_feat, \n",
    "                                                           domain=\"A\", netD_code=netD_code, real_code=code_A)\n",
    "    loss_DB, loss_DB2, loss_GB, loss_DB_code = define_loss(netDB, netDB2, netGB, real_B, fake_B, distorted_B, \n",
    "                                                           zs=None, vggface_feat=vggface_feat, \n",
    "                                                           domain=\"B\", netD_code=netD_code, real_code=code_B)\n",
    "elif base_model == \"VAE-GAN\":  \n",
    "    loss_DA, loss_DA2, loss_GA = define_loss(netDA, netDA2, netGA, real_A, fake_A, \n",
    "                                             (distorted_x_A+distorted_y_A)/2, \n",
    "                                             zs=zs_A, vggface_feat=vggface_feat)\n",
    "    loss_DB, loss_DB2, loss_GB = define_loss(netDB, netDB2, netGB, real_B, fake_B, \n",
    "                                             (distorted_x_B+distorted_y_B)/2, \n",
    "                                             zs=zs_B, vggface_feat=vggface_feat)\n",
    "elif base_model == \"GAN\":\n",
    "    loss_DA, loss_DA2, loss_GA = define_loss(netDA, netDA2, netGA, real_A, fake_A, \n",
    "                                             (distorted_x_A+distorted_y_A)/2, \n",
    "                                             zs=None, vggface_feat=vggface_feat)\n",
    "    loss_DB, loss_DB2, loss_GB = define_loss(netDB, netDB2, netGB, real_B, fake_B, \n",
    "                                             (distorted_x_B+distorted_y_B)/2, \n",
    "                                             zs=None, vggface_feat=vggface_feat)\n",
    "\n",
    "# Alpha mask regularization\n",
    "loss_GA += 1e-2 * K.mean(K.abs(mask_A))\n",
    "loss_GB += 1e-2 * K.mean(K.abs(mask_B))\n",
    "loss_GA += w_mask_fo * K.mean(first_order(mask_A, axis=1))\n",
    "loss_GA += w_mask_fo * K.mean(first_order(mask_A, axis=2))\n",
    "loss_GB += w_mask_fo * K.mean(first_order(mask_B, axis=1))\n",
    "loss_GB += w_mask_fo * K.mean(first_order(mask_B, axis=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "C3nVnCBw_6yc"
   },
   "outputs": [],
   "source": [
    "weightsDA = netDA.trainable_weights\n",
    "weightsDA2 = netDA2.trainable_weights\n",
    "weightsGA = netGA.trainable_weights\n",
    "weightsDB = netDB.trainable_weights\n",
    "weightsDB2 = netDB2.trainable_weights\n",
    "weightsGB = netGB.trainable_weights\n",
    "\n",
    "\n",
    "# Adam(..).get_updates(...)\n",
    "\n",
    "\"\"\"\n",
    "# Using the following update function spped up training time (per iter.) by ~15%.\n",
    "training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDA+weightsDA2,[],loss_DA+loss_DA2)\n",
    "netDA_train = K.function([distorted_A, real_A],[loss_DA+loss_DA2], training_updates)\n",
    "\"\"\"\n",
    "if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDA,[],loss_DA)\n",
    "    netDA_train = K.function([distorted_x_A, distorted_y_A, real_A],[loss_DA], training_updates)\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDA2,[],loss_DA2)\n",
    "    netDA2_train = K.function([distorted_x_A, distorted_y_A, real_A],[loss_DA2], training_updates)\n",
    "    training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(weightsGA,[], loss_GA)\n",
    "    netGA_train = K.function([distorted_x_A, distorted_y_A, real_A], [loss_GA], training_updates)\n",
    "\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDB,[],loss_DB)\n",
    "    netDB_train = K.function([distorted_x_B, distorted_y_B, real_B],[loss_DB], training_updates)\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDB2,[],loss_DB2)\n",
    "    netDB2_train = K.function([distorted_x_B, distorted_y_B, real_B],[loss_DB2], training_updates)\n",
    "    training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(weightsGB,[], loss_GB)\n",
    "    netGB_train = K.function([distorted_x_B, distorted_y_B, real_B], [loss_GB], training_updates)\n",
    "    \n",
    "elif base_model ==\"XGAN\":\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDA,[],loss_DA)\n",
    "    netDA_train = K.function([distorted_A, real_A],[loss_DA], training_updates)\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDA2,[],loss_DA2)\n",
    "    netDA2_train = K.function([distorted_A, real_A],[loss_DA2], training_updates)\n",
    "    training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(weightsGA,[], loss_GA)\n",
    "    netGA_train = K.function([distorted_A, real_A], [loss_GA], training_updates)\n",
    "\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDB,[],loss_DB)\n",
    "    netDB_train = K.function([distorted_B, real_B],[loss_DB], training_updates)\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsDB2,[],loss_DB2)\n",
    "    netDB2_train = K.function([distorted_B, real_B],[loss_DB2], training_updates)\n",
    "    training_updates = Adam(lr=lrG, beta_1=0.5).get_updates(weightsGB,[], loss_GB)\n",
    "    netGB_train = K.function([distorted_B, real_B], [loss_GB], training_updates)\n",
    "  \n",
    "    weightsD_code = netD_code.trainable_weights\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsD_code,[], loss_DA_code)\n",
    "    netDA_code_train = K.function([distorted_A, real_A],[loss_DA_code], training_updates)\n",
    "    training_updates = Adam(lr=lrD, beta_1=0.5).get_updates(weightsD_code,[], loss_DB_code)\n",
    "    netDB_code_train = K.function([distorted_B, real_B],[loss_DB_code], training_updates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pnPh4xDB_6yr"
   },
   "source": [
    "<a id='9'></a>\n",
    "# 9. Utils For Loading/Displaying Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "6h0sgJkOCMOn"
   },
   "outputs": [],
   "source": [
    "from scipy import ndimage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "-ODLAc5SCMOr"
   },
   "outputs": [],
   "source": [
    "def get_motion_blur_kernel(sz=7):\n",
    "    rot_angle = np.random.uniform(-180,180)\n",
    "    kernel = np.zeros((sz,sz))\n",
    "    kernel[int((sz-1)//2), :] = np.ones(sz)\n",
    "    kernel = ndimage.interpolation.rotate(kernel, rot_angle, reshape=False)\n",
    "    kernel = np.clip(kernel, 0, 1)\n",
    "    normalize_factor = 1 / np.sum(kernel)\n",
    "    kernel = kernel * normalize_factor\n",
    "    return kernel\n",
    "\n",
    "def motion_blur(images, sz=7):\n",
    "  # images is a list [image2, image2, ...]\n",
    "    blur_sz = np.random.choice([5, 7, 9, 11])\n",
    "    kernel_motion_blur = get_motion_blur_kernel(blur_sz)\n",
    "    for i, image in enumerate(images):\n",
    "        images[i] = cv2.filter2D(image, -1, kernel_motion_blur).astype(np.float64)\n",
    "    return images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "oi9EU6rt_6yu"
   },
   "outputs": [],
   "source": [
    "def load_data(file_pattern):\n",
    "    return glob.glob(file_pattern)\n",
    "\n",
    "def random_channel_shift(x, intensity=10):\n",
    "    x = cv2.cvtColor(x, cv2.COLOR_BGR2HSV).astype('float32')\n",
    "    x += np.array([np.random.uniform(-intensity, intensity), \n",
    "                   np.random.uniform(-3*intensity, 3*intensity), \n",
    "                   np.random.uniform(-3*intensity, 3*intensity)])\n",
    "    x = np.clip(x, 0, 255).astype('uint8')\n",
    "    x = cv2.cvtColor(x, cv2.COLOR_HSV2BGR)\n",
    "    return x\n",
    "  \n",
    "def random_cutout(x, intensity=0.2):\n",
    "    h, w, c = x.shape\n",
    "    cutout_range = np.min([h//2, w//2])\n",
    "    cutout_x0 = np.random.randint(h)\n",
    "    cutout_y0 = np.random.randint(w)\n",
    "    cutout_x1 = np.min([cutout_x0+np.random.randint(cutout_range)+10, h])\n",
    "    cutout_y1 = np.min([cutout_y0+np.random.randint(cutout_range)+10, w])\n",
    "    cutout_h = cutout_x1 - cutout_x0\n",
    "    cutout_w = cutout_y1 - cutout_y0\n",
    "    x[cutout_x0:cutout_x1, cutout_y0:cutout_y1, :] = np.random.normal(0., intensity, size=(cutout_h, cutout_w, c))\n",
    "    return x\n",
    "\n",
    "def random_warp_rev(image):\n",
    "    assert image.shape == (256,256,3)\n",
    "    rand_coverage = np.random.randint(16) + 80\n",
    "    range_ = np.linspace(128-rand_coverage, 128+rand_coverage, 5)\n",
    "    mapx = np.broadcast_to(range_, (5,5))\n",
    "    mapy = mapx.T\n",
    "    mapx = mapx + np.random.normal( size=(5,5), scale=6)\n",
    "    mapy = mapy + np.random.normal( size=(5,5), scale=6)\n",
    "    interp_mapx = cv2.resize( mapx, (80,80) )[8:72,8:72].astype('float32')\n",
    "    interp_mapy = cv2.resize( mapy, (80,80) )[8:72,8:72].astype('float32')\n",
    "    warped_image = cv2.remap(image, interp_mapx, interp_mapy, cv2.INTER_LINEAR)\n",
    "    src_points = np.stack( [ mapx.ravel(), mapy.ravel() ], axis=-1)\n",
    "    dst_points = np.mgrid[0:65:16,0:65:16].T.reshape(-1,2)\n",
    "    mat = umeyama(src_points, dst_points, True)[0:2]\n",
    "    target_image = cv2.warpAffine(image, mat, (64,64))\n",
    "    return warped_image, target_image\n",
    "\n",
    "random_transform_args = {\n",
    "    'rotation_range': 20,\n",
    "    'zoom_range': 0.15,\n",
    "    'shift_range': 0.05,\n",
    "    'random_flip': 0.5,\n",
    "    }\n",
    "def read_image(fn, random_transform_args=random_transform_args):\n",
    "    image = cv2.imread(fn)\n",
    "    image = cv2.resize(image, (256,256)) / 255 * 2 - 1\n",
    "    image = random_transform( image, **random_transform_args)\n",
    "    warped_img, target_img = random_warp_rev(image)    \n",
    "    \n",
    "    #if np.random.uniform() < 0.1: # random downscaling\n",
    "    #    rand_size = np.random.randint(48) + 16\n",
    "    #    warped_img = cv2.resize(warped_img, (rand_size,rand_size))\n",
    "    #    warped_img = cv2.resize(warped_img, (64,64))\n",
    "    #elif np.random.uniform() < 0.3:\n",
    "    #    warped_img = random_cutout(warped_img)\n",
    "    \n",
    "    # Motion blur data augmentation:\n",
    "    # we want the model to learn to preserve motion blurs of input images\n",
    "    #if np.random.uniform() < 0.4: \n",
    "    #    warped_img, target_img = motion_blur([warped_img, target_img])\n",
    "    \n",
    "    return warped_img, target_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "J7Y__RLv_6zD"
   },
   "outputs": [],
   "source": [
    "# A generator function that yields epoch, batchsize of warped_img and batchsize of target_img\n",
    "def minibatch(data, batchsize):\n",
    "    length = len(data)\n",
    "    epoch = i = 0\n",
    "    tmpsize = None  \n",
    "    shuffle(data)\n",
    "    while True:\n",
    "        size = tmpsize if tmpsize else batchsize\n",
    "        if i+size > length:\n",
    "            shuffle(data)\n",
    "            i = 0\n",
    "            epoch+=1        \n",
    "        rtn = np.float32([read_image(data[j]) for j in range(i,i+size)])\n",
    "        i+=size\n",
    "        tmpsize = yield epoch, rtn[:,0,:,:,:], rtn[:,1,:,:,:]       \n",
    "\n",
    "def minibatchAB(dataA, batchsize):\n",
    "    batchA = minibatch(dataA, batchsize)\n",
    "    tmpsize = None    \n",
    "    while True:        \n",
    "        ep1, warped_img, target_img = batchA.send(tmpsize)\n",
    "        tmpsize = yield ep1, warped_img, target_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "zjnYAMsx_6zQ"
   },
   "outputs": [],
   "source": [
    "def showG(test_A, test_B, path_A, path_B):\n",
    "    if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "        figure_A = np.stack([\n",
    "            test_A,\n",
    "            np.squeeze(np.array([path_A([test_A[i:i+1], test_A[i:i+1]]) for i in range(test_A.shape[0])])),\n",
    "            np.squeeze(np.array([path_B([test_A[i:i+1], test_A[i:i+1]]) for i in range(test_A.shape[0])])),\n",
    "            ], axis=1 )\n",
    "        figure_B = np.stack([\n",
    "            test_B,\n",
    "            np.squeeze(np.array([path_B([test_B[i:i+1], test_B[i:i+1]]) for i in range(test_B.shape[0])])),\n",
    "            np.squeeze(np.array([path_A([test_B[i:i+1], test_B[i:i+1]]) for i in range(test_B.shape[0])])),\n",
    "            ], axis=1 )\n",
    "    elif base_model == \"XGAN\":\n",
    "        figure_A = np.stack([\n",
    "            test_A,\n",
    "            np.squeeze(np.array([path_A([test_A[i:i+1]]) for i in range(test_A.shape[0])])),\n",
    "            np.squeeze(np.array([path_B([test_A[i:i+1]]) for i in range(test_A.shape[0])])),\n",
    "            ], axis=1 )\n",
    "        figure_B = np.stack([\n",
    "            test_B,\n",
    "            np.squeeze(np.array([path_B([test_B[i:i+1]]) for i in range(test_B.shape[0])])),\n",
    "            np.squeeze(np.array([path_A([test_B[i:i+1]]) for i in range(test_B.shape[0])])),\n",
    "            ], axis=1 )\n",
    "\n",
    "    figure = np.concatenate([figure_A, figure_B], axis=0 )\n",
    "    figure = figure.reshape((4,7) + figure.shape[1:])\n",
    "    figure = stack_images(figure)\n",
    "    figure = np.clip((figure + 1) * 255 / 2, 0, 255).astype('uint8')\n",
    "    figure = cv2.cvtColor(figure, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "    display(Image.fromarray(figure))\n",
    "    return figure\n",
    "    \n",
    "def showG_mask(test_A, test_B, path_A, path_B):\n",
    "    if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "        figure_A = np.stack([\n",
    "            test_A,\n",
    "            (np.squeeze(np.array([path_A([test_A[i:i+1], test_A[i:i+1]]) for i in range(test_A.shape[0])])))*2-1,\n",
    "            (np.squeeze(np.array([path_B([test_A[i:i+1], test_A[i:i+1]]) for i in range(test_A.shape[0])])))*2-1,\n",
    "            ], axis=1 )\n",
    "        figure_B = np.stack([\n",
    "            test_B,\n",
    "            (np.squeeze(np.array([path_B([test_B[i:i+1], test_B[i:i+1]]) for i in range(test_B.shape[0])])))*2-1,\n",
    "            (np.squeeze(np.array([path_A([test_B[i:i+1], test_B[i:i+1]]) for i in range(test_B.shape[0])])))*2-1,\n",
    "            ], axis=1 )\n",
    "    elif base_model == \"XGAN\":\n",
    "        figure_A = np.stack([\n",
    "            test_A,\n",
    "            (np.squeeze(np.array([path_A([test_A[i:i+1]]) for i in range(test_A.shape[0])])))*2-1,\n",
    "            (np.squeeze(np.array([path_B([test_A[i:i+1]]) for i in range(test_A.shape[0])])))*2-1,\n",
    "            ], axis=1 )\n",
    "        figure_B = np.stack([\n",
    "            test_B,\n",
    "            (np.squeeze(np.array([path_B([test_B[i:i+1]]) for i in range(test_B.shape[0])])))*2-1,\n",
    "            (np.squeeze(np.array([path_A([test_B[i:i+1]]) for i in range(test_B.shape[0])])))*2-1,\n",
    "            ], axis=1 )\n",
    "\n",
    "    figure = np.concatenate([figure_A, figure_B], axis=0 )\n",
    "    figure = figure.reshape((4,7) + figure.shape[1:])\n",
    "    figure = stack_images(figure)\n",
    "    figure = np.clip((figure + 1) * 255 / 2, 0, 255).astype('uint8')\n",
    "    figure = cv2.cvtColor(figure, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "    display(Image.fromarray(figure))\n",
    "    return figure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Avm226IZ_6zq"
   },
   "source": [
    "<a id='10'></a>\n",
    "# 10. Start Training\n",
    "\n",
    "Show results and save model weights every `display_iters` iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "KdxqnEsh_6zt"
   },
   "outputs": [],
   "source": [
    "!mkdir models # create ./models directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "YUEccJNe_60C"
   },
   "outputs": [],
   "source": [
    "!mkdir previews"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 53,
     "output_extras": [
      {
       "item_id": 1
      }
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 605,
     "status": "ok",
     "timestamp": 1520927888859,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -480
    },
    "id": "LLknxYfa_60R",
    "outputId": "045b5909-4201-4e19-c386-8d51d5424346"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of images in folder of face A: 538\n",
      "Number of images in folder of face B: 661\n"
     ]
    }
   ],
   "source": [
    "# Get filenames\n",
    "train_A = load_data(img_dirA)\n",
    "train_B = load_data(img_dirB)\n",
    "\n",
    "assert len(train_A), \"No image found in \" + str(img_dirA)\n",
    "assert len(train_B), \"No image found in \" + str(img_dirB)\n",
    "\n",
    "print (\"Number of images in folder of face A: \" + str(len(train_A)))\n",
    "print (\"Number of images in folder of face B: \" + str(len(train_B)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "JZaysMUlCMPe"
   },
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "niter = 300\n",
    "gen_iterations = 0\n",
    "epoch = 0\n",
    "errGA_sum = errGB_sum = errDA_sum = errDB_sum = errDA2_sum = errDB2_sum = errDA_code_sum = errDB_code_sum = 0\n",
    "\n",
    "display_iters = 300\n",
    "train_batchA = minibatchAB(train_A, batchSize)\n",
    "train_batchB = minibatchAB(train_B, batchSize)\n",
    "\n",
    "\n",
    "while gen_iterations < 30000: \n",
    "    epoch, warped_A, target_A = next(train_batchA) \n",
    "    epoch, warped_B, target_B = next(train_batchB) \n",
    "    \n",
    "    # Train dicriminators for one batch\n",
    "    if gen_iterations % 1 == 0:\n",
    "        if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "            errDA = netDA_train([warped_A, warped_A, target_A])\n",
    "            errDB = netDB_train([warped_B, warped_B, target_B])\n",
    "            errDA2 = netDA2_train([warped_A, warped_A, target_A])\n",
    "            errDB2 = netDB2_train([warped_B, warped_B, target_B])\n",
    "        elif base_model == \"XGAN\":\n",
    "            errDA = netDA_train([warped_A, target_A])\n",
    "            errDB = netDB_train([warped_B, target_B])\n",
    "            errDA2 = netDA2_train([warped_A, target_A])\n",
    "            errDB2 = netDB2_train([warped_B, target_B])\n",
    "            errDA_code = netDA_code_train([warped_A, target_A])\n",
    "            errDB_code = netDB_code_train([warped_B, target_B])\n",
    "    errDA_sum += errDA[0]\n",
    "    errDB_sum += errDB[0]\n",
    "    errDA2_sum += errDA2[0]\n",
    "    errDB2_sum += errDB2[0]\n",
    "    if base_model == \"XGAN\":\n",
    "        errDA_code_sum += errDA_code[0]\n",
    "        errDB_code_sum += errDB_code[0]\n",
    "    \n",
    "    if gen_iterations == 5:\n",
    "        print (\"working.\")\n",
    "        print('[epoch %d][iter %d] Loss_DA: %f Loss_DB: %f Loss_DA2: %f Loss_DB2: %f Loss_GA: %f Loss_GB: %f time: %f'\n",
    "        % (epoch, gen_iterations, errDA_sum/5, errDB_sum/5, errDA2_sum/5, errDB2_sum/5,\n",
    "           errGA_sum/5, errGB_sum/5, time.time()-t0))\n",
    "        if base_model == \"XGAN\":\n",
    "            print('Loss_DA_code: %f Loss_DB_code: %f' % (errDA_code_sum/5, errDB_code_sum/5))\n",
    "        _, wA, tA = train_batchA.send(14)  \n",
    "        _, wB, tB = train_batchB.send(14)\n",
    "        preview_im = showG(tA, tB, path_A, path_B)   \n",
    "        print (\" \")\n",
    "        preview_im = showG(wA, wB, path_bgr_A, path_bgr_B)   \n",
    "        print (\" \")    \n",
    "        preview_im = showG_mask(tA, tB, path_mask_A, path_mask_B) \n",
    "    \n",
    "    # Train generators for one batch\n",
    "    if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "        errGA = netGA_train([warped_A, warped_A, target_A])\n",
    "        errGB = netGB_train([warped_B, warped_B, target_B])\n",
    "    elif base_model == \"XGAN\":\n",
    "        errGA = netGA_train([warped_A, target_A])\n",
    "        errGB = netGB_train([warped_B, target_B])\n",
    "    errGA_sum += errGA[0]\n",
    "    errGB_sum += errGB[0]\n",
    "    gen_iterations += 1\n",
    "    \n",
    "    # Display previews & save model weights\n",
    "    if gen_iterations % display_iters == 0:\n",
    "        clear_output()\n",
    "        # Show losses\n",
    "        print('[epoch %d][iter %d] Loss_DA: %f Loss_DB: %f Loss_DA2: %f Loss_DB2: %f Loss_GA: %f Loss_GB: %f time: %f'\n",
    "        % (epoch, gen_iterations, errDA_sum/display_iters, errDB_sum/display_iters, errDA2_sum/display_iters, errDB2_sum/display_iters,\n",
    "           errGA_sum/display_iters, errGB_sum/display_iters, time.time()-t0)) \n",
    "        if base_model == \"XGAN\":\n",
    "            print('Loss_DA_code: %f Loss_DB_code: %f' % (errDA_code_sum/display_iters, errDB_code_sum/display_iters)) \n",
    "        \n",
    "        # Show previews\n",
    "        _, wA, tA = train_batchA.send(14)  \n",
    "        _, wB, tB = train_batchB.send(14)\n",
    "        preview_im = showG(tA, tB, path_A, path_B)   \n",
    "        plt.imsave(\"./previews/\"+str(gen_iterations)+\"_masked.jpg\", preview_im)\n",
    "        print (\" \")\n",
    "        preview_im = showG(wA, wB, path_bgr_A, path_bgr_B)   \n",
    "        plt.imsave(\"./previews/\"+str(gen_iterations)+\"_raw.jpg\", preview_im)  \n",
    "        print (\" \")    \n",
    "        preview_im = showG_mask(tA, tB, path_mask_A, path_mask_B)  \n",
    "        plt.imsave(\"./previews/\"+str(gen_iterations)+\"_mask.jpg\", preview_im)   \n",
    "        errGA_sum = errGB_sum = errDA_sum = errDB_sum = 0\n",
    "        errDA2_sum = errDB2_sum = 0\n",
    "        if base_model == \"XGAN\":\n",
    "            errDA_code_sum = errDB_code_sum = 0\n",
    "        \n",
    "        # Save models\n",
    "        encoder.save_weights(\"models/encoder.h5\")\n",
    "        decoder_A.save_weights(\"models/decoder_A.h5\")\n",
    "        decoder_B.save_weights(\"models/decoder_B.h5\")  \n",
    "        if base_model == \"VAE-GAN\" or base_model == \"GAN\":\n",
    "            decoder_AnB.save_weights(\"models/decoder_AnB.h5\")     \n",
    "            if decoder_type == 1:\n",
    "                decoder_common.save_weights(\"models/decoder_common.h5\")\n",
    "            elif decoder_type == 2:\n",
    "                decoder_A2.save_weights(\"models/decoder_A2.h5\")\n",
    "                decoder_B2.save_weights(\"models/decoder_B2.h5\")\n",
    "        netDA.save_weights(\"models/netDA.h5\")\n",
    "        netDB.save_weights(\"models/netDB.h5\")\n",
    "        netDA2.save_weights(\"models/netDA2.h5\")\n",
    "        netDB2.save_weights(\"models/netDB2.h5\")\n",
    "        if base_model == \"XGAN\":\n",
    "            netD_code.save_weights(\"models/netD_code.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "omaMY9WMCMPg"
   },
   "source": [
    "# 11. Video Conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "VKFfmCOHk-tQ"
   },
   "outputs": [],
   "source": [
    "from moviepy.editor import VideoFileClip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "ddy7_DJfgaTj"
   },
   "outputs": [],
   "source": [
    "path_func = path_abgr_B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dBaOTHZCCMPr"
   },
   "source": [
    "## MTCNN "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "NgPmXa1rgaKu"
   },
   "outputs": [],
   "source": [
    "import mtcnn_detect_face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "UhluNIGGiGo0"
   },
   "outputs": [],
   "source": [
    "def create_mtcnn(sess, model_path):\n",
    "    if not model_path:\n",
    "        model_path,_ = os.path.split(os.path.realpath(__file__))\n",
    "\n",
    "    with tf.variable_scope('pnet2'):\n",
    "        data = tf.placeholder(tf.float32, (None,None,None,3), 'input')\n",
    "        pnet = mtcnn_detect_face.PNet({'data':data})\n",
    "        pnet.load(os.path.join(model_path, 'det1.npy'), sess)\n",
    "    with tf.variable_scope('rnet2'):\n",
    "        data = tf.placeholder(tf.float32, (None,24,24,3), 'input')\n",
    "        rnet = mtcnn_detect_face.RNet({'data':data})\n",
    "        rnet.load(os.path.join(model_path, 'det2.npy'), sess)\n",
    "    with tf.variable_scope('onet2'):\n",
    "        data = tf.placeholder(tf.float32, (None,48,48,3), 'input')\n",
    "        onet = mtcnn_detect_face.ONet({'data':data})\n",
    "        onet.load(os.path.join(model_path, 'det3.npy'), sess)\n",
    "    return pnet, rnet, onet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 121,
     "output_extras": [
      {}
     ]
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 9057,
     "status": "ok",
     "timestamp": 1519398382252,
     "user": {
      "displayName": "Lu SA",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "109275333798683015269"
     },
     "user_tz": -540
    },
    "id": "6VX8pt5FiGfd",
    "outputId": "155fc4c4-f407-4ce7-b1af-36a46a6bc82a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /content/mtcnn_detect_face.py:212: calling reduce_max (from tensorflow.python.ops.math_ops) with keep_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "keep_dims is deprecated, use keepdims instead\n",
      "WARNING:tensorflow:From /content/mtcnn_detect_face.py:214: calling reduce_sum (from tensorflow.python.ops.math_ops) with keep_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "keep_dims is deprecated, use keepdims instead\n"
     ]
    }
   ],
   "source": [
    "WEIGHTS_PATH = \"./mtcnn_weights/\"\n",
    "\n",
    "sess = K.get_session()\n",
    "with sess.as_default():\n",
    "    global pnet, rnet, onet \n",
    "    pnet2, rnet2, onet2 = create_mtcnn(sess, WEIGHTS_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "pbepaCyxKEJ9"
   },
   "outputs": [],
   "source": [
    "global pnet, rnet, onet\n",
    "pnet_fun = K.function([pnet2.layers['data']],[pnet2.layers['conv4-2'], pnet2.layers['prob1']])\n",
    "rnet_fun = K.function([rnet2.layers['data']],[rnet2.layers['conv5-2'], rnet2.layers['prob1']])\n",
    "onet_fun = K.function([onet2.layers['data']], [onet2.layers['conv6-2'], onet2.layers['conv6-3'], onet2.layers['prob1']])\n",
    "\n",
    "with tf.variable_scope('pnet2', reuse=True):\n",
    "    data = tf.placeholder(tf.float32, (None,None,None,3), 'input')\n",
    "    pnet2 = mtcnn_detect_face.PNet({'data':data})\n",
    "    pnet2.load(os.path.join(\"./mtcnn_weights/\", 'det1.npy'), sess)\n",
    "with tf.variable_scope('rnet2', reuse=True):\n",
    "    data = tf.placeholder(tf.float32, (None,24,24,3), 'input')\n",
    "    rnet2 = mtcnn_detect_face.RNet({'data':data})\n",
    "    rnet2.load(os.path.join(\"./mtcnn_weights/\", 'det2.npy'), sess)\n",
    "with tf.variable_scope('onet2', reuse=True):\n",
    "    data = tf.placeholder(tf.float32, (None,48,48,3), 'input')\n",
    "    onet2 = mtcnn_detect_face.ONet({'data':data})\n",
    "    onet2.load(os.path.join(\"./mtcnn_weights/\", 'det3.npy'), sess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "9L4TFcSKKuHE"
   },
   "outputs": [],
   "source": [
    "pnet = K.function([pnet2.layers['data']],[pnet2.layers['conv4-2'], pnet2.layers['prob1']])\n",
    "rnet = K.function([rnet2.layers['data']],[rnet2.layers['conv5-2'], rnet2.layers['prob1']])\n",
    "onet = K.function([onet2.layers['data']], [onet2.layers['conv6-2'], onet2.layers['conv6-3'], onet2.layers['prob1']])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "GENCQrMtNTO4"
   },
   "source": [
    "## FCN8s for face segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "1Tl0tw7mCMQF"
   },
   "outputs": [],
   "source": [
    "from FCN8s_keras import FCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "IcrIMkwHNTDQ"
   },
   "outputs": [],
   "source": [
    "fcn = FCN()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "f1kFe3zUNS3u"
   },
   "outputs": [],
   "source": [
    "fcn.load_weights(\"Keras_FCN8s_face_seg_YuvalNirkin.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "ebC5nHg4Nacl"
   },
   "outputs": [],
   "source": [
    "def fcn_mask_seg(x):\n",
    "    def vgg_preprocess(im):\n",
    "        im = cv2.resize(im, (500, 500))\n",
    "        in_ = np.array(im, dtype=np.float32)\n",
    "        #in_ = in_[:,:,::-1]\n",
    "        in_ -= np.array((104.00698793,116.66876762,122.67891434))\n",
    "        in_ = in_[np.newaxis,:]\n",
    "        #in_ = in_.transpose((2,0,1))\n",
    "        return in_\n",
    "    \n",
    "    inp_fcn = vgg_preprocess(x)\n",
    "    out_mask_fcn = fcn.predict([inp_fcn])\n",
    "    out_mask_fcn = cv2.resize(np.squeeze(out_mask_fcn), (x.shape[1],x.shape[0]))\n",
    "    out_mask_fcn = np.clip(out_mask_fcn.argmax(axis=2), 0, 1).astype(np.float64)\n",
    "    return out_mask_fcn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "t9wHKYE3NaOu"
   },
   "source": [
    "## Video conversion functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Qiu3TdB4j0cj"
   },
   "outputs": [],
   "source": [
    "use_smoothed_mask = True\n",
    "use_smoothed_bbox = True\n",
    "\n",
    "def is_overlap(box1, box2):\n",
    "    overlap_x0 = np.max([box1[0], box2[0]]).astype(np.float32)\n",
    "    overlap_y1 = np.min([box1[1], box2[1]]).astype(np.float32)\n",
    "    overlap_x1 = np.min([box1[2], box2[2]]).astype(np.float32)\n",
    "    overlap_y0 = np.max([box1[3], box2[3]]).astype(np.float32)\n",
    "    area_iou = (overlap_x1-overlap_x0) * (overlap_y1-overlap_y0)\n",
    "    area_box1 = (box1[2]-box1[0]) * (box1[1]-box1[3])\n",
    "    area_box2 = (box2[2]-box2[0]) * (box2[1]-box2[3])    \n",
    "    return (area_iou / area_box1) >= 0.2\n",
    "    \n",
    "def remove_overlaps(faces):    \n",
    "    main_face = get_most_conf_face(faces)\n",
    "    main_face_bbox = main_face[0]\n",
    "    result_faces = []\n",
    "    result_faces.append(main_face_bbox)\n",
    "    for (x0, y1, x1, y0, conf_score) in faces:\n",
    "        if not is_overlap(main_face_bbox, (x0, y1, x1, y0)):\n",
    "            result_faces.append((x0, y1, x1, y0, conf_score))\n",
    "    return result_faces\n",
    "\n",
    "def get_most_conf_face(faces):\n",
    "    # Return the bbox w/ the highest confidence score\n",
    "    best_conf_score = 0\n",
    "    conf_face = None\n",
    "    for (x0, y1, x1, y0, conf_score) in faces: \n",
    "        if conf_score >= best_conf_score:\n",
    "            best_conf_score = conf_score\n",
    "            conf_face = [(x0, y1, x1, y0, conf_score)]\n",
    "    return conf_face\n",
    "\n",
    "def kalmanfilter_init(noise_coef):\n",
    "    kf = cv2.KalmanFilter(4,2)\n",
    "    kf.measurementMatrix = np.array([[1,0,0,0],[0,1,0,0]], np.float32)\n",
    "    kf.transitionMatrix = np.array([[1,0,1,0],[0,1,0,1],[0,0,1,0],[0,0,0,1]], np.float32)\n",
    "    kf.processNoiseCov = noise_coef * np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], np.float32)\n",
    "    return kf\n",
    "\n",
    "def is_higher_than_480p(x):\n",
    "    return (x.shape[0] * x.shape[1]) >= (858*480)\n",
    "\n",
    "def is_higher_than_720p(x):\n",
    "    return (x.shape[0] * x.shape[1]) >= (1280*720)\n",
    "\n",
    "def is_higher_than_1080p(x):\n",
    "    return (x.shape[0] * x.shape[1]) >= (1920*1080)\n",
    "\n",
    "def calibrate_coord(faces, video_scaling_factor):\n",
    "    for i, (x0, y1, x1, y0, _) in enumerate(faces):\n",
    "        faces[i] = (x0*video_scaling_factor, y1*video_scaling_factor, \n",
    "                    x1*video_scaling_factor, y0*video_scaling_factor, _)\n",
    "    return faces\n",
    "\n",
    "def process_mtcnn_bbox(bboxes, im_shape):\n",
    "    # outuut bbox coordinate of MTCNN is (y0, x0, y1, x1)\n",
    "    # Process the bbox coord. to a square bbox with ordering (x0, y1, x1, y0)\n",
    "    for i, bbox in enumerate(bboxes):\n",
    "        y0, x0, y1, x1 = bboxes[i,0:4]\n",
    "        w = int(y1 - y0)\n",
    "        h = int(x1 - x0)\n",
    "        length = (w + h) / 2\n",
    "        center = (int((x1+x0)/2),int((y1+y0)/2))\n",
    "        new_x0 = np.max([0, (center[0]-length//2)])#.astype(np.int32)\n",
    "        new_x1 = np.min([im_shape[0], (center[0]+length//2)])#.astype(np.int32)\n",
    "        new_y0 = np.max([0, (center[1]-length//2)])#.astype(np.int32)\n",
    "        new_y1 = np.min([im_shape[1], (center[1]+length//2)])#.astype(np.int32)\n",
    "        bboxes[i,0:4] = new_x0, new_y1, new_x1, new_y0\n",
    "    return bboxes\n",
    "\n",
    "def get_downscale_factor(image):\n",
    "    if is_higher_than_1080p(image):\n",
    "        return  4 + video_scaling_offset\n",
    "    elif is_higher_than_720p(image):\n",
    "        return 3 + video_scaling_offset\n",
    "    elif is_higher_than_480p(image):\n",
    "        return 2 + video_scaling_offset\n",
    "    else:\n",
    "        return 1\n",
    "\n",
    "def get_faces_bbox(image):  \n",
    "    global pnet, rnet, onet \n",
    "    global detec_threshold\n",
    "    minsize = 20 # minimum size of face\n",
    "    threshold = [ 0.6, 0.7, detec_threshold ]  # three steps's threshold\n",
    "    factor = 0.709 # scale factor\n",
    "    if manually_downscale:\n",
    "        video_scaling_factor = manual_downscale_factor\n",
    "        resized_image = cv2.resize(image, \n",
    "                                   (image.shape[1]//video_scaling_factor, \n",
    "                                    image.shape[0]//video_scaling_factor))\n",
    "        faces, pnts = mtcnn_detect_face.detect_face(resized_image, minsize, pnet, rnet, onet, threshold, factor)\n",
    "        faces = process_mtcnn_bbox(faces, resized_image.shape)\n",
    "        faces = calibrate_coord(faces, video_scaling_factor)\n",
    "    else:\n",
    "        video_scaling_factor = get_downscale_factor(image)\n",
    "        resized_image = cv2.resize(image, \n",
    "                                   (image.shape[1]//video_scaling_factor, \n",
    "                                    image.shape[0]//video_scaling_factor))\n",
    "        faces, pnts = mtcnn_detect_face.detect_face(resized_image, minsize, pnet, rnet, onet, threshold, factor)\n",
    "        faces = process_mtcnn_bbox(faces, resized_image.shape)\n",
    "        faces = calibrate_coord(faces, video_scaling_factor)\n",
    "    return faces\n",
    "\n",
    "def get_smoothed_coord(x0, x1, y0, y1, shape, ratio=0.65):\n",
    "    global prev_x0, prev_x1, prev_y0, prev_y1\n",
    "    if not use_kalman_filter:\n",
    "        x0 = int(ratio * prev_x0 + (1-ratio) * x0)\n",
    "        x1 = int(ratio * prev_x1 + (1-ratio) * x1)\n",
    "        y1 = int(ratio * prev_y1 + (1-ratio) * y1)\n",
    "        y0 = int(ratio * prev_y0 + (1-ratio) * y0)\n",
    "    else:\n",
    "        x0y0 = np.array([x0, y0]).astype(np.float32)\n",
    "        x1y1 = np.array([x1, y1]).astype(np.float32)\n",
    "        kf0.correct(x0y0)\n",
    "        pred_x0y0 = kf0.predict()\n",
    "        kf1.correct(x1y1)\n",
    "        pred_x1y1 = kf1.predict()\n",
    "        x0 = np.max([0, pred_x0y0[0][0]]).astype(np.int)\n",
    "        x1 = np.min([shape[0], pred_x1y1[0][0]]).astype(np.int)\n",
    "        y0 = np.max([0, pred_x0y0[1][0]]).astype(np.int)\n",
    "        y1 = np.min([shape[1], pred_x1y1[1][0]]).astype(np.int)\n",
    "        if x0 == x1 or y0 == y1:\n",
    "            x0, y0, x1, y1 = prev_x0, prev_y0, prev_x1, prev_y1\n",
    "    return x0, x1, y0, y1     \n",
    "    \n",
    "def set_global_coord(x0, x1, y0, y1):\n",
    "    global prev_x0, prev_x1, prev_y0, prev_y1\n",
    "    prev_x0 = x0\n",
    "    prev_x1 = x1\n",
    "    prev_y1 = y1\n",
    "    prev_y0 = y0\n",
    "    \n",
    "def merge_mask(mask1, mask2,mode=\"avg\"):\n",
    "    mask2 = cv2.resize(mask2, (mask1.shape[1],mask1.shape[0]))\n",
    "    if mode == \"avg\":\n",
    "        return (mask1 + mask2).astype(np.float32) / 2\n",
    "    elif mode == \"max\":\n",
    "        return np.maximum(mask1, mask2)\n",
    "    elif mode == \"adaptive\":\n",
    "        sum1 = np.sum(mask1)/3\n",
    "        sum2 = np.sum(mask2)\n",
    "        ratio = np.clip(sum1/(sum1+sum2), 0.2, 0.8)\n",
    "        return ratio*mask1 + (1-ratio)*mask2\n",
    "\n",
    "def get_tta_masked_result(result_bgr, result_a, roi_image):\n",
    "    result_bgr = cv2.resize(result_bgr, (64,64)).astype(np.float64)\n",
    "    result_a = cv2.resize(result_a, (64,64)).astype(np.float64)\n",
    "    roi_image_resized = cv2.resize(roi_image, (64,64)).astype(np.float64)\n",
    "    result_a = cv2.GaussianBlur(result_a, (7,7), 6)[:,:,np.newaxis]\n",
    "    result = (result_a/255 * result_bgr.astype(np.float64) + (1-result_a/255) * roi_image_resized).astype(np.uint8)\n",
    "    result = cv2.resize(result, (roi_image.shape[1],roi_image.shape[0]))\n",
    "    if use_color_correction:\n",
    "        result = color_hist_match(result, roi_image).astype('uint8')\n",
    "    result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB).astype('uint8')\n",
    "    return result    \n",
    "\n",
    "def generate_face(ae_input, path_abgr, roi_size, roi_image):\n",
    "    if base_model == \"XGAN\":\n",
    "        result = np.squeeze(np.array([path_abgr([[ae_input]])]))\n",
    "    elif base_model == \"VAE-GAN\" or base_model ==\"GAN\":\n",
    "        result = np.squeeze(np.array([path_abgr([[ae_input], [ae_input]])]))\n",
    "    result_a = result[:,:,0] * 255\n",
    "    result_bgr = np.clip( (result[:,:,1:] + 1) * 255 / 2, 0, 255 )\n",
    "    result_a_clear = np.copy(result_a)\n",
    "    result_a = cv2.GaussianBlur(result_a ,(7,7),6)\n",
    "    if use_landmark_match and False:\n",
    "        resized_roi = cv2.resize(roi_image, (64,64))\n",
    "        result_bgr, result_a = landmarks_match_mtcnn(resized_roi, result_bgr, result_a)\n",
    "    #if use_color_correction:\n",
    "    #    result_bgr = color_hist_match(result_bgr, roi_image)\n",
    "    \n",
    "    if use_FCN_mask:\n",
    "        fcn_mask = fcn_mask_seg(roi_image)\n",
    "        if merge_mask_mode == \"avg\":\n",
    "            fcn_mask = merge_mask(fcn_mask, result[:,:,0], mode=\"avg\")\n",
    "        elif merge_mask_mode == \"max\":\n",
    "            fcn_mask = merge_mask(fcn_mask, result[:,:,0], mode=\"max\")\n",
    "        elif merge_mask_mode == \"adaptive\":\n",
    "            fcn_mask = merge_mask(fcn_mask, result[:,:,0], mode=\"adaptive\")\n",
    "        result_a = fcn_mask[:,:,np.newaxis] * 255\n",
    "        result_a_clear = np.copy(fcn_mask * 255)\n",
    "        fcn_mask_sz64 = cv2.resize(fcn_mask, (64,64))\n",
    "        fcn_mask_sz64 = cv2.GaussianBlur(fcn_mask_sz64, (7,7), 6)[:,:,np.newaxis]\n",
    "        result = (fcn_mask_sz64 * result_bgr.astype(np.float64) + (1-fcn_mask_sz64) * ((ae_input+1)*255/2)).astype(np.uint8)\n",
    "    else:\n",
    "        result_a = np.expand_dims(result_a, axis=2)\n",
    "        result = (result_a/255 * result_bgr + (1 - result_a/255) * ((ae_input + 1) * 255 / 2)).astype('uint8')\n",
    "        result = result_bgr.astype('uint8')\n",
    "    if use_color_correction:\n",
    "        result = color_hist_match(result, roi_image).astype('uint8')\n",
    "    result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)\n",
    "    result = cv2.resize(result, (roi_size[1],roi_size[0]))\n",
    "    result_a_clear = np.expand_dims(cv2.resize(result_a_clear, (roi_size[1],roi_size[0])), axis=2)\n",
    "    return result, result_a_clear, result_bgr\n",
    "\n",
    "def get_init_mask_map(image):\n",
    "    return np.zeros_like(image)\n",
    "\n",
    "def get_init_comb_img(input_img):\n",
    "    comb_img = np.zeros([input_img.shape[0], input_img.shape[1]*2,input_img.shape[2]])\n",
    "    comb_img[:, :input_img.shape[1], :] = input_img\n",
    "    comb_img[:, input_img.shape[1]:, :] = input_img\n",
    "    return comb_img    \n",
    "\n",
    "def get_init_triple_img(input_img, no_face=False):\n",
    "    if no_face:\n",
    "        triple_img = np.zeros([input_img.shape[0], input_img.shape[1]*3,input_img.shape[2]])\n",
    "        triple_img[:, :input_img.shape[1], :] = input_img\n",
    "        triple_img[:, input_img.shape[1]:input_img.shape[1]*2, :] = input_img      \n",
    "        triple_img[:, input_img.shape[1]*2:, :] = (input_img * .15).astype('uint8')  \n",
    "        return triple_img\n",
    "    else:\n",
    "        triple_img = np.zeros([input_img.shape[0], input_img.shape[1]*3,input_img.shape[2]])\n",
    "        return triple_img\n",
    "\n",
    "def get_mask(roi_image, h, w):\n",
    "    mask = np.zeros_like(roi_image)\n",
    "    mask[h//15:-h//15,w//15:-w//15,:] = 255\n",
    "    mask = cv2.GaussianBlur(mask,(15,15),10)\n",
    "    return mask\n",
    "\n",
    "def hist_match(source, template):\n",
    "    # Code borrow from:\n",
    "    # https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x\n",
    "    oldshape = source.shape\n",
    "    source = source.ravel()\n",
    "    template = template.ravel()\n",
    "    s_values, bin_idx, s_counts = np.unique(source, return_inverse=True,\n",
    "                                            return_counts=True)\n",
    "    t_values, t_counts = np.unique(template, return_counts=True)\n",
    "\n",
    "    s_quantiles = np.cumsum(s_counts).astype(np.float64)\n",
    "    s_quantiles /= s_quantiles[-1]\n",
    "    t_quantiles = np.cumsum(t_counts).astype(np.float64)\n",
    "    t_quantiles /= t_quantiles[-1]\n",
    "    interp_t_values = np.interp(s_quantiles, t_quantiles, t_values)\n",
    "\n",
    "    return interp_t_values[bin_idx].reshape(oldshape)\n",
    "\n",
    "def color_hist_match(src_im, tar_im):\n",
    "    #src_im = cv2.cvtColor(src_im, cv2.COLOR_BGR2HSV)\n",
    "    #tar_im = cv2.cvtColor(tar_im, cv2.COLOR_BGR2HSV)\n",
    "    matched_0 = hist_match(src_im[:,:,0], tar_im[:,:,0])\n",
    "    matched_1 = hist_match(src_im[:,:,1], tar_im[:,:,1])\n",
    "    matched_2 = hist_match(src_im[:,:,2], tar_im[:,:,2])\n",
    "    matched = np.stack((matched_1, matched_1, matched_2), axis=2).astype(\"uint8\")\n",
    "    #matched = cv2.cvtColor(matched, cv2.COLOR_HSV2BGR)\n",
    "    return matched\n",
    "\n",
    "def landmarks_match_mtcnn(source, target, alpha):\n",
    "    global prev_pnts1, prev_pnts2\n",
    "    ratio = 0.2\n",
    "    \"\"\"\n",
    "    TODO: Reuse the landmarks of source image. Conceivable bug: coordinate mismatch.\n",
    "    \"\"\"\n",
    "    minsize = 20 # minimum size of face\n",
    "    threshold = [ 0.6, 0.7, 0.93 ]  # three steps's threshold\n",
    "    factor = 0.709 # scale factor\n",
    "    _, pnts1 = mtcnn_detect_face.detect_face(source, minsize, pnet, rnet, onet, threshold, factor) # redundant detection\n",
    "    _, pnts2 = mtcnn_detect_face.detect_face(target, minsize, pnet, rnet, onet, threshold, factor)  \n",
    "    \n",
    "    if len(prev_pnts1) == 0 and len(prev_pnts2) == 0:\n",
    "        if pnts1.shape[0] == 10 and pnts2.shape[0] == 10:\n",
    "            prev_pnts1, prev_pnts2 = pnts1, pnts2        \n",
    "    try:\n",
    "        landmarks_XY1 = []\n",
    "        landmarks_XY2 = []\n",
    "        for i in range(5):\n",
    "            landmarks_XY1.extend([((1-ratio)*pnts1[i+5][0] + ratio*prev_pnts1[i+5][0], \n",
    "                                   (1-ratio)*pnts1[i][0] + ratio*prev_pnts1[i][0])])\n",
    "            landmarks_XY2.extend([((1-ratio)*pnts2[i+5][0] + ratio*prev_pnts2[i+5][0], \n",
    "                                   (1-ratio)*pnts2[i][0] + ratio*prev_pnts2[i][0])])\n",
    "        M = umeyama(np.array(landmarks_XY1), np.array(landmarks_XY2), True)[0:2]\n",
    "        result = cv2.warpAffine(source, M, (64, 64), borderMode=cv2.BORDER_REPLICATE)  \n",
    "        mask = np.stack([alpha, alpha, alpha], axis=2)\n",
    "        assert len(mask.shape) == 3, \"len(mask.shape) is \" + str(len(mask.shape))\n",
    "        mask = cv2.warpAffine(mask, M, (64, 64), borderMode=cv2.BORDER_REPLICATE) \n",
    "        prev_landmarks_XY1, prev_landmarks_XY2 = landmarks_XY1, landmarks_XY2\n",
    "        return result, mask[:,:,0].astype(np.float32)\n",
    "    except:\n",
    "        return source, alpha\n",
    "\n",
    "def process_video(input_img): \n",
    "    global prev_x0, prev_x1, prev_y0, prev_y1\n",
    "    global frames      \n",
    "    global pnet, rnet, onet\n",
    "    \"\"\"The following if statement is meant to solve a bug that has an unknow cause.\"\"\"\n",
    "    if frames <= 2:\n",
    "        with tf.variable_scope('pnet2', reuse=True):\n",
    "            pnet2 = None\n",
    "            data = tf.placeholder(tf.float32, (None,None,None,3), 'input')\n",
    "            pnet2 = mtcnn_detect_face.PNet({'data':data})\n",
    "            pnet2.load(os.path.join(\"./mtcnn_weights/\", 'det1.npy'), sess)\n",
    "        with tf.variable_scope('rnet2', reuse=True):\n",
    "            rnet2 = None\n",
    "            data = tf.placeholder(tf.float32, (None,24,24,3), 'input')\n",
    "            rnet2 = mtcnn_detect_face.RNet({'data':data})\n",
    "            rnet2.load(os.path.join(\"./mtcnn_weights/\", 'det2.npy'), sess)\n",
    "        with tf.variable_scope('onet2', reuse=True):\n",
    "            onet2 = None\n",
    "            data = tf.placeholder(tf.float32, (None,48,48,3), 'input')\n",
    "            onet2 = mtcnn_detect_face.ONet({'data':data})\n",
    "            onet2.load(os.path.join(\"./mtcnn_weights/\", 'det3.npy'), sess)\n",
    "        pnet = K.function([pnet2.layers['data']],\n",
    "                          [pnet2.layers['conv4-2'], \n",
    "                           pnet2.layers['prob1']])\n",
    "        rnet = K.function([rnet2.layers['data']],\n",
    "                          [rnet2.layers['conv5-2'], \n",
    "                           rnet2.layers['prob1']])\n",
    "        onet = K.function([onet2.layers['data']], \n",
    "                          [onet2.layers['conv6-2'], \n",
    "                           onet2.layers['conv6-3'], \n",
    "                           onet2.layers['prob1']])\n",
    "    \"\"\"Get face bboxes\"\"\"\n",
    "    image = input_img\n",
    "    faces = get_faces_bbox(image) # faces: face bbox coord\n",
    "    \n",
    "    \"\"\"Init.\"\"\"\n",
    "    if len(faces) == 0:\n",
    "        comb_img = get_init_comb_img(input_img)\n",
    "        triple_img = get_init_triple_img(input_img, no_face=True)\n",
    "    else:\n",
    "        _ = remove_overlaps(faces) # Has non-max suppression already been implemented in MTCNN?        \n",
    "    mask_map = get_init_mask_map(image)\n",
    "    comb_img = get_init_comb_img(input_img)\n",
    "    best_conf_score = 0\n",
    "    \n",
    "    \"\"\"Process detected faces\"\"\"\n",
    "    for (x0, y1, x1, y0, conf_score) in faces: \n",
    "        \"\"\"Apply moving average bounding box\"\"\"        \n",
    "        if use_smoothed_bbox:\n",
    "            if frames != 0 and conf_score >= best_conf_score:\n",
    "                x0, x1, y0, y1 = get_smoothed_coord(x0, x1, y0, y1, \n",
    "                                                    image.shape, \n",
    "                                                    ratio=0.65 if use_kalman_filter else bbox_moving_avg_coef)\n",
    "                set_global_coord(x0, x1, y0, y1)\n",
    "                best_conf_score = conf_score\n",
    "                frames += 1\n",
    "            elif conf_score <= best_conf_score:\n",
    "                frames += 1\n",
    "            else:\n",
    "                if conf_score >= best_conf_score:\n",
    "                    set_global_coord(x0, x1, y0, y1)\n",
    "                    best_conf_score = conf_score\n",
    "                if use_kalman_filter:\n",
    "                    for i in range(200):\n",
    "                        kf0.predict()\n",
    "                        kf1.predict()\n",
    "                frames += 1\n",
    "        \n",
    "        \"\"\"ROI params\"\"\"\n",
    "        h = int(x1 - x0)\n",
    "        w = int(y1 - y0)\n",
    "        roi_coef_h = 25\n",
    "        roi_coef_w = 25\n",
    "        roi_x0, roi_x1, roi_y0, roi_y1 = int(x0+h//roi_coef_h), int(x1-h//roi_coef_h), int(y0+w//roi_coef_w), int(y1-w//roi_coef_w)            \n",
    "        cv2_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n",
    "        roi_image = cv2_img[roi_x0:roi_x1,roi_y0:roi_y1,:]\n",
    "        roi_size = roi_image.shape  \n",
    "        \n",
    "        \"\"\"Feed face image into generator\"\"\"\n",
    "        ae_input = cv2.resize(roi_image, (64,64))/255. * 2 - 1  \n",
    "        result, result_a, result_bgr = generate_face(ae_input, path_func, roi_size, roi_image)\n",
    "        \n",
    "        \"\"\"Apply test time augmentation\"\"\"\n",
    "        if use_TTA:\n",
    "            result_flip, result_a_flip, result_bgr_flip = generate_face(ae_input[:,::-1,:], \n",
    "                                                                        path_func, roi_size, \n",
    "                                                                        roi_image[:,::-1,:])\n",
    "            result = (result.astype(np.float32) + result_flip[:,::-1,:].astype(np.float32))/2\n",
    "            result_a = (result_a.astype(np.float32) + result_a_flip[:,::-1,:].astype(np.float32))/2\n",
    "            result_bgr = (result_bgr.astype(np.float32) + result_bgr_flip[:,::-1,:].astype(np.float32))/2          \n",
    "            result = result.astype('uint8')\n",
    "            result_a = result_a.astype('uint8')\n",
    "            result = get_tta_masked_result(result_bgr, result_a, roi_image)\n",
    "            \n",
    "        \"\"\"Post processing\"\"\"\n",
    "        if conf_score >= best_conf_score:\n",
    "            mask_map[roi_x0:roi_x1,roi_y0:roi_y1,:] = result_a\n",
    "            mask_map = np.clip(mask_map + .15 * input_img, 0, 255)     \n",
    "        else:\n",
    "            mask_map[roi_x0:roi_x1,roi_y0:roi_y1,:] += result_a\n",
    "            mask_map = np.clip(mask_map, 0, 255)        \n",
    "        if use_smoothed_mask:\n",
    "            mask = get_mask(roi_image, h, w)\n",
    "            roi_rgb = cv2.cvtColor(roi_image, cv2.COLOR_BGR2RGB)\n",
    "            smoothed_result = mask/255 * result + (1-mask/255) * roi_rgb\n",
    "            comb_img[roi_x0:roi_x1, input_img.shape[1]+roi_y0:input_img.shape[1]+roi_y1,:] = smoothed_result\n",
    "        else:\n",
    "            comb_img[roi_x0:roi_x1, input_img.shape[1]+roi_y0:input_img.shape[1]+roi_y1,:] = result\n",
    "            \n",
    "        triple_img = get_init_triple_img(input_img)\n",
    "        triple_img[:, :input_img.shape[1]*2, :] = comb_img\n",
    "        triple_img[:, input_img.shape[1]*2:, :] = mask_map\n",
    "    \n",
    "    \"\"\"Return frame result\"\"\"\n",
    "    global output_type\n",
    "    if output_type == 1:\n",
    "        return comb_img[:, input_img.shape[1]:, :]  # return only result image\n",
    "    elif output_type == 2:\n",
    "        return comb_img  # return input and result image combined as one\n",
    "    elif output_type == 3:\n",
    "        return triple_img #return input,result and mask heatmap image combined as one"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8N9OciuACMQk"
   },
   "source": [
    "## Video conversion config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "yBSdHX1mi9ZA"
   },
   "outputs": [],
   "source": [
    "use_kalman_filter = True\n",
    "if use_kalman_filter:\n",
    "    noise_coef = 8e-2 # Increase by 10x if tracking is slow. \n",
    "    kf0 = kalmanfilter_init(noise_coef)\n",
    "    kf1 = kalmanfilter_init(noise_coef)\n",
    "else:\n",
    "    bbox_moving_avg_coef = 0.65\n",
    "    \n",
    "video_scaling_offset = 0 \n",
    "manually_downscale = False\n",
    "manual_downscale_factor = int(2) # should be an positive integer\n",
    "use_color_correction = False\n",
    "use_landmark_match = False # Under developement, This is not functioning.\n",
    "use_TTA = False # test time augmentation\n",
    "use_FCN_mask = True\n",
    "merge_mask_mode = \"adaptive\" # avg, max, adaptive\n",
    "\n",
    "# ========== Change the following line for different output type==========\n",
    "# Output type: \n",
    "#    1. [ result ] \n",
    "#    2. [ source | result ] \n",
    "#    3. [ source | result | mask ]\n",
    "global output_type\n",
    "output_type = 3\n",
    "\n",
    "# Detection threshold:  a float point between 0 and 1. Decrease this value if faces are missed.\n",
    "global detec_threshold\n",
    "detec_threshold = 0.7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "u94-s5_yCMQp"
   },
   "source": [
    "## Start video conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "_IhPUzAPCMQq"
   },
   "outputs": [],
   "source": [
    "# Variables for smoothing bounding box\n",
    "global prev_x0, prev_x1, prev_y0, prev_y1\n",
    "global frames\n",
    "global prev_pnts1, prev_pnts2\n",
    "prev_x0 = prev_x1 = prev_y0 = prev_y1 = 0\n",
    "frames = 0\n",
    "prev_pnts1 = prev_pnts2 = np.array([])\n",
    "\n",
    "output = 'OUTPUT_VIDEO.mp4'\n",
    "clip1 = VideoFileClip(\"INPUT_VIDEO.mp4\")\n",
    "clip = clip1.fl_image(process_video)#.subclip(1, 10) #NOTE: this function expects color images!!\n",
    "%time clip.write_videofile(output, audio=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "dhXbNcXxLN9K"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Xb_5RmKdJFfm"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "kj8S-YAO_653"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "WesFh8C-_65_"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "default_view": {},
   "name": "FaceSwap_GAN_v2.1_train.ipynb",
   "provenance": [],
   "version": "0.3.2",
   "views": {}
  },
  "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}