main_alignmif.py

import torch
import configargparse
import os
import numpy as np

from lidarnerf.nerf.utils import (seed_everything, RMSEMeter, MAEMeter,
                                  PSNRMeter, LPIPSMeter, DepthMeter,
                                  PointsMeter, SSIMMeter)


def get_arg_parser():
    parser = configargparse.ArgumentParser()

    parser.add_argument("--config",
                        is_config_file=True,
                        default="configs/kitti360_1908.txt",
                        help="config file path")
    parser.add_argument("--path", type=str, default="data/kitti360")
    parser.add_argument("-L",
                        action="store_true",
                        help="equals --fp16 --preload")
    parser.add_argument("--test", action="store_true", help="test mode")
    parser.add_argument("--test_eval",
                        action="store_true",
                        help="test and eval mode")
    parser.add_argument("--workspace", type=str, default="workspace")
    parser.add_argument("--cluster_summary_path",
                        type=str,
                        default="/summary",
                        help="Overwrite default summary path if on cluster")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--dataloader",
                        type=str,
                        choices=("kitti360", "nerf_mvl", "waymo",
                                 "aiodrive"),
                        default="kitti360")
    parser.add_argument("--network",
                        type=str,
                        choices=("tcnn", "mif", "alignmif"),
                        default="tcnn")
    parser.add_argument("--sequence_id", type=str, default="1908")


    # multi-nerf
    parser.add_argument("--activate_levels", type=int, default=0)

    parser.add_argument('--enable_rgb', action='store_true', help="Enable rgb.")
    parser.add_argument('--alpha_rgb', type=float, default=1)
    parser.add_argument('--rgb_loss',
                        type=str,
                        default='mse',
                        help="l1, bce, mse, huber")

    parser.add_argument("--enable_lidar",
                        action="store_true",
                        help="Enable lidar.")
    parser.add_argument("--alpha_d", type=float, default=1e3)
    parser.add_argument("--alpha_r", type=float, default=1)
    parser.add_argument("--alpha_rd", type=float, default=1)
    parser.add_argument("--alpha_i", type=float, default=1)
    parser.add_argument("--alpha_grad_norm", type=float, default=1)
    parser.add_argument("--alpha_spatial", type=float, default=0.1)
    parser.add_argument("--alpha_tv", type=float, default=1)
    parser.add_argument("--alpha_grad", type=float, default=1e2)

    parser.add_argument("--intensity_inv_scale", type=float, default=1)

    parser.add_argument("--spatial_smooth", action="store_true")
    parser.add_argument("--grad_norm_smooth", action="store_true")
    parser.add_argument("--tv_loss", action="store_true")
    parser.add_argument("--grad_loss", action="store_true")
    parser.add_argument("--sobel_grad", action="store_true")

    parser.add_argument("--desired_resolution",
                        type=int,
                        default=2048,
                        help="TCN finest resolution at the smallest scale")
    parser.add_argument("--log2_hashmap_size", type=int, default=19)
    parser.add_argument("--n_features_per_level", type=int, default=2)
    parser.add_argument("--num_layers",
                        type=int,
                        default=2,
                        help="num_layers of sigmanet")
    parser.add_argument("--hidden_dim",
                        type=int,
                        default=64,
                        help="hidden_dim of sigmanet")
    parser.add_argument("--geo_feat_dim",
                        type=int,
                        default=15,
                        help="geo_feat_dim of sigmanet")
    parser.add_argument("--eval_interval", type=int, default=50)
    parser.add_argument(
        "--num_rays_lidar",
        type=int,
        default=4096,
        help="num rays sampled per image for each training step")
    parser.add_argument("--min_near_lidar",
                        type=float,
                        default=0.01,
                        help="minimum near distance for camera")
    parser.add_argument("--lidar_max_depth",
                        type=float,
                        default=81.0,
                        help="max distance for lidar")
    parser.add_argument("--depth_loss",
                        type=str,
                        default="l1",
                        help="l1, bce, mse, huber")
    parser.add_argument("--rgb_depth_loss",
                        type=str,
                        default="l1",
                        help="l1, bce, mse, huber")
    parser.add_argument("--depth_grad_loss",
                        type=str,
                        default="l1",
                        help="l1, bce, mse, huber")
    parser.add_argument("--intensity_loss",
                        type=str,
                        default="mse",
                        help="l1, bce, mse, huber")
    parser.add_argument("--raydrop_loss",
                        type=str,
                        default="mse",
                        help="l1, bce, mse, huber")
    parser.add_argument("--patch_size_lidar",
                        type=int,
                        default=1,
                        help="[experimental] render patches in training. "
                        "1 means disabled, use [64, 32, 16] to enable")
    parser.add_argument(
        "--change_patch_size_lidar",
        nargs="+",
        type=int,
        default=[1, 1],
        help="[experimental] render patches in training. "
        "1 means disabled, use [64, 32, 16] to enable, change during training")
    parser.add_argument("--change_patch_size_epoch",
                        type=int,
                        default=2,
                        help="change patch_size intenvel")

    ### 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="best")
    parser.add_argument(
        "--num_rays",
        type=int,
        default=4096,
        help="num rays sampled per image for each training step")
    parser.add_argument("--num_steps",
                        type=int,
                        default=768,
                        help="num steps sampled per ray")
    parser.add_argument("--upsample_steps",
                        type=int,
                        default=64,
                        help="num steps up-sampled per ray")
    parser.add_argument("--max_ray_batch",
                        type=int,
                        default=4096,
                        help="batch size of rays at inference to avoid OOM)")
    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("--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)")

    return parser


def main():
    parser = get_arg_parser()
    opt = parser.parse_args()

    # Check sequence id.
    aiodrive_sequence_ids = ["64"]
    kitti360_sequence_ids = [
        "1538",
        "1728",
        "1908",
        "3353",
    ]
    nerf_mvl_sequence_ids = [
        "bollard",
        "car",
        "pedestrian",
        "pier",
        "plant",
        "tire",
        "traffic_cone",
        "warning_sign",
        "water_safety_barrier",
    ]

    # Specify dataloader class
    if opt.dataloader == "kitti360":
        from lidarnerf.dataset.kitti360_dataset import KITTI360Dataset as NeRFDataset
        if opt.sequence_id not in kitti360_sequence_ids:
            raise ValueError(
                f"Unknown sequence id {opt.sequence_id} for {opt.dataloader}")
    elif opt.dataloader == "aiodrive":
        from lidarnerf.dataset.aiodrive_dataset import AIODriveDataset as NeRFDataset
        if opt.sequence_id not in aiodrive_sequence_ids:
            raise ValueError(
                f"Unknown sequence id {opt.sequence_id} for {opt.dataloader}")
    elif opt.dataloader == "nerf_mvl":
        from lidarnerf.dataset.nerfmvl_dataset import NeRFMVLDataset as NeRFDataset
        if opt.sequence_id not in nerf_mvl_sequence_ids:
            raise ValueError(
                f"Unknown sequence id {opt.sequence_id} for {opt.dataloader}")
    elif opt.dataloader == "waymo":
        from lidarnerf.dataset.waymo_dataset import WaymoDataset as NeRFDataset
    else:
        raise RuntimeError("Should not reach here.")

    os.makedirs(opt.workspace, exist_ok=True)
    f = os.path.join(opt.workspace, "args.txt")
    with open(f, "w") as file:
        for arg in vars(opt):
            attr = getattr(opt, arg)
            file.write("{} = {}\n".format(arg, attr))

    if opt.L:
        opt.fp16 = True
        opt.tcnn = True
        opt.preload = True

    if opt.patch_size > 1:
        # 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."

    opt.min_near = opt.scale  # hard-code, set min_near ori 1m
    opt.min_near_lidar = opt.scale


    if opt.network == 'tcnn':
        opt.fp16 = True
        assert opt.bg_radius <= 0, "background model is not implemented for --tcnn"
        assert opt.enable_lidar or opt.enable_rgb, "single modality training"
        if opt.enable_lidar:
            from lidarnerf.nerf.network_tcnn_lidar import NeRFNetwork
        elif opt.enable_rgb:
            from lidarnerf.nerf.network_tcnn_rgb import NeRFNetwork
        model = NeRFNetwork(
            encoding="hashgrid",
            desired_resolution=opt.desired_resolution,
            log2_hashmap_size=opt.log2_hashmap_size,
            n_features_per_level=opt.n_features_per_level,
            num_layers=opt.num_layers,
            hidden_dim=opt.hidden_dim,
            geo_feat_dim=opt.geo_feat_dim,
            bound=opt.bound,
            density_scale=1,
            min_near=opt.min_near,
            min_near_lidar=opt.min_near_lidar,
            density_thresh=opt.density_thresh,
            bg_radius=opt.bg_radius,
        )
    elif opt.network == 'mif':
        opt.fp16 = True
        assert opt.bg_radius <= 0, "background model is not implemented for --tcnn"
        from lidarnerf.nerf.network_tcnn_mif import NeRFNetwork
        model = NeRFNetwork(
            encoding="hashgrid",
            desired_resolution=opt.desired_resolution,
            log2_hashmap_size=opt.log2_hashmap_size,
            n_features_per_level=opt.n_features_per_level,
            num_layers=opt.num_layers,
            hidden_dim=opt.hidden_dim,
            geo_feat_dim=opt.geo_feat_dim,
            bound=opt.bound,
            density_scale=1,
            min_near=opt.min_near,
            min_near_lidar=opt.min_near_lidar,
            density_thresh=opt.density_thresh,
            bg_radius=opt.bg_radius,
        )
    elif opt.network == 'alignmif':
        opt.fp16 = True
        assert opt.bg_radius <= 0, "background model is not implemented for --tcnn"
        from lidarnerf.nerf.network_tcnn_alignmif import NeRFNetwork
        model = NeRFNetwork(
            encoding="hashgrid",
            desired_resolution=opt.desired_resolution,
            log2_hashmap_size=opt.log2_hashmap_size,
            n_features_per_level=opt.n_features_per_level,
            num_layers=opt.num_layers,
            hidden_dim=opt.hidden_dim,
            geo_feat_dim=opt.geo_feat_dim,
            bound=opt.bound,
            density_scale=1,
            min_near=opt.min_near,
            min_near_lidar=opt.min_near_lidar,
            density_thresh=opt.density_thresh,
            bg_radius=opt.bg_radius,
        )
    print(opt)
    seed_everything(opt.seed)
    print(model)

    from lidarnerf.nerf.utils import Trainer

    loss_dict = {
        "mse": torch.nn.MSELoss(reduction="none"),
        "l1": torch.nn.L1Loss(reduction="none"),
        "bce": torch.nn.BCEWithLogitsLoss(reduction="none"),
        "huber": torch.nn.HuberLoss(reduction="none", delta=0.2 * opt.scale),
        "cos": torch.nn.CosineSimilarity()
    }
    criterion = {
        'rgb': loss_dict[opt.rgb_loss],
        "depth": loss_dict[opt.depth_loss],
        "rgb_depth": loss_dict[opt.rgb_depth_loss],
        "raydrop": loss_dict[opt.raydrop_loss],
        "intensity": loss_dict[opt.intensity_loss],
        "grad": loss_dict[opt.depth_grad_loss]
    }

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

    if opt.test or opt.test_eval:
        test_loader = NeRFDataset(
            device=device,
            split="test",
            root_path=opt.path,
            sequence_id=opt.sequence_id,
            preload=opt.preload,
            scale=opt.scale,
            offset=opt.offset,
            fp16=opt.fp16,
            patch_size_lidar=opt.patch_size_lidar,
            enable_lidar=opt.enable_lidar,
            enable_rgb=opt.enable_rgb,
            color_space=opt.color_space,
            num_rays=opt.num_rays,
            num_rays_lidar=opt.num_rays_lidar).dataloader()
        if opt.enable_lidar:
            depth_metrics = [
                MAEMeter(intensity_inv_scale=opt.intensity_inv_scale),
                RMSEMeter(),
                DepthMeter(scale=opt.scale),
                PointsMeter(
                    scale=opt.scale,
                    intrinsics=test_loader._data.intrinsics_lidar,
                    beam_inclinations=test_loader._data.beam_inclinations
                    if opt.dataloader in ["waymo"] else None)
            ]
        else:
            depth_metrics = []
        if opt.enable_rgb:
            metrics = [
                RMSEMeter(rgb_metric=True),
                PSNRMeter(),
                LPIPSMeter(device=device),
                SSIMMeter(),
            ]
        else:
            metrics = []
        trainer = Trainer("alignmif",
                          opt,
                          model,
                          device=device,
                          workspace=opt.workspace,
                          criterion=criterion,
                          fp16=opt.fp16,
                          metrics=metrics,
                          depth_metrics=depth_metrics,
                          use_checkpoint=opt.ckpt)
        trainer.show_hash_feat()
        if test_loader.has_gt and opt.test_eval:
            trainer.evaluate(test_loader)  # blender has gt, so evaluate it.
        trainer.test(test_loader, write_video=False)  # test and save video

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

        train_loader = NeRFDataset(device=device,
                                   split="train",
                                   root_path=opt.path,
                                   sequence_id=opt.sequence_id,
                                   preload=opt.preload,
                                   scale=opt.scale,
                                   offset=opt.offset,
                                   fp16=opt.fp16,
                                   patch_size_lidar=opt.patch_size_lidar,
                                   enable_lidar=opt.enable_lidar,
                                   enable_rgb=opt.enable_rgb,
                                   color_space=opt.color_space,
                                   num_rays=opt.num_rays,
                                   num_rays_lidar=opt.num_rays_lidar,).dataloader()

        # decay to 0.1 * init_lr at last iter step
        scheduler = lambda optimizer: torch.optim.lr_scheduler.LambdaLR(
            optimizer, lambda iter: 0.1**min(iter / opt.iters, 1))
        if opt.enable_lidar:
            depth_metrics = [
                MAEMeter(intensity_inv_scale=opt.intensity_inv_scale),
                RMSEMeter(),
                DepthMeter(scale=opt.scale),
                PointsMeter(
                    scale=opt.scale,
                    intrinsics=train_loader._data.intrinsics_lidar,
                    beam_inclinations=train_loader._data.beam_inclinations
                    if opt.dataloader in ["waymo"] else None)
            ]
        else:
            depth_metrics = []
        if opt.enable_rgb:
            metrics = [
                RMSEMeter(rgb_metric=True),
                PSNRMeter(),
                LPIPSMeter(device=device),
                SSIMMeter(),
            ]
        else:
            metrics = []

        trainer = Trainer("alignmif",
                          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,
                          depth_metrics=depth_metrics,
                          metrics=metrics,
                          use_checkpoint=opt.ckpt,
                          eval_interval=opt.eval_interval)

        valid_loader = NeRFDataset(
            device=device,
            split="val",
            root_path=opt.path,
            sequence_id=opt.sequence_id,
            preload=opt.preload,
            scale=opt.scale,
            offset=opt.offset,
            fp16=opt.fp16,
            patch_size_lidar=opt.patch_size_lidar,
            enable_lidar=opt.enable_lidar,
            enable_rgb=opt.enable_rgb,
            color_space=opt.color_space,
            num_rays=opt.num_rays,
            num_rays_lidar=opt.num_rays_lidar).dataloader()

        max_epoch = np.ceil(opt.iters / len(train_loader)).astype(np.int32)
        print(f"max_epoch: {max_epoch}")
        trainer.train(train_loader, valid_loader, max_epoch)

        # also test
        test_loader = NeRFDataset(
            device=device,
            split="test",
            root_path=opt.path,
            sequence_id=opt.sequence_id,
            preload=opt.preload,
            scale=opt.scale,
            offset=opt.offset,
            fp16=opt.fp16,
            patch_size_lidar=opt.patch_size_lidar,
            enable_lidar=opt.enable_lidar,
            enable_rgb=opt.enable_rgb,
            color_space=opt.color_space,
            num_rays=opt.num_rays,
            num_rays_lidar=opt.num_rays_lidar).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=128, threshold=10)


if __name__ == "__main__":
    main()