text: # Copyright (c) 2021 VISTEC - Vidyasirimedhi Institute of Science and Technology
  # Distribute under MIT License
  # Authors:
  #  - Suttisak Wizadwongsa <suttisak.w_s19[-at-]vistec.ac.th>
  #  - Pakkapon Phongthawee <pakkapon.p_s19[-at-]vistec.ac.th>
  #  - Jiraphon Yenphraphai <jiraphony_pro[-at-]vistec.ac.th>
  #  - Supasorn Suwajanakorn <supasorn.s[-at-]vistec.ac.th>
  
  from __future__ import division
  from __future__ import print_function
  
  import argparse
  import getpass
  import shutil
  
  import torch as pt
  import torch.nn.functional as F
  import torch.nn as nn
  from torch.utils.data import Dataset, DataLoader
  from torchvision.utils import save_image, make_grid
  from torch.utils.tensorboard import SummaryWriter
  
  import os, sys, json
  import numpy as np
  from skimage import io, transform
  from datetime import datetime
  
  from utils.utils import *
  from utils.mpi_utils import *
  from utils.mlp import *
  from utils.colmap_runner import colmapGenPoses
  
  parser = argparse.ArgumentParser()
  
  #training schedule
  parser.add_argument('-epochs', type=int, default=4000, help='total epochs to train')
  parser.add_argument('-steps', type=int, default=-1, help='total steps to train. In our paper, we proposed to use epoch instead.')
  parser.add_argument('-tb_saveimage', type=int, default=50, help='write an output image to tensorboard for every <tb_saveimage> epochs')
  parser.add_argument('-tb_savempi', type=int, default=200, help='generate MPI (WebGL) and measure PSNR/SSIM of validation image for every <tb_savempi> epochs')
  parser.add_argument('-checkpoint', type=int, default=100, help='save checkpoint for every <checkpoint> epochs. Be aware that! It will replace the previous checkpoint.')
  parser.add_argument('-tb_toc',type=int, default=500, help="print output to terminal for every tb_toc epochs")
  
  #lr schedule
  parser.add_argument('-lrc', type=float, default=10, help='the number of times of lr using for learning rate of explicit basis (k0).')
  parser.add_argument('-lr', type=float, default=1e-3, help='learning rate of a multi-layer perceptron')
  parser.add_argument('-decay_epoch', type=int, default=1333, help='the number of epochs for decay learning rate')
  parser.add_argument('-decay_rate', type=float, default=0.1, help='ratio of decay rate at every <decay_epoch> epochs')
  
  #network (First MLP)
  parser.add_argument('-ray', type=int, default=8000, help='the number of sampled ray that is used to train in each step')
  parser.add_argument('-hidden', type=int, default=384, help='the number of hidden node of the main MLP')
  parser.add_argument('-mlp', type=int, default=4, help='the number of hidden layer of the main MLP')
  parser.add_argument('-pos_level', type=int, default=10, help='the number of positional encoding in terms of image size. We recommend to set 2^(pos_level) > image_height and image_width')
  parser.add_argument('-depth_level', type=int, default=8,help='the number of positional encoding in terms number of plane. We recommend to set 2^(depth_level) > layers * subplayers')
  parser.add_argument('-lrelu_slope', type=float, default=0.01, help='slope of leaky relu')
  parser.add_argument('-sigmoid_offset', type=float, default=5, help='sigmoid offset that is applied to alpha before sigmoid')
  
  #basis (Second MLP)
  parser.add_argument('-basis_hidden', type=int, default=64, help='the number of hidden node in the learned basis MLP')
  parser.add_argument('-basis_mlp', type=int, default=1, help='the number of hidden layer in the learned basis MLP')
  parser.add_argument('-basis_order', type=int, default=3, help='the number of  positional encoding in terms of viewing angle')
  parser.add_argument('-basis_out', type=int, default=8, help='the number of coeffcient output (N in equation 3 under seftion 3.1)')
  
  #loss
  parser.add_argument('-gradloss', type=float, default=0.05, help='hyperparameter for grad loss')
  parser.add_argument('-tvc', type=float, default=0.03, help='hyperparameter for total variation regularizer')
  
  #training and eval data
  parser.add_argument('-scene', type=str, default="", help='directory to the scene')
  parser.add_argument('-ref_img', type=str, default="",  help='reference image, camera parameter of reference image is use to create MPI')
  parser.add_argument('-dmin', type=float, default=-1, help='first plane depth')
  parser.add_argument('-dmax', type=float, default=-1, help='last plane depth')
  parser.add_argument('-invz', action='store_true', help='place MPI with inverse depth')
  parser.add_argument('-scale', type=float, default=-1, help='scale the MPI size')
  parser.add_argument('-llff_width', type=int, default=1008, help='if input dataset is LLFF it will resize the image to <llff_width>')
  parser.add_argument('-deepview_width', type=int, default=800, help='if input dataset is deepview dataset, it will resize the image to <deepview_width>')
  parser.add_argument('-train_ratio', type=float, default=0.875, help='ratio to split number of train/test (in case dataset doesn\'t specify how to split)')
  parser.add_argument('-random_split', action='store_true', help='random split the train/test set. (in case dataset doesn\'t specify how to split)')
  parser.add_argument('-num_workers', type=int, default=8, help='number of pytorch\'s dataloader worker')
  parser.add_argument('-cv2resize', action='store_true', help='apply cv2.resize instead of skimage.transform.resize to match the score in our paper (see note in github readme for more detail) ')
  
  #MPI
  parser.add_argument('-offset', type=int, default=200, help='the offset (padding) of the MPI.')
  parser.add_argument('-layers', type=int, default=16, help='the number of plane that stores base color')
  parser.add_argument('-sublayers', type=int, default=12, help='the number of plane that share the same texture. (please refer to coefficient sharing under section 3.4 in the paper)')
  
  #predict
  parser.add_argument('-no_eval', action='store_true', help='do not measurement the score (PSNR/SSIM/LPIPS) ')
  parser.add_argument('-no_csv', action='store_true', help="do not write CSV on evaluation")
  parser.add_argument('-no_video', action='store_true', help="do not write the video on prediction")
  parser.add_argument('-no_webgl', action='store_true', help='do not predict webgl (realtime demo) related content.')
  parser.add_argument('-predict', action='store_true', help='predict validation images')
  parser.add_argument('-eval_path', type=str, default='runs/evaluation/', help='path to save validation image')
  parser.add_argument('-web_path', type=str, default='runs/html/', help='path to output real time demo')
  parser.add_argument('-web_width', type=int, default=16000, help='max texture size (pixel) of realtime demo. WebGL on Highend PC is support up to 16384px, while mobile phone support only 4096px')
  parser.add_argument('-http', action='store_true', help='serve real-time demo on http server')
  parser.add_argument('-render_viewing', action='store_true', help='genereate view-dependent-effect video')
  parser.add_argument('-render_nearest', action='store_true', help='genereate nearest input video')
  parser.add_argument('-render_depth', action='store_true', help='generate depth')
  
  # render path
  parser.add_argument('-nice_llff', action='store_true', help="generate video that its rendering path matches real-forward facing dataset")
  parser.add_argument('-nice_shiny', action='store_true', help="generate video that its rendering path matches shiny dataset")
  
  
  #training utility
  parser.add_argument('-model_dir', type=str, default="scene", help='model (scene) directory which store in runs/<model_dir>/')
  parser.add_argument('-pretrained', type=str, default="", help='location of checkpoint file, if not provide will use model_dir instead')
  parser.add_argument('-restart', action='store_true', help='delete old weight and retrain')
  parser.add_argument('-clean', action='store_true', help='delete old weight without start training process')
  
  
  #miscellaneous
  parser.add_argument('-all_gpu',action='store_true',help="In multiple GPU training, We don't train MLP (data parallel) on the first GPU. This make training slower but we can utilize more VRAM on other GPU.")
  parser.add_argument('-transformation', type=int, default=0, help='transformation to be applied to the coordinates')
  
  args = parser.parse_args()
  
  def computeHomographies(sfm, feature, planes):
    fx = feature['fx'][0]
    fy = feature['fy'][0]
    px = feature['px'][0]
    py = feature['py'][0]
  
    new_r = feature['r'][0] @ sfm.ref_rT
    new_t = (-new_r @ sfm.ref_t) + feature['t'][0]
  
    n = pt.tensor([[0.0, 0.0, 1.0]])
    Ha = new_r.t()
    Hb = Ha @ new_t @ n @ Ha
    Hc = (n @ Ha @ new_t)[0]
  
    ki = pt.tensor([[fx, 0, px],
                    [0, fy, py],
                    [0, 0, 1]], dtype=pt.float).inverse()
  
    tt = sfm.ref_cam
    ref_k = pt.tensor( [[tt['fx'], 0, tt['px']],
                        [0, tt['fy'], tt['py']],
                        [0,        0,       1]])
  
    planes_mat = pt.Tensor(planes).view(-1, 1, 1)
    return (ref_k @ (Ha + Hb / (-planes_mat - Hc))) @ ki
  
  def computeHomoWarp(sfm, input_shape, input_offset,
                      output_shape, selection,
                      feature, planes, inv=False, inv_offset = False):
  
    selection = selection.cuda()
    # coords: (sel, 3)
    coords = pt.stack([selection % output_shape[1], selection // output_shape[1],
                      pt.ones_like(selection)], -1).float()
  
    # Hs: (n, 3, 3)
  
    Hs = computeHomographies(sfm, feature, planes)
    if inv: Hs = Hs.inverse()
    if inv_offset:
      coords[:, :2] += input_offset
    prod = coords @ pt.transpose(Hs, 1, 2).cuda()
    scale = pt.tensor([input_shape[1] - 1, input_shape[0] - 1]).cuda()
  
    ref_coords = prod[:, :, :2] / prod[:, :, 2:]
    if not inv_offset:
      warp = ((ref_coords + input_offset) / scale.view(1, 1, 2)) * 2 - 1
    else:
      warp = ((ref_coords) / scale.view(1, 1, 2)) * 2 - 1
    warp = warp[:, :, None]
  
  
    return warp, ref_coords
  
  def totalVariation(images):
    pixel_dif1 = images[:, :, 1:, :] - images[:, :, :-1, :]
    pixel_dif2 = images[:, :, :, 1:] - images[:, :, :, :-1]
    sum_axis = [1, 2, 3]
  
    tot_var = (
        pt.sum(pt.abs(pixel_dif1), dim=sum_axis) +
        pt.sum(pt.abs(pixel_dif2), dim=sum_axis))
  
    return tot_var / (images.shape[2]-1) / (images.shape[3]-1)
  
  def cumprod_exclusive(x):
    cp = pt.cumprod(x, 0)
    cp = pt.roll(cp, 1, 0)
    cp[0] = 1.0
    return cp
  
  def getWarp3d(warp, interpolate = False):
    if not interpolate:
      depths = pt.repeat_interleave(pt.linspace(-1, 1, args.layers), args.sublayers).view(1, -1, 1, 1, 1).cuda()
    else:
      depths = pt.linspace(-1, 1, args.layers * args.sublayers).view(1, -1, 1, 1, 1).cuda()
    warp3d = warp[None] # 1, n, sel, 1, 2
    warp3d = pt.cat([warp3d, pt.ones_like(warp3d[:, :, :, :, :1]) * depths], -1)
    return warp3d
  
  def normalized(v, dim):
    return v / (pt.norm(v, dim=dim, keepdim=True) + 1e-7)
  
  class Basis(nn.Module):
    def __init__(self, shape, out_view):
      super().__init__()
      #choosing illumination model
      self.order = args.basis_order
  
      # network for learn basis
      self.seq_basis = nn.DataParallel(
        ReluMLP(
          args.basis_mlp, #basis_mlp
          args.basis_hidden, #basis_hidden
          self.order * 4,
          args.lrelu_slope,
          out_node = args.basis_out, #basis_out
        )
      )
      print('Basis Network:',self.seq_basis)
  
      # positional encoding pre compute
      self.pos_freq_viewing = pt.Tensor([(2 ** i) for i in range(self.order)]).view(1, 1, 1, 1, -1).cuda()
  
    def forward(self, sfm, feature, ref_coords, warp, planes, coeff = None):
      vi, xy = get_viewing_angle(sfm, feature, ref_coords, planes)
      n, sel = vi.shape[:2]
  
      # positional encoding for learn basis
      hinv_xy = vi[...,  :2, None] * self.pos_freq_viewing
      big = pt.reshape(hinv_xy, [n, sel, 1, hinv_xy.shape[-2] * hinv_xy.shape[-1]])
      vi = pt.cat([pt.sin(0.5*np.pi*big), pt.cos(0.5*np.pi*big)], -1)
  
      out2 = self.seq_basis(vi)
      out2 = pt.tanh(out2)
  
      vi = out2.view(n, sel, 1, 1, -1)
  
      coeff = coeff.view(coeff.shape[0], coeff.shape[1], coeff.shape[2], 3,  -1)
      coeff = pt.tanh(coeff)
  
      illumination = pt.sum(coeff * vi,-1).permute([0, 3, 1, 2])
  
      return illumination
  
  def get_viewing_angle(sfm, feature, ref_coords, planes):
    camera = sfm.ref_rT.t() @ feature["center"][0] + sfm.ref_t
  
    # (n, rays, 2) -> (n, 2, rays)
    coords = ref_coords.permute([0, 2, 1])
    # (n, 2, rays) -> (n, 3, rays)
    coords = pt.cat([coords, pt.ones_like(coords[:, :1])], 1)
  
    # coords: (n, 3, rays)
    # viewed planes: (n, 1, 1)
    # xyz: (n, 3, rays)
    xyz = coords * pt.Tensor(planes).view(-1, 1, 1).cuda()
  
    ki = pt.tensor([[feature['fx'][0], 0, feature['px'][0]],
                    [0, feature['fy'][0], feature['py'][0]],
                    [0, 0, 1]], dtype=pt.float).inverse().cuda()
  
    xyz = ki @ xyz
  
    # camera: (3, 1) -> (1, 3, 1)
    # xyz: (n, 3, rays)
    # viewing_angle: (n, 3, rays)
    # viewing_angle = normalized(camera[None].cuda() - xyz, 1)
    inv_viewing_angle = normalized(xyz - camera[None].cuda(), 1)
  
    view = inv_viewing_angle.permute([0, 2, 1])
    xyz = xyz.permute([0, 2, 1])
    return view[:,:,None], xyz[:,:,None]
  
  class Network(nn.Module):
    def __init__(self, shape, sfm):
      super(Network, self).__init__()
      self.shape = [shape[2], shape[3]]
      total_cuda = pt.cuda.device_count()
      mlp_first_device = 1 if (not args.all_gpu) and total_cuda > 1 else 0
      mlp_devices = list(range(mlp_first_device, total_cuda))
      #mpi_c (k0) as an explicit
      mpi_c = pt.empty((shape[0], 3, shape[2], shape[3]), device='cuda:0').uniform_(-1, 1)
      self.mpi_c = nn.Parameter(mpi_c)
      self.specular = Basis(shape, args.basis_out * 3).cuda()
      self.seq1 = nn.DataParallel(
        VanillaMLP(
          args.mlp,
          args.hidden,
          args.pos_level,
          args.depth_level,
          args.lrelu_slope,
          out_node = 1 + args.basis_out * 3,
          first_gpu = mlp_first_device
        ),
        device_ids = mlp_devices
      )
  
      self.seq1 = self.seq1.cuda("cuda:{}".format(mlp_first_device))
      self.pos_freq = pt.Tensor([0.5 * np.pi * (2 ** i) for i in range(args.pos_level)] * 2).view(1, 1, 1, 2, -1).cuda()
      self.depth_freq = pt.Tensor([0.5 * np.pi * (2 ** i) for i in range(args.depth_level)]).view(1, 1, 1, -1).cuda()
  
      self.z_coords = pt.linspace(-1, 1, args.layers * args.sublayers).view(-1, 1, 1, 1).cuda()
      if args.render_depth:
        self.rainbow_mpi = np.zeros((shape[0], 3, shape[2], shape[3]), dtype=np.float32)
        for i,s in enumerate(np.linspace(1, 0, shape[0])):
          color = Rainbow(s)
          for c in range(3):
            self.rainbow_mpi[i,c] = color[c]
        self.rainbow_mpi = pt.from_numpy(self.rainbow_mpi).to('cuda:0')
      else:
        self.rainbow_mpi = None
  
      if sfm.dmin < 0 or sfm.dmax < 0:
        raise ValueError("invalid dmin dmax")
  
      self.planes = getPlanes(sfm, args.layers * args.sublayers)
      print('Mpi Size: {}'.format(self.mpi_c.shape))
      print('All combined layers: {}'.format(args.layers * args.sublayers))
      print(self.planes)
      print('Using inverse depth: {}, Min depth: {}, Max depth: {}'.format(sfm.invz == 1, self.planes[0],self.planes[-1]))
      print('Layer of MLP: {}'.format(args.mlp + 2))
      print('Hidden Channel of MLP: {}'.format(args.hidden))
      print('Main Network',self.seq1)
  
  
    def forward(self, sfm, feature, output_shape, selection):
      ''' Rendering
      Args:
        sfm: reference camera parameter
        feature: target camera parameter
        output_shape: [h, w]. Desired rendered image
        selection: [ray]. pixel to train
      Returns:
        output: [1, 3, rays, 1] rendered image
      '''
      # (n, sel, 1, 2), (n, sel, 1, 1), (n, sel, 2)
      warp, ref_coords = computeHomoWarp(sfm,
                                   self.shape,
                                   sfm.offset,
                                   output_shape, selection,
                                   feature, self.planes)
  
      n = warp.shape[0]
      sel = warp.shape[1]
      # vxy: (n, sel, 1, 2, pos_level)
      vxy = warp[:, :, :, :, None] * self.pos_freq
      vxy = vxy.view(n, sel, 1, -1) # (n, sel, 1, pos_level*2)
  
      # vz: (n, sel, 1, depth_level)
      vz = pt.ones_like(warp[:, :, :, :1]) * self.z_coords * self.depth_freq
  
      vxyz = pt.cat([vxy, vz], -1)
      bigcoords = pt.cat([pt.sin(vxyz), pt.cos(vxyz)], -1)
      # (n, sel, 1, out_node)
      out = self.seq1(bigcoords).cuda()
  
      node = 0
      self.mpi_a = out[..., node:node + 1]
      node += 1
      # n, 1, sel, 1
      self.mpi_a = self.mpi_a.view(self.mpi_a.shape[0], 1, self.mpi_a.shape[1], self.mpi_a.shape[2])
      mpi_a_sig  = pt.sigmoid(self.mpi_a - args.sigmoid_offset)
  
      if args.render_depth:
        # generate Rainbow MPI instead of real mpi to visualize the depth
        # self.rainbow_mpi: n, 3, h, w   warp: (n, sel, 1, 2)
        # Need: N, C, Din, Hin, Win;  N, Dout, Hout, Wout, 3
        rainbow_3d = self.rainbow_mpi.permute([1, 0, 2, 3])[None]
        warp3d = getWarp3d(warp)
        #samples: N, C, Dout, Hout, Wout
        samples = F.grid_sample(rainbow_3d, warp3d, align_corners=True)
        # (layers, out_node, rays, 1)
        rgb = samples[0].permute([1, 0, 2, 3])
      else:
        mpi_sig = pt.sigmoid(self.mpi_c)
        # mpi_sig: n, 3, h, w   warp: (n, sel, 1, 2)
        # Need: N, C, Din, Hin, Win;    N, Dout, Hout, Wout, 3
        mpi_sig3d = mpi_sig.permute([1, 0, 2, 3])[None]
        warp3d = getWarp3d(warp)
        #samples: N, C, Dout, Hout, Wout
        samples = F.grid_sample(mpi_sig3d, warp3d, align_corners=True)
        # (layers, out_node, rays, 1)
        rgb = samples[0].permute([1, 0, 2, 3])
        cof = out[::args.sublayers, ..., node:]
        cof = pt.repeat_interleave(cof, args.sublayers, 0)
        self.illumination = self.specular(sfm, feature, ref_coords, warp, self.planes, coeff = cof)
        # rgb: (layers, 3, rays, 1)
        rgb = pt.clamp(rgb + self.illumination, 0.0, 1.0)
  
      weight = cumprod_exclusive(1 - mpi_a_sig)
  
      output = pt.sum(weight * rgb * mpi_a_sig, dim=0, keepdim=True)
  
      return output
  
  def getMPI(model, sfm, m = 1, dataloader = None):
    ''' convert from neural network to MPI planes
      Args:
        model: Neural net model
        sfm: reference camera parameter
        m: target camera parameter
        dataloader:
      Returns:
        output: dict({
          'mpi_c':'explicit coefficient k0',
          'mpi_a':'alpha transparentcy',
          'mpi_b':'basis',
          'mpi_v':'Kn coefficient'
        })
      '''
    sh = sfm.ref_cam['height'] + sfm.offset * 2
    sw = sfm.ref_cam['width'] + sfm.offset * 2
    #print((sh, sw))
  
    y, x = pt.meshgrid([
      (pt.arange(0, sh, dtype=pt.float)) / (sh-1) * 2 - 1,
      (pt.arange(0, sw, dtype=pt.float)) / (sw-1) * 2 - 1])
  
    coords = pt.cat([x[:,:,None].cuda(), y[:,:,None].cuda()], -1)
    model.eval()
    sh_v = 400
    sw_v = 400
    rangex, rangey =  0.7, 0.6
    y_v, x_v = pt.meshgrid([
      (pt.linspace(-rangey, rangey, sh_v)),
      (pt.linspace(-rangex, rangex , sw_v))])
    #viewing [sh_v, sh_w, 2]
    viewing = pt.cat([x_v[:,:,None].cuda(), y_v[:,:,None].cuda()], -1)
    #hinv_xy [1, sh_v, sh_w, 2 * pos_lev]
    hinv_xy = viewing.view(1, sh_v, sw_v, 2, 1) * model.specular.pos_freq_viewing
    hinv_xy = hinv_xy.view(1, sh_v, sw_v, -1)
    #pe_view [1, sh_v, sw_v, 2 * 2 * pos_lev]
    pe_view = pt.cat([pt.sin(0.5*np.pi *hinv_xy), pt.cos(0.5*np.pi *hinv_xy)], -1)
    #out2 [1, sh, sw, num_basis]
    out2 = model.specular.seq_basis(pe_view)
    #imgs_b [num_basis, 1, sh_v, sw_v]
    imgs_b = pt.tanh(out2.permute([3, 0, 1, 2])).cpu().detach()
  
    n = args.layers * args.sublayers
    imgs_c, imgs_a, imgs_v = [], [], []
  
    with pt.no_grad():
      for i in range(0, n, m):
  
        #coords [sh, sw, 2] --> [1, sh, sw, 2, 1]
        #vxy [1, sh, sw, 2, pos_lev] -->  [1, sh, sw, 2*pos_lev]
        vxy = coords.view(1, sh, sw, 2, 1) * model.pos_freq
        vxy = vxy.view(1, sh, sw, -1)
        #vz [1, sh, sw, depth_lev]
        vz = pt.ones_like(coords.view(1, sh, sw, -1)[..., :1]) * model.z_coords[i:i+1] * model.depth_freq
        #vxyz [1, sh, sw, 2*pos_lev + depth_lev]
        vxyz = pt.cat([vxy, vz], -1)
        bigcoords = pt.cat([pt.sin(vxyz), pt.cos(vxyz)], -1)
        if sfm.offset > 270:
          out =  [model.seq1(bigy) for bigy in [bigcoords[:, :int(sh/2)], bigcoords[:, int(sh/2):]]]
          out = pt.cat(out, 1)
        else:
          out = model.seq1(bigcoords)
        node = 0
  
        mpi_a = out[..., node:node + 1].cpu()
        node +=1
        mpi_a = mpi_a.view(mpi_a.shape[0], 1, mpi_a.shape[1], mpi_a.shape[2])
        imgs_a.append(pt.sigmoid(mpi_a[0] - args.sigmoid_offset))
  
        if i % args.sublayers == 0:
          out = out[..., node:].cpu()
          mpi_v = out.view(out.shape[0], out.shape[1], out.shape[2], 3, -1)
          mpi_v = mpi_v.permute([0, 3, 1, 2, 4])
          mpi_v = mpi_v[0]
          mpi_v =  pt.tanh(mpi_v)
          imgs_v.append(mpi_v)
  
      mpi_c_sig = pt.sigmoid(model.mpi_c)
      mpi_a_sig = pt.stack(imgs_a, 0)
      mpi_v_tanh = pt.stack(imgs_v, 0)
      info = {
        'mpi_c': mpi_c_sig.cpu(),
        'mpi_a': mpi_a_sig.cpu(),
        'mpi_v' : mpi_v_tanh.cpu(),
        'mpi_b':  imgs_b.cpu()
      }
  
    pt.cuda.empty_cache()
    return info
  
  def generateAlpha(model, dataset, dataloader, writer, runpath, suffix="", dataloader_train = None):
    ''' Prediction
      Args.
        model.   --> trained model
        dataset. --> valiade dataset
        writer.  --> tensorboard
    '''
    suffix_str = "/%06d" % suffix if isinstance(suffix, int) else "/"+str(suffix)
    # create webgl only when using -predict or finish training
    if not args.no_webgl and suffix =="":
      info = getMPI(model, dataset.sfm, dataloader = dataloader_train)
  
      outputCompactMPI(info,
                     dataset.sfm,
                     model.planes,
                     runpath + args.model_dir + suffix_str,
                     args.layers,
                     args.sublayers,
                     dataset.sfm.offset,
                     args.invz,
                     webpath=args.web_path,
                     web_width= args.web_width)
  
    if not args.no_eval and len(dataloader) > 0:
      out = evaluation(model,
                       dataset,
                       dataloader,
                       args.ray,
                       runpath + args.model_dir + suffix_str,
                       webpath=args.eval_path,
                       write_csv = not args.no_csv)
      if writer is not None and isinstance(suffix, int):
        for metrics, score in out.items():
          writer.add_scalar('METRICS/{}'.format(metrics), score, suffix)
  
  
  def setLearningRate(optimizer, epoch):
    ds = int(epoch / args.decay_epoch)
    lr = args.lr * (args.decay_rate ** ds)
  
    optimizer.param_groups[0]['lr'] = lr
    if args.lrc > 0:
      optimizer.param_groups[1]['lr'] = lr * args.lrc
  
  def train():
    pt.manual_seed(1)
    np.random.seed(1)
  
    if args.restart or args.clean:
      os.system("rm -rf " + "runs/" + args.model_dir)
    if args.clean:
      exit()
  
    dpath = args.scene
  
    dataset = loadDataset(dpath)
    sampler_train, sampler_val, dataloader_train, dataloader_val = prepareDataloaders(
      dataset,
      dpath,
      random_split = args.random_split,
      train_ratio = args.train_ratio,
      num_workers = args.num_workers
    )
  
    mpi_h = int(dataset.sfm.ref_cam['height'] + dataset.sfm.offset * 2)
    mpi_w = int(dataset.sfm.ref_cam['width'] + dataset.sfm.offset * 2)
    model = Network((args.layers,
                   4,
                   mpi_h,
                   mpi_w,
                   ), dataset.sfm)
  
    if args.lrc > 0:
      #mlp use lower lr, while mpi_c, light dir intensity get higher lr
      my_list = [name for name, params in model.named_parameters() if 'seq1' in name]
      mlp_params = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] in my_list, model.named_parameters()))))
      other_params = list(map(lambda x: x[1],list(filter(lambda kv: kv[0] not in my_list, model.named_parameters()))))
      optimizer = pt.optim.Adam([
        {'params': mlp_params, 'lr': 0},
        {'params': other_params, 'lr': 0}])
    else:
      optimizer = pt.optim.Adam(model.parameters(), lr=0)
  
    start_epoch = 0
    runpath = "runs/"
    ckpt = runpath + args.model_dir + "/ckpt.pt"
    if os.path.exists(ckpt):
      start_epoch = loadFromCheckpoint(ckpt, model, optimizer)
    elif args.pretrained != "":
      start_epoch = loadFromCheckpoint(runpath + args.pretrained + "/ckpt.pt", model, optimizer)
  
    step = start_epoch * len(sampler_train)
  
    if args.epochs < 0 and args.steps < 0:
      raise Exception("Need to specify epochs or steps")
  
    if args.epochs < 0:
      args.epochs = int(np.ceil(args.steps / len(sampler_train)))
  
    if args.predict:
      generateAlpha(model, dataset, dataloader_val, None, runpath, dataloader_train = dataloader_train)
      if not args.no_video:
        if args.render_nearest:
          vid_path = 'video_nearest'
          render_type = 'nearest'
        elif args.render_viewing:
          vid_path = 'viewing_output'
          render_type = 'viewing'
        elif args.render_depth:
          vid_path = 'video_depth'
          render_type = 'depth'
        else:
          vid_path = 'video_output'
          render_type = 'default'
        pt.cuda.empty_cache()
        render_video(model, dataset, args.ray, os.path.join(runpath, vid_path, args.model_dir),
                    render_type = render_type, dataloader = dataloader_train)
      if args.http:
        serve_files(args.model_dir, args.web_path)
      exit()
  
  
    backupConfigAndCode(runpath)
    ts = TrainingStatus(num_steps=args.epochs * len(sampler_train))
    writer = SummaryWriter(runpath + args.model_dir)
    writer.add_text('command',' '.join(sys.argv), 0)
    ts.tic()
  
     # shift by 1 epoch to save last epoch to tensorboard
    for epoch in range(start_epoch, args.epochs+1):
      epoch_loss_total = 0
      epoch_mse = 0
  
      model.train()
  
      for i, feature in enumerate(dataloader_train):
        #print("step: {}".format(i))
        setLearningRate(optimizer, epoch)
        optimizer.zero_grad()
  
        output_shape = feature['image'].shape[-2:]
  
        #sample L-shaped rays
        sel = Lsel(output_shape, args.ray)
  
  
        gt = feature['image']
        gt = gt.view(gt.shape[0], gt.shape[1], gt.shape[2] * gt.shape[3])
        gt = gt[:, :, sel, None].cuda()
        output = model(dataset.sfm, feature, output_shape, sel)
  
        mse = pt.mean((output - gt) ** 2)
  
        loss_total = mse
  
        #tvc regularizer
        tvc = args.tvc * pt.mean(totalVariation(pt.sigmoid(model.mpi_c[:, :3])))
        loss_total = loss_total + tvc
  
        # grad loss
        ox = output[:, :, 1::3,  :] - output[:, :, 0::3, :]
        oy = output[:, :, 2::3,  :] - output[:, :, 0::3, :]
        gx = gt[:, :, 1::3,  :] - gt[:, :, 0::3, :]
        gy = gt[:, :, 2::3, :] - gt[:, :, 0::3, :]
        loss_total = loss_total + args.gradloss * (pt.mean(pt.abs(ox - gx)) + pt.mean(pt.abs(oy - gy)))
  
        epoch_loss_total += loss_total
        epoch_mse += mse
  
        loss_total.backward()
        optimizer.step()
  
        step += 1
        toc_msg = ts.toc(step, loss_total.item())
        if step % args.tb_toc == 0:  print(toc_msg)
        ts.tic()
  
      writer.add_scalar('loss/total', epoch_loss_total/len(sampler_train), epoch)
      writer.add_scalar('loss/mse', epoch_mse/len(sampler_train), epoch)
  
      if epoch % args.tb_saveimage == 0 and args.tb_saveimage > 0:
        with pt.no_grad():
          render = patch_render(model, dataset.sfm, feature, args.ray)
          Spec = getMPI(model, dataset.sfm, m = args.sublayers, dataloader = None)
  
          writer.add_image('images/render', pt.cat([feature['image'].cuda(), render], 2)[0], epoch)
          writer.add_image('images/2_mpia', make_grid(F.interpolate(Spec['mpi_a'],
            (int(mpi_h * 0.3), int(mpi_w * 0.3)),
            mode='area'), 4), epoch)
          writer.add_image('images/2_mpic', make_grid(F.interpolate(Spec['mpi_c'],
            (int(mpi_h * 0.3), int(mpi_w * 0.3)),
            mode='area'), 4), epoch)
          pt.cuda.empty_cache()
  
      if epoch % args.tb_savempi == 0 and args.tb_savempi > 0 and epoch > 0:
        generateAlpha(model, dataset, dataloader_val, writer, runpath, epoch)
        pt.cuda.empty_cache()
  
      var = pt.mean(pt.std(model.illumination, 2)** 2)
      mean = pt.mean(model.illumination)
      writer.add_scalar('loss/illumination_mean', mean, epoch)
      writer.add_scalar('loss/illumination_var', var, epoch)
  
      if (epoch+1) % args.checkpoint == 0 or epoch == args.epochs-1:
        if np.isnan(loss_total.item()):
          exit()
        checkpoint(ckpt, model, optimizer, epoch+1)
  
    print('Finished Training')
    generateAlpha(model, dataset, dataloader_val, None, runpath, dataloader_train = dataloader_train)
    if not args.no_video:
      render_video(model, dataset, args.ray, os.path.join(runpath, 'video_output', args.model_dir))
    if args.http:
      serve_files(args.model_dir, args.web_path)
  
  def checkpoint(file, model, optimizer, epoch):
    print("Checkpointing Model @ Epoch %d ..." % epoch)
    pt.save({
      'epoch': epoch,
      'model_state_dict': model.state_dict(),
      'optimizer_state_dict': optimizer.state_dict(),
      }, file)
  
  def loadFromCheckpoint(file, model, optimizer):
    checkpoint = pt.load(file)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch']
    print("Loading %s Model @ Epoch %d" % (args.pretrained, start_epoch))
    return start_epoch
  
  def backupConfigAndCode(runpath):
    if args.predict or args.clean:
      return
    model_path = os.path.join(runpath, args.model_dir)
    os.makedirs(model_path, exist_ok = True)
    now = datetime.now()
    t = now.strftime("_%Y_%m_%d_%H:%M:%S")
    with open(model_path + "/args.json", 'w') as out:
      json.dump(vars(args), out, indent=2, sort_keys=True)
    os.system("cp " + os.path.abspath(__file__) + " " + model_path + "/")
    os.system("cp " + os.path.abspath(__file__) + " " + model_path + "/" + os.path.basename(__file__).replace(".py", t + ".py"))
    os.system("cp " + model_path + "/args.json " + model_path + "/args" + t + ".json")
  
  
  def loadDataset(dpath):
    # if dataset directory has only image, create LLFF poses
    colmapGenPoses(dpath, args.transformation)
  
    if args.scale == -1:
      args.scale = getDatasetScale(dpath, args.deepview_width, args.llff_width)
  
    if is_deepview(dpath) and args.ref_img == '':
      with open(dpath + "/ref_image.txt", "r") as fi:
        args.ref_img = str(fi.readline().strip())
    render_style = 'llff' if args.nice_llff else 'shiny' if args.nice_shiny else ''
    return OrbiterDataset(dpath, ref_img=args.ref_img, scale=args.scale,
                             dmin=args.dmin,
                             dmax=args.dmax,
                             invz=args.invz,
                             render_style=render_style,
                             offset=args.offset,
                             cv2resize=args.cv2resize)
  
  
  if __name__ == "__main__":
    sys.excepthook = colored_hook(os.path.dirname(os.path.realpath(__file__)))
  
    # delete the colmap product so we force the execution
    keep = ['images', 'database.db']
  
    for file_name in os.listdir(args.scene):
      if file_name in keep:
        continue
      file_name = os.path.join(args.scene, file_name)
      if os.path.isdir(file_name):
        shutil.rmtree(file_name)
      elif os.path.isfile(file_name):
        os.remove(file_name)
  
    # save the arguments
    with open('arguments.json', 'w') as f:
      json.dump(vars(args), f)
  
    # train the model
    train()
  