DepthMapReconstruction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "sys.path.append(os.getcwd())\n",
    "\n",
    "import random\n",
    "import time\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import sklearn.datasets\n",
    "\n",
    "MODE = 'wgan-gp' # wgan or wgan-gp\n",
    "DATASET = 'colins_private_gallery' #'8gaussians' # 8gaussians, 25gaussians, swissroll\n",
    "DIM = 512 # Model dimensionality\n",
    "FIXED_GENERATOR = False # whether to hold the generator fixed at real data plus\n",
    "                        # Gaussian noise, as in the plots in the paper\n",
    "LAMBDA = .1 # Smaller lambda makes things faster for toy tasks, but isn't\n",
    "            # necessary if you increase CRITIC_ITERS enough\n",
    "CRITIC_ITERS = 5 # How many critic iterations per generator iteration\n",
    "BATCH_SIZE = 256 # Batch size\n",
    "ITERS = 100000 # how many generator iterations to train for\n",
    "\n",
    "lib.print_model_settings(locals().copy())\n",
    "\n",
    "def ReLULayer(name, n_in, n_out, inputs):\n",
    "    output = lib.ops.linear.Linear(\n",
    "        name+'.Linear',\n",
    "        n_in,\n",
    "        n_out,\n",
    "        inputs,\n",
    "        initialization='he'\n",
    "    )\n",
    "    output = tf.nn.relu(output)\n",
    "    return output\n",
    "\n",
    "def Generator(n_samples, real_data):\n",
    "    if FIXED_GENERATOR:\n",
    "        return real_data + (1.*tf.random_normal(tf.shape(real_data)))\n",
    "    else:\n",
    "        noise = tf.random_normal([n_samples, 2])\n",
    "        output = ReLULayer('Generator.1', 2, DIM, noise)\n",
    "        output = ReLULayer('Generator.2', DIM, DIM, output)\n",
    "        output = ReLULayer('Generator.3', DIM, DIM, output)\n",
    "        output = lib.ops.linear.Linear('Generator.4', DIM, 2, output)\n",
    "        return output\n",
    "\n",
    "def Discriminator(inputs):\n",
    "    output = ReLULayer('Discriminator.1', 2, DIM, inputs)\n",
    "    output = ReLULayer('Discriminator.2', DIM, DIM, output)\n",
    "    output = ReLULayer('Discriminator.3', DIM, DIM, output)\n",
    "    output = lib.ops.linear.Linear('Discriminator.4', DIM, 1, output)\n",
    "    return tf.reshape(output, [-1])\n",
    "\n",
    "real_data = tf.placeholder(tf.float32, shape=[None, 2])\n",
    "fake_data = Generator(BATCH_SIZE, real_data)\n",
    "\n",
    "disc_real = Discriminator(real_data)\n",
    "disc_fake = Discriminator(fake_data)\n",
    "\n",
    "# WGAN loss\n",
    "disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)\n",
    "gen_cost = -tf.reduce_mean(disc_fake)\n",
    "\n",
    "# WGAN gradient penalty\n",
    "if MODE == 'wgan-gp':\n",
    "    alpha = tf.random_uniform(\n",
    "        shape=[BATCH_SIZE,1], \n",
    "        minval=0.,\n",
    "        maxval=1.\n",
    "    )\n",
    "    interpolates = alpha*real_data + ((1-alpha)*fake_data)\n",
    "    disc_interpolates = Discriminator(interpolates)\n",
    "    gradients = tf.gradients(disc_interpolates, [interpolates])[0]\n",
    "    slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))\n",
    "    gradient_penalty = tf.reduce_mean((slopes-1)**2)\n",
    " \n",
    "    disc_cost += LAMBDA*gradient_penalty\n",
    "\n",
    "disc_params = lib.params_with_name('Discriminator')\n",
    "gen_params = lib.params_with_name('Generator')\n",
    "\n",
    "if MODE == 'wgan-gp':\n",
    "    disc_train_op = tf.train.AdamOptimizer(\n",
    "        learning_rate=1e-4, \n",
    "        beta1=0.5, \n",
    "        beta2=0.9\n",
    "    ).minimize(\n",
    "        disc_cost, \n",
    "        var_list=disc_params\n",
    "    )\n",
    "    if len(gen_params) > 0:\n",
    "        gen_train_op = tf.train.AdamOptimizer(\n",
    "            learning_rate=1e-4, \n",
    "            beta1=0.5, \n",
    "            beta2=0.9\n",
    "        ).minimize(\n",
    "            gen_cost, \n",
    "            var_list=gen_params\n",
    "        )\n",
    "    else:\n",
    "        gen_train_op = tf.no_op()\n",
    "\n",
    "else:\n",
    "    disc_train_op = tf.train.RMSPropOptimizer(learning_rate=5e-5).minimize(\n",
    "        disc_cost, \n",
    "        var_list=disc_params\n",
    "    )\n",
    "    if len(gen_params) > 0:\n",
    "        gen_train_op = tf.train.RMSPropOptimizer(learning_rate=5e-5).minimize(\n",
    "            gen_cost, \n",
    "            var_list=gen_params\n",
    "        )\n",
    "    else:\n",
    "        gen_train_op = tf.no_op()\n",
    "\n",
    "\n",
    "    # Build an op to do the weight clipping\n",
    "    clip_ops = []\n",
    "    for var in disc_params:\n",
    "        clip_bounds = [-.01, .01]\n",
    "        clip_ops.append(\n",
    "            tf.assign(\n",
    "                var, \n",
    "                tf.clip_by_value(var, clip_bounds[0], clip_bounds[1])\n",
    "            )\n",
    "        )\n",
    "    clip_disc_weights = tf.group(*clip_ops)\n",
    "\n",
    "print \"Generator params:\"\n",
    "for var in lib.params_with_name('Generator'):\n",
    "    print \"\\t{}\\t{}\".format(var.name, var.get_shape())\n",
    "print \"Discriminator params:\"\n",
    "for var in lib.params_with_name('Discriminator'):\n",
    "    print \"\\t{}\\t{}\".format(var.name, var.get_shape())\n",
    "\n",
    "frame_index = [0]\n",
    "def generate_image(true_dist):\n",
    "    \"\"\"\n",
    "    Generates and saves a plot of the true distribution, the generator, and the\n",
    "    critic.\n",
    "    \"\"\"\n",
    "    N_POINTS = 128\n",
    "    RANGE = 3\n",
    "\n",
    "    points = np.zeros((N_POINTS, N_POINTS, 2), dtype='float32')\n",
    "    points[:,:,0] = np.linspace(-RANGE, RANGE, N_POINTS)[:,None]\n",
    "    points[:,:,1] = np.linspace(-RANGE, RANGE, N_POINTS)[None,:]\n",
    "    points = points.reshape((-1,2))\n",
    "    samples, disc_map = session.run(\n",
    "        [fake_data, disc_real], \n",
    "        feed_dict={real_data:points}\n",
    "    )\n",
    "    disc_map = session.run(disc_real, feed_dict={real_data:points})\n",
    "\n",
    "    plt.clf()\n",
    "\n",
    "    x = y = np.linspace(-RANGE, RANGE, N_POINTS)\n",
    "    plt.contour(x,y,disc_map.reshape((len(x), len(y))).transpose())\n",
    "\n",
    "    plt.scatter(true_dist[:, 0], true_dist[:, 1], c='orange',  marker='+')\n",
    "    plt.scatter(samples[:, 0],    samples[:, 1],    c='green', marker='+')\n",
    "\n",
    "    plt.savefig('frame'+str(frame_index[0])+'.jpg')\n",
    "    frame_index[0] += 1\n",
    "\n",
    "# Dataset iterator\n",
    "def inf_train_gen():\n",
    "    if DATASET == '25gaussians':\n",
    "    \n",
    "        dataset = []\n",
    "        for i in xrange(100000/25):\n",
    "            for x in xrange(-2, 3):\n",
    "                for y in xrange(-2, 3):\n",
    "                    point = np.random.randn(2)*0.05\n",
    "                    point[0] += 2*x\n",
    "                    point[1] += 2*y\n",
    "                    dataset.append(point)\n",
    "        dataset = np.array(dataset, dtype='float32')\n",
    "        np.random.shuffle(dataset)\n",
    "        dataset /= 2.828 # stdev\n",
    "        while True:\n",
    "            for i in xrange(len(dataset)/BATCH_SIZE):\n",
    "                yield dataset[i*BATCH_SIZE:(i+1)*BATCH_SIZE]\n",
    "\n",
    "    elif DATASET == 'swissroll':\n",
    "\n",
    "        while True:\n",
    "            data = sklearn.datasets.make_swiss_roll(\n",
    "                n_samples=BATCH_SIZE, \n",
    "                noise=0.25\n",
    "            )[0]\n",
    "            data = data.astype('float32')[:, [0, 2]]\n",
    "            data /= 7.5 # stdev plus a little\n",
    "            yield data\n",
    "\n",
    "    elif DATASET == '8gaussians':\n",
    "    \n",
    "        scale = 2.\n",
    "        centers = [\n",
    "            (1,0),\n",
    "            (-1,0),\n",
    "            (0,1),\n",
    "            (0,-1),\n",
    "            (1./np.sqrt(2), 1./np.sqrt(2)),\n",
    "            (1./np.sqrt(2), -1./np.sqrt(2)),\n",
    "            (-1./np.sqrt(2), 1./np.sqrt(2)),\n",
    "            (-1./np.sqrt(2), -1./np.sqrt(2))\n",
    "        ]\n",
    "        centers = [(scale*x,scale*y) for x,y in centers]\n",
    "        while True:\n",
    "            dataset = []\n",
    "            for i in xrange(BATCH_SIZE):\n",
    "                point = np.random.randn(2)*.02\n",
    "                center = random.choice(centers)\n",
    "                point[0] += center[0]\n",
    "                point[1] += center[1]\n",
    "                dataset.append(point)\n",
    "            dataset = np.array(dataset, dtype='float32')\n",
    "            dataset /= 1.414 # stdev\n",
    "            yield dataset\n",
    "\n",
    "# Train loop!\n",
    "with tf.Session() as session:\n",
    "    session.run(tf.initialize_all_variables())\n",
    "    gen = inf_train_gen()\n",
    "    for iteration in xrange(ITERS):\n",
    "        # Train generator\n",
    "        if iteration > 0:\n",
    "            _ = session.run(gen_train_op)\n",
    "        # Train critic\n",
    "        for i in xrange(CRITIC_ITERS):\n",
    "            _data = gen.next()\n",
    "            _disc_cost, _ = session.run(\n",
    "                [disc_cost, disc_train_op],\n",
    "                feed_dict={real_data: _data}\n",
    "            )\n",
    "            if MODE == 'wgan':\n",
    "                _ = session.run([clip_disc_weights])\n",
    "        # Write logs and save samples\n",
    "        lib.plot.plot('disc cost', _disc_cost)\n",
    "        if iteration % 100 == 99:\n",
    "            lib.plot.flush()\n",
    "            generate_image(_data)\n",
    "        lib.plot.tick()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import argparse\n",
    "import os\n",
    "import json\n",
    "import glob\n",
    "import random\n",
    "import collections\n",
    "import math\n",
    "import time\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--input_dir\", help=\"path to folder containing images\")\n",
    "parser.add_argument(\"--mode\", required=True, choices=[\"train\", \"test\", \"export\"])\n",
    "parser.add_argument(\"--output_dir\", required=True, help=\"where to put output files\")\n",
    "parser.add_argument(\"--seed\", type=int)\n",
    "parser.add_argument(\"--checkpoint\", default=None, help=\"directory with checkpoint to resume training from or use for testing\")\n",
    "\n",
    "parser.add_argument(\"--max_steps\", type=int, help=\"number of training steps (0 to disable)\")\n",
    "parser.add_argument(\"--max_epochs\", type=int, help=\"number of training epochs\")\n",
    "parser.add_argument(\"--summary_freq\", type=int, default=100, help=\"update summaries every summary_freq steps\")\n",
    "parser.add_argument(\"--progress_freq\", type=int, default=50, help=\"display progress every progress_freq steps\")\n",
    "parser.add_argument(\"--trace_freq\", type=int, default=0, help=\"trace execution every trace_freq steps\")\n",
    "parser.add_argument(\"--display_freq\", type=int, default=0, help=\"write current training images every display_freq steps\")\n",
    "parser.add_argument(\"--save_freq\", type=int, default=5000, help=\"save model every save_freq steps, 0 to disable\")\n",
    "\n",
    "parser.add_argument(\"--separable_conv\", action=\"store_true\", help=\"use separable convolutions in the generator\")\n",
    "parser.add_argument(\"--aspect_ratio\", type=float, default=1.0, help=\"aspect ratio of output images (width/height)\")\n",
    "parser.add_argument(\"--lab_colorization\", action=\"store_true\", help=\"split input image into brightness (A) and color (B)\")\n",
    "parser.add_argument(\"--batch_size\", type=int, default=1, help=\"number of images in batch\")\n",
    "parser.add_argument(\"--which_direction\", type=str, default=\"AtoB\", choices=[\"AtoB\", \"BtoA\"])\n",
    "parser.add_argument(\"--ngf\", type=int, default=64, help=\"number of generator filters in first conv layer\")\n",
    "parser.add_argument(\"--ndf\", type=int, default=64, help=\"number of discriminator filters in first conv layer\")\n",
    "parser.add_argument(\"--scale_size\", type=int, default=286, help=\"scale images to this size before cropping to 256x256\")\n",
    "parser.add_argument(\"--flip\", dest=\"flip\", action=\"store_true\", help=\"flip images horizontally\")\n",
    "parser.add_argument(\"--no_flip\", dest=\"flip\", action=\"store_false\", help=\"don't flip images horizontally\")\n",
    "parser.set_defaults(flip=True)\n",
    "parser.add_argument(\"--lr\", type=float, default=0.0002, help=\"initial learning rate for adam\")\n",
    "parser.add_argument(\"--beta1\", type=float, default=0.5, help=\"momentum term of adam\")\n",
    "parser.add_argument(\"--l1_weight\", type=float, default=100.0, help=\"weight on L1 term for generator gradient\")\n",
    "parser.add_argument(\"--gan_weight\", type=float, default=1.0, help=\"weight on GAN term for generator gradient\")\n",
    "\n",
    "# export options\n",
    "parser.add_argument(\"--output_filetype\", default=\"png\", choices=[\"png\", \"jpeg\"])\n",
    "a = parser.parse_args()\n",
    "\n",
    "EPS = 1e-12\n",
    "CROP_SIZE = 256\n",
    "\n",
    "Examples = collections.namedtuple(\"Examples\", \"paths, inputs, targets, count, steps_per_epoch\")\n",
    "Model = collections.namedtuple(\"Model\", \"outputs, predict_real, predict_fake, discrim_loss, discrim_grads_and_vars, gen_loss_GAN, gen_loss_L1, gen_grads_and_vars, train\")\n",
    "\n",
    "\n",
    "def preprocess(image):\n",
    "    with tf.name_scope(\"preprocess\"):\n",
    "        # [0, 1] => [-1, 1]\n",
    "        return image * 2 - 1\n",
    "\n",
    "\n",
    "def deprocess(image):\n",
    "    with tf.name_scope(\"deprocess\"):\n",
    "        # [-1, 1] => [0, 1]\n",
    "        return (image + 1) / 2\n",
    "\n",
    "\n",
    "def preprocess_lab(lab):\n",
    "    with tf.name_scope(\"preprocess_lab\"):\n",
    "        L_chan, a_chan, b_chan = tf.unstack(lab, axis=2)\n",
    "        # L_chan: black and white with input range [0, 100]\n",
    "        # a_chan/b_chan: color channels with input range ~[-110, 110], not exact\n",
    "        # [0, 100] => [-1, 1],  ~[-110, 110] => [-1, 1]\n",
    "        return [L_chan / 50 - 1, a_chan / 110, b_chan / 110]\n",
    "\n",
    "\n",
    "def deprocess_lab(L_chan, a_chan, b_chan):\n",
    "    with tf.name_scope(\"deprocess_lab\"):\n",
    "        # this is axis=3 instead of axis=2 because we process individual images but deprocess batches\n",
    "        return tf.stack([(L_chan + 1) / 2 * 100, a_chan * 110, b_chan * 110], axis=3)\n",
    "\n",
    "\n",
    "def augment(image, brightness):\n",
    "    # (a, b) color channels, combine with L channel and convert to rgb\n",
    "    a_chan, b_chan = tf.unstack(image, axis=3)\n",
    "    L_chan = tf.squeeze(brightness, axis=3)\n",
    "    lab = deprocess_lab(L_chan, a_chan, b_chan)\n",
    "    rgb = lab_to_rgb(lab)\n",
    "    return rgb\n",
    "\n",
    "\n",
    "def discrim_conv(batch_input, out_channels, stride):\n",
    "    padded_input = tf.pad(batch_input, [[0, 0], [1, 1], [1, 1], [0, 0]], mode=\"CONSTANT\")\n",
    "    return tf.layers.conv2d(padded_input, out_channels, kernel_size=4, strides=(stride, stride), padding=\"valid\", kernel_initializer=tf.random_normal_initializer(0, 0.02))\n",
    "\n",
    "\n",
    "def gen_conv(batch_input, out_channels):\n",
    "    # [batch, in_height, in_width, in_channels] => [batch, out_height, out_width, out_channels]\n",
    "    initializer = tf.random_normal_initializer(0, 0.02)\n",
    "    if a.separable_conv:\n",
    "        return tf.layers.separable_conv2d(batch_input, out_channels, kernel_size=4, strides=(2, 2), padding=\"same\", depthwise_initializer=initializer, pointwise_initializer=initializer)\n",
    "    else:\n",
    "        return tf.layers.conv2d(batch_input, out_channels, kernel_size=4, strides=(2, 2), padding=\"same\", kernel_initializer=initializer)\n",
    "\n",
    "\n",
    "def gen_deconv(batch_input, out_channels):\n",
    "    # [batch, in_height, in_width, in_channels] => [batch, out_height, out_width, out_channels]\n",
    "    initializer = tf.random_normal_initializer(0, 0.02)\n",
    "    if a.separable_conv:\n",
    "        _b, h, w, _c = batch_input.shape\n",
    "        resized_input = tf.image.resize_images(batch_input, [h * 2, w * 2], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
    "        return tf.layers.separable_conv2d(resized_input, out_channels, kernel_size=4, strides=(1, 1), padding=\"same\", depthwise_initializer=initializer, pointwise_initializer=initializer)\n",
    "    else:\n",
    "        return tf.layers.conv2d_transpose(batch_input, out_channels, kernel_size=4, strides=(2, 2), padding=\"same\", kernel_initializer=initializer)\n",
    "\n",
    "\n",
    "def lrelu(x, a):\n",
    "    with tf.name_scope(\"lrelu\"):\n",
    "        # adding these together creates the leak part and linear part\n",
    "        # then cancels them out by subtracting/adding an absolute value term\n",
    "        # leak: a*x/2 - a*abs(x)/2\n",
    "        # linear: x/2 + abs(x)/2\n",
    "\n",
    "        # this block looks like it has 2 inputs on the graph unless we do this\n",
    "        x = tf.identity(x)\n",
    "        return (0.5 * (1 + a)) * x + (0.5 * (1 - a)) * tf.abs(x)\n",
    "\n",
    "\n",
    "def batchnorm(inputs):\n",
    "    return tf.layers.batch_normalization(inputs, axis=3, epsilon=1e-5, momentum=0.1, training=True, gamma_initializer=tf.random_normal_initializer(1.0, 0.02))\n",
    "\n",
    "\n",
    "def check_image(image):\n",
    "    assertion = tf.assert_equal(tf.shape(image)[-1], 3, message=\"image must have 3 color channels\")\n",
    "    with tf.control_dependencies([assertion]):\n",
    "        image = tf.identity(image)\n",
    "\n",
    "    if image.get_shape().ndims not in (3, 4):\n",
    "        raise ValueError(\"image must be either 3 or 4 dimensions\")\n",
    "\n",
    "    # make the last dimension 3 so that you can unstack the colors\n",
    "    shape = list(image.get_shape())\n",
    "    shape[-1] = 3\n",
    "    image.set_shape(shape)\n",
    "    return image\n",
    "\n",
    "# based on https://github.com/torch/image/blob/9f65c30167b2048ecbe8b7befdc6b2d6d12baee9/generic/image.c\n",
    "def rgb_to_lab(srgb):\n",
    "    with tf.name_scope(\"rgb_to_lab\"):\n",
    "        srgb = check_image(srgb)\n",
    "        srgb_pixels = tf.reshape(srgb, [-1, 3])\n",
    "\n",
    "        with tf.name_scope(\"srgb_to_xyz\"):\n",
    "            linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)\n",
    "            exponential_mask = tf.cast(srgb_pixels > 0.04045, dtype=tf.float32)\n",
    "            rgb_pixels = (srgb_pixels / 12.92 * linear_mask) + (((srgb_pixels + 0.055) / 1.055) ** 2.4) * exponential_mask\n",
    "            rgb_to_xyz = tf.constant([\n",
    "                #    X        Y          Z\n",
    "                [0.412453, 0.212671, 0.019334], # R\n",
    "                [0.357580, 0.715160, 0.119193], # G\n",
    "                [0.180423, 0.072169, 0.950227], # B\n",
    "            ])\n",
    "            xyz_pixels = tf.matmul(rgb_pixels, rgb_to_xyz)\n",
    "\n",
    "        # https://en.wikipedia.org/wiki/Lab_color_space#CIELAB-CIEXYZ_conversions\n",
    "        with tf.name_scope(\"xyz_to_cielab\"):\n",
    "            # convert to fx = f(X/Xn), fy = f(Y/Yn), fz = f(Z/Zn)\n",
    "\n",
    "            # normalize for D65 white point\n",
    "            xyz_normalized_pixels = tf.multiply(xyz_pixels, [1/0.950456, 1.0, 1/1.088754])\n",
    "\n",
    "            epsilon = 6/29\n",
    "            linear_mask = tf.cast(xyz_normalized_pixels <= (epsilon**3), dtype=tf.float32)\n",
    "            exponential_mask = tf.cast(xyz_normalized_pixels > (epsilon**3), dtype=tf.float32)\n",
    "            fxfyfz_pixels = (xyz_normalized_pixels / (3 * epsilon**2) + 4/29) * linear_mask + (xyz_normalized_pixels ** (1/3)) * exponential_mask\n",
    "\n",
    "            # convert to lab\n",
    "            fxfyfz_to_lab = tf.constant([\n",
    "                #  l       a       b\n",
    "                [  0.0,  500.0,    0.0], # fx\n",
    "                [116.0, -500.0,  200.0], # fy\n",
    "                [  0.0,    0.0, -200.0], # fz\n",
    "            ])\n",
    "            lab_pixels = tf.matmul(fxfyfz_pixels, fxfyfz_to_lab) + tf.constant([-16.0, 0.0, 0.0])\n",
    "\n",
    "        return tf.reshape(lab_pixels, tf.shape(srgb))\n",
    "\n",
    "\n",
    "def lab_to_rgb(lab):\n",
    "    with tf.name_scope(\"lab_to_rgb\"):\n",
    "        lab = check_image(lab)\n",
    "        lab_pixels = tf.reshape(lab, [-1, 3])\n",
    "\n",
    "        # https://en.wikipedia.org/wiki/Lab_color_space#CIELAB-CIEXYZ_conversions\n",
    "        with tf.name_scope(\"cielab_to_xyz\"):\n",
    "            # convert to fxfyfz\n",
    "            lab_to_fxfyfz = tf.constant([\n",
    "                #   fx      fy        fz\n",
    "                [1/116.0, 1/116.0,  1/116.0], # l\n",
    "                [1/500.0,     0.0,      0.0], # a\n",
    "                [    0.0,     0.0, -1/200.0], # b\n",
    "            ])\n",
    "            fxfyfz_pixels = tf.matmul(lab_pixels + tf.constant([16.0, 0.0, 0.0]), lab_to_fxfyfz)\n",
    "\n",
    "            # convert to xyz\n",
    "            epsilon = 6/29\n",
    "            linear_mask = tf.cast(fxfyfz_pixels <= epsilon, dtype=tf.float32)\n",
    "            exponential_mask = tf.cast(fxfyfz_pixels > epsilon, dtype=tf.float32)\n",
    "            xyz_pixels = (3 * epsilon**2 * (fxfyfz_pixels - 4/29)) * linear_mask + (fxfyfz_pixels ** 3) * exponential_mask\n",
    "\n",
    "            # denormalize for D65 white point\n",
    "            xyz_pixels = tf.multiply(xyz_pixels, [0.950456, 1.0, 1.088754])\n",
    "\n",
    "        with tf.name_scope(\"xyz_to_srgb\"):\n",
    "            xyz_to_rgb = tf.constant([\n",
    "                #     r           g          b\n",
    "                [ 3.2404542, -0.9692660,  0.0556434], # x\n",
    "                [-1.5371385,  1.8760108, -0.2040259], # y\n",
    "                [-0.4985314,  0.0415560,  1.0572252], # z\n",
    "            ])\n",
    "            rgb_pixels = tf.matmul(xyz_pixels, xyz_to_rgb)\n",
    "            # avoid a slightly negative number messing up the conversion\n",
    "            rgb_pixels = tf.clip_by_value(rgb_pixels, 0.0, 1.0)\n",
    "            linear_mask = tf.cast(rgb_pixels <= 0.0031308, dtype=tf.float32)\n",
    "            exponential_mask = tf.cast(rgb_pixels > 0.0031308, dtype=tf.float32)\n",
    "            srgb_pixels = (rgb_pixels * 12.92 * linear_mask) + ((rgb_pixels ** (1/2.4) * 1.055) - 0.055) * exponential_mask\n",
    "\n",
    "        return tf.reshape(srgb_pixels, tf.shape(lab))\n",
    "\n",
    "\n",
    "def load_examples():\n",
    "    if a.input_dir is None or not os.path.exists(a.input_dir):\n",
    "        raise Exception(\"input_dir does not exist\")\n",
    "\n",
    "    input_paths = glob.glob(os.path.join(a.input_dir, \"*.jpg\"))\n",
    "    decode = tf.image.decode_jpeg\n",
    "    if len(input_paths) == 0:\n",
    "        input_paths = glob.glob(os.path.join(a.input_dir, \"*.png\"))\n",
    "        decode = tf.image.decode_png\n",
    "\n",
    "    if len(input_paths) == 0:\n",
    "        raise Exception(\"input_dir contains no image files\")\n",
    "\n",
    "    def get_name(path):\n",
    "        name, _ = os.path.splitext(os.path.basename(path))\n",
    "        return name\n",
    "\n",
    "    # if the image names are numbers, sort by the value rather than asciibetically\n",
    "    # having sorted inputs means that the outputs are sorted in test mode\n",
    "    if all(get_name(path).isdigit() for path in input_paths):\n",
    "        input_paths = sorted(input_paths, key=lambda path: int(get_name(path)))\n",
    "    else:\n",
    "        input_paths = sorted(input_paths)\n",
    "\n",
    "    with tf.name_scope(\"load_images\"):\n",
    "        path_queue = tf.train.string_input_producer(input_paths, shuffle=a.mode == \"train\")\n",
    "        reader = tf.WholeFileReader()\n",
    "        paths, contents = reader.read(path_queue)\n",
    "        raw_input = decode(contents)\n",
    "        raw_input = tf.image.convert_image_dtype(raw_input, dtype=tf.float32)\n",
    "\n",
    "        assertion = tf.assert_equal(tf.shape(raw_input)[2], 3, message=\"image does not have 3 channels\")\n",
    "        with tf.control_dependencies([assertion]):\n",
    "            raw_input = tf.identity(raw_input)\n",
    "\n",
    "        raw_input.set_shape([None, None, 3])\n",
    "\n",
    "        if a.lab_colorization:\n",
    "            # load color and brightness from image, no B image exists here\n",
    "            lab = rgb_to_lab(raw_input)\n",
    "            L_chan, a_chan, b_chan = preprocess_lab(lab)\n",
    "            a_images = tf.expand_dims(L_chan, axis=2)\n",
    "            b_images = tf.stack([a_chan, b_chan], axis=2)\n",
    "        else:\n",
    "            # break apart image pair and move to range [-1, 1]\n",
    "            width = tf.shape(raw_input)[1] # [height, width, channels]\n",
    "            a_images = preprocess(raw_input[:,:width//2,:])\n",
    "            b_images = preprocess(raw_input[:,width//2:,:])\n",
    "\n",
    "    if a.which_direction == \"AtoB\":\n",
    "        inputs, targets = [a_images, b_images]\n",
    "    elif a.which_direction == \"BtoA\":\n",
    "        inputs, targets = [b_images, a_images]\n",
    "    else:\n",
    "        raise Exception(\"invalid direction\")\n",
    "\n",
    "    # synchronize seed for image operations so that we do the same operations to both\n",
    "    # input and output images\n",
    "    seed = random.randint(0, 2**31 - 1)\n",
    "    def transform(image):\n",
    "        r = image\n",
    "        if a.flip:\n",
    "            r = tf.image.random_flip_left_right(r, seed=seed)\n",
    "\n",
    "        # area produces a nice downscaling, but does nearest neighbor for upscaling\n",
    "        # assume we're going to be doing downscaling here\n",
    "        r = tf.image.resize_images(r, [a.scale_size, a.scale_size], method=tf.image.ResizeMethod.AREA)\n",
    "\n",
    "        offset = tf.cast(tf.floor(tf.random_uniform([2], 0, a.scale_size - CROP_SIZE + 1, seed=seed)), dtype=tf.int32)\n",
    "        if a.scale_size > CROP_SIZE:\n",
    "            r = tf.image.crop_to_bounding_box(r, offset[0], offset[1], CROP_SIZE, CROP_SIZE)\n",
    "        elif a.scale_size < CROP_SIZE:\n",
    "            raise Exception(\"scale size cannot be less than crop size\")\n",
    "        return r\n",
    "\n",
    "    with tf.name_scope(\"input_images\"):\n",
    "        input_images = transform(inputs)\n",
    "\n",
    "    with tf.name_scope(\"target_images\"):\n",
    "        target_images = transform(targets)\n",
    "\n",
    "    paths_batch, inputs_batch, targets_batch = tf.train.batch([paths, input_images, target_images], batch_size=a.batch_size)\n",
    "    steps_per_epoch = int(math.ceil(len(input_paths) / a.batch_size))\n",
    "\n",
    "    return Examples(\n",
    "        paths=paths_batch,\n",
    "        inputs=inputs_batch,\n",
    "        targets=targets_batch,\n",
    "        count=len(input_paths),\n",
    "        steps_per_epoch=steps_per_epoch,\n",
    "    )\n",
    "\n",
    "\n",
    "def create_generator(generator_inputs, generator_outputs_channels):\n",
    "    layers = []\n",
    "\n",
    "    # encoder_1: [batch, 256, 256, in_channels] => [batch, 128, 128, ngf]\n",
    "    with tf.variable_scope(\"encoder_1\"):\n",
    "        output = gen_conv(generator_inputs, a.ngf)\n",
    "        layers.append(output)\n",
    "\n",
    "    layer_specs = [\n",
    "        a.ngf * 2, # encoder_2: [batch, 128, 128, ngf] => [batch, 64, 64, ngf * 2]\n",
    "        a.ngf * 4, # encoder_3: [batch, 64, 64, ngf * 2] => [batch, 32, 32, ngf * 4]\n",
    "        a.ngf * 8, # encoder_4: [batch, 32, 32, ngf * 4] => [batch, 16, 16, ngf * 8]\n",
    "        a.ngf * 8, # encoder_5: [batch, 16, 16, ngf * 8] => [batch, 8, 8, ngf * 8]\n",
    "        a.ngf * 8, # encoder_6: [batch, 8, 8, ngf * 8] => [batch, 4, 4, ngf * 8]\n",
    "        a.ngf * 8, # encoder_7: [batch, 4, 4, ngf * 8] => [batch, 2, 2, ngf * 8]\n",
    "        a.ngf * 8, # encoder_8: [batch, 2, 2, ngf * 8] => [batch, 1, 1, ngf * 8]\n",
    "    ]\n",
    "\n",
    "    for out_channels in layer_specs:\n",
    "        with tf.variable_scope(\"encoder_%d\" % (len(layers) + 1)):\n",
    "            rectified = lrelu(layers[-1], 0.2)\n",
    "            # [batch, in_height, in_width, in_channels] => [batch, in_height/2, in_width/2, out_channels]\n",
    "            convolved = gen_conv(rectified, out_channels)\n",
    "            output = batchnorm(convolved)\n",
    "            layers.append(output)\n",
    "\n",
    "    layer_specs = [\n",
    "        (a.ngf * 8, 0.5),   # decoder_8: [batch, 1, 1, ngf * 8] => [batch, 2, 2, ngf * 8 * 2]\n",
    "        (a.ngf * 8, 0.5),   # decoder_7: [batch, 2, 2, ngf * 8 * 2] => [batch, 4, 4, ngf * 8 * 2]\n",
    "        (a.ngf * 8, 0.5),   # decoder_6: [batch, 4, 4, ngf * 8 * 2] => [batch, 8, 8, ngf * 8 * 2]\n",
    "        (a.ngf * 8, 0.0),   # decoder_5: [batch, 8, 8, ngf * 8 * 2] => [batch, 16, 16, ngf * 8 * 2]\n",
    "        (a.ngf * 4, 0.0),   # decoder_4: [batch, 16, 16, ngf * 8 * 2] => [batch, 32, 32, ngf * 4 * 2]\n",
    "        (a.ngf * 2, 0.0),   # decoder_3: [batch, 32, 32, ngf * 4 * 2] => [batch, 64, 64, ngf * 2 * 2]\n",
    "        (a.ngf, 0.0),       # decoder_2: [batch, 64, 64, ngf * 2 * 2] => [batch, 128, 128, ngf * 2]\n",
    "    ]\n",
    "\n",
    "    num_encoder_layers = len(layers)\n",
    "    for decoder_layer, (out_channels, dropout) in enumerate(layer_specs):\n",
    "        skip_layer = num_encoder_layers - decoder_layer - 1\n",
    "        with tf.variable_scope(\"decoder_%d\" % (skip_layer + 1)):\n",
    "            if decoder_layer == 0:\n",
    "                # first decoder layer doesn't have skip connections\n",
    "                # since it is directly connected to the skip_layer\n",
    "                input = layers[-1]\n",
    "            else:\n",
    "                input = tf.concat([layers[-1], layers[skip_layer]], axis=3)\n",
    "\n",
    "            rectified = tf.nn.relu(input)\n",
    "            # [batch, in_height, in_width, in_channels] => [batch, in_height*2, in_width*2, out_channels]\n",
    "            output = gen_deconv(rectified, out_channels)\n",
    "            output = batchnorm(output)\n",
    "\n",
    "            if dropout > 0.0:\n",
    "                output = tf.nn.dropout(output, keep_prob=1 - dropout)\n",
    "\n",
    "            layers.append(output)\n",
    "\n",
    "    # decoder_1: [batch, 128, 128, ngf * 2] => [batch, 256, 256, generator_outputs_channels]\n",
    "    with tf.variable_scope(\"decoder_1\"):\n",
    "        input = tf.concat([layers[-1], layers[0]], axis=3)\n",
    "        rectified = tf.nn.relu(input)\n",
    "        output = gen_deconv(rectified, generator_outputs_channels)\n",
    "        output = tf.tanh(output)\n",
    "        layers.append(output)\n",
    "\n",
    "    return layers[-1]\n",
    "\n",
    "\n",
    "def create_model(inputs, targets):\n",
    "    def create_discriminator(discrim_inputs, discrim_targets):\n",
    "        n_layers = 3\n",
    "        layers = []\n",
    "\n",
    "        # 2x [batch, height, width, in_channels] => [batch, height, width, in_channels * 2]\n",
    "        input = tf.concat([discrim_inputs, discrim_targets], axis=3)\n",
    "\n",
    "        # layer_1: [batch, 256, 256, in_channels * 2] => [batch, 128, 128, ndf]\n",
    "        with tf.variable_scope(\"layer_1\"):\n",
    "            convolved = discrim_conv(input, a.ndf, stride=2)\n",
    "            rectified = lrelu(convolved, 0.2)\n",
    "            layers.append(rectified)\n",
    "\n",
    "        # layer_2: [batch, 128, 128, ndf] => [batch, 64, 64, ndf * 2]\n",
    "        # layer_3: [batch, 64, 64, ndf * 2] => [batch, 32, 32, ndf * 4]\n",
    "        # layer_4: [batch, 32, 32, ndf * 4] => [batch, 31, 31, ndf * 8]\n",
    "        for i in range(n_layers):\n",
    "            with tf.variable_scope(\"layer_%d\" % (len(layers) + 1)):\n",
    "                out_channels = a.ndf * min(2**(i+1), 8)\n",
    "                stride = 1 if i == n_layers - 1 else 2  # last layer here has stride 1\n",
    "                convolved = discrim_conv(layers[-1], out_channels, stride=stride)\n",
    "                normalized = batchnorm(convolved)\n",
    "                rectified = lrelu(normalized, 0.2)\n",
    "                layers.append(rectified)\n",
    "\n",
    "        # layer_5: [batch, 31, 31, ndf * 8] => [batch, 30, 30, 1]\n",
    "        with tf.variable_scope(\"layer_%d\" % (len(layers) + 1)):\n",
    "            convolved = discrim_conv(rectified, out_channels=1, stride=1)\n",
    "            output = tf.sigmoid(convolved)\n",
    "            layers.append(output)\n",
    "\n",
    "        return layers[-1]\n",
    "\n",
    "    with tf.variable_scope(\"generator\"):\n",
    "        out_channels = int(targets.get_shape()[-1])\n",
    "        outputs = create_generator(inputs, out_channels)\n",
    "\n",
    "    # create two copies of discriminator, one for real pairs and one for fake pairs\n",
    "    # they share the same underlying variables\n",
    "    with tf.name_scope(\"real_discriminator\"):\n",
    "        with tf.variable_scope(\"discriminator\"):\n",
    "            # 2x [batch, height, width, channels] => [batch, 30, 30, 1]\n",
    "            predict_real = create_discriminator(inputs, targets)\n",
    "\n",
    "    with tf.name_scope(\"fake_discriminator\"):\n",
    "        with tf.variable_scope(\"discriminator\", reuse=True):\n",
    "            # 2x [batch, height, width, channels] => [batch, 30, 30, 1]\n",
    "            predict_fake = create_discriminator(inputs, outputs)\n",
    "\n",
    "    with tf.name_scope(\"discriminator_loss\"):\n",
    "        # minimizing -tf.log will try to get inputs to 1\n",
    "        # predict_real => 1\n",
    "        # predict_fake => 0\n",
    "        discrim_loss = tf.reduce_mean(-(tf.log(predict_real + EPS) + tf.log(1 - predict_fake + EPS)))\n",
    "\n",
    "    with tf.name_scope(\"generator_loss\"):\n",
    "        # predict_fake => 1\n",
    "        # abs(targets - outputs) => 0\n",
    "        gen_loss_GAN = tf.reduce_mean(-tf.log(predict_fake + EPS))\n",
    "        gen_loss_L1 = tf.reduce_mean(tf.abs(targets - outputs))\n",
    "        gen_loss = gen_loss_GAN * a.gan_weight + gen_loss_L1 * a.l1_weight\n",
    "\n",
    "    with tf.name_scope(\"discriminator_train\"):\n",
    "        discrim_tvars = [var for var in tf.trainable_variables() if var.name.startswith(\"discriminator\")]\n",
    "        discrim_optim = tf.train.AdamOptimizer(a.lr, a.beta1)\n",
    "        discrim_grads_and_vars = discrim_optim.compute_gradients(discrim_loss, var_list=discrim_tvars)\n",
    "        discrim_train = discrim_optim.apply_gradients(discrim_grads_and_vars)\n",
    "\n",
    "    with tf.name_scope(\"generator_train\"):\n",
    "        with tf.control_dependencies([discrim_train]):\n",
    "            gen_tvars = [var for var in tf.trainable_variables() if var.name.startswith(\"generator\")]\n",
    "            gen_optim = tf.train.AdamOptimizer(a.lr, a.beta1)\n",
    "            gen_grads_and_vars = gen_optim.compute_gradients(gen_loss, var_list=gen_tvars)\n",
    "            gen_train = gen_optim.apply_gradients(gen_grads_and_vars)\n",
    "\n",
    "    ema = tf.train.ExponentialMovingAverage(decay=0.99)\n",
    "    update_losses = ema.apply([discrim_loss, gen_loss_GAN, gen_loss_L1])\n",
    "\n",
    "    global_step = tf.train.get_or_create_global_step()\n",
    "    incr_global_step = tf.assign(global_step, global_step+1)\n",
    "\n",
    "    return Model(\n",
    "        predict_real=predict_real,\n",
    "        predict_fake=predict_fake,\n",
    "        discrim_loss=ema.average(discrim_loss),\n",
    "        discrim_grads_and_vars=discrim_grads_and_vars,\n",
    "        gen_loss_GAN=ema.average(gen_loss_GAN),\n",
    "        gen_loss_L1=ema.average(gen_loss_L1),\n",
    "        gen_grads_and_vars=gen_grads_and_vars,\n",
    "        outputs=outputs,\n",
    "        train=tf.group(update_losses, incr_global_step, gen_train),\n",
    "    )\n",
    "\n",
    "\n",
    "def save_images(fetches, step=None):\n",
    "    image_dir = os.path.join(a.output_dir, \"images\")\n",
    "    if not os.path.exists(image_dir):\n",
    "        os.makedirs(image_dir)\n",
    "\n",
    "    filesets = []\n",
    "    for i, in_path in enumerate(fetches[\"paths\"]):\n",
    "        name, _ = os.path.splitext(os.path.basename(in_path.decode(\"utf8\")))\n",
    "        fileset = {\"name\": name, \"step\": step}\n",
    "        for kind in [\"inputs\", \"outputs\", \"targets\"]:\n",
    "            filename = name + \"-\" + kind + \".png\"\n",
    "            if step is not None:\n",
    "                filename = \"%08d-%s\" % (step, filename)\n",
    "            fileset[kind] = filename\n",
    "            out_path = os.path.join(image_dir, filename)\n",
    "            contents = fetches[kind][i]\n",
    "            with open(out_path, \"wb\") as f:\n",
    "                f.write(contents)\n",
    "        filesets.append(fileset)\n",
    "    return filesets\n",
    "\n",
    "\n",
    "def append_index(filesets, step=False):\n",
    "    index_path = os.path.join(a.output_dir, \"index.html\")\n",
    "    if os.path.exists(index_path):\n",
    "        index = open(index_path, \"a\")\n",
    "    else:\n",
    "        index = open(index_path, \"w\")\n",
    "        index.write(\"<html><body><table><tr>\")\n",
    "        if step:\n",
    "            index.write(\"<th>step</th>\")\n",
    "        index.write(\"<th>name</th><th>input</th><th>output</th><th>target</th></tr>\")\n",
    "\n",
    "    for fileset in filesets:\n",
    "        index.write(\"<tr>\")\n",
    "\n",
    "        if step:\n",
    "            index.write(\"<td>%d</td>\" % fileset[\"step\"])\n",
    "        index.write(\"<td>%s</td>\" % fileset[\"name\"])\n",
    "\n",
    "        for kind in [\"inputs\", \"outputs\", \"targets\"]:\n",
    "            index.write(\"<td><img src='images/%s'></td>\" % fileset[kind])\n",
    "\n",
    "        index.write(\"</tr>\")\n",
    "    return index_path\n",
    "\n",
    "\n",
    "def main():\n",
    "    if a.seed is None:\n",
    "        a.seed = random.randint(0, 2**31 - 1)\n",
    "\n",
    "    tf.set_random_seed(a.seed)\n",
    "    np.random.seed(a.seed)\n",
    "    random.seed(a.seed)\n",
    "\n",
    "    if not os.path.exists(a.output_dir):\n",
    "        os.makedirs(a.output_dir)\n",
    "\n",
    "    if a.mode == \"test\" or a.mode == \"export\":\n",
    "        if a.checkpoint is None:\n",
    "            raise Exception(\"checkpoint required for test mode\")\n",
    "\n",
    "        # load some options from the checkpoint\n",
    "        options = {\"which_direction\", \"ngf\", \"ndf\", \"lab_colorization\"}\n",
    "        with open(os.path.join(a.checkpoint, \"options.json\")) as f:\n",
    "            for key, val in json.loads(f.read()).items():\n",
    "                if key in options:\n",
    "                    print(\"loaded\", key, \"=\", val)\n",
    "                    setattr(a, key, val)\n",
    "        # disable these features in test mode\n",
    "        a.scale_size = CROP_SIZE\n",
    "        a.flip = False\n",
    "\n",
    "    for k, v in a._get_kwargs():\n",
    "        print(k, \"=\", v)\n",
    "\n",
    "    with open(os.path.join(a.output_dir, \"options.json\"), \"w\") as f:\n",
    "        f.write(json.dumps(vars(a), sort_keys=True, indent=4))\n",
    "\n",
    "    if a.mode == \"export\":\n",
    "        # export the generator to a meta graph that can be imported later for standalone generation\n",
    "        if a.lab_colorization:\n",
    "            raise Exception(\"export not supported for lab_colorization\")\n",
    "\n",
    "        input = tf.placeholder(tf.string, shape=[1])\n",
    "        input_data = tf.decode_base64(input[0])\n",
    "        input_image = tf.image.decode_png(input_data)\n",
    "\n",
    "        # remove alpha channel if present\n",
    "        input_image = tf.cond(tf.equal(tf.shape(input_image)[2], 4), lambda: input_image[:,:,:3], lambda: input_image)\n",
    "        # convert grayscale to RGB\n",
    "        input_image = tf.cond(tf.equal(tf.shape(input_image)[2], 1), lambda: tf.image.grayscale_to_rgb(input_image), lambda: input_image)\n",
    "\n",
    "        input_image = tf.image.convert_image_dtype(input_image, dtype=tf.float32)\n",
    "        input_image.set_shape([CROP_SIZE, CROP_SIZE, 3])\n",
    "        batch_input = tf.expand_dims(input_image, axis=0)\n",
    "\n",
    "        with tf.variable_scope(\"generator\"):\n",
    "            batch_output = deprocess(create_generator(preprocess(batch_input), 3))\n",
    "\n",
    "        output_image = tf.image.convert_image_dtype(batch_output, dtype=tf.uint8)[0]\n",
    "        if a.output_filetype == \"png\":\n",
    "            output_data = tf.image.encode_png(output_image)\n",
    "        elif a.output_filetype == \"jpeg\":\n",
    "            output_data = tf.image.encode_jpeg(output_image, quality=80)\n",
    "        else:\n",
    "            raise Exception(\"invalid filetype\")\n",
    "        output = tf.convert_to_tensor([tf.encode_base64(output_data)])\n",
    "\n",
    "        key = tf.placeholder(tf.string, shape=[1])\n",
    "        inputs = {\n",
    "            \"key\": key.name,\n",
    "            \"input\": input.name\n",
    "        }\n",
    "        tf.add_to_collection(\"inputs\", json.dumps(inputs))\n",
    "        outputs = {\n",
    "            \"key\":  tf.identity(key).name,\n",
    "            \"output\": output.name,\n",
    "        }\n",
    "        tf.add_to_collection(\"outputs\", json.dumps(outputs))\n",
    "\n",
    "        init_op = tf.global_variables_initializer()\n",
    "        restore_saver = tf.train.Saver()\n",
    "        export_saver = tf.train.Saver()\n",
    "\n",
    "        with tf.Session() as sess:\n",
    "            sess.run(init_op)\n",
    "            print(\"loading model from checkpoint\")\n",
    "            checkpoint = tf.train.latest_checkpoint(a.checkpoint)\n",
    "            restore_saver.restore(sess, checkpoint)\n",
    "            print(\"exporting model\")\n",
    "            export_saver.export_meta_graph(filename=os.path.join(a.output_dir, \"export.meta\"))\n",
    "            export_saver.save(sess, os.path.join(a.output_dir, \"export\"), write_meta_graph=False)\n",
    "\n",
    "        return\n",
    "\n",
    "    examples = load_examples()\n",
    "    print(\"examples count = %d\" % examples.count)\n",
    "\n",
    "    # inputs and targets are [batch_size, height, width, channels]\n",
    "    model = create_model(examples.inputs, examples.targets)\n",
    "\n",
    "    # undo colorization splitting on images that we use for display/output\n",
    "    if a.lab_colorization:\n",
    "        if a.which_direction == \"AtoB\":\n",
    "            # inputs is brightness, this will be handled fine as a grayscale image\n",
    "            # need to augment targets and outputs with brightness\n",
    "            targets = augment(examples.targets, examples.inputs)\n",
    "            outputs = augment(model.outputs, examples.inputs)\n",
    "            # inputs can be deprocessed normally and handled as if they are single channel\n",
    "            # grayscale images\n",
    "            inputs = deprocess(examples.inputs)\n",
    "        elif a.which_direction == \"BtoA\":\n",
    "            # inputs will be color channels only, get brightness from targets\n",
    "            inputs = augment(examples.inputs, examples.targets)\n",
    "            targets = deprocess(examples.targets)\n",
    "            outputs = deprocess(model.outputs)\n",
    "        else:\n",
    "            raise Exception(\"invalid direction\")\n",
    "    else:\n",
    "        inputs = deprocess(examples.inputs)\n",
    "        targets = deprocess(examples.targets)\n",
    "        outputs = deprocess(model.outputs)\n",
    "\n",
    "    def convert(image):\n",
    "        if a.aspect_ratio != 1.0:\n",
    "            # upscale to correct aspect ratio\n",
    "            size = [CROP_SIZE, int(round(CROP_SIZE * a.aspect_ratio))]\n",
    "            image = tf.image.resize_images(image, size=size, method=tf.image.ResizeMethod.BICUBIC)\n",
    "\n",
    "        return tf.image.convert_image_dtype(image, dtype=tf.uint8, saturate=True)\n",
    "\n",
    "    # reverse any processing on images so they can be written to disk or displayed to user\n",
    "    with tf.name_scope(\"convert_inputs\"):\n",
    "        converted_inputs = convert(inputs)\n",
    "\n",
    "    with tf.name_scope(\"convert_targets\"):\n",
    "        converted_targets = convert(targets)\n",
    "\n",
    "    with tf.name_scope(\"convert_outputs\"):\n",
    "        converted_outputs = convert(outputs)\n",
    "\n",
    "    with tf.name_scope(\"encode_images\"):\n",
    "        display_fetches = {\n",
    "            \"paths\": examples.paths,\n",
    "            \"inputs\": tf.map_fn(tf.image.encode_png, converted_inputs, dtype=tf.string, name=\"input_pngs\"),\n",
    "            \"targets\": tf.map_fn(tf.image.encode_png, converted_targets, dtype=tf.string, name=\"target_pngs\"),\n",
    "            \"outputs\": tf.map_fn(tf.image.encode_png, converted_outputs, dtype=tf.string, name=\"output_pngs\"),\n",
    "        }\n",
    "\n",
    "    # summaries\n",
    "    with tf.name_scope(\"inputs_summary\"):\n",
    "        tf.summary.image(\"inputs\", converted_inputs)\n",
    "\n",
    "    with tf.name_scope(\"targets_summary\"):\n",
    "        tf.summary.image(\"targets\", converted_targets)\n",
    "\n",
    "    with tf.name_scope(\"outputs_summary\"):\n",
    "        tf.summary.image(\"outputs\", converted_outputs)\n",
    "\n",
    "    with tf.name_scope(\"predict_real_summary\"):\n",
    "        tf.summary.image(\"predict_real\", tf.image.convert_image_dtype(model.predict_real, dtype=tf.uint8))\n",
    "\n",
    "    with tf.name_scope(\"predict_fake_summary\"):\n",
    "        tf.summary.image(\"predict_fake\", tf.image.convert_image_dtype(model.predict_fake, dtype=tf.uint8))\n",
    "\n",
    "    tf.summary.scalar(\"discriminator_loss\", model.discrim_loss)\n",
    "    tf.summary.scalar(\"generator_loss_GAN\", model.gen_loss_GAN)\n",
    "    tf.summary.scalar(\"generator_loss_L1\", model.gen_loss_L1)\n",
    "\n",
    "    for var in tf.trainable_variables():\n",
    "        tf.summary.histogram(var.op.name + \"/values\", var)\n",
    "\n",
    "    for grad, var in model.discrim_grads_and_vars + model.gen_grads_and_vars:\n",
    "        tf.summary.histogram(var.op.name + \"/gradients\", grad)\n",
    "\n",
    "    with tf.name_scope(\"parameter_count\"):\n",
    "        parameter_count = tf.reduce_sum([tf.reduce_prod(tf.shape(v)) for v in tf.trainable_variables()])\n",
    "\n",
    "    saver = tf.train.Saver(max_to_keep=1)\n",
    "\n",
    "    logdir = a.output_dir if (a.trace_freq > 0 or a.summary_freq > 0) else None\n",
    "    sv = tf.train.Supervisor(logdir=logdir, save_summaries_secs=0, saver=None)\n",
    "    with sv.managed_session() as sess:\n",
    "        print(\"parameter_count =\", sess.run(parameter_count))\n",
    "\n",
    "        if a.checkpoint is not None:\n",
    "            print(\"loading model from checkpoint\")\n",
    "            checkpoint = tf.train.latest_checkpoint(a.checkpoint)\n",
    "            saver.restore(sess, checkpoint)\n",
    "\n",
    "        max_steps = 2**32\n",
    "        if a.max_epochs is not None:\n",
    "            max_steps = examples.steps_per_epoch * a.max_epochs\n",
    "        if a.max_steps is not None:\n",
    "            max_steps = a.max_steps\n",
    "\n",
    "        if a.mode == \"test\":\n",
    "            # testing\n",
    "            # at most, process the test data once\n",
    "            start = time.time()\n",
    "            max_steps = min(examples.steps_per_epoch, max_steps)\n",
    "            for step in range(max_steps):\n",
    "                results = sess.run(display_fetches)\n",
    "                filesets = save_images(results)\n",
    "                for i, f in enumerate(filesets):\n",
    "                    print(\"evaluated image\", f[\"name\"])\n",
    "                index_path = append_index(filesets)\n",
    "            print(\"wrote index at\", index_path)\n",
    "            print(\"rate\", (time.time() - start) / max_steps)\n",
    "        else:\n",
    "            # training\n",
    "            start = time.time()\n",
    "\n",
    "            for step in range(max_steps):\n",
    "                def should(freq):\n",
    "                    return freq > 0 and ((step + 1) % freq == 0 or step == max_steps - 1)\n",
    "\n",
    "                options = None\n",
    "                run_metadata = None\n",
    "                if should(a.trace_freq):\n",
    "                    options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)\n",
    "                    run_metadata = tf.RunMetadata()\n",
    "\n",
    "                fetches = {\n",
    "                    \"train\": model.train,\n",
    "                    \"global_step\": sv.global_step,\n",
    "                }\n",
    "\n",
    "                if should(a.progress_freq):\n",
    "                    fetches[\"discrim_loss\"] = model.discrim_loss\n",
    "                    fetches[\"gen_loss_GAN\"] = model.gen_loss_GAN\n",
    "                    fetches[\"gen_loss_L1\"] = model.gen_loss_L1\n",
    "\n",
    "                if should(a.summary_freq):\n",
    "                    fetches[\"summary\"] = sv.summary_op\n",
    "\n",
    "                if should(a.display_freq):\n",
    "                    fetches[\"display\"] = display_fetches\n",
    "\n",
    "                results = sess.run(fetches, options=options, run_metadata=run_metadata)\n",
    "\n",
    "                if should(a.summary_freq):\n",
    "                    print(\"recording summary\")\n",
    "                    sv.summary_writer.add_summary(results[\"summary\"], results[\"global_step\"])\n",
    "\n",
    "                if should(a.display_freq):\n",
    "                    print(\"saving display images\")\n",
    "                    filesets = save_images(results[\"display\"], step=results[\"global_step\"])\n",
    "                    append_index(filesets, step=True)\n",
    "\n",
    "                if should(a.trace_freq):\n",
    "                    print(\"recording trace\")\n",
    "                    sv.summary_writer.add_run_metadata(run_metadata, \"step_%d\" % results[\"global_step\"])\n",
    "\n",
    "                if should(a.progress_freq):\n",
    "                    # global_step will have the correct step count if we resume from a checkpoint\n",
    "                    train_epoch = math.ceil(results[\"global_step\"] / examples.steps_per_epoch)\n",
    "                    train_step = (results[\"global_step\"] - 1) % examples.steps_per_epoch + 1\n",
    "                    rate = (step + 1) * a.batch_size / (time.time() - start)\n",
    "                    remaining = (max_steps - step) * a.batch_size / rate\n",
    "                    print(\"progress  epoch %d  step %d  image/sec %0.1f  remaining %dm\" % (train_epoch, train_step, rate, remaining / 60))\n",
    "                    print(\"discrim_loss\", results[\"discrim_loss\"])\n",
    "                    print(\"gen_loss_GAN\", results[\"gen_loss_GAN\"])\n",
    "                    print(\"gen_loss_L1\", results[\"gen_loss_L1\"])\n",
    "\n",
    "                if should(a.save_freq):\n",
    "                    print(\"saving model\")\n",
    "                    saver.save(sess, os.path.join(a.output_dir, \"model\"), global_step=sv.global_step)\n",
    "\n",
    "                if sv.should_stop():\n",
    "                    break\n",
    "\n",
    "\n",
    "main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (tensorflow)",
   "language": "python",
   "name": "tensorflow"
  },
  "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}