nerfstudio-gnt 0.0.1__tar.gz

nerfstudio_gnt-0.0.1/GNT/config.py

@@ -0,0 +1,197 @@
+import configargparse
+
+
+def config_parser():
+    parser = configargparse.ArgumentParser()
+    # general
+    parser.add_argument("--config", is_config_file=True, help="config file path")
+    parser.add_argument(
+        "--rootdir",
+        type=str,
+        default="./",
+        help="the path to the project root directory. Replace this path with yours!",
+    )
+    parser.add_argument("--expname", type=str, help="experiment name")
+    parser.add_argument("--distributed", action="store_true", help="if use distributed training")
+    parser.add_argument("--local_rank", type=int, default=0, help="rank for distributed training")
+    parser.add_argument(
+        "-j",
+        "--workers",
+        default=8,
+        type=int,
+        metavar="N",
+        help="number of data loading workers (default: 8)",
+    )
+
+    ########## dataset options ##########
+    ## train and eval dataset
+    parser.add_argument(
+        "--train_dataset",
+        type=str,
+        default="ibrnet_collected",
+        help="the training dataset, should either be a single dataset, "
+        'or multiple datasets connected with "+", for example, ibrnet_collected+llff+spaces',
+    )
+    parser.add_argument(
+        "--dataset_weights",
+        nargs="+",
+        type=float,
+        default=[],
+        help="the weights for training datasets, valid when multiple datasets are used.",
+    )
+    parser.add_argument(
+        "--train_scenes",
+        nargs="+",
+        default=[],
+        help="optional, specify a subset of training scenes from training dataset",
+    )
+    parser.add_argument(
+        "--eval_dataset", type=str, default="llff_test", help="the dataset to evaluate"
+    )
+    parser.add_argument(
+        "--eval_scenes",
+        nargs="+",
+        default=[],
+        help="optional, specify a subset of scenes from eval_dataset to evaluate",
+    )
+    ## others
+    parser.add_argument(
+        "--testskip",
+        type=int,
+        default=8,
+        help="will load 1/N images from test/val sets, "
+        "useful for large datasets like deepvoxels or nerf_synthetic",
+    )
+
+    ########## model options ##########
+    ## ray sampling options
+    parser.add_argument(
+        "--sample_mode",
+        type=str,
+        default="uniform",
+        help="how to sample pixels from images for training:" "uniform|center",
+    )
+    parser.add_argument(
+        "--center_ratio", type=float, default=0.8, help="the ratio of center crop to keep"
+    )
+    parser.add_argument(
+        "--N_rand",
+        type=int,
+        default=32 * 16,
+        help="batch size (number of random rays per gradient step)",
+    )
+    parser.add_argument(
+        "--chunk_size",
+        type=int,
+        default=1024 * 4,
+        help="number of rays processed in parallel, decrease if running out of memory",
+    )
+
+    ## model options
+    parser.add_argument(
+        "--coarse_feat_dim", type=int, default=32, help="2D feature dimension for coarse level"
+    )
+    parser.add_argument(
+        "--fine_feat_dim", type=int, default=32, help="2D feature dimension for fine level"
+    )
+    parser.add_argument(
+        "--num_source_views",
+        type=int,
+        default=10,
+        help="the number of input source views for each target view",
+    )
+    parser.add_argument(
+        "--rectify_inplane_rotation", action="store_true", help="if rectify inplane rotation"
+    )
+    parser.add_argument("--coarse_only", action="store_true", help="use coarse network only")
+    parser.add_argument(
+        "--anti_alias_pooling", type=int, default=1, help="if use anti-alias pooling"
+    )
+    parser.add_argument("--trans_depth", type=int, default=4, help="number of transformer layers")
+    parser.add_argument("--netwidth", type=int, default=64, help="network intermediate dimension")
+    parser.add_argument(
+        "--single_net",
+        type=bool,
+        default=True,
+        help="use single network for both coarse and/or fine sampling",
+    )
+
+    ########## checkpoints ##########
+    parser.add_argument(
+        "--no_reload", action="store_true", help="do not reload weights from saved ckpt"
+    )
+    parser.add_argument(
+        "--ckpt_path",
+        type=str,
+        default="",
+        help="specific weights npy file to reload for coarse network",
+    )
+    parser.add_argument(
+        "--no_load_opt", action="store_true", help="do not load optimizer when reloading"
+    )
+    parser.add_argument(
+        "--no_load_scheduler", action="store_true", help="do not load scheduler when reloading"
+    )
+
+    ########### iterations & learning rate options ##########
+    parser.add_argument("--n_iters", type=int, default=250000, help="num of iterations")
+    parser.add_argument(
+        "--lrate_feature", type=float, default=1e-3, help="learning rate for feature extractor"
+    )
+    parser.add_argument("--lrate_gnt", type=float, default=5e-4, help="learning rate for gnt")
+    parser.add_argument(
+        "--lrate_decay_factor",
+        type=float,
+        default=0.5,
+        help="decay learning rate by a factor every specified number of steps",
+    )
+    parser.add_argument(
+        "--lrate_decay_steps",
+        type=int,
+        default=50000,
+        help="decay learning rate by a factor every specified number of steps",
+    )
+
+    ########## rendering options ##########
+    parser.add_argument(
+        "--N_samples", type=int, default=64, help="number of coarse samples per ray"
+    )
+    parser.add_argument(
+        "--N_importance", type=int, default=64, help="number of important samples per ray"
+    )
+    parser.add_argument(
+        "--inv_uniform", action="store_true", help="if True, will uniformly sample inverse depths"
+    )
+    parser.add_argument(
+        "--det", action="store_true", help="deterministic sampling for coarse and fine samples"
+    )
+    parser.add_argument(
+        "--white_bkgd",
+        action="store_true",
+        help="apply the trick to avoid fitting to white background",
+    )
+    parser.add_argument(
+        "--render_stride",
+        type=int,
+        default=1,
+        help="render with large stride for validation to save time",
+    )
+
+    ########## logging/saving options ##########
+    parser.add_argument("--i_print", type=int, default=100, help="frequency of terminal printout")
+    parser.add_argument(
+        "--i_img", type=int, default=500, help="frequency of tensorboard image logging"
+    )
+    parser.add_argument(
+        "--i_weights", type=int, default=10000, help="frequency of weight ckpt saving"
+    )
+
+    ########## evaluation options ##########
+    parser.add_argument(
+        "--llffhold",
+        type=int,
+        default=8,
+        help="will take every 1/N images as LLFF test set, paper uses 8",
+    )
+
+    return parser

nerfstudio_gnt-0.0.1/GNT/eval.py

@@ -0,0 +1,236 @@
+import os
+import numpy as np
+import shutil
+import torch
+import torch.utils.data.distributed
+
+from torch.utils.data import DataLoader
+
+from gnt.data_loaders import dataset_dict
+from gnt.render_image import render_single_image
+from gnt.model import GNTModel
+from gnt.sample_ray import RaySamplerSingleImage
+from utils import img_HWC2CHW, colorize, img2psnr, lpips, ssim
+import config
+import torch.distributed as dist
+from gnt.projection import Projector
+from gnt.data_loaders.create_training_dataset import create_training_dataset
+import imageio
+
+
+def worker_init_fn(worker_id):
+    np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+
+def synchronize():
+    """
+    Helper function to synchronize (barrier) among all processes when
+    using distributed training
+    """
+    if not dist.is_available():
+        return
+    if not dist.is_initialized():
+        return
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return
+    dist.barrier()
+
+
+@torch.no_grad()
+def eval(args):
+
+    device = "cuda:{}".format(args.local_rank)
+    out_folder = os.path.join(args.rootdir, "out", args.expname)
+    print("outputs will be saved to {}".format(out_folder))
+    os.makedirs(out_folder, exist_ok=True)
+
+    # save the args and config files
+    f = os.path.join(out_folder, "args.txt")
+    with open(f, "w") as file:
+        for arg in sorted(vars(args)):
+            attr = getattr(args, arg)
+            file.write("{} = {}\n".format(arg, attr))
+
+    if args.config is not None:
+        f = os.path.join(out_folder, "config.txt")
+        if not os.path.isfile(f):
+            shutil.copy(args.config, f)
+
+    if args.run_val == False:
+        # create training dataset
+        dataset, sampler = create_training_dataset(args)
+        # currently only support batch_size=1 (i.e., one set of target and source views) for each GPU node
+        # please use distributed parallel on multiple GPUs to train multiple target views per batch
+        loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=1,
+            worker_init_fn=lambda _: np.random.seed(),
+            num_workers=args.workers,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=True if sampler is None else False,
+        )
+        iterator = iter(loader)
+    else:
+        # create validation dataset
+        dataset = dataset_dict[args.eval_dataset](args, "validation", scenes=args.eval_scenes)
+        loader = DataLoader(dataset, batch_size=1)
+        iterator = iter(loader)
+
+    # Create GNT model
+    model = GNTModel(
+        args, load_opt=not args.no_load_opt, load_scheduler=not args.no_load_scheduler
+    )
+    # create projector
+    projector = Projector(device=device)
+
+    indx = 0
+    psnr_scores = []
+    lpips_scores = []
+    ssim_scores = []
+    while True:
+        try:
+            data = next(iterator)
+        except:
+            break
+        if args.local_rank == 0:
+            tmp_ray_sampler = RaySamplerSingleImage(data, device, render_stride=args.render_stride)
+            H, W = tmp_ray_sampler.H, tmp_ray_sampler.W
+            gt_img = tmp_ray_sampler.rgb.reshape(H, W, 3)
+            psnr_curr_img, lpips_curr_img, ssim_curr_img = log_view(
+                indx,
+                args,
+                model,
+                tmp_ray_sampler,
+                projector,
+                gt_img,
+                render_stride=args.render_stride,
+                prefix="val/" if args.run_val else "train/",
+                out_folder=out_folder,
+                ret_alpha=args.N_importance > 0,
+                single_net=args.single_net,
+            )
+            psnr_scores.append(psnr_curr_img)
+            lpips_scores.append(lpips_curr_img)
+            ssim_scores.append(ssim_curr_img)
+            torch.cuda.empty_cache()
+            indx += 1
+    print("Average PSNR: ", np.mean(psnr_scores))
+    print("Average LPIPS: ", np.mean(lpips_scores))
+    print("Average SSIM: ", np.mean(ssim_scores))
+
+
+@torch.no_grad()
+def log_view(
+    global_step,
+    args,
+    model,
+    ray_sampler,
+    projector,
+    gt_img,
+    render_stride=1,
+    prefix="",
+    out_folder="",
+    ret_alpha=False,
+    single_net=True,
+):
+    model.switch_to_eval()
+    with torch.no_grad():
+        ray_batch = ray_sampler.get_all()
+        if model.feature_net is not None:
+            featmaps = model.feature_net(ray_batch["src_rgbs"].squeeze(0).permute(0, 3, 1, 2))
+        else:
+            featmaps = [None, None]
+        ret = render_single_image(
+            ray_sampler=ray_sampler,
+            ray_batch=ray_batch,
+            model=model,
+            projector=projector,
+            chunk_size=args.chunk_size,
+            N_samples=args.N_samples,
+            inv_uniform=args.inv_uniform,
+            det=True,
+            N_importance=args.N_importance,
+            white_bkgd=args.white_bkgd,
+            render_stride=render_stride,
+            featmaps=featmaps,
+            ret_alpha=ret_alpha,
+            single_net=single_net,
+        )
+
+    average_im = ray_sampler.src_rgbs.cpu().mean(dim=(0, 1))
+
+    if args.render_stride != 1:
+        gt_img = gt_img[::render_stride, ::render_stride]
+        average_im = average_im[::render_stride, ::render_stride]
+
+    rgb_gt = img_HWC2CHW(gt_img)
+    average_im = img_HWC2CHW(average_im)
+
+    rgb_coarse = img_HWC2CHW(ret["outputs_coarse"]["rgb"].detach().cpu())
+    if "depth" in ret["outputs_coarse"].keys():
+        depth_pred = ret["outputs_coarse"]["depth"].detach().cpu()
+        depth_coarse = img_HWC2CHW(colorize(depth_pred, cmap_name="jet"))
+    else:
+        depth_coarse = None
+
+    if ret["outputs_fine"] is not None:
+        rgb_fine = img_HWC2CHW(ret["outputs_fine"]["rgb"].detach().cpu())
+        if "depth" in ret["outputs_fine"].keys():
+            depth_pred = ret["outputs_fine"]["depth"].detach().cpu()
+            depth_fine = img_HWC2CHW(colorize(depth_pred, cmap_name="jet"))
+    else:
+        rgb_fine = None
+        depth_fine = None
+
+    rgb_coarse = rgb_coarse.permute(1, 2, 0).detach().cpu().numpy()
+    filename = os.path.join(out_folder, prefix[:-1] + "_{:03d}_coarse.png".format(global_step))
+    imageio.imwrite(filename, rgb_coarse)
+
+    if depth_coarse is not None:
+        depth_coarse = depth_coarse.permute(1, 2, 0).detach().cpu().numpy()
+        filename = os.path.join(
+            out_folder, prefix[:-1] + "_{:03d}_coarse_depth.png".format(global_step)
+        )
+        imageio.imwrite(filename, depth_coarse)
+
+    if rgb_fine is not None:
+        rgb_fine = rgb_fine.permute(1, 2, 0).detach().cpu().numpy()
+        filename = os.path.join(out_folder, prefix[:-1] + "_{:03d}_fine.png".format(global_step))
+        imageio.imwrite(filename, rgb_fine)
+
+    if depth_fine is not None:
+        depth_fine = depth_fine.permute(1, 2, 0).detach().cpu().numpy()
+        filename = os.path.join(
+            out_folder, prefix[:-1] + "_{:03d}_fine_depth.png".format(global_step)
+        )
+        imageio.imwrite(filename, depth_fine)
+
+    # write scalar
+    pred_rgb = (
+        ret["outputs_fine"]["rgb"]
+        if ret["outputs_fine"] is not None
+        else ret["outputs_coarse"]["rgb"]
+    )
+    pred_rgb = torch.clip(pred_rgb, 0.0, 1.0)
+    lpips_curr_img = lpips(pred_rgb, gt_img, format="HWC").item()
+    ssim_curr_img = ssim(pred_rgb, gt_img, format="HWC").item()
+    psnr_curr_img = img2psnr(pred_rgb.detach().cpu(), gt_img)
+    print(prefix + "psnr_image: ", psnr_curr_img)
+    print(prefix + "lpips_image: ", lpips_curr_img)
+    print(prefix + "ssim_image: ", ssim_curr_img)
+    return psnr_curr_img, lpips_curr_img, ssim_curr_img
+
+
+if __name__ == "__main__":
+    parser = config.config_parser()
+    parser.add_argument("--run_val", action="store_true", help="run on val set")
+    args = parser.parse_args()
+
+    if args.distributed:
+        torch.cuda.set_device(args.local_rank)
+        torch.distributed.init_process_group(backend="nccl", init_method="env://")
+        synchronize()
+
+    eval(args)

nerfstudio_gnt-0.0.1/GNT/gnt/criterion.py

@@ -0,0 +1,22 @@
+import torch.nn as nn
+from utils import img2mse
+
+
+class Criterion(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, outputs, ray_batch, scalars_to_log):
+        """
+        training criterion
+        """
+        pred_rgb = outputs["rgb"]
+        if "mask" in outputs:
+            pred_mask = outputs["mask"].float()
+        else:
+            pred_mask = None
+        gt_rgb = ray_batch["rgb"]
+
+        loss = img2mse(pred_rgb, gt_rgb, pred_mask)
+
+        return loss, scalars_to_log

nerfstudio_gnt-0.0.1/GNT/gnt/data_loaders/__init__.py

@@ -0,0 +1,31 @@
+from .google_scanned_objects import *
+from .realestate import *
+from .deepvoxels import *
+from .realestate import *
+from .llff import *
+from .llff_test import *
+from .ibrnet_collected import *
+from .realestate import *
+from .spaces_dataset import *
+from .nerf_synthetic import *
+from .shiny import *
+from .llff_render import *
+from .shiny_render import *
+from .nerf_synthetic_render import *
+from .nmr_dataset import *
+
+dataset_dict = {
+    "spaces": SpacesFreeDataset,
+    "google_scanned": GoogleScannedDataset,
+    "realestate": RealEstateDataset,
+    "deepvoxels": DeepVoxelsDataset,
+    "nerf_synthetic": NerfSyntheticDataset,
+    "llff": LLFFDataset,
+    "ibrnet_collected": IBRNetCollectedDataset,
+    "llff_test": LLFFTestDataset,
+    "shiny": ShinyDataset,
+    "llff_render": LLFFRenderDataset,
+    "shiny_render": ShinyRenderDataset,
+    "nerf_synthetic_render": NerfSyntheticRenderDataset,
+    "nmr": NMRDataset,
+}