Local_Light_Field_Synthesis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "from scipy import io\n",
    "from scipy import interpolate\n",
    "from scipy import ndimage\n",
    "import os\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#parameters\n",
    "\n",
    "lfsize = [372, 540, 8, 8] #dimensions of Lytro light fields\n",
    "batchsize = 1 #modify based on user's GPU memory\n",
    "patchsize = [192, 192] #spatial dimensions of training light fields\n",
    "disp_mult = 4.0 #max disparity between adjacent veiws\n",
    "num_crops = 4 #number of random spatial crops per light field for each input queue thread to push\n",
    "learning_rate = 0.001\n",
    "train_iters = 120000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#functions for CNN layers\n",
    "\n",
    "def weight_variable(w_shape):\n",
    "    return tf.get_variable('weights', w_shape, initializer=tf.contrib.layers.xavier_initializer_conv2d())\n",
    "\n",
    "def bias_variable(b_shape, init_bias=0.0):\n",
    "    return tf.get_variable('bias', b_shape, initializer=tf.constant_initializer(init_bias))\n",
    "\n",
    "def cnn_layer(input_tensor, w_shape, b_shape, layer_name, rate=1, ds=1):\n",
    "    with tf.variable_scope(layer_name):\n",
    "        W = weight_variable(w_shape)\n",
    "        pad_amt_0 = rate * (w_shape[0] - 1)//2\n",
    "        pad_amt_1 = rate * (w_shape[1] - 1)//2\n",
    "        input_tensor = tf.pad(input_tensor, [[0,0],[pad_amt_0,pad_amt_0],[pad_amt_1,pad_amt_1],[0,0]], mode='SYMMETRIC')\n",
    "        h = tf.nn.convolution(input_tensor, W, strides=[ds, ds], padding='VALID', dilation_rate=[rate, rate], name=layer_name + '_conv')\n",
    "        h = tf.contrib.layers.instance_norm(h + bias_variable(b_shape))\n",
    "        h = tf.nn.leaky_relu(h)\n",
    "        return h\n",
    "    \n",
    "def cnn_layer_plain(input_tensor, w_shape, b_shape, layer_name, rate=1, ds=1):\n",
    "    with tf.variable_scope(layer_name):\n",
    "        W = weight_variable(w_shape)\n",
    "        pad_amt_0 = rate * (w_shape[0] - 1)//2\n",
    "        pad_amt_1 = rate * (w_shape[1] - 1)//2\n",
    "        input_tensor = tf.pad(input_tensor, [[0,0],[pad_amt_0,pad_amt_0],[pad_amt_1,pad_amt_1],[0,0]], mode='SYMMETRIC')\n",
    "        h = tf.nn.convolution(input_tensor, W, strides=[ds, ds], padding='VALID', dilation_rate=[rate, rate], name=layer_name + '_conv')\n",
    "        h = h + bias_variable(b_shape)\n",
    "        return h\n",
    "    \n",
    "def cnn_layer_3D(input_tensor, w_shape, b_shape, layer_name, rate=1, ds=1):\n",
    "    with tf.variable_scope(layer_name):\n",
    "        W = weight_variable(w_shape)\n",
    "        pad_amt_0 = rate * (w_shape[0] - 1)//2\n",
    "        pad_amt_1 = rate * (w_shape[1] - 1)//2\n",
    "        pad_amt_2 = rate * (w_shape[2] - 1)//2\n",
    "        input_tensor = tf.pad(input_tensor, [[0,0],[pad_amt_0,pad_amt_0],[pad_amt_1,pad_amt_1],[pad_amt_2,pad_amt_2],[0,0]], mode='SYMMETRIC')\n",
    "        h = tf.nn.convolution(input_tensor, W, strides=[ds, ds, ds], padding='VALID', dilation_rate=[rate, rate, rate], name=layer_name + '_conv')\n",
    "        h = tf.contrib.layers.instance_norm(h + bias_variable(b_shape))\n",
    "        h = tf.nn.leaky_relu(h)\n",
    "        return h\n",
    "    \n",
    "def cnn_layer_3D_plain(input_tensor, w_shape, b_shape, layer_name, rate=1, ds=1):\n",
    "    with tf.variable_scope(layer_name):\n",
    "        W = weight_variable(w_shape)\n",
    "        pad_amt_0 = rate * (w_shape[0] - 1)//2\n",
    "        pad_amt_1 = rate * (w_shape[1] - 1)//2\n",
    "        pad_amt_2 = rate * (w_shape[2] - 1)//2\n",
    "        input_tensor = tf.pad(input_tensor, [[0,0],[pad_amt_0,pad_amt_0],[pad_amt_1,pad_amt_1],[pad_amt_2,pad_amt_2],[0,0]], mode='SYMMETRIC')\n",
    "        h = tf.nn.convolution(input_tensor, W, strides=[ds, ds, ds], padding='VALID', dilation_rate=[rate, rate, rate], name=layer_name + '_conv')\n",
    "        h = h + bias_variable(b_shape)\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#network to predict ray depths from input image\n",
    "\n",
    "def depth_network(x, lfsize, disp_mult, name):\n",
    "    with tf.variable_scope(name):\n",
    "        \n",
    "        b_sz = tf.shape(x)[0]\n",
    "        y_sz = tf.shape(x)[1]\n",
    "        x_sz = tf.shape(x)[2]\n",
    "        v_sz = lfsize[2]\n",
    "        u_sz = lfsize[3]\n",
    "        \n",
    "        c1 = cnn_layer(x, [3, 3, 3, 16], [16], 'c1')\n",
    "        c2 = cnn_layer(c1, [3, 3, 16, 64], [64], 'c2')\n",
    "        c3 = cnn_layer(c2, [3, 3, 64, 128], [128], 'c3')\n",
    "        c4 = cnn_layer(c3, [3, 3, 128, 128], [128], 'c4', rate=2)\n",
    "        c5 = cnn_layer(c4, [3, 3, 128, 128], [128], 'c5', rate=4)\n",
    "        c6 = cnn_layer(c5, [3, 3, 128, 128], [128], 'c6', rate=8)\n",
    "        c7 = cnn_layer(c6, [3, 3, 128, 128], [128], 'c7', rate=16)\n",
    "        c8 = cnn_layer(c7, [3, 3, 128, 128], [128], 'c8')\n",
    "        c9 = cnn_layer(c8, [3, 3, 128, lfsize[2]*lfsize[3]], [lfsize[2]*lfsize[3]], 'c9')\n",
    "        c10 = disp_mult*tf.tanh(cnn_layer_plain(c9, [3, 3, lfsize[2]*lfsize[3], lfsize[2]*lfsize[3]], \\\n",
    "                                                [lfsize[2]*lfsize[3]], 'c10'))\n",
    "        \n",
    "        return tf.reshape(c10, [b_sz, y_sz, x_sz, v_sz, u_sz])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#network for refining Lambertian light field (predict occluded rays and non-Lambertian effects)\n",
    "\n",
    "def occlusions_network(x, shear, lfsize, name):\n",
    "    with tf.variable_scope(name):\n",
    "        \n",
    "        b_sz = tf.shape(x)[0]\n",
    "        y_sz = tf.shape(x)[1]\n",
    "        x_sz = tf.shape(x)[2]\n",
    "        v_sz = lfsize[2]\n",
    "        u_sz = lfsize[3]\n",
    "        \n",
    "        x = tf.transpose(tf.reshape(tf.transpose(x, perm=[0, 5, 1, 2, 3, 4]), \\\n",
    "                                    [b_sz, 4, y_sz, x_sz, u_sz*v_sz]), perm=[0, 4, 2, 3, 1])\n",
    "        \n",
    "        c1 = cnn_layer_3D(x, [3, 3, 3, 4, 8], [8], 'c1')\n",
    "        c2 = cnn_layer_3D(c1, [3, 3, 3, 8, 8], [8], 'c2')\n",
    "        c3 = cnn_layer_3D(c2, [3, 3, 3, 8, 8], [8], 'c3')\n",
    "        c4 = cnn_layer_3D(c3, [3, 3, 3, 8, 8], [8], 'c4')\n",
    "        c5 = tf.tanh(cnn_layer_3D_plain(c4, [3, 3, 3, 8, 3], [3], 'c5'))\n",
    "        \n",
    "        output = tf.transpose(tf.reshape(tf.transpose(c5, perm=[0, 4, 2, 3, 1]), \\\n",
    "                                         [b_sz, 3, y_sz, x_sz, v_sz, u_sz]), perm=[0, 2, 3, 4, 5, 1]) + shear\n",
    "                \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#full forward model\n",
    "\n",
    "def forward_model(x, lfsize, disp_mult):\n",
    "    with tf.variable_scope('forward_model') as scope:\n",
    "        #predict ray depths from input image\n",
    "        ray_depths = depth_network(x, lfsize, disp_mult, 'ray_depths')\n",
    "        #shear input image by predicted ray depths to render Lambertian light field\n",
    "        lf_shear_r = depth_rendering(x[:, :, :, 0], ray_depths, lfsize)\n",
    "        lf_shear_g = depth_rendering(x[:, :, :, 1], ray_depths, lfsize)\n",
    "        lf_shear_b = depth_rendering(x[:, :, :, 2], ray_depths, lfsize)\n",
    "        lf_shear = tf.stack([lf_shear_r, lf_shear_g, lf_shear_b], axis=5)\n",
    "        #occlusion/non-Lambertian prediction network\n",
    "        shear_and_depth = tf.stack([lf_shear_r, lf_shear_g, lf_shear_b, tf.stop_gradient(ray_depths)], axis=5)\n",
    "        y = occlusions_network(shear_and_depth, lf_shear, lfsize, 'occlusions')\n",
    "        return ray_depths, lf_shear, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#render light field from input image and ray depths\n",
    "\n",
    "def depth_rendering(central, ray_depths, lfsize):\n",
    "    with tf.variable_scope('depth_rendering') as scope:\n",
    "        b_sz = tf.shape(central)[0]\n",
    "        y_sz = tf.shape(central)[1]\n",
    "        x_sz = tf.shape(central)[2]\n",
    "        u_sz = lfsize[2]\n",
    "        v_sz = lfsize[3]\n",
    "        \n",
    "        central = tf.expand_dims(tf.expand_dims(central, 3), 4)\n",
    "                                                \n",
    "        #create and reparameterize light field grid\n",
    "        b_vals = tf.to_float(tf.range(b_sz))\n",
    "        v_vals = tf.to_float(tf.range(v_sz)) - tf.to_float(v_sz)/2.0\n",
    "        u_vals = tf.to_float(tf.range(u_sz)) - tf.to_float(u_sz)/2.0\n",
    "        y_vals = tf.to_float(tf.range(y_sz))\n",
    "        x_vals = tf.to_float(tf.range(x_sz))\n",
    "    \n",
    "        b, y, x, v, u = tf.meshgrid(b_vals, y_vals, x_vals, v_vals, u_vals, indexing='ij')\n",
    "               \n",
    "        #warp coordinates by ray depths\n",
    "        y_t = y + v * ray_depths\n",
    "        x_t = x + u * ray_depths\n",
    "        \n",
    "        v_r = tf.zeros_like(b)\n",
    "        u_r = tf.zeros_like(b)\n",
    "        \n",
    "        #indices for linear interpolation\n",
    "        b_1 = tf.to_int32(b)\n",
    "        y_1 = tf.to_int32(tf.floor(y_t))\n",
    "        y_2 = y_1 + 1\n",
    "        x_1 = tf.to_int32(tf.floor(x_t))\n",
    "        x_2 = x_1 + 1\n",
    "        v_1 = tf.to_int32(v_r)\n",
    "        u_1 = tf.to_int32(u_r)\n",
    "        \n",
    "        y_1 = tf.clip_by_value(y_1, 0, y_sz-1)\n",
    "        y_2 = tf.clip_by_value(y_2, 0, y_sz-1)\n",
    "        x_1 = tf.clip_by_value(x_1, 0, x_sz-1)\n",
    "        x_2 = tf.clip_by_value(x_2, 0, x_sz-1)\n",
    "        \n",
    "        #assemble interpolation indices\n",
    "        interp_pts_1 = tf.stack([b_1, y_1, x_1, v_1, u_1], -1)\n",
    "        interp_pts_2 = tf.stack([b_1, y_2, x_1, v_1, u_1], -1)\n",
    "        interp_pts_3 = tf.stack([b_1, y_1, x_2, v_1, u_1], -1)\n",
    "        interp_pts_4 = tf.stack([b_1, y_2, x_2, v_1, u_1], -1)\n",
    "        \n",
    "        #gather light fields to be interpolated\n",
    "        lf_1 = tf.gather_nd(central, interp_pts_1)\n",
    "        lf_2 = tf.gather_nd(central, interp_pts_2)\n",
    "        lf_3 = tf.gather_nd(central, interp_pts_3)\n",
    "        lf_4 = tf.gather_nd(central, interp_pts_4)\n",
    "        \n",
    "        #calculate interpolation weights\n",
    "        y_1_f = tf.to_float(y_1)\n",
    "        x_1_f = tf.to_float(x_1)\n",
    "        d_y_1 = 1.0 - (y_t - y_1_f)\n",
    "        d_y_2 = 1.0 - d_y_1\n",
    "        d_x_1 = 1.0 - (x_t - x_1_f)\n",
    "        d_x_2 = 1.0 - d_x_1\n",
    "        \n",
    "        w1 = d_y_1 * d_x_1\n",
    "        w2 = d_y_2 * d_x_1\n",
    "        w3 = d_y_1 * d_x_2\n",
    "        w4 = d_y_2 * d_x_2\n",
    "        \n",
    "        lf = tf.add_n([w1*lf_1, w2*lf_2, w3*lf_3, w4*lf_4])\n",
    "                        \n",
    "    return lf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#resample ray depths for depth consistency regularization\n",
    "\n",
    "def transform_ray_depths(ray_depths, u_step, v_step, lfsize):\n",
    "    with tf.variable_scope('transform_ray_depths') as scope:\n",
    "        b_sz = tf.shape(ray_depths)[0]\n",
    "        y_sz = tf.shape(ray_depths)[1]\n",
    "        x_sz = tf.shape(ray_depths)[2]\n",
    "        u_sz = lfsize[2]\n",
    "        v_sz = lfsize[3]\n",
    "                                                        \n",
    "        #create and reparameterize light field grid\n",
    "        b_vals = tf.to_float(tf.range(b_sz))\n",
    "        v_vals = tf.to_float(tf.range(v_sz)) - tf.to_float(v_sz)/2.0\n",
    "        u_vals = tf.to_float(tf.range(u_sz)) - tf.to_float(u_sz)/2.0\n",
    "        y_vals = tf.to_float(tf.range(y_sz))\n",
    "        x_vals = tf.to_float(tf.range(x_sz))\n",
    "    \n",
    "        b, y, x, v, u = tf.meshgrid(b_vals, y_vals, x_vals, v_vals, u_vals, indexing='ij')\n",
    "               \n",
    "        #warp coordinates by ray depths\n",
    "        y_t = y + v_step * ray_depths\n",
    "        x_t = x + u_step * ray_depths\n",
    "        \n",
    "        v_t = v - v_step + tf.to_float(v_sz)/2.0\n",
    "        u_t = u - u_step + tf.to_float(u_sz)/2.0\n",
    "        \n",
    "        #indices for linear interpolation\n",
    "        b_1 = tf.to_int32(b)\n",
    "        y_1 = tf.to_int32(tf.floor(y_t))\n",
    "        y_2 = y_1 + 1\n",
    "        x_1 = tf.to_int32(tf.floor(x_t))\n",
    "        x_2 = x_1 + 1\n",
    "        v_1 = tf.to_int32(v_t)\n",
    "        u_1 = tf.to_int32(u_t)\n",
    "        \n",
    "        y_1 = tf.clip_by_value(y_1, 0, y_sz-1)\n",
    "        y_2 = tf.clip_by_value(y_2, 0, y_sz-1)\n",
    "        x_1 = tf.clip_by_value(x_1, 0, x_sz-1)\n",
    "        x_2 = tf.clip_by_value(x_2, 0, x_sz-1)\n",
    "        v_1 = tf.clip_by_value(v_1, 0, v_sz-1)\n",
    "        u_1 = tf.clip_by_value(u_1, 0, u_sz-1)\n",
    "        \n",
    "        #assemble interpolation indices\n",
    "        interp_pts_1 = tf.stack([b_1, y_1, x_1, v_1, u_1], -1)\n",
    "        interp_pts_2 = tf.stack([b_1, y_2, x_1, v_1, u_1], -1)\n",
    "        interp_pts_3 = tf.stack([b_1, y_1, x_2, v_1, u_1], -1)\n",
    "        interp_pts_4 = tf.stack([b_1, y_2, x_2, v_1, u_1], -1)\n",
    "        \n",
    "        #gather light fields to be interpolated\n",
    "        lf_1 = tf.gather_nd(ray_depths, interp_pts_1)\n",
    "        lf_2 = tf.gather_nd(ray_depths, interp_pts_2)\n",
    "        lf_3 = tf.gather_nd(ray_depths, interp_pts_3)\n",
    "        lf_4 = tf.gather_nd(ray_depths, interp_pts_4)\n",
    "        \n",
    "        #calculate interpolation weights\n",
    "        y_1_f = tf.to_float(y_1)\n",
    "        x_1_f = tf.to_float(x_1)\n",
    "        d_y_1 = 1.0 - (y_t - y_1_f)\n",
    "        d_y_2 = 1.0 - d_y_1\n",
    "        d_x_1 = 1.0 - (x_t - x_1_f)\n",
    "        d_x_2 = 1.0 - d_x_1\n",
    "        \n",
    "        w1 = d_y_1 * d_x_1\n",
    "        w2 = d_y_2 * d_x_1\n",
    "        w3 = d_y_1 * d_x_2\n",
    "        w4 = d_y_2 * d_x_2\n",
    "        \n",
    "        lf = tf.add_n([w1*lf_1, w2*lf_2, w3*lf_3, w4*lf_4])\n",
    "                        \n",
    "    return lf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#loss to encourage consistency of ray depths corresponding to same scene point\n",
    "\n",
    "def depth_consistency_loss(x, lfsize):\n",
    "    x_u = transform_ray_depths(x, 1.0, 0.0, lfsize)\n",
    "    x_v = transform_ray_depths(x, 0.0, 1.0, lfsize)\n",
    "    x_uv = transform_ray_depths(x, 1.0, 1.0, lfsize)\n",
    "    d1 = (x[:,:,:,1:,1:]-x_u[:,:,:,1:,1:])\n",
    "    d2 = (x[:,:,:,1:,1:]-x_v[:,:,:,1:,1:])\n",
    "    d3 = (x[:,:,:,1:,1:]-x_uv[:,:,:,1:,1:])\n",
    "    l1 = tf.reduce_mean(tf.abs(d1)+tf.abs(d2)+tf.abs(d3))\n",
    "    return l1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#spatial TV loss (l1 of spatial derivatives)\n",
    "\n",
    "def image_derivs(x, nc):\n",
    "    dy = tf.nn.depthwise_conv2d(x, tf.tile(tf.expand_dims(tf.expand_dims([[1.0, 2.0, 1.0], [0.0, 0.0, 0.0], [-1.0, -2.0, -1.0]], 2), 3), [1, 1, nc, 1]), strides=[1, 1, 1, 1], padding='VALID')\n",
    "    dx = tf.nn.depthwise_conv2d(x, tf.tile(tf.expand_dims(tf.expand_dims([[1.0, 0.0, -1.0], [2.0, 0.0, -2.0], [1.0, 0.0, -1.0]], 2), 3), [1, 1, nc, 1]), strides=[1, 1, 1, 1], padding='VALID')\n",
    "    return dy, dx\n",
    "\n",
    "def tv_loss(x):\n",
    "    b_sz = tf.shape(x)[0]\n",
    "    y_sz = tf.shape(x)[1]\n",
    "    x_sz = tf.shape(x)[2]\n",
    "    u_sz = lfsize[2]\n",
    "    v_sz = lfsize[3]\n",
    "    temp = tf.reshape(x, [b_sz, y_sz, x_sz, u_sz*v_sz])\n",
    "    dy, dx = image_derivs(temp, u_sz*v_sz)\n",
    "    l1 = tf.reduce_mean(tf.abs(dy)+tf.abs(dx))\n",
    "    return l1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#normalize to between -1 and 1, given input between 0 and 1\n",
    "\n",
    "def normalize_lf(lf):\n",
    "    return 2.0*(lf-0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#input pipeline\n",
    "\n",
    "def process_lf(lf, num_crops, lfsize, patchsize):\n",
    "    gamma_val = tf.random_uniform(shape=[], minval=0.4, maxval=1.0) #random gamma for data augmentation (change at test time, I suggest 0.4-0.5)\n",
    "    lf = normalize_lf(tf.image.adjust_gamma(tf.to_float(lf[:lfsize[0]*14, :lfsize[1]*14, :])/255.0, gamma=gamma_val))\n",
    "    lf = tf.transpose(tf.reshape(lf, [lfsize[0], 14, lfsize[1], 14, 3]), [0, 2, 1, 3, 4])\n",
    "    lf = lf[:, :, (14//2)-(lfsize[2]//2):(14//2)+(lfsize[2]//2), (14//2)-(lfsize[3]//2):(14//2)+(lfsize[3]//2), :]\n",
    "    aif = lf[:, :, lfsize[2]//2, lfsize[3]//2, :]\n",
    "    aif_list = []\n",
    "    lf_list = []\n",
    "    for i in range(num_crops):\n",
    "        r = tf.random_uniform(shape=[], minval=0, maxval=tf.shape(lf)[0]-patchsize[0], dtype=tf.int32)\n",
    "        c = tf.random_uniform(shape=[], minval=0, maxval=tf.shape(lf)[1]-patchsize[1], dtype=tf.int32)\n",
    "        aif_list.append(aif[r:r+patchsize[0], c:c+patchsize[1], :])\n",
    "        lf_list.append(lf[r:r+patchsize[0], c:c+patchsize[1], :, :, :])\n",
    "    return aif_list, lf_list\n",
    "\n",
    "def read_lf(filename_queue, num_crops, lfsize, patchsize):\n",
    "    value = tf.read_file(filename_queue[0])\n",
    "    lf = tf.image.decode_image(value, channels=3)\n",
    "    aif_list, lf_list = process_lf(lf, num_crops, lfsize, patchsize)\n",
    "    return aif_list, lf_list\n",
    "\n",
    "def input_pipeline(filenames, lfsize, patchsize, batchsize, num_crops):\n",
    "    filename_queue = tf.train.slice_input_producer([filenames], shuffle=True)\n",
    "    example_list = [read_lf(filename_queue, num_crops, lfsize, patchsize) for _ in range(2)] #number of threads for populating queue\n",
    "    min_after_dequeue = 0\n",
    "    capacity = 8\n",
    "    aif_batch, lf_batch = tf.train.shuffle_batch_join(example_list, batch_size=batchsize, capacity=capacity, \n",
    "                                                      min_after_dequeue=min_after_dequeue, enqueue_many=True,\n",
    "                                                      shapes=[[patchsize[0], patchsize[1], 3], \n",
    "                                                              [patchsize[0], patchsize[1], lfsize[2], lfsize[3], 3]])\n",
    "    return aif_batch, lf_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_path = '/Users/pratulsrinivasan/Dropbox/LF_Flowers' #path to training examples\n",
    "train_filenames = [os.path.join(train_path, f) for f in os.listdir(train_path) if not f.startswith('.')]\n",
    "\n",
    "aif_batch, lf_batch = input_pipeline(train_filenames, lfsize, patchsize, batchsize, num_crops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#forward model\n",
    "ray_depths, lf_shear, y = forward_model(aif_batch, lfsize, disp_mult)\n",
    "\n",
    "#training losses to minimize\n",
    "lam_tv = 0.01\n",
    "lam_dc = 0.005\n",
    "with tf.name_scope('loss'):\n",
    "    shear_loss = tf.reduce_mean(tf.abs(lf_shear-lf_batch))\n",
    "    output_loss = tf.reduce_mean(tf.abs(y-lf_batch)) \n",
    "    tv_loss = lam_tv * tv_loss(ray_depths)\n",
    "    depth_consistency_loss = lam_dc * depth_consistency_loss(ray_depths, lfsize)\n",
    "    train_loss = shear_loss + output_loss + tv_loss + depth_consistency_loss\n",
    "    \n",
    "with tf.name_scope('train'):\n",
    "    train_step = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(train_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#tensorboard summaries\n",
    "tf.summary.scalar('shear_loss', shear_loss)\n",
    "tf.summary.scalar('output_loss', output_loss)\n",
    "tf.summary.scalar('tv_loss', tv_loss)\n",
    "tf.summary.scalar('depth_consistency_loss', depth_consistency_loss)\n",
    "tf.summary.scalar('train_loss', train_loss)\n",
    "\n",
    "tf.summary.histogram('ray_depths', ray_depths)\n",
    "\n",
    "tf.summary.image('input_image', aif_batch)\n",
    "tf.summary.image('lf_shear', tf.reshape(tf.transpose(lf_shear, perm=[0, 3, 1, 4, 2, 5]), \n",
    "                                        [batchsize, patchsize[0]*lfsize[2], patchsize[1]*lfsize[3], 3]))\n",
    "tf.summary.image('lf_output', tf.reshape(tf.transpose(y, perm=[0, 3, 1, 4, 2, 5]), \n",
    "                                        [batchsize, patchsize[0]*lfsize[2], patchsize[1]*lfsize[3], 3]))\n",
    "tf.summary.image('ray_depths', tf.reshape(tf.transpose(ray_depths, perm=[0, 3, 1, 4, 2]), \n",
    "                                        [batchsize, patchsize[0]*lfsize[2], patchsize[1]*lfsize[3], 1]))\n",
    "\n",
    "merged = tf.summary.merge_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logdir = 'logs/train/' #path to store logs\n",
    "checkpointdir = 'checkpoints/' #path to store checkpoints\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    train_writer = tf.summary.FileWriter(logdir, sess.graph)\n",
    "    saver = tf.train.Saver()\n",
    "    sess.run(tf.global_variables_initializer()) #initialize variables \n",
    "    \n",
    "    coord = tf.train.Coordinator() #coordinator for input queue threads\n",
    "    threads = tf.train.start_queue_runners(sess=sess, coord=coord) #start input queue threads\n",
    "    \n",
    "    for i in range(train_iters):\n",
    "        #training training step\n",
    "        _ = sess.run(train_step)\n",
    "        #save training summaries\n",
    "        if (i+1) % 1 == 0: #can change the frequency of writing summaries if desired\n",
    "            print('training step: ', i)\n",
    "            trainsummary = sess.run(merged)\n",
    "            train_writer.add_summary(trainsummary, i)  \n",
    "        #save checkpoint\n",
    "        if (i+1) % 4000 == 0:\n",
    "            saver.save(sess, checkpointdir + 'model.ckpt', global_step=i)\n",
    "            \n",
    "    #cleanup\n",
    "    train_writer.close()\n",
    "    coord.request_stop()\n",
    "    coord.join(threads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}