import os
import logging
logger = logging.getLogger(__name__)
import warnings
warnings.simplefilter('ignore')

import torch
import torch.optim as optim
import matplotlib.pyplot as plt

from lib import rendering
from lib.utils import metrics, utils
from lib.nerf_mlp import MLP


def main(args):
    if torch.cuda.is_available():
        device = 'cuda'
        torch.backends.cudnn.benchmark = True
    else:
        device = 'cpu'

    model = MLP(args.pos_emb_L, args.ray_emb_L, args.hidden_dim,
                args.n_layers, args.n_aux_layers, args.skips).to(device)
    params = [{'params': model.parameters()}]

    if args.n_fine_sample > 0:
        model_coarse = MLP(args.pos_emb_L, args.ray_emb_L, args.c_hidden_dim,
                           args.c_n_layers, args.c_n_aux_layers, args.skips).to(device)
        params.append({'params': model_coarse.parameters()})

    optimizer = optim.Adam(params, args.lr)
    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lambda i: 0.1 ** (i / (args.lr_decay * 1000)))

    if args.dataset_type == 'blender':
        from lib.data import blender
        train_data = blender.Blender(args.root, 'train', near=args.near, far=args.far,
                                    pixel_as_sample=args.pixel_as_sample, factor=args.factor, ndc=args.ndc)
    elif args.dataset_type == 'llff':
        from lib.data import llff
        train_data = llff.LLFF(args.root, args.factor, True, 0.75,
                               args.ndc, args.pixel_as_sample)

    start_iter = 1
    if args.resume or args.phase != 'train':
        checkpoint = torch.load(os.path.join(args.log_dir, 'checkpoint.pth'))
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        start_iter += checkpoint['iter']
        if args.n_fine_sample > 0:
            model_coarse.load_state_dict(checkpoint['model_coarse'])
        logger.info(f'resume training from {start_iter} iteration')

    meter = utils.AvgMeter()

    if args.phase == 'train':
        logger.info('start training')
        for i in range(start_iter, args.n_iters + 1):
            # sample training data
            xyz, ray_d, ray_o, view_d, z, target_rgb = train_data.random_batch_sampling(args.batch_size,
                                                                                        args.n_sample,
                                                                                        center_crop=i < args.crop_iter,
                                                                                        device=device)
            # rendering image
            if args.n_fine_sample > 0:
                # coarse to fine rendering in Sec. 5.2
                (c_rgb, _), (rgb, disparity) = rendering.coarse2fine_rendering(model, model_coarse,
                                                                               xyz, ray_o, ray_d, view_d,
                                                                               z, args.n_fine_sample,
                                                                               args.density_noise_std,
                                                                               True, white_bkgd=args.dataset_type=='blender')
            else:
                rgb, disparity, _ = rendering.rendering(model, xyz, ray_d, view_d,
                                                        z, args.density_noise_std,
                                                        white_bkgd=args.dataset_type=='blender')
            # update gradient
            optimizer.zero_grad()
            psnr, mse = metrics.calc_psnr(rgb, target_rgb)
            if args.n_fine_sample > 0:
                c_mse = metrics.calc_mse(c_rgb, target_rgb)
                mse = c_mse + mse
            mse.backward()
            optimizer.step()
            scheduler.step()

            meter.add({'mse': mse.item(), 'psnr': psnr.item()})
            # display mse and psnr
            if i % args.disp_freq == 0:
                state = meter.state(f'[{i}/{args.n_iters}]',
                                    f'lr {optimizer.param_groups[0]["lr"]:.6g}')
                logger.info(state)
            # visualize rendered image
            if i % args.val_freq == 0:
                logger.info(f'synthesize novel view')
                with torch.no_grad():
                    height, width = train_data.img_size
                    # make directory for saving results
                    rgb_dir = os.path.join(args.log_dir, 'rgb', f'{i}iteration')
                    disp_dir = os.path.join(args.log_dir, 'disparity', f'{i}iteration')
                    os.makedirs(rgb_dir, exist_ok=True)
                    os.makedirs(disp_dir, exist_ok=True)
                    for idx in range(train_data.render_poses.shape[0]):
                        # sampling ray
                        xyz, ray_d, ray_o, view_d, z = train_data.sample_ray_for_novel_view_rendering(idx, args.n_sample, device)
                        rgb_list, disparity_list = [], []
                        # chunk data for saving GPU memory
                        iterator = zip(xyz.chunk(args.test_chunk),
                                       ray_d.chunk(args.test_chunk),
                                       ray_o.chunk(args.test_chunk),
                                       view_d.chunk(args.test_chunk),
                                       z.chunk(args.test_chunk))
                        for xyz_c, ray_d_c, ray_o_c, view_d_c, z_c in iterator:
                            # rendering
                            if args.n_fine_sample > 0:
                                _, (rgb, disparity) = rendering.coarse2fine_rendering(model, model_coarse,
                                                                                      xyz_c, ray_o_c, ray_d_c,
                                                                                      view_d_c, z_c, args.n_fine_sample,
                                                                                    white_bkgd=args.dataset_type=='blender')
                            else:
                                rgb, disparity, _ = rendering.rendering(model, xyz_c, ray_d_c, view_d_c, z_c,
                                                                        white_bkgd=args.dataset_type=='blender')
                            rgb_list.append(rgb)
                            disparity_list.append(disparity)
                        # visualize rendered image and disparity
                        rgb = torch.cat(rgb_list, 0)
                        disparity = torch.cat(disparity_list, 0)
                        rgb = rgb.reshape(height, width, 3).clamp(max=1, min=0).to('cpu').numpy()
                        disparity = disparity.reshape(height, width).to('cpu').numpy()
                        disparity = disparity / disparity.max()
                        plt.imsave(os.path.join(rgb_dir, f'{str(idx).zfill(5)}.png'), rgb)
                        plt.imsave(os.path.join(disp_dir, f'{str(idx).zfill(5)}.png'), disparity, cmap=plt.cm.gray)
            # checkpoint
            if i  % args.save_freq == 0:
                logger.info('checkpoint')
                checkpoint = {'model': model.state_dict(),
                            'optimizer': optimizer.state_dict(),
                            'scheduler': scheduler.state_dict(),
                            'iter': i}
                if args.n_fine_sample > 0:
                    checkpoint['model_coarse'] = model_coarse.state_dict()
                torch.save(checkpoint, os.path.join(args.log_dir, 'checkpoint.pth'))

    else: # evaluation
        logger.info(f'evaluation on {args.phase} set')
        if args.dataset_type == 'blender':
            from lib.data import blender
            data = blender.Blender(args.root, args.phase, near=args.near, far=args.far, factor=args.factor, ndc=args.ndc)
        else:
            raise NotImplementedError
        with torch.no_grad():
            height, width = data.img_size
            # make directory for saving results
            rgb_dir = os.path.join(args.log_dir, f'{args.phase}_rgb')
            disp_dir = os.path.join(args.log_dir, f'{args.phase}_disparity')
            os.makedirs(rgb_dir, exist_ok=True)
            os.makedirs(disp_dir, exist_ok=True)
            mean_psnr = 0
            for idx in range(len(data)):
                # sampling ray
                xyz, ray_d, ray_o, view_d, z, target_rgb = data.sample_image(idx, args.n_sample, False, device)
                rgb_list, disparity_list = [], []
                iterator = zip(xyz.chunk(args.test_chunk),
                                ray_d.chunk(args.test_chunk),
                                ray_o.chunk(args.test_chunk),
                                view_d.chunk(args.test_chunk),
                                z.chunk(args.test_chunk))
                for xyz_c, ray_d_c, ray_o_c, view_d_c, z_c in iterator:
                    # rendering
                    if args.n_fine_sample > 0:
                        _, (rgb, disparity) = rendering.coarse2fine_rendering(model, model_coarse,
                                                                              xyz_c, ray_o_c, ray_d_c,
                                                                              view_d_c, z_c, args.n_fine_sample,
                                                                              white_bkgd=args.dataset_type=='blender')
                    else:
                        rgb, disparity, _ = rendering.rendering(model, xyz_c, ray_d_c, view_d_c, z_c,
                                                                white_bkgd=args.dataset_type=='blender')
                    rgb_list.append(rgb)
                    disparity_list.append(disparity)
                rgb = torch.cat(rgb_list, 0)
                disparity = torch.cat(disparity_list, 0)
                # calculate PSNR
                psnr, _ = metrics.calc_psnr(rgb, target_rgb)
                logger.info(f'image {idx} PSNR {psnr.item()}')
                mean_psnr += psnr / len(data)
                # saving rendered image and disparity
                rgb = rgb.reshape(height, width, 3).clamp(max=1, min=0).to('cpu').numpy()
                disparity = disparity.reshape(height, width).to('cpu').numpy()
                disparity = disparity / disparity.max()
                plt.imsave(os.path.join(rgb_dir, f'{str(idx).zfill(5)}.png'), rgb)
                plt.imsave(os.path.join(disp_dir, f'{str(idx).zfill(5)}.png'), disparity, cmap=plt.cm.gray)
            logger.info(f'mean PSNR {mean_psnr}')


if __name__=='__main__':
    import json
    import argparse

    parser = argparse.ArgumentParser()
    # dataset
    parser.add_argument('--root', type=str, required=True,
                        help='/path/to/dataset')
    parser.add_argument('--dataset_type', default='blender', choices=['blender', 'llff'], type=str)
    parser.add_argument('--log_dir', default='log', type=str)
    parser.add_argument('--ndc', action='store_true')
    parser.add_argument('--factor', default=None, type=int,
                        help='downsampling factor')
    # nerf
    ## position embedding
    parser.add_argument('--pos_emb_L', default=10, type=int)
    parser.add_argument('--ray_emb_L', default=4, type=int)
    ## model parameters
    parser.add_argument('--hidden_dim', default=256, type=int)
    parser.add_argument('--n_layers', default=8, type=int)
    parser.add_argument('--n_aux_layers', default=1, type=int)
    parser.add_argument('--skips', default='4', type=str)
    parser.add_argument('--n_fine_sample', default=128, type=int)
    ## model parameters for coarse model
    parser.add_argument('--c_hidden_dim', default=256, type=int)
    parser.add_argument('--c_n_layers', default=8, type=int)
    parser.add_argument('--c_n_aux_layers', default=1, type=int)
    ## misc
    parser.add_argument('--crop_iter', default=0, type=int)
    parser.add_argument('--near', default=None, type=float)
    parser.add_argument('--far', default=None, type=float)
    parser.add_argument('--density_noise_std', default=0, type=float)
    # optimizer
    parser.add_argument('--lr', default=5e-4, type=float)
    parser.add_argument('--n_iters', default=200000, type=int)
    parser.add_argument('--lr_decay', default=250, type=float)
    parser.add_argument('--batch_size', default=4096, type=int)
    parser.add_argument('--n_sample', default=64, type=int)
    # misc
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--save_freq', default=20000, type=int)
    parser.add_argument('--disp_freq', default=500, type=int)
    parser.add_argument('--val_freq', default=50000, type=int)
    parser.add_argument('--yaml', default=None, type=str)
    parser.add_argument('--json', default=None, type=str)
    parser.add_argument('--test_chunk', default=32768, type=int)
    parser.add_argument('--pixel_as_sample', action='store_true')
    parser.add_argument('--phase', default='train', choices=['train', 'val', 'test'])

    args = parser.parse_args()

    if args.yaml is not None:
        yaml_cfg = utils.load_yaml(args.yaml)
        args = utils.override_config(args, yaml_cfg)
    if args.json is not None:
        json_cfg = utils.load_yaml(args.json)
        args = utils.override_config(args, json_cfg)

    os.makedirs(args.log_dir, exist_ok=True)
    os.makedirs(os.path.join(args.log_dir, 'rgb'), exist_ok=True)
    os.makedirs(os.path.join(args.log_dir, 'disparity'), exist_ok=True)
    with open(os.path.join(args.log_dir, 'args.json'), 'w') as f:
        json.dump(args.__dict__, f)

    args.skips = list(map(int, args.skips.split(',')))

    utils.set_seed(args.seed)
    utils.setup_logger(args.log_dir, args.resume)

    logger.info(args)

    main(args)

import os
import logging
logger = logging.getLogger(__name__)
import warnings
warnings.simplefilter('ignore')

import torch
import torch.optim as optim
import matplotlib.pyplot as plt

from lib import rendering
from lib.utils import metrics, utils
from lib.nerf_mlp import MLP


def main(args):
    if torch.cuda.is_available():
        device = 'cuda'
        torch.backends.cudnn.benchmark = True
    else:
        device = 'cpu'

    model = MLP(args.pos_emb_L, args.ray_emb_L, args.hidden_dim,
                args.n_layers, args.n_aux_layers, args.skips).to(device)
    params = [{'params': model.parameters()}]

    if args.n_fine_sample > 0:
        model_coarse = MLP(args.pos_emb_L, args.ray_emb_L, args.c_hidden_dim,
                           args.c_n_layers, args.c_n_aux_layers, args.skips).to(device)
        params.append({'params': model_coarse.parameters()})

    optimizer = optim.Adam(params, args.lr)
    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lambda i: 0.1 ** (i / (args.lr_decay * 1000)))

    if args.dataset_type == 'blender':
        from lib.data import blender
        train_data = blender.Blender(args.root, 'train', near=args.near, far=args.far,
                                    pixel_as_sample=args.pixel_as_sample, factor=args.factor, ndc=args.ndc)
    elif args.dataset_type == 'llff':
        from lib.data import llff
        train_data = llff.LLFF(args.root, args.factor, True, 0.75,
                               args.ndc, args.pixel_as_sample)

    start_iter = 1
    if args.resume or args.phase != 'train':
        checkpoint = torch.load(os.path.join(args.log_dir, 'checkpoint.pth'))
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        start_iter += checkpoint['iter']
        if args.n_fine_sample > 0:
            model_coarse.load_state_dict(checkpoint['model_coarse'])
        logger.info(f'resume training from {start_iter} iteration')

    meter = utils.AvgMeter()

    if args.phase == 'train':
        logger.info('start training')
        for i in range(start_iter, args.n_iters + 1):
            # sample training data
            xyz, ray_d, ray_o, view_d, z, target_rgb = train_data.random_batch_sampling(args.batch_size,
                                                                                        args.n_sample,
                                                                                        center_crop=i < args.crop_iter,
                                                                                        device=device)
            # rendering image
            if args.n_fine_sample > 0:
                # coarse to fine rendering in Sec. 5.2
                (c_rgb, _), (rgb, disparity) = rendering.coarse2fine_rendering(model, model_coarse,
                                                                               xyz, ray_o, ray_d, view_d,
                                                                               z, args.n_fine_sample,
                                                                               args.density_noise_std,
                                                                               True, white_bkgd=args.dataset_type=='blender')
            else:
                rgb, disparity, _ = rendering.rendering(model, xyz, ray_d, view_d,
                                                        z, args.density_noise_std,
                                                        white_bkgd=args.dataset_type=='blender')
            # update gradient
            optimizer.zero_grad()
            psnr, mse = metrics.calc_psnr(rgb, target_rgb)
            if args.n_fine_sample > 0:
                c_mse = metrics.calc_mse(c_rgb, target_rgb)
                mse = c_mse + mse
            mse.backward()
            optimizer.step()
            scheduler.step()

            meter.add({'mse': mse.item(), 'psnr': psnr.item()})
            # display mse and psnr
            if i % args.disp_freq == 0:
                state = meter.state(f'[{i}/{args.n_iters}]',
                                    f'lr {optimizer.param_groups[0]["lr"]:.6g}')
                logger.info(state)
            # visualize rendered image
            if i % args.val_freq == 0:
                logger.info(f'synthesize novel view')
                with torch.no_grad():
                    height, width = train_data.img_size
                    # make directory for saving results
                    rgb_dir = os.path.join(args.log_dir, 'rgb', f'{i}iteration')
                    disp_dir = os.path.join(args.log_dir, 'disparity', f'{i}iteration')
                    os.makedirs(rgb_dir, exist_ok=True)
                    os.makedirs(disp_dir, exist_ok=True)
                    for idx in range(train_data.render_poses.shape[0]):
                        # sampling ray
                        xyz, ray_d, ray_o, view_d, z = train_data.sample_ray_for_novel_view_rendering(idx, args.n_sample, device)
                        rgb_list, disparity_list = [], []
                        # chunk data for saving GPU memory
                        iterator = zip(xyz.chunk(args.test_chunk),
                                       ray_d.chunk(args.test_chunk),
                                       ray_o.chunk(args.test_chunk),
                                       view_d.chunk(args.test_chunk),
                                       z.chunk(args.test_chunk))
                        for xyz_c, ray_d_c, ray_o_c, view_d_c, z_c in iterator:
                            # rendering
                            if args.n_fine_sample > 0:
                                _, (rgb, disparity) = rendering.coarse2fine_rendering(model, model_coarse,
                                                                                      xyz_c, ray_o_c, ray_d_c,
                                                                                      view_d_c, z_c, args.n_fine_sample,
                                                                                    white_bkgd=args.dataset_type=='blender')
                            else:
                                rgb, disparity, _ = rendering.rendering(model, xyz_c, ray_d_c, view_d_c, z_c,
                                                                        white_bkgd=args.dataset_type=='blender')
                            rgb_list.append(rgb)
                            disparity_list.append(disparity)
                        # visualize rendered image and disparity
                        rgb = torch.cat(rgb_list, 0)
                        disparity = torch.cat(disparity_list, 0)
                        rgb = rgb.reshape(height, width, 3).clamp(max=1, min=0).to('cpu').numpy()
                        disparity = disparity.reshape(height, width).to('cpu').numpy()
                        disparity = disparity / disparity.max()
                        plt.imsave(os.path.join(rgb_dir, f'{str(idx).zfill(5)}.png'), rgb)
                        plt.imsave(os.path.join(disp_dir, f'{str(idx).zfill(5)}.png'), disparity, cmap=plt.cm.gray)
            # checkpoint
            if i  % args.save_freq == 0:
                logger.info('checkpoint')
                checkpoint = {'model': model.state_dict(),
                            'optimizer': optimizer.state_dict(),
                            'scheduler': scheduler.state_dict(),
                            'iter': i}
                if args.n_fine_sample > 0:
                    checkpoint['model_coarse'] = model_coarse.state_dict()
                torch.save(checkpoint, os.path.join(args.log_dir, 'checkpoint.pth'))

    else: # evaluation
        logger.info(f'evaluation on {args.phase} set')
        if args.dataset_type == 'blender':
            from lib.data import blender
            data = blender.Blender(args.root, args.phase, near=args.near, far=args.far, factor=args.factor, ndc=args.ndc)
        else:
            raise NotImplementedError
        with torch.no_grad():
            height, width = data.img_size
            # make directory for saving results
            rgb_dir = os.path.join(args.log_dir, f'{args.phase}_rgb')
            disp_dir = os.path.join(args.log_dir, f'{args.phase}_disparity')
            os.makedirs(rgb_dir, exist_ok=True)
            os.makedirs(disp_dir, exist_ok=True)
            mean_psnr = 0
            for idx in range(len(data)):
                # sampling ray
                xyz, ray_d, ray_o, view_d, z, target_rgb = data.sample_image(idx, args.n_sample, False, device)
                rgb_list, disparity_list = [], []
                iterator = zip(xyz.chunk(args.test_chunk),
                                ray_d.chunk(args.test_chunk),
                                ray_o.chunk(args.test_chunk),
                                view_d.chunk(args.test_chunk),
                                z.chunk(args.test_chunk))
                for xyz_c, ray_d_c, ray_o_c, view_d_c, z_c in iterator:
                    # rendering
                    if args.n_fine_sample > 0:
                        _, (rgb, disparity) = rendering.coarse2fine_rendering(model, model_coarse,
                                                                              xyz_c, ray_o_c, ray_d_c,
                                                                              view_d_c, z_c, args.n_fine_sample,
                                                                              white_bkgd=args.dataset_type=='blender')
                    else:
                        rgb, disparity, _ = rendering.rendering(model, xyz_c, ray_d_c, view_d_c, z_c,
                                                                white_bkgd=args.dataset_type=='blender')
                    rgb_list.append(rgb)
                    disparity_list.append(disparity)
                rgb = torch.cat(rgb_list, 0)
                disparity = torch.cat(disparity_list, 0)
                # calculate PSNR
                psnr, _ = metrics.calc_psnr(rgb, target_rgb)
                logger.info(f'image {idx} PSNR {psnr.item()}')
                mean_psnr += psnr / len(data)
                # saving rendered image and disparity
                rgb = rgb.reshape(height, width, 3).clamp(max=1, min=0).to('cpu').numpy()
                disparity = disparity.reshape(height, width).to('cpu').numpy()
                disparity = disparity / disparity.max()
                plt.imsave(os.path.join(rgb_dir, f'{str(idx).zfill(5)}.png'), rgb)
                plt.imsave(os.path.join(disp_dir, f'{str(idx).zfill(5)}.png'), disparity, cmap=plt.cm.gray)
            logger.info(f'mean PSNR {mean_psnr}')


if __name__=='__main__':
    import json
    import argparse

    parser = argparse.ArgumentParser()
    # dataset
    parser.add_argument('--root', type=str, required=True,
                        help='/path/to/dataset')
    parser.add_argument('--dataset_type', default='blender', choices=['blender', 'llff'], type=str)
    parser.add_argument('--log_dir', default='log', type=str)
    parser.add_argument('--ndc', action='store_true')
    parser.add_argument('--factor', default=None, type=int,
                        help='downsampling factor')
    # nerf
    ## position embedding
    parser.add_argument('--pos_emb_L', default=10, type=int)
    parser.add_argument('--ray_emb_L', default=4, type=int)
    ## model parameters
    parser.add_argument('--hidden_dim', default=256, type=int)
    parser.add_argument('--n_layers', default=8, type=int)
    parser.add_argument('--n_aux_layers', default=1, type=int)
    parser.add_argument('--skips', default='4', type=str)
    parser.add_argument('--n_fine_sample', default=128, type=int)
    ## model parameters for coarse model
    parser.add_argument('--c_hidden_dim', default=256, type=int)
    parser.add_argument('--c_n_layers', default=8, type=int)
    parser.add_argument('--c_n_aux_layers', default=1, type=int)
    ## misc
    parser.add_argument('--crop_iter', default=0, type=int)
    parser.add_argument('--near', default=None, type=float)
    parser.add_argument('--far', default=None, type=float)
    parser.add_argument('--density_noise_std', default=0, type=float)
    # optimizer
    parser.add_argument('--lr', default=5e-4, type=float)
    parser.add_argument('--n_iters', default=200000, type=int)
    parser.add_argument('--lr_decay', default=250, type=float)
    parser.add_argument('--batch_size', default=4096, type=int)
    parser.add_argument('--n_sample', default=64, type=int)
    # misc
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--save_freq', default=20000, type=int)
    parser.add_argument('--disp_freq', default=500, type=int)
    parser.add_argument('--val_freq', default=50000, type=int)
    parser.add_argument('--yaml', default=None, type=str)
    parser.add_argument('--json', default=None, type=str)
    parser.add_argument('--test_chunk', default=32768, type=int)
    parser.add_argument('--pixel_as_sample', action='store_true')
    parser.add_argument('--phase', default='train', choices=['train', 'val', 'test'])

    args = parser.parse_args()

    if args.yaml is not None:
        yaml_cfg = utils.load_yaml(args.yaml)
        args = utils.override_config(args, yaml_cfg)
    if args.json is not None:
        json_cfg = utils.load_yaml(args.json)
        args = utils.override_config(args, json_cfg)

    os.makedirs(args.log_dir, exist_ok=True)
    os.makedirs(os.path.join(args.log_dir, 'rgb'), exist_ok=True)
    os.makedirs(os.path.join(args.log_dir, 'disparity'), exist_ok=True)
    with open(os.path.join(args.log_dir, 'args.json'), 'w') as f:
        json.dump(args.__dict__, f)

    args.skips = list(map(int, args.skips.split(',')))

    utils.set_seed(args.seed)
    utils.setup_logger(args.log_dir, args.resume)

    logger.info(args)

    main(args)

# sample training data # rendering image # coarse to fine rendering in Sec. 5.2 # update gradient # display mse and psnr # visualize rendered image # make directory for saving results # sampling ray # chunk data for saving GPU memory # rendering # visualize rendered image and disparity # checkpoint # evaluation # make directory for saving results # sampling ray # rendering # calculate PSNR # saving rendered image and disparity # dataset # nerf ## position embedding ## model parameters ## model parameters for coarse model ## misc # optimizer # misc