main_nerf.py

import torch
import argparse

from nerf.provider import NeRFDataset
from nerf.gui import NeRFGUI
from nerf.utils import *

from editing.editgrid import EditGrid

from functools import partial
from loss import huber_loss

# torch.autograd.set_detect_anomaly(True)

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('path', type=str)
    parser.add_argument('-O', action='store_true', help="equals --fp16 --cuda_ray --preload")
    parser.add_argument('--test', action='store_true', help="test mode")
    parser.add_argument('--workspace', type=str, default='workspace')
    parser.add_argument('--seed', type=int, default=0)

    ### training options
    parser.add_argument('--iters', type=int, default=30000, help="training iters")
    parser.add_argument('--lr', type=float, default=1e-2, help="initial learning rate")
    parser.add_argument('--ckpt', type=str, default='latest')
    parser.add_argument('--num_rays', type=int, default=4096, help="num rays sampled per image for each training step")
    parser.add_argument('--cuda_ray', action='store_true', help="use CUDA raymarching instead of pytorch")
    parser.add_argument('--max_steps', type=int, default=1024,
                        help="max num steps sampled per ray (only valid when using --cuda_ray)")
    parser.add_argument('--num_steps', type=int, default=512,
                        help="num steps sampled per ray (only valid when NOT using --cuda_ray)")
    parser.add_argument('--upsample_steps', type=int, default=0,
                        help="num steps up-sampled per ray (only valid when NOT using --cuda_ray)")
    parser.add_argument('--update_extra_interval', type=int, default=16,
                        help="iter interval to update extra status (only valid when using --cuda_ray)")
    parser.add_argument('--max_ray_batch', type=int, default=4096,
                        help="batch size of rays at inference to avoid OOM (only valid when NOT using --cuda_ray)")
    parser.add_argument('--patch_size', type=int, default=1,
                        help="[experimental] render patches in training, so as to apply LPIPS loss. 1 means disabled, use [64, 32, 16] to enable")

    ### network backbone options
    parser.add_argument('--fp16', action='store_true', help="use amp mixed precision training")
    parser.add_argument('--ff', action='store_true', help="use fully-fused MLP")
    parser.add_argument('--tcnn', action='store_true', help="use TCNN backend")

    ### dataset options
    parser.add_argument('--color_space', type=str, default='srgb', help="Color space, supports (linear, srgb)")
    parser.add_argument('--preload', action='store_true',
                        help="preload all data into GPU, accelerate training but use more GPU memory")
    # (the default value is for the fox dataset)
    parser.add_argument('--bound', type=float, default=2,
                        help="assume the scene is bounded in box[-bound, bound]^3, if > 1, will invoke adaptive ray marching.")
    parser.add_argument('--scale', type=float, default=0.33, help="scale camera location into box[-bound, bound]^3")
    parser.add_argument('--offset', type=float, nargs='*', default=[0, 0, 0], help="offset of camera location")
    parser.add_argument('--dt_gamma', type=float, default=1 / 128,
                        help="dt_gamma (>=0) for adaptive ray marching. set to 0 to disable, >0 to accelerate rendering (but usually with worse quality)")
    parser.add_argument('--min_near', type=float, default=0.2, help="minimum near distance for camera")
    parser.add_argument('--density_thresh', type=float, default=10, help="threshold for density grid to be occupied")
    parser.add_argument('--bg_radius', type=float, default=-1,
                        help="if positive, use a background model at sphere(bg_radius)")
    parser.add_argument('--no_bg', action='store_true',
                        help="for real world scenes such as LLFF, Mip360, results in different blending")

    ### GUI options
    parser.add_argument('--gui', action='store_true', help="start a GUI")
    parser.add_argument('--W', type=int, default=1920, help="GUI width")
    parser.add_argument('--H', type=int, default=1080, help="GUI height")
    parser.add_argument('--radius', type=float, default=5, help="default GUI camera radius from center")
    parser.add_argument('--fovy', type=float, default=50, help="default GUI camera fovy")
    parser.add_argument('--max_spp', type=int, default=24, help="GUI rendering max sample per pixel")

    ### experimental
    parser.add_argument('--error_map', action='store_true', help="use error map to sample rays")
    parser.add_argument('--clip_text', type=str, default='', help="text input for CLIP guidance")
    parser.add_argument('--rand_pose', type=int, default=-1,
                        help="<0 uses no rand pose, =0 only uses rand pose, >0 sample one rand pose every $ known poses")

    ### Editing and Stylization
    parser.add_argument('--ablation_folder', type=str, default="test",
                        help="folder name where everything is saved to")
    parser.add_argument('--tv_weight', type=float, default=0., help="weight for the tv loss")
    parser.add_argument('--depth_disc_weight', type=float, default=0.,
                        help="weight for the depth disc loss")
    parser.add_argument('--smooth_trans_weight', type=float, default=0.,
                        help="weight for the smooth continuity loss")
    parser.add_argument('--style_weight', type=float, default=0., help="weight for the style loss")
    parser.add_argument('--style_layers', action='append', type=int,
                        help="layers to use for style transfer")
    parser.add_argument('--tv_depth_guide', action='store_true',
                        help="whether to use depth discontinuities for guiding the TV loss")
    parser.add_argument('--intensity_weight', type=float, default=0.,
                        help="weight for the intensity-based regularization")
    parser.add_argument('--preserve_color', action='store_true',
                        help="whether to preserve the color of the NeRF model")
    parser.add_argument('--train_steps_style', type=int, default=3000,
                        help="number of steps to train the style encoder")
    parser.add_argument('--train_steps_distill', type=int, default=3000,
                        help="number of steps to distill the style transfer")
    parser.add_argument('--style_image', type=str, default=None,
                        help="style transfer image (in the style_images directory)")
    parser.add_argument('--offset_loss', type=float, default=0.,
                        help="loss weight for the offset")
    parser.add_argument('--weight_loss_non_uniform', type=float, default=0.,
                        help="non-uniform weight loss (each pixel should be affected by few pixels only)")
    parser.add_argument('--weight_loss_uniform', type=float, default=0.,
                        help="uniform weight loss (each palette should contribute equally)")
    parser.add_argument('--palette_loss_valid', type=float, default=0.,
                        help="valid palette loss (each palette shouldn't be out-of-gamut)")
    parser.add_argument('--palette_loss_distinct', type=float, default=0.,
                        help="distinct palette loss (each palette should be distinct)")
    parser.add_argument('--ablation_dir', type=str, default="ablation_",
                        help="directory to save ablation_folder in")
    parser.add_argument('--num_palette_bases', type=int, default=4, help="number of palette bases")
    parser.add_argument('--distill_palette_steps', type=int, default=1500,
                        help="modify num palette bases x iterations before termination")
    parser.add_argument('--run_all', action='store_true',
                        help="Run everything (StyleEnc train, distill)")
    parser.add_argument('--warmup_iterations', type=int, default=1000, help="number of warmup steps")
    parser.add_argument('--crop_size', type=int, default=256, help='size of the random crop for stylization')

    parser.add_argument('--style_enc_path', type=str, help='load a style encoder for automatic recoloring/stylization')
    parser.add_argument('--palette_path', type=str, help='load a color palette for automatic recoloring/stylization')
    parser.add_argument('--depth_diff', type=float, default=0.5)
    parser.add_argument('--use_error_maps', action='store_true')
    parser.add_argument('--load_edit_dataset', type=str, default=None)
    parser.add_argument('--filter_close_point', action='store_true')

    ### NPR-based rendering
    parser.add_argument('--ref_npr_config', type=str, default=None,
                        help='path for config file for ref-based stylization')
    parser.add_argument('--reg_max_dist', type=float, default=2e-2,
                        help="maximal distance for registration")
    parser.add_argument('--tv_min_dist', type=float, default=10e-2,
                        help="maximum distance for full tv_loss")
    parser.add_argument('--min_tv_factor', type=float, default=0.1, help="minimum factor for tv")
    parser.add_argument('--cos_loss_factor', type=float, default=2.5, help="cos loss factor")
    parser.add_argument('--mse_loss', type=float, default=6., help="ray reg loss")
    parser.add_argument('--color_patch_loss', type=float, default=3e1, help="color supervision loss")

    parser.add_argument('--style_weight_d', type=float, default=5e-1, help="color supervision loss")
    parser.add_argument('--depth_weight_d', type=float, default=1e-3, help="color supervision loss")
    parser.add_argument('--feature_size', type=int, default=256, help="color supervision loss")

    opt = parser.parse_args()

    # default args for the list parameter wont work...
    if opt.style_layers is None:
        opt.style_layers = [10, 12, 14]

    if opt.gui:
        if not os.path.exists(opt.ablation_dir):
            os.mkdir(opt.ablation_dir)
        if not os.path.exists(os.path.join(opt.ablation_dir, opt.ablation_folder)):
            os.mkdir(os.path.join(opt.ablation_dir, opt.ablation_folder))

    if opt.O:
        opt.fp16 = True
        opt.cuda_ray = True
        opt.preload = True

    if opt.patch_size > 1:
        opt.error_map = False  # do not use error_map if use patch-based training
        # assert opt.patch_size > 16, "patch_size should > 16 to run LPIPS loss."
        assert opt.num_rays % (opt.patch_size ** 2) == 0, "patch_size ** 2 should be dividable by num_rays."

    if opt.ff:
        opt.fp16 = True
        assert opt.bg_radius <= 0, "background model is not implemented for --ff"
        from nerf.network_ff import NeRFNetwork
    elif opt.tcnn:
        opt.fp16 = True
        assert opt.bg_radius <= 0, "background model is not implemented for --tcnn"
        from nerf.network_tcnn import NeRFNetwork
    else:
        from nerf.network import NeRFNetwork

    print(opt)

    seed_everything(opt.seed)

    model = NeRFNetwork(
        encoding="hashgrid",
        bound=opt.bound,
        cuda_ray=opt.cuda_ray,
        density_scale=1,
        min_near=opt.min_near,
        density_thresh=opt.density_thresh,
        bg_radius=opt.bg_radius,
    )

    print(model)

    criterion = torch.nn.MSELoss(reduction='none')
    # criterion = partial(huber_loss, reduction='none')
    # criterion = torch.nn.HuberLoss(reduction='none', beta=0.1) # only available after torch 1.10 ?

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    if opt.test:

        metrics = [PSNRMeter(), LPIPSMeter(device=device)]
        trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, criterion=criterion, fp16=opt.fp16,
                          metrics=metrics, use_checkpoint=opt.ckpt)

        if opt.gui:
            gui = NeRFGUI(opt, trainer)
            gui.render()

        else:
            test_loader = NeRFDataset(opt, device=device, type='test').dataloader()

            if test_loader.has_gt:
                trainer.evaluate(test_loader)  # blender has gt, so evaluate it.

            trainer.test(test_loader, write_video=True)  # test and save video

            trainer.save_mesh(resolution=256, threshold=10)

    else:

        optimizer = lambda model: torch.optim.Adam(model.get_params(opt.lr), betas=(0.9, 0.99), eps=1e-15)

        train_loader = NeRFDataset(opt, device=device, type='train').dataloader()
        test_loader = NeRFDataset(opt, device=device, type='val').dataloader()
        t_loader = NeRFDataset(opt, device=device, type='test').dataloader()
        try:
            vid_loader = NeRFDataset(opt, device=device, type='video').dataloader()
        except FileNotFoundError:
            vid_loader = None
            # FIX: use the NeRF test dataset for rendering a video
            if t_loader is not None:
                vid_loader = t_loader
        except AttributeError:
            vid_loader = None

        # decay to 0.1 * init_lr at last iter step
        scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer,
                                                                  lambda iter: 0.1 ** min(iter / opt.iters, 1))

        metrics = [PSNRMeter(), LPIPSMeter(device=device)]
        trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, optimizer=optimizer,
                          criterion=criterion, ema_decay=0.95, fp16=opt.fp16, lr_scheduler=scheduler,
                          scheduler_update_every_step=True, metrics=metrics, use_checkpoint=opt.ckpt, eval_interval=50)

        if opt.gui:
            gui = NeRFGUI(opt, trainer, train_loader, val_loader=test_loader, video_loader=vid_loader,
                          test_loader=t_loader)
            gui.render()

        else:
            valid_loader = NeRFDataset(opt, device=device, type='val', downscale=1).dataloader()

            max_epoch = np.ceil(opt.iters / len(train_loader)).astype(np.int32)
            trainer.train(train_loader, valid_loader, max_epoch)

            # also test
            test_loader = NeRFDataset(opt, device=device, type='test').dataloader()

            if test_loader.has_gt:
                trainer.evaluate(test_loader)  # blender has gt, so evaluate it.

            trainer.test(test_loader, write_video=False)  # test and save video

            # trainer.save_mesh(resolution=256, threshold=10)