Official Pytorch implementation of SHACIRA: Scalable HAsh-grid Compression for Implicit Neural Representations
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import os
import argparse
import logging
import numpy as np
import wandb
import torch
import wisp
import json
import shutil
import tempfile
import logging as log
from torch.utils.tensorboard import SummaryWriter
from wisp.app_utils import default_log_setup, args_to_log_format, copy_dir_msrsync, copy_dir, copy_file
import wisp.config_parser as config_parser
from wisp.framework import WispState, ImageState
from wisp.datasets import BaseImageDataset, SampleRays
from wisp.models.grids import BLASGrid, OctreeGrid, CodebookOctreeGrid, TriplanarGrid, HashGrid, LatentGrid
from wisp.models.latent_decoders.basic_latent_decoder import DecoderIdentity
from wisp.models.latent_decoders.multi_latent_decoder import MultiLatentDecoder
from wisp.tracers import BaseTracer, PackedRFTracer
from wisp.models.nefs import BaseNeuralField, NeuralImage
from wisp.models.pipeline import Pipeline
from wisp.trainers import BaseTrainer, ImageTrainer

def parse_args():
    """Wisp mains define args per app.
    Args are collected by priority: cli args > config yaml > argparse defaults
    For convenience, args are divided into groups.
    """
    parser = argparse.ArgumentParser(description='A script for training simple NeRF variants.',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config', type=str,
                        help='Path to config file to replace defaults.')
    parser.add_argument('--profile', action='store_true',
                        help='Enable NVTX profiling')
    parser.add_argument('--seed', type=int,
                        help='Random seed. (Not implemented)')

    log_group = parser.add_argument_group('logging')
    log_group.add_argument('--exp-name', type=str,
                           help='Experiment name, unique id for trainers, logs.')
    log_group.add_argument('--log-level', action='store', type=int, default=logging.INFO,
                           help='Logging level to use globally, DEBUG: 10, INFO: 20, WARN: 30, ERROR: 40.')
    log_group.add_argument('--perf', action='store_true', default=False,
                           help='Use high-level profiling for the trainer.')
    log_group.add_argument('--resume', action='store_true', default=False,
                           help='Resume training from the latest checkpoint.')
    log_group.add_argument('--log-every', type=int, default=100, 
                               help='Log data every N iterations')
    log_group.add_argument('--log-dir', type=str, default='_results/logs/',
                               help='Log file directory for checkpoints.')

    data_group = parser.add_argument_group('dataset')
    data_group.add_argument('--dataset-path', type=str,
                            help='Path to the dataset')
    data_group.add_argument('--copy-local', action='store_true', default=False,
                            help='Whether to copy the dataset to local disk.')
    data_group.add_argument('--local-path', type=str, default='',
                            help='Local path to copy the dataset to. If not specified, a temporary directory will be used.')
    data_group.add_argument('--msrsync-exec', type=str, default='/path/to/msrsync/executable',
                            help='Path to msrsync executable.')
    data_group.add_argument('--msrsync-workers', type=int, default=1,
                            help='Number of workers for msrsync copy')
    data_group.add_argument('--dataset-num-workers', type=int, default=-1,
                            help='Number of workers for dataset preprocessing, if it supports multiprocessing. '
                                 '-1 indicates no multiprocessing.')
    data_group.add_argument('--dataloader-num-workers', type=int, default=1,
                            help='Number of workers for dataloader.')
    data_group.add_argument('--num-samples', type=int,  
                               default=-1, help='Number of coordinate samples per batch (power of 2), -1 corresponds to full batch')
    data_group.add_argument('--sample-mode', type=str,  choices=['full','wreplace','woreplace','sequential'],
                               default='full', help='Different sample methods to apply for each batch of coords.')


    grid_group = parser.add_argument_group('grid')
    grid_group.add_argument('--grid-type', type=str, default='LatentGrid',
                            choices=config_parser.list_modules('grid'),
                            help='Type of to use, i.e.:'
                                 '"OctreeGrid", "CodebookOctreeGrid", "TriplanarGrid", "HashGrid", "LatentGrid".'
                                 'Grids are located in `wisp.models.grids`')
    grid_group.add_argument('--interpolation-type', type=str, default='linear', choices=['linear', 'closest'],
                            help='Interpolation type to use for samples within grids.'
                                 'For a 3D grid structure, linear uses trilinear interpolation of 8 cell nodes,'
                                 'closest uses the nearest neighbor.')
    grid_group.add_argument('--blas-type', type=str, default='octree', 
                            choices=['octree',],
                            help='Type of acceleration structure to use for fast raymarch occupancy queries.')
    grid_group.add_argument('--multiscale-type', type=str, default='cat', choices=['sum', 'cat'],
                            help='Aggregation of choice for multi-level grids, for features from different LODs.')
    grid_group.add_argument('--feature-dim', type=int, default=32,
                            help='Dimensionality for features stored within the grid nodes.')
    grid_group.add_argument('--feature-std', type=float, default=0.0,
                            help='Grid initialization: standard deviation used for randomly sampling initial features.')
    grid_group.add_argument('--feature-bias', type=float, default=0.0,
                            help='Grid initialization: bias used for randomly sampling initial features.')
    grid_group.add_argument('--base-lod', type=int, default=1,
                            help='Number of levels in grid, which book-keep occupancy but not features.'
                                 'The total number of levels in a grid is `base_lod + num_lod - 1`')
    grid_group.add_argument('--num-lods', type=int, default=1,
                            help='Number of levels in grid, which store concrete features.')
    grid_group.add_argument('--codebook-bitwidth', type=int, default=8,
                            help='For Codebook and HashGrids/LatentGrids only: determines the table size as 2**(bitwidth).')
    grid_group.add_argument('--tree-type', type=str, default='geometric', choices=['geometric', 'quad'],
                            help='For HashGrids/LatentGrids only: how the resolution of the grid is determined. '
                                 '"geometric" uses the geometric sequence initialization from InstantNGP,'
                                 'where "quad" uses an octree sampling pattern.')
    grid_group.add_argument('--min-grid-res', type=int, default=16,
                            help='For HashGrids/LatentGrids only: min grid resolution, used only in geometric initialization mode')
    grid_group.add_argument('--max-grid-res', type=int, default=2048,
                            help='For HashGrids/LatentGrids only: max grid resolution, used only in geometric initialization mode')
    grid_group.add_argument('--blas-level', type=float, default=7,
                            help='For HashGrids/LatentGrids only: Determines the number of levels in the acceleration structure '
                                 'used to track the occupancy status (bottom level acceleration structure).')
    grid_group.add_argument('--init-grid', type=str, default='uniform', choices=['normal', 'uniform'],
                            help='Grid initialization: distribution used for randomly sampling initial features.')

    nef_group = parser.add_argument_group('nef')
    nef_group.add_argument('--pos-embedder', type=str, choices=['none', 'identity', 'positional'],
                           default='positional',
                           help='MLP Decoder of neural field: Positional embedder used to encode input coordinates.')
    nef_group.add_argument('--position-input', type=bool, default=False,
                           help='If True, position coords will be concatenated to the '
                                'features / positional embeddings when fed into the decoder.')
    nef_group.add_argument('--pos-multires', type=int, default=10,
                           help='MLP Decoder of neural field: Number of frequencies to use for positional encoding'
                                'of input coordinates')
    nef_group.add_argument('--layer-type', type=str, default='none',
                           choices=['none', 'spectral_norm', 'frobenius_norm', 'l_1_norm', 'l_inf_norm'])
    nef_group.add_argument('--activation-type', type=str, default='relu',
                           choices=['relu', 'sin'])
    nef_group.add_argument('--hidden-dim', type=int, help='MLP Decoder of neural field: width of all hidden layers.')
    nef_group.add_argument('--num-layers', type=int, help='MLP Decoder of neural field: number of hidden layers.')

    trainer_group = parser.add_argument_group('trainer')
    trainer_group.add_argument('--epochs', type=int, default=10000,
                               help='Number of epochs to run the training.')
    trainer_group.add_argument('--batch-size', type=int, default=1,
                               help='Batch size for the training. Only included for compatibility for images')
    trainer_group.add_argument('--resample', action='store_true',
                               help='Resample the dataset after every epoch.')
    trainer_group.add_argument('--resample-every', type=int, default=1,
                               help='Resample every N epochs')
    trainer_group.add_argument('--model-format', type=str, default='full', choices=['full', 'state_dict'],
                               help='Format in which to save models.')
    trainer_group.add_argument('--pretrained', type=str,
                               help='Path to pretrained model weights.')
    trainer_group.add_argument('--save-as-new', action='store_true',
                               help='Save the model at every epoch (no overwrite).')
    trainer_group.add_argument('--save-every', type=int, default=-1,
                               help='Save the model at every N iterations.')
    trainer_group.add_argument('--render-tb-every', type=int, default=-1,
                               help='Render (log image) every N epochs') 
    trainer_group.add_argument('--valid-only', action='store_true',
                               help='Run validation only (and do not run training).')
    trainer_group.add_argument('--valid-only-load-path', type=str, default=None,
                               help='If valid only, load model.pth from this path.')
    trainer_group.add_argument('--valid-every', type=int, default=-1,
                               help='Frequency of running validation.')
    
    wandb_group = parser.add_argument_group('wandb')
    wandb_group.add_argument('--wandb-project', type=str, default=None,
                               help='Weights & Biases Project')
    wandb_group.add_argument('--wandb-mode', type=str, default='online',
                               help='Whether to run Weights & Biases in online or offline mode.')
    wandb_group.add_argument('--log-metrics-only', action='store_true',
                               help='Weights & Biases Log metrics only (no images or artifacts).')
    wandb_group.add_argument('--wandb-run-name', type=str, default=None,
                               help='Weights & Biases Run Name')
    wandb_group.add_argument('--wandb-entity', type=str, default=None,
                               help='Weights & Biases Entity')
    wandb_group.add_argument('--wandb-viz-nerf-angles', type=int, default=20,
                               help='Number of Angles to visualize a scene on Weights & Biases. '
                                    'Set this to 0 to disable 360 degree visualizations.')
    wandb_group.add_argument('--wandb-viz-nerf-distance', type=int, default=3,
                               help='Distance to visualize Scene from on Weights & Biases')

    optimizer_group = parser.add_argument_group('optimizer')
    optimizer_group.add_argument('--disable-amp', action='store_true',
                                 help='Disable Automatic Mixed Precision (AMP) training.')
    optimizer_group.add_argument('--disable-scaler', action='store_true',
                                 help='Disable scaler for training.')
    optimizer_group.add_argument('--optimizer-type', type=str, default='adam',
                                 choices=config_parser.list_modules('optim'),
                                 help='Optimizer to be used, includes optimizer modules available within `torch.optim` '
                                      'and fused optimizers from `apex`, if apex is installed.')
    optimizer_group.add_argument('--lr', type=float, default=0.001,
                                 help='Base optimizer learning rate.')
    optimizer_group.add_argument('--eps', type=float, default=1e-8,
                                 help='Eps value for numerical stability.')
    optimizer_group.add_argument('--weight-decay', type=float, default=0,
                                 help='Weight decay.')
    optimizer_group.add_argument('--weight-decay-decoder', type=float, default=0,
                                 help='Weight decay, applied only to decoder weights.')
    optimizer_group.add_argument('--grid-lr', type=float, default=1.0e-2,
                                 help='Learning rate applied only for the grid parameters'
                                      '(e.g. parameters which contain "grid" in their name)')
    optimizer_group.add_argument('--scale-grid-lr', type=str, default='none', choices=['none','mul','div'],
                                 help='Scale grid-lr by multiplying/dividing by decoder norm or use as is.')
    optimizer_group.add_argument('--ldec-lr', type=float, default=1.0e-2,
                                 help='Learning rate applied only for the latent decoder parameters')
    optimizer_group.add_argument('--ldec-lr-warmup', type=int, default=1,
                                 help='Warmup epochs for latent decoder learning rate')
    optimizer_group.add_argument('--rgb-loss', type=float, default=1.0,
                                 help='Weight of rgb loss')

    latent_decoder_group = parser.add_argument_group('latent_decoder')
    latent_decoder_group.add_argument('--ldecode-enabled', action='store_true', default=False,
                           help='Whether to perform quantization and latent decoding.')
    latent_decoder_group.add_argument('--ldecode-type', type=str, choices=['single', 'multi', 'hierarchical'], default='single',
                                        help='Type of latent decoder to use.')
    latent_decoder_group.add_argument('--use-sga', action='store_true', default=False,
                           help='Whether to use stochastic gumbel annealing for the quantized weights before decoding.')
    latent_decoder_group.add_argument('--diff-sampling', action='store_true', default=False,
                                        help='Whether to make sampling differentiable for SGA.')
    latent_decoder_group.add_argument('--ldecode-matrix', type=str, choices=['sq','dft'], default='sq',
                                        help='Whether to use a learnable decoder or a DFT basis decoder for the weight matrices.')
    latent_decoder_group.add_argument('--latent-dim', type=int, default=0,
                               help='Latent dimension. If 0, latent dimension is feature dimension.')
    latent_decoder_group.add_argument('--norm', type=str, choices=['none','max','std'], default='none',
                                        help='Whether to normalize latents.')
    latent_decoder_group.add_argument('--norm-every', type=int, default=1,
                                        help='Frequency of calculating norm of latents for normalizing')
    latent_decoder_group.add_argument('--use-shift', action='store_true', default=False,
                           help='Whether to use shift parameter in decoder.')
    latent_decoder_group.add_argument('--num-layers-dec', type=int, default=0,
                               help='Number of hidden layers in decoder. If 0, decoder is linear.')
    latent_decoder_group.add_argument('--hidden-dim-dec', type=int, default=0,
                               help='Hidden dimension in decoder. If 0, hidden dimension is feature dimension.')
    latent_decoder_group.add_argument('--activation', type=str, choices=['none','sigmoid','tanh','relu','sine'], 
                                      default='none', help='Activation function to use in decoder.')
    latent_decoder_group.add_argument('--final-activation', type=str, choices=['none','sigmoid','tanh','relu','sine'], 
                                      default='none', help='Final layer activation to use after decoder.')
    latent_decoder_group.add_argument('--clamp-weights', type=float, default=0.0,
                                        help='Clamp norm of decoder weights to this value. 0 disables clamping.')
    latent_decoder_group.add_argument('--ldec-std', type=float, default=1.0,
                                        help='Standard deviation for decoder weight initialization.')
    latent_decoder_group.add_argument('--num-decoders', type=int, default=2,
                                        help='Number of latent decoders. Used only when decode-type is multi.')
    latent_decoder_group.add_argument('--temperature', type=float, default=1.0,
                                        help='Temperature for decoder softmax scores and for SGA annealing. Used only when decode-type is multi or SGA enabled')
    latent_decoder_group.add_argument('--decay-period', type=float, default=0.8,
                                        help='Fraction of epochs for temperature decay. Set to value of temperature after this.')
    latent_decoder_group.add_argument('--alpha-std', type=float, default=10.0,
                                        help='Initialization of softmax logits for decoder. Used only when decode-type is multi.')
    
    entropy_reg_group = parser.add_argument_group('entropy_reg')
    entropy_reg_group.add_argument('--num-prob-layers', type=int, default=2,
                                        help='Number of hidden layers for probability model.')
    entropy_reg_group.add_argument('--entropy-reg', type=float, default=0.0,
                                        help='Starting entropy regularization lambda for latents.')
    entropy_reg_group.add_argument('--entropy-reg-end', type=float, default=0.0,
                                        help='Ending entropy regularization lambda for latents.')
    entropy_reg_group.add_argument('--entropy-reg-sched', type=str, default='cosine', choices=['fix', 'linear', 'exp','cosine'],
                                        help='Decay schedule for entropy regularization lambda for latents.')
    entropy_reg_group.add_argument('--noise-freq', type=int, default=50,
                                        help='Iteration frequency for generating noise for prob model.')

    # Parse CLI args & config files
    args = config_parser.parse_args(parser)

    # Also obtain args as grouped hierarchy, useful for, i.e., logging
    args_dict = config_parser.get_grouped_args(parser, args)
    return args, args_dict


def load_dataset(args) -> BaseImageDataset:
    """ Loads an image dataset.
    """
    train_dataset = wisp.datasets.load_image_dataset( 
                                                    dataset_path=args.dataset_path,
                                                    num_samples=args.num_samples,
                                                    sample_mode=args.sample_mode)
    return train_dataset


def load_grid(args, args_dict, dataset: BaseImageDataset) -> BLASGrid:
    """ Wisp's implementation of NeRF uses feature grids to improve the performance and quality (allowing therefore,
    interactivity).
    This function loads the feature grid to use within the neural pipeline.
    Grid choices are interesting to explore, so we leave the exact backbone type configurable,
    and show how grid instances may be explicitly constructed.
    Grids choices, for example, are: OctreeGrid, TriplanarGrid, HashGrid,LatentGrid, CodebookOctreeGrid
    See corresponding grid constructors for each of their arg details.
    """
    grid = None

    # Optimization: For octrees based grids, if dataset contains depth info, initialize only cells known to be occupied
    # WARNING: Not implemented/tested for grids other than HashGrid/LatentGrid 
    if args.grid_type == "OctreeGrid":
        if dataset.supports_depth():
            grid = OctreeGrid.from_pointcloud(
                pointcloud=dataset.as_pointcloud(),
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
            )
        else:
            grid = OctreeGrid.make_dense(
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
            )
    elif args.grid_type == "CodebookOctreeGrid":
        if dataset.supports_depth():
            grid = CodebookOctreeGrid.from_pointcloud(
                pointcloud=dataset.as_pointcloud(),
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth
            )
        else:
            grid = CodebookOctreeGrid.make_dense(
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth
            )
    elif args.grid_type == "TriplanarGrid":
        grid = TriplanarGrid(
            feature_dim=args.feature_dim,
            base_lod=args.base_lod,
            num_lods=args.num_lods,
            interpolation_type=args.interpolation_type,
            multiscale_type=args.multiscale_type,
            feature_std=args.feature_std,
            feature_bias=args.feature_bias,
        )
    elif args.grid_type == "HashGrid":
        # "geometric" - determines the resolution of the grid using geometric sequence initialization from InstantNGP,
        if args.tree_type == "geometric":
            grid = HashGrid.from_geometric(
                feature_dim=args.feature_dim,
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                min_grid_res=args.min_grid_res,
                max_grid_res=args.max_grid_res,
                blas_level=args.blas_level
            )
        # "quad" - determines the resolution of the grid using an octree sampling pattern.
        elif args.tree_type == "octree":
            grid = HashGrid.from_octree(
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                blas_level=args.blas_level
            )
    elif args.grid_type == "LatentGrid":
        conf_latent_decoder = args_dict["latent_decoder"]
        # "geometric" - determines the resolution of the grid using geometric sequence initialization from InstantNGP,
        if args.tree_type == "geometric":
            grid = LatentGrid.from_geometric(
                feature_dim=args.feature_dim,
                latent_dim=conf_latent_decoder["latent_dim"],
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                resolution_dim=2,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                min_grid_res=args.min_grid_res,
                max_grid_res=args.max_grid_res,
                blas_level=args.blas_level,
                init_grid=args.init_grid,
                conf_latent_decoder=conf_latent_decoder,
                conf_entropy_reg=args_dict["entropy_reg"]
            )
        # "quad" - determines the resolution of the grid using an octree sampling pattern.
        elif args.tree_type == "octree":
            grid = LatentGrid.from_octree(
                feature_dim=args.feature_dim,
                latent_dim=conf_latent_decoder["latent_dim"],
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                resolution_dim=2,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                blas_level=args.blas_level,
                init_grid=args.init_grid,
                conf_latent_decoder=conf_latent_decoder,
                conf_entropy_reg=args_dict["entropy_reg"]
            )
    else:
        raise ValueError(f"Unknown grid_type argument: {args.grid_type}")
    return grid


def load_neural_field(args, args_dict, dataset: BaseImageDataset) -> NeuralImage:
    """ Creates a "Neural Field" instance which converts input coordinates to some output signal.
    Here a NeuralRadianceField is created, which maps 2D coordinates -> RGB.
    The NeuralRadianceField uses spatial feature grids internally for faster feature interpolation and raymarching.
    """
    grid = load_grid(args=args, args_dict=args_dict, dataset=dataset)
    nef = NeuralImage(
        grid=grid,
        pos_embedder=args.pos_embedder,
        position_input=args.position_input,
        pos_multires=args.pos_multires,
        activation_type=args.activation_type,
        layer_type=args.layer_type,
        hidden_dim=args.hidden_dim,
        num_layers=args.num_layers,
    )
    return nef

def load_neural_pipeline(args, args_dict, dataset, device) -> Pipeline:
    """ In Wisp, a Pipeline comprises of a neural field + a tracer (the latter is optional in this image case).
    Together, they form the complete pipeline required to render a neural primitive from input rays / coordinates.
    """
    nef = load_neural_field(args=args, args_dict=args_dict, dataset=dataset)
    pipeline = Pipeline(nef=nef, tracer=None)
    if args.pretrained:
        if args.model_format == "full":
            pipeline = torch.load(args.pretrained)
        else:
            pipeline.load_state_dict(torch.load(args.pretrained))
    pipeline.to(device)
    return pipeline


def load_trainer(pipeline, train_dataset, validation_dataset, device, scene_state, args, args_dict, writer) -> BaseTrainer:
    """ Loads the NeRF trainer.
    The trainer is responsible for managing the optimization life-cycles which will run 
    the train() function until all training steps are exhausted.
    """
    # args.optimizer_type is the name of some optimizer class (from torch.optim or apex),
    # Wisp's config_parser is able to pick this app's args with corresponding names to the optimizer constructor args.
    # The actual construction of the optimizer instance happens within the trainer.
    optimizer_cls = config_parser.get_module(name=args.optimizer_type)
    optimizer_params = config_parser.get_args_for_function(args, optimizer_cls)

    trainer = ImageTrainer(pipeline=pipeline,
                               train_dataset=train_dataset,
                               validation_dataset=validation_dataset,
                               num_epochs=args.epochs,
                               batch_size=args.batch_size,
                               optim_cls=optimizer_cls,
                               lr=args.lr,
                               weight_decay=args.weight_decay,
                               weight_decay_decoder=args.weight_decay_decoder,
                               grid_lr=args.grid_lr,
                               ldec_lr=args.ldec_lr,
                               optim_params=optimizer_params,
                               log_dir=args.log_dir,
                               device=device,
                               exp_name=args.exp_name,
                               info=args_to_log_format(args_dict),
                               extra_args=vars(args),
                               render_tb_every=args.render_tb_every,
                               save_every=args.save_every,
                               scene_state=scene_state,
                               writer=writer,
                               trainer_mode='validate' if args.valid_only else 'train',
                               using_wandb=args.wandb_project is not None,
                               metrics_only=args.log_metrics_only,
                               enable_amp=not args.disable_amp,
                               use_scaler=not args.disable_scaler)
    return trainer

def set_seed(seed: int):
    import random, os
    import numpy as np
    import torch
    
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

if __name__ == "__main__":
    args, args_dict = parse_args()  # Obtain args by priority: cli args > config yaml > argparse defaults
    default_log_setup(args.log_level)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if args.num_samples>-1:
        args.num_samples = 2 ** args.num_samples

    set_seed(args.seed)
    # Copy data to local path for faster reading
    if args.copy_local:
        # # If dataset is a single file such as an image
        # if not os.path.isdir(args.dataset_path):
        #     args.dataset_path = copy_file(args.dataset_path, args.local_path if args.local_path else None)
        # Use msrsync executable if available for fast copying
        if os.path.exists(args.msrsync_exec):
            args.dataset_path = copy_dir_msrsync(args.dataset_path, args.msrsync_num_workers, args.msrsync_exec,
                                                    args.local_path if args.local_path else None)
        else:
            args.dataset_path = copy_dir(args.dataset_path, args.local_path if args.local_path else None)
        args_dict['dataset_path'] = args.dataset_path

    # Load dataset for training
    train_dataset = load_dataset(args=args)
    validation_dataset = train_dataset

    # If resuming, load image number in dataset to resume training from
    start_image_idx = 0
    if args.resume and os.path.exists(os.path.join(args.log_dir, args.exp_name, "resume_state.pth")):
        state = torch.load(os.path.join(args.log_dir, args.exp_name, "resume_state.pth"))
        start_image_idx = state["image_idx"]
    train_dataset.image_idx = start_image_idx

    # Default TensorBoard Logging
    writer = SummaryWriter(os.path.join(args.log_dir,args.exp_name), purge_step=0)
    writer.add_text('Info', args_to_log_format(args_dict))

    if args.wandb_project is not None:
        wandb_project = args.wandb_project
        wandb_run_name = args.wandb_run_name
        wandb_entity = args.wandb_entity
        wandb_mode = args.wandb_mode
        wandb.init(
            project=wandb_project,
            name=args.exp_name if wandb_run_name is None else wandb_run_name,
            entity=wandb_entity,
            job_type='validate' if args.valid_only else 'train',
            config=args_dict,
            sync_tensorboard=not args.log_metrics_only,
            dir=args.log_dir,
            mode=wandb_mode,
            resume=args.resume,
        )

    avg_metrics = None
    # Perform training for each image separately
    for image_idx in range(start_image_idx, train_dataset.num_images):
        # Load the training coordinates, pixels and initialize model pipeline for current image
        train_dataset.load_next()
        pipeline = load_neural_pipeline(args=args, args_dict=args_dict, dataset=train_dataset, device=device)
        scene_state = ImageState()   # Joint trainer / app state
        # Create trainer for current image
        trainer = load_trainer(pipeline=pipeline, train_dataset=train_dataset, validation_dataset=validation_dataset,
                               device=device, scene_state=scene_state, args=args, args_dict=args_dict,
                               writer=writer)
        if args.resume and os.path.exists(os.path.join(args.log_dir, args.exp_name, "resume_state.pth")):
            trainer.resume_state()
            
        if args.valid_only:
            if args.valid_only_load_path:
                ckpt = torch.load(os.path.join(
                                                args.valid_only_load_path,
                                                os.path.splitext(os.path.basename(train_dataset.image_path))[0],
                                                'model_best.pth'))
                trainer.pipeline.load_state_dict(ckpt.state_dict())
                
                # Decode codebook only once for validation for fast inference
                if isinstance(trainer.pipeline.nef.grid, LatentGrid): 
                    grid = trainer.pipeline.nef.grid
                    assert isinstance(grid, LatentGrid)
                    if args.use_sga:
                        grid.latent_dec.use_sga = False
                    if isinstance(grid.latent_dec, MultiLatentDecoder):
                        grid.latent_dec.straight_through = True
                    decoded_codebook = grid.latent_dec(grid.codebook)
                    grid.codebook = decoded_codebook
                    trainer.pipeline.nef.latent_dec = DecoderIdentity()

            trainer.validate()
        else:
            trainer.train()
            if avg_metrics is None:
                avg_metrics = {k:[v] for k,v in trainer.best_state.items()}
            else:
                avg_metrics = {k:avg_metrics[k]+[v] for k,v in trainer.best_state.items()}

            avg_metrics = {k:sum(avg_metrics[k])/len(avg_metrics[k]) for k in avg_metrics}
            log_text = 'Metrics (avg): '
            for k,v in avg_metrics.items():
                log_text += f'{k}:{v} |'
            log.info(log_text)

            with open(os.path.join(args.log_dir, args.exp_name, 'metrics.json'), 'w') as f:
                json.dump(avg_metrics, f)
        
            if args.wandb_project is not None:
                wandb.run.summary({f"Average/{k}": v for k,v in avg_metrics.items()})

            open(os.path.join(trainer.log_dir, 'complete'), "a").close()
    if args.copy_local:
        shutil.rmtree(args.dataset_path)

    writer.close()
    if args.resume:
        if os.path.exists(os.path.join(args.log_dir, args.exp_name, "resume_state.pth")):
            os.remove(os.path.join(args.log_dir, args.exp_name, "resume_state.pth"))
            print("Finished training. Removed resume state file at"
                  +f"{os.path.join(args.log_dir, args.exp_name, 'resume_state.pth')}")
    if args.wandb_project is not None:
        wandb.finish()

app/image/main_image.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import os
import argparse
import logging
import numpy as np
import torch
import wisp
import tempfile
from wisp.app_utils import default_log_setup, args_to_log_format
import wisp.config_parser as config_parser
from wisp.framework import WispState
from wisp.datasets import MultiviewDataset, SampleRays
from wisp.models.grids import BLASGrid, OctreeGrid, CodebookOctreeGrid, TriplanarGrid, HashGrid, LatentGrid
from wisp.models.latent_decoders.basic_latent_decoder import DecoderIdentity
from wisp.tracers import BaseTracer, PackedRFTracer
from wisp.models.nefs import BaseNeuralField, NeuralRadianceField
from wisp.models.pipeline import Pipeline
from wisp.trainers import BaseTrainer, MultiviewTrainer

def copy_dir_msrsync(input_dir, num_procs, msrsync_exec, dest_dir=None):
    if dest_dir is None:
        dest_dir = tempfile.mkdtemp()
    else:
        dest_dir = dest_dir.rstrip("/")
        os.makedirs(dest_dir,exist_ok=True)
    data_name = input_dir.rstrip("/").split("/")[-1]

    destination_dir = f"{dest_dir}/{data_name}"
    complete_flag = f"{destination_dir}/copy_complete"

    if os.path.exists(complete_flag):
        print(f'Found data already copied to {destination_dir}')
        return destination_dir
    else:
        print(
            f"Copying {input_dir} to dir {destination_dir} using {num_procs} processes"
        )
        # We have to do multi-threaded rsync to speed up copy.
        cmd = (
            f"{msrsync_exec} -p {num_procs} {input_dir.rstrip('/')} {dest_dir}"
        )
        os.system(cmd)
        open(complete_flag, "a").close()
        print("Copied to local directory")
        return destination_dir

def parse_args():
    """Wisp mains define args per app.
    Args are collected by priority: cli args > config yaml > argparse defaults
    For convenience, args are divided into groups.
    """
    parser = argparse.ArgumentParser(description='A script for training simple NeRF variants.',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--config', type=str,
                        help='Path to config file to replace defaults.')
    parser.add_argument('--profile', action='store_true',
                        help='Enable NVTX profiling')
    parser.add_argument('--seed', type=int,
                        help='Random seed. (Not implemented)')
    parser.add_argument('--headless', action='store_true',
                        help='Enable NVTX profiling')

    log_group = parser.add_argument_group('logging')
    log_group.add_argument('--exp-name', type=str,
                           help='Experiment name, unique id for trainers, logs.')
    log_group.add_argument('--log-level', action='store', type=int, default=logging.INFO,
                           help='Logging level to use globally, DEBUG: 10, INFO: 20, WARN: 30, ERROR: 40.')
    log_group.add_argument('--perf', action='store_true', default=False,
                           help='Use high-level profiling for the trainer.')
    log_group.add_argument('--resume', action='store_true', default=False,
                           help='Resume training from the latest checkpoint.')

    data_group = parser.add_argument_group('dataset')
    data_group.add_argument('--dataset-path', type=str,
                            help='Path to the dataset')
    data_group.add_argument('--copy-local', action='store_true', default=False,
                            help='Whether to copy the dataset to local disk.')
    data_group.add_argument('--local-path', type=str, default='',
                            help='Local path to copy the dataset to. If not specified, a temporary directory will be used.')
    data_group.add_argument('--msrsync-exec', type=str, default='/path/to/msrsync/executable',
                            help='Path to msrsync executable.')
    data_group.add_argument('--dataset-num-workers', type=int, default=-1,
                            help='Number of workers for dataset preprocessing, if it supports multiprocessing. '
                                 '-1 indicates no multiprocessing.')
    data_group.add_argument('--dataloader-num-workers', type=int, default=0,
                            help='Number of workers for dataloader.')
    data_group.add_argument('--bg-color', default='black' if is_interactive() else 'white',
                            choices=['white', 'black'], help='Background color')
    data_group.add_argument('--multiview-dataset-format', default='standard', choices=['standard', 'rtmv'],
                            help='Data format for the transforms')
    data_group.add_argument('--num-rays-sampled-per-img', type=int, default=4096,
                            help='Number of rays to sample per image')
    data_group.add_argument('--mip', type=int, default=None,
                            help='MIP level of ground truth image')

    grid_group = parser.add_argument_group('grid')
    grid_group.add_argument('--grid-type', type=str, default='OctreeGrid',
                            choices=config_parser.list_modules('grid'),
                            help='Type of to use, i.e.:'
                                 '"OctreeGrid", "CodebookOctreeGrid", "TriplanarGrid", "HashGrid", "LatentGrid".'
                                 'Grids are located in `wisp.models.grids`')
    grid_group.add_argument('--interpolation-type', type=str, default='linear', choices=['linear', 'closest'],
                            help='Interpolation type to use for samples within grids.'
                                 'For a 3D grid structure, linear uses trilinear interpolation of 8 cell nodes,'
                                 'closest uses the nearest neighbor.')
    grid_group.add_argument('--blas-type', type=str, default='octree', 
                            choices=['octree',],
                            help='Type of acceleration structure to use for fast raymarch occupancy queries.')
    grid_group.add_argument('--multiscale-type', type=str, default='sum', choices=['sum', 'cat'],
                            help='Aggregation of choice for multi-level grids, for features from different LODs.')
    grid_group.add_argument('--feature-dim', type=int, default=32,
                            help='Dimensionality for features stored within the grid nodes.')
    grid_group.add_argument('--feature-std', type=float, default=0.0,
                            help='Grid initialization: standard deviation used for randomly sampling initial features.')
    grid_group.add_argument('--feature-bias', type=float, default=0.0,
                            help='Grid initialization: bias used for randomly sampling initial features.')
    grid_group.add_argument('--base-lod', type=int, default=2,
                            help='Number of levels in grid, which book-keep occupancy but not features.'
                                 'The total number of levels in a grid is `base_lod + num_lod - 1`')
    grid_group.add_argument('--num-lods', type=int, default=1,
                            help='Number of levels in grid, which store concrete features.')
    grid_group.add_argument('--codebook-bitwidth', type=int, default=8,
                            help='For Codebook and HashGrids/LatentGrids only: determines the table size as 2**(bitwidth).')
    grid_group.add_argument('--tree-type', type=str, default='geometric', choices=['geometric', 'quad'],
                            help='For HashGrids/LatentGrids only: how the resolution of the grid is determined. '
                                 '"geometric" uses the geometric sequence initialization from InstantNGP,'
                                 'where "quad" uses an octree sampling pattern.')
    grid_group.add_argument('--min-grid-res', type=int, default=16,
                            help='For HashGrids/LatentGrids only: min grid resolution, used only in geometric initialization mode')
    grid_group.add_argument('--max-grid-res', type=int, default=2048,
                            help='For HashGrids/LatentGrids only: max grid resolution, used only in geometric initialization mode')
    grid_group.add_argument('--prune-min-density', type=float, default=(0.01 * 512) / np.sqrt(3),
                            help='For HashGrids/LatentGrids only: Minimum density value for pruning')
    grid_group.add_argument('--prune-density-decay', type=float, default=0.6,
                            help='For HashGrids/LatentGrids only: The decay applied on the density every pruning')
    grid_group.add_argument('--blas-level', type=float, default=7,
                            help='For HashGrids/LatentGrids only: Determines the number of levels in the acceleration structure '
                                 'used to track the occupancy status (bottom level acceleration structure).')
    grid_group.add_argument('--init-grid', type=str, default='normal', choices=['normal', 'uniform'],
                            help='Grid initialization: distribution used for randomly sampling initial features.')

    nef_group = parser.add_argument_group('nef')
    nef_group.add_argument('--pos-embedder', type=str, choices=['none', 'identity', 'positional'],
                           default='positional',
                           help='MLP Decoder of neural field: Positional embedder used to encode input coordinates'
                                'or view directions.')
    nef_group.add_argument('--view-embedder', type=str, choices=['none', 'identity', 'positional'],
                           default='positional',
                           help='MLP Decoder of neural field: Positional embedder used to encode view direction')
    nef_group.add_argument('--position-input', type=bool, default=False,
                           help='If True, position coords will be concatenated to the '
                                'features / positional embeddings when fed into the decoder.')
    nef_group.add_argument('--pos-multires', type=int, default=10,
                           help='MLP Decoder of neural field: Number of frequencies to use for positional encoding'
                                'of input coordinates')
    nef_group.add_argument('--view-multires', type=int, default=4,
                           help='MLP Decoder of neural field: Number of frequencies to use for positional encoding'
                                'of view direction')
    nef_group.add_argument('--layer-type', type=str, default='none',
                           choices=['none', 'spectral_norm', 'frobenius_norm', 'l_1_norm', 'l_inf_norm'])
    nef_group.add_argument('--activation-type', type=str, default='relu',
                           choices=['relu', 'sin'])
    nef_group.add_argument('--hidden-dim', type=int, help='MLP Decoder of neural field: width of all hidden layers.')
    nef_group.add_argument('--num-layers', type=int, help='MLP Decoder of neural field: number of hidden layers.')

    trainer_group = parser.add_argument_group('trainer')
    trainer_group.add_argument('--epochs', type=int, default=250,
                               help='Number of epochs to run the training.')
    trainer_group.add_argument('--batch-size', type=int, default=512,
                               help='Batch size for the training.')
    trainer_group.add_argument('--resample', action='store_true',
                               help='Resample the dataset after every epoch.')
    trainer_group.add_argument('--only-last', action='store_true',
                               help='Train only last LOD.')
    trainer_group.add_argument('--resample-every', type=int, default=1,
                               help='Resample every N epochs')
    trainer_group.add_argument('--model-format', type=str, default='full', choices=['full', 'state_dict'],
                               help='Format in which to save models.')
    trainer_group.add_argument('--pretrained', type=str,
                               help='Path to pretrained model weights.')
    trainer_group.add_argument('--save-as-new', action='store_true',
                               help='Save the model at every epoch (no overwrite).')
    trainer_group.add_argument('--save-every', type=int, default=(-1 if is_interactive() else 5),
                               help='Save the model at every N epoch.')
    trainer_group.add_argument('--render-tb-every', type=int, default=(-1 if is_interactive() else 5),
                               help='Render every N epochs')
    trainer_group.add_argument('--log-tb-every', type=int, default=5, 
                               help='Render to tensorboard every N epochs')
    trainer_group.add_argument('--log-dir', type=str, default='_results/logs/runs/',
                               help='Log file directory for checkpoints.')
    trainer_group.add_argument('--prune-every', type=int, default=-1,
                               help='Prune every N epochs')
    trainer_group.add_argument('--grow-every', type=int, default=-1,
                               help='Grow network every X epochs')
    trainer_group.add_argument('--growth-strategy', type=str, default='increase',
                               choices=['onebyone',      # One by one trains one level at a time.
                                        'increase',      # Increase starts from [0] and ends up at [0,...,N]
                                        'shrink',        # Shrink strats from [0,...,N] and ends up at [N]
                                        'finetocoarse',  # Fine to coarse starts from [N] and ends up at [0,...,N]
                                        'onlylast'],     # Only last starts and ends at [N]
                               help='Strategy for coarse-to-fine training')
    trainer_group.add_argument('--valid-only', action='store_true',
                               help='Run validation only (and do not run training).')
    trainer_group.add_argument('--valid-only-load-path', type=str, default=None,
                               help='If valid only, load model.pth from this path.')
    trainer_group.add_argument('--valid-every', type=int, default=-1,
                               help='Frequency of running validation.')
    trainer_group.add_argument('--random-lod', action='store_true',
                               help='Use random lods to train.')
    
    tracer_group = parser.add_argument_group('tracer')
    tracer_group.add_argument('--raymarch-type', type=str, choices=['ray', 'voxel'], default='ray',
                              help='Marching algorithm to use when generating samples along rays in tracers.'
                                   '`ray` samples fixed amount of randomized `num_steps` along the ray.'
                                   '`voxel` samples `num_steps` samples in each cell the ray intersects.')
    tracer_group.add_argument('--num-steps', type=int, default=1024,
                              help='Number of samples to generate along traced rays. See --raymarch-type for '
                                   'algorithm used to generate the samples.')

    wandb_group = parser.add_argument_group('wandb')
    wandb_group.add_argument('--wandb-project', type=str, default=None,
                               help='Weights & Biases Project')
    wandb_group.add_argument('--wandb-mode', type=str, default='online',
                               help='Whether to run Weights & Biases in online or offline mode.')
    wandb_group.add_argument('--log-metrics-only', action='store_true',
                               help='Weights & Biases Log metrics only (no images or artifacts).')
    wandb_group.add_argument('--wandb-run-name', type=str, default=None,
                               help='Weights & Biases Run Name')
    wandb_group.add_argument('--wandb-entity', type=str, default=None,
                               help='Weights & Biases Entity')
    wandb_group.add_argument('--wandb-viz-nerf-angles', type=int, default=20,
                               help='Number of Angles to visualize a scene on Weights & Biases. '
                                    'Set this to 0 to disable 360 degree visualizations.')
    wandb_group.add_argument('--wandb-viz-nerf-distance', type=int, default=3,
                               help='Distance to visualize Scene from on Weights & Biases')

    optimizer_group = parser.add_argument_group('optimizer')
    optimizer_group.add_argument('--disable-amp', action='store_true',
                                 help='Disable Automatic Mixed Precision (AMP) training.')
    optimizer_group.add_argument('--disable-scaler', action='store_true',
                                 help='Disable scaler for training.')
    optimizer_group.add_argument('--optimizer-type', type=str, default='adam',
                                 choices=config_parser.list_modules('optim'),
                                 help='Optimizer to be used, includes optimizer modules available within `torch.optim` '
                                      'and fused optimizers from `apex`, if apex is installed.')
    optimizer_group.add_argument('--lr', type=float, default=0.001,
                                 help='Base optimizer learning rate.')
    optimizer_group.add_argument('--eps', type=float, default=1e-8,
                                 help='Eps value for numerical stability.')
    optimizer_group.add_argument('--weight-decay', type=float, default=0,
                                 help='Weight decay, applied only to decoder weights.')
    optimizer_group.add_argument('--weight-decay-decoder', type=float, default=0,
                                 help='Weight decay, applied only to decoder weights.')
    optimizer_group.add_argument('--grid-lr', type=float, default=1.0e-2,
                                 help='Learning rate applied only for the grid parameters'
                                      '(e.g. parameters which contain "grid" in their name)')
    optimizer_group.add_argument('--scale-grid-lr', type=str, default='none', choices=['none','mul','div'],
                                 help='Scale grid-lr by multiplying/dividing by decoder norm or use as is.')
    optimizer_group.add_argument('--ldec-lr', type=float, default=1.0e-2,
                                 help='Learning rate applied only for the latent decoder parameters')
    optimizer_group.add_argument('--ldec-lr-warmup', type=int, default=1,
                                 help='Warmup epochs for latent decoder learning rate')
    optimizer_group.add_argument('--rgb-loss', type=float, default=1.0,
                                 help='Weight of rgb loss')

    # Evaluation renderer (definitions do not affect interactive renderer)
    offline_renderer_group = parser.add_argument_group('renderer')
    offline_renderer_group.add_argument('--render-res', type=int, nargs=2, default=[512, 512],
                                        help='Width/height to render at.')
    offline_renderer_group.add_argument('--render-batch', type=int, default=0,
                                        help='Batch size (in number of rays) for batched rendering.')
    offline_renderer_group.add_argument('--camera-origin', type=float, nargs=3, default=[-2.8, 2.8, -2.8],
                                        help='Camera origin.')
    offline_renderer_group.add_argument('--camera-lookat', type=float, nargs=3, default=[0, 0, 0],
                                        help='Camera look-at/target point.')
    offline_renderer_group.add_argument('--camera-fov', type=float, default=30,
                                        help='Camera field of view (FOV).')
    offline_renderer_group.add_argument('--camera-proj', type=str, choices=['ortho', 'persp'], default='persp',
                                        help='Camera projection.')
    offline_renderer_group.add_argument('--camera-clamp', nargs=2, type=float, default=[0, 10],
                                        help='Camera clipping bounds.')

    latent_decoder_group = parser.add_argument_group('latent_decoder')
    latent_decoder_group.add_argument('--ldecode-enabled', action='store_true', default=False,
                           help='Whether to perform quantization and latent decoding.')
    latent_decoder_group.add_argument('--ldecode-type', type=str, choices=['single', 'multi', 'hierarchical'], default='single',
                                        help='Type of latent decoder to use.')
    latent_decoder_group.add_argument('--use-sga', action='store_true', default=False,
                           help='Whether to use stochastic gumbel annealing for the quantized weights before decoding.')
    latent_decoder_group.add_argument('--diff-sampling', action='store_true', default=False,
                                        help='Whether to make sampling differentiable for SGA.')
    latent_decoder_group.add_argument('--ldecode-matrix', type=str, choices=['sq','dft'], default='sq',
                                        help='Whether to use a learnable decoder or a DFT basis decoder for the weight matrices.')
    latent_decoder_group.add_argument('--latent-dim', type=int, default=0,
                               help='Latent dimension. If 0, latent dimension is feature dimension.')
    latent_decoder_group.add_argument('--norm', type=str, choices=['none','min_max'], default='none',
                                        help='Whether to normalize latents.')
    latent_decoder_group.add_argument('--use-shift', action='store_true', default=False,
                           help='Whether to use shift parameter in decoder.')
    latent_decoder_group.add_argument('--num-layers-dec', type=int, default=0,
                               help='Number of hidden layers in decoder. If 0, decoder is linear.')
    latent_decoder_group.add_argument('--hidden-dim-dec', type=int, default=0,
                               help='Hidden dimension in decoder. If 0, hidden dimension is feature dimension.')
    latent_decoder_group.add_argument('--activation', type=str, choices=['none','sigmoid','tanh','relu','sine'], 
                                      default='none', help='Activation function to use in decoder.')
    latent_decoder_group.add_argument('--final-activation', type=str, choices=['none','sigmoid','tanh','relu','sine'], 
                                      default='none', help='Final layer activation to use in decoder.')
    latent_decoder_group.add_argument('--clamp-weights', type=float, default=0.0,
                                        help='Clamp norm of decoder weights to this value. 0 disables clamping.')
    latent_decoder_group.add_argument('--ldec-std', type=float, default=1.0,
                                        help='Standard deviation for decoder weight initialization.')
    latent_decoder_group.add_argument('--num-decoders', type=int, default=2,
                                        help='Number of latent decoders. Used only when decode-type is multi.')
    latent_decoder_group.add_argument('--temperature', type=float, default=1.0,
                                        help='Temperature for decoder softmax scores. Used only when decode-type is multi.')
    latent_decoder_group.add_argument('--decay-period', type=float, default=0.8,
                                        help='Fraction of epochs for temperature decay. Set to value of temperature after this.')
    latent_decoder_group.add_argument('--alpha-std', type=float, default=10.0,
                                        help='Initialization of softmax logits for decoder. Used only when decode-type is multi.')
    
    entropy_reg_group = parser.add_argument_group('entropy_reg')
    entropy_reg_group.add_argument('--num-prob-layers', type=int, default=2,
                                        help='Number of hidden layers for probability model.')
    entropy_reg_group.add_argument('--entropy-reg', type=float, default=0.0,
                                        help='Starting entropy regularization lambda for latents.')
    entropy_reg_group.add_argument('--entropy-reg-end', type=float, default=0.0,
                                        help='Ending entropy regularization lambda for latents.')
    entropy_reg_group.add_argument('--entropy-reg-sched', type=str, default='fix', choices=['fix', 'linear', 'exp','cosine'],
                                        help='Decay schedule for entropy regularization lambda for latents.')
    entropy_reg_group.add_argument('--noise-freq', type=int, default=50,
                                        help='Iteration frequency for generating noise for prob model.')

    # Parse CLI args & config files
    args = config_parser.parse_args(parser)

    # Also obtain args as grouped hierarchy, useful for, i.e., logging
    args_dict = config_parser.get_grouped_args(parser, args)
    return args, args_dict


def load_dataset(args) -> MultiviewDataset:
    """ Loads a multiview dataset comprising of pairs of images and calibrated cameras.
    The types of supported datasets are defined by multiview_dataset_format:
    'standard' - refers to the standard NeRF format popularized by Mildenhall et al. 2020,
                 including additions to the metadata format added by Muller et al. 2022.
    'rtmv' - refers to the dataset published by Tremblay et. al 2022,
            "RTMV: A Ray-Traced Multi-View Synthetic Dataset for Novel View Synthesis".
            This dataset includes depth information which allows for performance improving optimizations in some cases.
    """
    transform = SampleRays(num_samples=args.num_rays_sampled_per_img)
    train_dataset = wisp.datasets.load_multiview_dataset(dataset_path=args.dataset_path,
                                                         split='train',
                                                         mip=args.mip,
                                                         bg_color=args.bg_color,
                                                         dataset_num_workers=args.dataset_num_workers,
                                                         transform=transform)
    validation_dataset = None
    if args.valid_every > -1 or args.valid_only:
        validation_dataset = train_dataset.create_split(split='val', transform=None)
    return train_dataset, validation_dataset


def load_grid(args, args_dict, dataset: MultiviewDataset) -> BLASGrid:
    """ Wisp's implementation of NeRF uses feature grids to improve the performance and quality (allowing therefore,
    interactivity).
    This function loads the feature grid to use within the neural pipeline.
    Grid choices are interesting to explore, so we leave the exact backbone type configurable,
    and show how grid instances may be explicitly constructed.
    Grids choices, for example, are: OctreeGrid, TriplanarGrid, HashGrid,LatentGrid, CodebookOctreeGrid
    See corresponding grid constructors for each of their arg details.
    """
    grid = None

    # Optimization: For octrees based grids, if dataset contains depth info, initialize only cells known to be occupied
    if args.grid_type == "OctreeGrid":
        if dataset.supports_depth():
            grid = OctreeGrid.from_pointcloud(
                pointcloud=dataset.as_pointcloud(),
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
            )
        else:
            grid = OctreeGrid.make_dense(
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
            )
    elif args.grid_type == "CodebookOctreeGrid":
        if dataset.supports_depth():
            grid = CodebookOctreeGrid.from_pointcloud(
                pointcloud=dataset.as_pointcloud(),
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth
            )
        else:
            grid = CodebookOctreeGrid.make_dense(
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                interpolation_type=args.interpolation_type,
                multiscale_type=args.multiscale_type,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth
            )
    elif args.grid_type == "TriplanarGrid":
        grid = TriplanarGrid(
            feature_dim=args.feature_dim,
            base_lod=args.base_lod,
            num_lods=args.num_lods,
            interpolation_type=args.interpolation_type,
            multiscale_type=args.multiscale_type,
            feature_std=args.feature_std,
            feature_bias=args.feature_bias,
        )
    elif args.grid_type == "HashGrid":
        # "geometric" - determines the resolution of the grid using geometric sequence initialization from InstantNGP,
        if args.tree_type == "geometric":
            grid = HashGrid.from_geometric(
                feature_dim=args.feature_dim,
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                resolution_dim=3,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                min_grid_res=args.min_grid_res,
                max_grid_res=args.max_grid_res,
                blas_level=args.blas_level
            )
        # "quad" - determines the resolution of the grid using an octree sampling pattern.
        elif args.tree_type == "octree":
            grid = HashGrid.from_octree(
                feature_dim=args.feature_dim,
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                resolution_dim=3,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                blas_level=args.blas_level
            )
    elif args.grid_type == "LatentGrid":
        conf_latent_decoder = args_dict["latent_decoder"]
        # "geometric" - determines the resolution of the grid using geometric sequence initialization from InstantNGP,
        if args.tree_type == "geometric":
            grid = LatentGrid.from_geometric(
                feature_dim=args.feature_dim,
                latent_dim=conf_latent_decoder["latent_dim"],
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                resolution_dim=3,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                min_grid_res=args.min_grid_res,
                max_grid_res=args.max_grid_res,
                blas_level=args.blas_level,
                init_grid=args.init_grid,
                conf_latent_decoder=conf_latent_decoder,
                conf_entropy_reg=args_dict["entropy_reg"]
            )
        # "quad" - determines the resolution of the grid using an octree sampling pattern.
        elif args.tree_type == "octree":
            grid = LatentGrid.from_octree(
                feature_dim=args.feature_dim,
                latent_dim=conf_latent_decoder["latent_dim"],
                base_lod=args.base_lod,
                num_lods=args.num_lods,
                multiscale_type=args.multiscale_type,
                resolution_dim=3,
                feature_std=args.feature_std,
                feature_bias=args.feature_bias,
                codebook_bitwidth=args.codebook_bitwidth,
                blas_level=args.blas_level,
                init_grid=args.init_grid,
                conf_latent_decoder=conf_latent_decoder,
                conf_entropy_reg=args_dict["entropy_reg"]
            )
    else:
        raise ValueError(f"Unknown grid_type argument: {args.grid_type}")
    return grid


def load_neural_field(args, args_dict, dataset: MultiviewDataset) -> BaseNeuralField:
    """ Creates a "Neural Field" instance which converts input coordinates to some output signal.
    Here a NeuralRadianceField is created, which maps 3D coordinates (+ 2D view direction) -> RGB + density.
    The NeuralRadianceField uses spatial feature grids internally for faster feature interpolation and raymarching.
    """
    grid = load_grid(args=args, args_dict=args_dict, dataset=dataset)
    nef = NeuralRadianceField(
        grid=grid,
        pos_embedder=args.pos_embedder,
        view_embedder=args.view_embedder,
        position_input=args.position_input,
        pos_multires=args.pos_multires,
        view_multires=args.view_multires,
        activation_type=args.activation_type,
        layer_type=args.layer_type,
        hidden_dim=args.hidden_dim,
        num_layers=args.num_layers,
        prune_density_decay=args.prune_density_decay,   # Used only for grid types which support pruning
        prune_min_density=args.prune_min_density        # Used only for grid types which support pruning
    )
    return nef


def load_tracer(args) -> BaseTracer:
    """ Wisp "Tracers" are responsible for taking input rays, marching them through the neural field to render
    an output RenderBuffer.
    Wisp's implementation of NeRF uses the PackedRFTracer to trace the neural field:
    - Packed: each ray yields a custom number of samples, which are therefore packed in a flat form within a tensor,
     see: https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#packed
    - RF: Radiance Field
    PackedRFTracer is employed within the training loop, and is responsible for making use of the neural field's
    grid to generate samples and decode them to pixel values.
    """
    tracer = PackedRFTracer(
        raymarch_type=args.raymarch_type,   # Chooses the ray-marching algorithm
        num_steps=args.num_steps,           # Number of steps depends on raymarch_type
        bg_color=args.bg_color
    )
    return tracer


def load_neural_pipeline(args, args_dict, dataset, device) -> Pipeline:
    """ In Wisp, a Pipeline comprises of a neural field + a tracer (the latter is optional in some cases).
    Together, they form the complete pipeline required to render a neural primitive from input rays / coordinates.
    """
    nef = load_neural_field(args=args, args_dict=args_dict, dataset=dataset)
    tracer = load_tracer(args=args)
    pipeline = Pipeline(nef=nef, tracer=tracer)
    if args.pretrained:
        if args.model_format == "full":
            pipeline = torch.load(args.pretrained)
        else:
            pipeline.load_state_dict(torch.load(args.pretrained))
    pipeline.to(device)
    return pipeline


def load_trainer(pipeline, train_dataset, validation_dataset, device, scene_state, args, args_dict) -> BaseTrainer:
    """ Loads the NeRF trainer.
    The trainer is responsible for managing the optimization life-cycles and can be operated in 2 modes:
    - Headless, which will run the train() function until all training steps are exhausted.
    - Interactive mode, which uses the gui. In this case, an OptimizationApp uses events to prompt the trainer to
      take training steps, while also taking care to render output to users (see: iterate()).
      In interactive mode, trainers can also share information with the app through the scene_state (WispState object).
    """
    # args.optimizer_type is the name of some optimizer class (from torch.optim or apex),
    # Wisp's config_parser is able to pick this app's args with corresponding names to the optimizer constructor args.
    # The actual construction of the optimizer instance happens within the trainer.
    optimizer_cls = config_parser.get_module(name=args.optimizer_type)
    optimizer_params = config_parser.get_args_for_function(args, optimizer_cls)

    trainer = MultiviewTrainer(pipeline=pipeline,
                               train_dataset=train_dataset,
                               validation_dataset=validation_dataset,
                               num_epochs=args.epochs,
                               batch_size=args.batch_size,
                               optim_cls=optimizer_cls,
                               lr=args.lr,
                               weight_decay=args.weight_decay,
                               weight_decay_decoder=args.weight_decay_decoder,
                               grid_lr=args.grid_lr,
                               ldec_lr=args.ldec_lr,
                               optim_params=optimizer_params,
                               log_dir=args.log_dir,
                               device=device,
                               exp_name=args.exp_name,
                               info=args_to_log_format(args_dict),
                               extra_args=vars(args),
                               render_tb_every=args.render_tb_every,
                               save_every=args.save_every,
                               scene_state=scene_state,
                               trainer_mode='validate' if args.valid_only else 'train',
                               using_wandb=args.wandb_project is not None,
                               metrics_only=args.log_metrics_only,
                               enable_amp=not args.disable_amp,
                               use_scaler=not args.disable_scaler)
    return trainer


def load_app(args, scene_state, trainer):
    """ Used only in interactive mode. Creates an interactive app, which employs a renderer which displays
    the latest information from the trainer (see: OptimizationApp).
    The OptimizationApp can be customized or further extend to support even more functionality.
    """
    if not is_interactive():
        logging.info("Running headless. For the app, set $WISP_HEADLESS=0.")
        return None  # Interactive mode is disabled
    else:
        from wisp.renderer.app.optimization_app import OptimizationApp
        scene_state.renderer.device = trainer.device  # Use same device for trainer and app renderer
        app = OptimizationApp(wisp_state=scene_state,
                              trainer_step_func=trainer.iterate,
                              experiment_name="wisp trainer")
        return app

def set_seed(seed: int):
    import random, os
    import numpy as np
    import torch
    
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

def is_interactive() -> bool:
    """ Returns True if interactive mode with gui is on, False is HEADLESS mode is forced """
    return os.environ.get('WISP_HEADLESS') != '1'


args, args_dict = parse_args()  # Obtain args by priority: cli args > config yaml > argparse defaults
default_log_setup(args.log_level)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

set_seed(args.seed)
if args.headless:
    os.environ['WISP_HEADLESS'] = '1'  # Force headless mode
    
if args.copy_local:
    args.dataset_path = copy_dir_msrsync(args.dataset_path, args.dataset_num_workers, args.msrsync_exec,
                                            args.local_path if args.local_path else None)
    args_dict['dataset_path'] = args.dataset_path

train_dataset, validation_dataset = load_dataset(args=args)
pipeline = load_neural_pipeline(args=args, args_dict=args_dict, dataset=train_dataset, device=device)
scene_state = WispState()   # Joint trainer / app state
trainer = load_trainer(pipeline=pipeline,
                       train_dataset=train_dataset, validation_dataset=validation_dataset,
                       device=device, scene_state=scene_state,
                       args=args, args_dict=args_dict)
app = load_app(args=args, scene_state=scene_state, trainer=trainer)

if app is not None:
    app.run()  # Run in interactive mode
else:
    if args.valid_only:
        if args.valid_only_load_path:
            ckpt = torch.load(os.path.join(args.valid_only_load_path,'model.pth'))
            trainer.pipeline.load_state_dict(ckpt.state_dict())
            if isinstance(trainer.pipeline.nef.grid, LatentGrid): # Decode codebook only once for validation for fast inference
                decoded_codebook = trainer.pipeline.nef.grid.latent_dec(trainer.pipeline.nef.grid.codebook)
                trainer.pipeline.nef.grid.codebook = decoded_codebook
                trainer.pipeline.nef.latent_dec = DecoderIdentity()
        trainer.validate()
    else:
        trainer.train()

app/nerf/main_nerf.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import os
import sys
from setuptools import setup, find_packages, dist
import glob
import logging
import subprocess

import torch
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME

PACKAGE_NAME = 'wisp'
DESCRIPTION = 'Kaolin-Wisp: A PyTorch library for performing research on neural fields'
URL = 'https://github.com/NVIDIAGameWorks/kaolin-wisp'
AUTHOR = 'Towaki Takikawa'
LICENSE = 'NVIDIA Source Code License'
version = '0.1.0'

def get_cuda_bare_metal_version(cuda_dir):
    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
    output = raw_output.split()
    release_idx = output.index("release") + 1
    release = output[release_idx].split(".")
    bare_metal_major = release[0]
    bare_metal_minor = release[1][0]

    return raw_output, bare_metal_major, bare_metal_minor

if not torch.cuda.is_available():
    if os.getenv('FORCE_CUDA', '0') == '1':
        # From: https://github.com/NVIDIA/apex/blob/c4e85f7bf144cb0e368da96d339a6cbd9882cea5/setup.py
        # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query torch.cuda.get_device_capability(),
        # which will fail if you are compiling in an environment without visible GPUs (e.g. during an nvidia-docker build command).
        logging.warning(
            "Torch did not find available GPUs on this system.\n"
            "If your intention is to cross-compile, this is not an error.\n"
            "By default, Apex will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\n"
            "Volta (compute capability 7.0), Turing (compute capability 7.5),\n"
            "and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\n"
            "If you wish to cross-compile for a single specific architecture,\n"
            'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n'
        )
        if os.getenv("TORCH_CUDA_ARCH_LIST", None) is None:
            _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
            if int(bare_metal_major) == 11:
                if int(bare_metal_minor) == 0:
                    os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0"
                else:
                    os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6"
            else:
                os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
    else:
        logging.warning(
            "Torch did not find available GPUs on this system.\n"
            "This script will install only with CPU support and will have very limited features.\n"
            'If your wish to cross-compile for GPU `export FORCE_CUDA=1` before running setup.py\n'
            "By default, Apex will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\n"
            "Volta (compute capability 7.0), Turing (compute capability 7.5),\n"
            "and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\n"
            "If you wish to cross-compile for a single specific architecture,\n"
            'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n'
        )

def get_extensions():
    extra_compile_args = {'cxx': ['-O3']} 
    define_macros = []
    include_dirs = []
    extensions = []
    sources = glob.glob('wisp/csrc/**/*.cpp', recursive=True)
 
    if len(sources) == 0:
        print("No source files found for extension, skipping extension compilation")
        return None


    if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
        define_macros += [("WITH_CUDA", None), ("THRUST_IGNORE_CUB_VERSION_CHECK", None)]
        sources += glob.glob('wisp/csrc/**/*.cu', recursive=True)
        extension = CUDAExtension
        extra_compile_args.update({'nvcc': ['-O3']})
        #include_dirs = get_include_dirs()
    else:
        assert(False, "CUDA is not available. Set FORCE_CUDA=1 for Docker builds")

    extensions.append(
        extension(
            name='wisp._C',
            sources=sources,
            define_macros=define_macros,
            extra_compile_args=extra_compile_args,
            #include_dirs=include_dirs
        )
    )

    for ext in extensions:
        ext.libraries = ['cudart_static' if x == 'cudart' else x
                         for x in ext.libraries]
 
    return extensions

if __name__ == '__main__':
    setup(
        # Metadata
        name=PACKAGE_NAME,
        version=version,
        author=AUTHOR,
        description=DESCRIPTION,
        url=URL,
        license=LICENSE,
        python_requires='>=3.8',

        # Package info
        packages=['wisp'] + find_packages(),
        #package_dir={'':'wisp'},
        include_package_data=True,
        zip_safe=True,
        ext_modules=get_extensions(),
        cmdclass={
            'build_ext': BuildExtension.with_options(no_python_abi_suffix=True)    
        }

    )

setup.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

wisp/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .base_as import BaseAS, ASQueryResults, ASRaytraceResults, ASRaymarchResults
from .octree_as import OctreeAS
from .aabb_as import AxisAlignedBBoxAS

wisp/accelstructs/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from wisp.accelstructs.octree_as import OctreeAS
import wisp.ops.spc as wisp_spc_ops


class AxisAlignedBBoxAS(OctreeAS):
    """Axis Aligned Bounding Box, as a bottom-level acceleration structure class.
    Can be used to to quickly query cells occupancy, and trace rays against the volume.
    """
    
    def __init__(self):
        """Initializes a simple acceleration structure of an AABB (axis aligned bounding box). """
        # Builds a root-only octree, of one level, which is essentially a bounding box.
        # Useful for hacking together a quick AABB (axis aligned bounding box) tracer.
        octree = wisp_spc_ops.create_dense_octree(1)
        super().__init__(octree)

    def name(self) -> str:
        return "AABB"

wisp/accelstructs/aabb_as.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
import torch
from wisp.core import WispModule


@dataclass
class ASQueryResults:
    """ A data holder for keeping the results of a single acceleration structure query() call.
    A query receives a set of input coordinates and returns the cell indices of the acceleration structure where
    the query coordinates fall.
    """

    pidx: torch.LongTensor
    """ Holds the query results.
    - If the query is invoked with `with_parents=False`, this field is a tensor of shape [num_coords],
      containing indices of cells of the acceleration structure, where the query coordinates match.
    - If the query is invoked with `with_parents=True`, this field is a tensor of shape [num_coords, level+1],
      containing indices of the cells of the acceleration structure + the full parent hierarchy of each 
      cell query result.
    """


@dataclass
class ASRaytraceResults:
    """ A data holder for keeping the results of a single acceleration structure raytrace() call.
    A raytrace operation returns all intersections of the ray set with the acceleration structure cells.
    Ray/cell intersections are also referred to in Kaolin & Wisp as "nuggets".
    """

    ridx: torch.LongTensor
    """ A tensor containing the ray index of the ray that intersected each cell [num_nuggets]. 
    (can be used to index into rays.origins and rays.dirs)
    """

    pidx: torch.LongTensor
    """ Point indices into the cells of the acceleration structure, where the ray intersected a cell [num_nuggets] """

    depth: torch.FloatTensor
    """ Depths of each nugget, representing:
      - The first intersection point of the ray with the cell (entry), and 
      - Optionally also a second intersection point of the ray with the cell (exit).
      A tensor of [num_intersections, 1 or 2]. 
    """


@dataclass
class ASRaymarchResults:
    """ A data holder for keeping the results of a single acceleration structure raymarching() call.
    A raymarch operation advances the ray set within the acceleration structure and generates samples along the rays.
    """

    samples: torch.FloatTensor
    """ Sample coordinates of shape [num_hit_samples, 3]"""

    ridx: torch.LongTensor
    """ A tensor containing the ray index of the ray that generated each sample [num_hit_samples]. 
    Can be used to index into rays.origins and rays.dirs.
    """

    depth_samples: Optional[torch.FloatTensor]
    """ Depths of each sample, a tensor of shape [num_hit_samples, 1] """

    deltas: Optional[torch.FloatTensor]
    """ Depth diffs between each sample and the previous one, a tensor of shape [num_hit_samples, 1] """

    boundary: Optional[torch.BoolTensor]
    """ Boundary tensor which marks the beginning of each variable-sized sample pack of shape [num_hit_samples].
        That is: [True, False, False, False, True, False False] represents two rays of 4 and 3 samples respectively. 
    """


class BaseAS(WispModule, ABC):
    """
    A base interface for all acceleration structures within Wisp.
    """

    def __init__(self):
        super().__init__()

    def query(self, coords, level=None, with_parents=False) -> ASQueryResults:
        """Returns the ``cell index`` corresponding to the sample coordinates
        (indices of acceleration structure cells the query coords match).

        Args:
            coords (torch.FloatTensor) : 3D coordinates of shape [num_coords, 3].
                Certain acceleration structures may require the coords to be normalized within some range.
            level (int) : For acceleration structures which support multiple level-of-details,
                level specifies the desired LOD to query. If None, queries the highest level.
            with_parents (bool) : For acceleration structures which support multiple level-of-details,
                If `with_parents=True`, this call will return the hierarchical parent indices as well.

        Returns:
            (ASQueryResults): A plain data object carrying the results of the results of the query operation,
            containing the indices of cells matching the query coords.
        """
        raise NotImplementedError(f"{self.name} acceleration structure does not support the 'query' method.")

    def raytrace(self, rays, level=None, with_exit=False) -> ASRaytraceResults:
        """Traces rays against the SPC structure, returning all intersections along the ray with the SPC points
        (SPC points are quantized, and can be interpreted as octree cell centers or corners).

        Args:
            rays (wisp.core.Rays): Ray origins and directions of shape [batch, 3].
            level (int) : The level of the octree to raytrace. If None, traces the highest level.
            with_exit (bool) : If True, also returns exit depth.

        Returns:
            (ASRaytraceResults): A plain data object carrying the results of the raytrace operation,
             containing the indices of ray and cell intersections.
        """
        raise NotImplementedError(f"{self.name} acceleration structure does not support the 'raytrace' method.")

    def raymarch(self, rays, *args, **kwargs) -> ASRaymarchResults:
        """Samples points along the ray within the acceleration structure.
        The exact algorithm used is under the responsibility of the acceleration structure.

        Args:
            rays (wisp.core.Rays): Ray origins and directions of shape [batch, 3].
            *args, **kwargs: Optional args this specific acceleration structure requires for marching.
        Returns:
            (RaymarchResults): a plain data object carrying the results of raymarching,
            containing samples generated by the ray marching.
        """
        raise NotImplementedError(f"{self.name} acceleration structure does not support the 'raymarch' method.")

    def occupancy(self) -> List[int]:
        """ Returns a list of length [LODs], where each element contains the number of cells occupied in that LOD.
        Used for logging, debugging and visualization purposes, for acceleration structures with
        a defined concept of occupancy.
        """
        return list()

    def capacity(self) -> List[int]:
        """ Returns a list of length [LODs], where each element contains the total cell capacity in that LOD.
        Used for logging, debugging and visualization purposes, for acceleration structures with
        a defined concept of capacity.
        """
        return list()

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.

        Acceleration structures are encouraged to report:
        - 'occupancy' - a list of length [LODs], where each element contains the number of cells occupied in that LOD
        - 'capacity' - a list of length [LODs], where each element contains the total cell capacity in that LOD
        """
        return {
            '#Used Cells (LOD)': self.occupancy(),
            '#Capacity (LOD)': self.capacity()
        }

wisp/accelstructs/base_as.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import List, Tuple
import torch
import kaolin.ops.spc as spc_ops
import kaolin.render.spc as spc_render
import wisp.ops.mesh as mesh_ops
import wisp.ops.spc as wisp_spc_ops
from wisp.accelstructs.base_as import BaseAS, ASQueryResults, ASRaytraceResults, ASRaymarchResults


@torch.jit.script
def fast_filter_method(mask_idx: torch.Tensor, depth: torch.Tensor, deltas: torch.Tensor, samples: torch.Tensor,
                       num_samples: int, num_rays: int, device: torch.device) -> \
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:


    depth_samples = depth[mask_idx[:, 0], mask_idx[:, 1]][:, None]
    deltas = deltas[mask_idx[:, 0], mask_idx[:, 1]].reshape(-1, 1)
    samples = samples[mask_idx[:, 0], mask_idx[:, 1], :]

    ridx = torch.arange(0, num_rays, device=device)
    ridx = ridx[..., None].repeat(1, num_samples)[mask_idx[:, 0], mask_idx[:, 1]]

    return depth_samples, deltas, samples, ridx


class OctreeAS(BaseAS):
    """Octree bottom-level acceleration structure class implemented using Kaolin SPC.
    Can be used to to quickly query cells occupancy, and trace rays against the volume.
    """
    
    def __init__(self, octree):
        """Initializes the acceleration structure from the topology of a sparse octree (Structured Point Cloud).
        Structured Point Cloud (SPC) is a compact data structure for organizing and efficiently pack sparse 3D geometric
        information.
        Intuitively, SPCs can also be described as sparse voxel-grids, quantized point clouds, or
        voxelized point clouds.

        Args:
            octree (torch.ByteTensor): SPC octree tensor, containing the acceleration structure topology.
            For more details about this format, see:
             https://kaolin.readthedocs.io/en/latest/notes/spc_summary.html
        """
        super().__init__()
        self.octree = octree
        self.points, self.pyramid, self.prefix = wisp_spc_ops.octree_to_spc(octree)
        self.max_level = self.pyramid.shape[-1] - 2
        self.extent = dict()    # Additional optional information which may be stored in this struct

    @classmethod
    def from_mesh(cls, mesh_path, level, sample_tex=False, num_samples=100000000) -> OctreeAS:
        """ Builds the acceleration structure and initializes occupancy of cells from samples over mesh faces.
        Assumes a path to the mesh, only OBJ is supported for now.

        Args:
            mesh_path (str): Path to the OBJ file.
            level (int): The depth of the octree.
            sample_tex (bool): If True, will also sample textures and store it in the accelstruct within
                self.texv, self.texf, self.mats fields.
                This feature is currently unused.
            num_samples (int): The number of samples to be generated on the mesh surface.
                More samples require additional processing time but are more likely to produce faithful occupancy
                output without "holes".
        """
        if sample_tex:
            out = mesh_ops.load_obj(mesh_path, load_materials=True)
            vertices, faces, texture_vertices, texture_faces, materials = out
        else:
            vertices, faces = mesh_ops.load_obj(mesh_path)

        # For now only supports sphere normalization, which is a bit more robust for SDF type workloads
        # (although it will underutilize the voxels)
        vertices, faces = mesh_ops.normalize(vertices, faces, 'sphere')

        # Note: This function is not deterministic since it relies on sampling.
        # Eventually this will be replaced by 3D rasterization.
        octree = wisp_spc_ops.mesh_to_octree(vertices, faces, level, num_samples)
        accel_struct = OctreeAS(octree)
        accel_struct.extent['vertices'] = vertices
        accel_struct.extent['faces'] = faces
        if sample_tex:
            accel_struct.extent['texv'] = texture_vertices
            accel_struct.extent['texf'] = texture_faces
            accel_struct.extent['mats'] = materials
        return accel_struct

    @classmethod
    def from_pointcloud(cls, pointcloud, level) -> OctreeAS:
        """ Builds the acceleration structure and initializes occupancy of cells from a pointcloud.

        Args:
            pointcloud (torch.FloatTensor): 3D coordinates of shape [num_coords, 3] in 
                                            normalized space [-1, 1].
            level (int): The depth of the octree.
        """
        octree = wisp_spc_ops.pointcloud_to_octree(pointcloud, level, dilate=0)
        return OctreeAS(octree)

    @classmethod
    def from_quantized_points(cls, quantized_points, level) -> OctreeAS:
        """ Builds the acceleration structure from quantized (integer) point coordinates.

        Args:
            quantized_points (torch.LongTensor): 3D coordinates of shape [num_coords, 3] in
                                                 integer coordinate space [0, 2**level]
            level (int): The depth of the octree.
        """
        octree = spc_ops.unbatched_points_to_octree(quantized_points, level, sorted=False)
        return OctreeAS(octree)

    @classmethod
    def make_dense(cls, level) -> OctreeAS:
        """ Builds the acceleration structure and initializes full occupancy of all cells.

        Args:
            level (int): The depth of the octree.
        """
        octree = wisp_spc_ops.create_dense_octree(level)
        return OctreeAS(octree)

    def query(self, coords, level=None, with_parents=False) -> ASQueryResults:
        """Returns the ``pidx`` for the sample coordinates (indices of acceleration structure cells returned by
        this query).

        Args:
            coords (torch.FloatTensor) : 3D coordinates of shape [num_coords, 3] in normalized [-1, 1] space.
            level (int) : The depth of the octree to query. If None, queries the highest level.
            with_parents (bool) : If True, also returns hierarchical parent indices.
        
        Returns:
            (ASQueryResults): containing the indices into the point hierarchy of shape [num_query].
                If with_parents is True, then the query result will be of shape [num_query, level+1].
        """
        if level is None:
            level = self.max_level
        
        pidx = spc_ops.unbatched_query(self.octree, self.prefix, coords, level, with_parents)
        return ASQueryResults(pidx=pidx)

    def raytrace(self, rays, level=None, with_exit=False) -> ASRaytraceResults:
        """Traces rays against the SPC structure, returning all intersections along the ray with the SPC points
        (SPC points are quantized, and can be interpreted as octree cell centers or corners).

        Args:
            rays (wisp.core.Rays): Ray origins and directions of shape [batch, 3].
            level (int) : The level of the octree to raytrace. If None, traces the highest level.
            with_exit (bool) : If True, also returns exit depth.

        Returns:
            (ASRaytraceResults): with fields containing -
                - Indices into rays.origins and rays.dirs of shape [num_intersections]
                - Indices into the point_hierarchy of shape [num_intersections]
                - Depths of shape [num_intersections, 1 or 2]
        """
        if level is None:
            level = self.max_level

        ridx, pidx, depth = spc_render.unbatched_raytrace(
            self.octree, self.points, self.pyramid, self.prefix,
            rays.origins, rays.dirs, level, return_depth=True, with_exit=with_exit)
        return ASRaytraceResults(ridx=ridx, pidx=pidx, depth=depth)

    def _raymarch_voxel(self, rays, num_samples, level=None) -> ASRaymarchResults:
        """Samples points along the ray inside the SPC structure.
        Raymarch is achieved by intersecting the rays with the SPC cells.
        Then among the intersected cells, each cell is sampled num_samples times.
        In this scheme, num_hit_samples <= num_intersections*num_samples

        Args:
            rays (wisp.core.Rays): Ray origins and directions of shape [batch, 3].
            num_samples (int) : Number of samples generated per voxel. The total number of samples generated will
                also depend on the number of cells a ray have intersected.
            level (int) : The level of the octree to raytrace. If None, traces the highest level.

        Returns:
            (ASRaymarchResults) with fields containing:
                - Indices into rays.origins and rays.dirs of shape [num_hit_samples]
                - Sample coordinates of shape [num_hit_samples, 3]
                - Sample depths of shape [num_hit_samples, 1]
                - Sample depth diffs of shape [num_hit_samples, 1]
                - Boundary tensor which marks the beginning of each variable-sized
                  sample pack of shape [num_hit_samples]
        """
        # NUM_INTERSECTIONS = number of nuggets: ray / cell intersections
        # NUM_INTERSECTIONS can be 0!
        # ridx, pidx ~ (NUM_INTERSECTIONS,)
        # depth ~ (NUM_INTERSECTIONS, 2)
        raytrace_results = self.raytrace(rays, level, with_exit=True)
        ridx = raytrace_results.ridx.long()
        num_intersections = ridx.shape[0]

        # depth_samples ~ (NUM_INTERSECTIONS, NUM_SAMPLES, 1)
        depth = raytrace_results.depth
        depth_samples = wisp_spc_ops.sample_from_depth_intervals(depth, num_samples)[..., None]
        deltas = depth_samples[..., 0].diff(dim=-1, prepend=depth[..., 0:1]).reshape(num_intersections * num_samples, 1)

        # samples ~ (NUM_INTERSECTIONS, NUM_SAMPLES, 1)
        samples = torch.addcmul(rays.origins.index_select(0, ridx)[:, None],
                                rays.dirs.index_select(0, ridx)[:, None], depth_samples)

        # boundary ~ (NUM_INTERSECTIONS * NUM_SAMPLES,)
        # (each intersected cell is sampled NUM_SAMPLES times)
        boundary = wisp_spc_ops.expand_pack_boundary(spc_render.mark_first_hit(ridx.int()), num_samples)
        boundary = boundary.bool()

        # ridx ~ (NUM_INTERSECTIONS * NUM_SAMPLES,)
        # samples ~ (NUM_INTERSECTIONS * NUM_SAMPLES, 3)
        # depth_samples ~ (NUM_INTERSECTIONS * NUM_SAMPLES, 1)
        ridx = ridx[:,None].expand(num_intersections, num_samples).reshape(num_intersections*num_samples)
        samples = samples.reshape(num_intersections * num_samples, 3)
        depth_samples = depth_samples.reshape(num_intersections * num_samples, 1)

        return ASRaymarchResults(
            ridx=ridx,
            samples=samples,
            depth_samples=depth_samples,
            deltas=deltas,
            boundary=boundary
        )

    def _raymarch_ray(self, rays, num_samples, level=None) -> ASRaymarchResults:
        """Samples points along the ray inside the SPC structure.
        Raymarch is achieved by sampling num_samples along each ray,
        and then filtering out samples which falls outside of occupied cells.
        In this scheme, num_hit_samples <= num_rays * num_samples.
        Ray boundaries are determined by the ray dist_min / dist_max values
        (which could, for example, be set by the near / far planes).

        Args:
            rays (wisp.core.Rays): Ray origins and directions of shape [batch, 3].
            num_samples (int) : Number of samples generated per ray. The actual number of generated samples may be lower
                due to samples intersecting empty cells.
            level (int) : The level of the octree to raytrace. If None, traces the highest level.

        Returns:
            (ASRaymarchResults) with fields containing:
                - Indices into rays.origins and rays.dirs of shape [num_hit_samples]
                - Sample coordinates of shape [num_hit_samples, 3]
                - Sample depths of shape [num_hit_samples, 1]
                - Sample depth diffs of shape [num_hit_samples, 1]
                - Boundary tensor which marks the beginning of each variable-sized
                  sample pack of shape [num_hit_samples]
        """
        # Sample points along 1D line
        # depth ~ (NUM_RAYS, NUM_SAMPLES)
        depth = torch.linspace(0, 1.0, num_samples, device=rays.origins.device)[None] + \
                (torch.rand(rays.origins.shape[0], num_samples, device=rays.origins.device) / num_samples)

        # Normalize between near and far plane
        depth *= (rays.dist_max - rays.dist_min)
        depth += rays.dist_min

        # Batched generation of samples
        # samples ~ (NUM_RAYS, NUM_SAMPLES, 3)
        # deltas, pidx, mask ~ (NUM_RAYS, NUM_SAMPLES)
        num_rays = rays.shape[0]
        samples = torch.addcmul(rays.origins[:, None], rays.dirs[:, None], depth[..., None])
        query_results = self.query(samples.reshape(num_rays * num_samples, 3), level=level)
        pidx = query_results.pidx
        pidx = pidx.reshape(num_rays, num_samples)
        mask = pidx > -1
        non_masked_idx = torch.nonzero(mask)
        
        deltas = depth.diff(dim=-1,
                            prepend=(torch.zeros(rays.origins.shape[0], 1, device=depth.device) + rays.dist_min))

        # NUM_HIT_SAMPLES: number of samples sampled within occupied cells
        # NUM_HIT_SAMPLES can be 0!
        # depth_samples, deltas, ridx, boundary ~ (NUM_HIT_SAMPLES,)
        # samples ~ (NUM_HIT_SAMPLES, 3)

        depth_samples, deltas, samples, ridx = fast_filter_method(non_masked_idx, depth, deltas, samples, num_samples, num_rays, pidx.device)

        boundary = spc_render.mark_pack_boundaries(ridx)

        return ASRaymarchResults(
            ridx=ridx.long(),
            samples=samples.float(),
            depth_samples=depth_samples.float(),
            deltas=deltas.float(),
            boundary=boundary
        )

    def raymarch(self, rays, raymarch_type, num_samples, level=None) -> ASRaymarchResults:
        """Samples points along the ray inside the SPC structure.
        The exact algorithm employed for raymarching is determined by `raymarch_type`.

        Args:
            rays (wisp.core.Rays): Ray origins and directions of shape [batch, 3].
            raymarch_type (str): Sampling strategy to use for raymarch.
                'voxel' - intersects the rays with the SPC cells. Then among the intersected cells, each cell
                    is sampled num_samples times.
                    In this scheme, num_hit_samples <= num_intersections*num_samples
                'ray' - samples num_samples along each ray, and then filters out samples which falls outside of occupied
                    cells.
                    In this scheme, num_hit_samples <= num_rays * num_samples
            num_samples (int) : Number of samples generated per voxel or ray. The exact meaning of this arg depends on
                the value of `raymarch_type`.
            level (int) : The level of the octree to raytrace. If None, traces the highest level.
        
        Returns:
            (ASRaymarchResults) with fields containing:
                - Indices into rays.origins and rays.dirs of shape [num_hit_samples]
                - Sample coordinates of shape [num_hit_samples, 3]
                - Sample depths of shape [num_hit_samples, 1]
                - Sample depth diffs of shape [num_hit_samples, 1]
                - Boundary tensor which marks the beginning of each variable-sized
                  sample pack of shape [num_hit_samples]
        """
        if level is None:
            level = self.max_level

        # Samples points along the rays by first tracing it against the SPC object.
        # Then, given each SPC voxel hit, will sample some number of samples in each voxel.
        # This setting is pretty nice for getting decent outputs from outside-looking-in scenes, 
        # but in general it's not very robust or proper since the ray samples will be weirdly distributed
        # and or aliased. 
        if raymarch_type == 'voxel':
            raymarch_results = self._raymarch_voxel(rays=rays, num_samples=num_samples, level=level)

        # Samples points along the rays, and then uses the SPC object the filter out samples that don't hit
        # the SPC objects. This is a much more well-spaced-out sampling scheme and will work well for 
        # inside-looking-out scenes. The camera near and far planes will have to be adjusted carefully, however.
        elif raymarch_type == 'ray':
            raymarch_results = self._raymarch_ray(rays=rays, num_samples=num_samples, level=level)

        else:
            raise TypeError(f"Raymarch sampler type: {raymarch_type} is not supported by OctreeAS.")

        return raymarch_results

    def occupancy(self) -> List[int]:
        """ Returns a list of length [LODs], where each element contains the number of cells occupied in that LOD """
        return self.pyramid[0, :-2].cpu().numpy().tolist()

    def capacity(self) -> List[int]:
        """ Returns a list of length [LODs], where each element contains the total cell capacity in that LOD """
        return [8**lod for lod in range(self.max_level)]

    def name(self) -> str:
        return "Octree"

wisp/accelstructs/octree_as.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import tempfile
import logging
import os
import sys
import pprint
import shutil
from distutils.dir_util import copy_tree


def default_log_setup(level=logging.INFO):
    """
    Sets up default logging, always logging to stdout.

    :param level: logging level, e.g. logging.INFO
    """
    handlers = [logging.StreamHandler(sys.stdout)]
    
    logging.basicConfig(level=level,
                        format='%(asctime)s|%(levelname)8s| %(message)s',
                        handlers=handlers)


def args_to_log_format(args_dict) -> str:
    """Convert args hierarchy to string representation suitable for logging (i.e. with Tensorboard).

    Args:
        args_dict : The parsed arguments, grouped within a dictionary.

    Returns:
        arg_str : The args encoded in a string format.
    """
    pp = pprint.PrettyPrinter(indent=2)
    args_str = pp.pformat(args_dict)
    args_str = f'```{args_str}```'
    return args_str

def copy_file(input_file, dest_dir=None) -> str:
    """Copy file from input directory to a temporary directory.

    Args:
        input_file (str): Input file containing data to be copied.
        dest_dir (str): Path to directory to be copied. Creates temporary directory
                        if None.

    Returns:
        dest_dir (str): Destination directory path
    """
    # Create destination directory
    if dest_dir is None:
        dest_dir = tempfile.mkdtemp()
    else:
        dest_dir = dest_dir.rstrip("/")
        os.makedirs(dest_dir,exist_ok=True)

    filename = os.path.basename(input_file)
    shutil.copyfile(input_file, os.path.join(dest_dir,filename))
    return os.path.join(dest_dir,filename)

def copy_dir(input_dir, dest_dir=None) -> str:
    """Copy data from input directory to a temporary directory.

    Args:
        input_dir (str): Input directory containing data to be copied.
        dest_dir (str): Path to directory to be copied. Creates temporary directory
                        if None.

    Returns:
        dest_dir (str): Destination directory path
    """
    # Create destination directory
    if dest_dir is None:
        dest_dir = tempfile.mkdtemp()
    else:
        dest_dir = dest_dir.rstrip("/")
        os.makedirs(dest_dir,exist_ok=True)

    copy_tree(input_dir, dest_dir)
    return dest_dir

def copy_dir_msrsync(input_dir, num_procs, msrsync_exec, dest_dir=None) -> str:
    """Copy data from input directory to a temporary directory for faster data loading. 
       Uses msrsync (https://github.com/jbd/msrsync) for fast parallelized copy.

    Args:
        input_dir (str): Input directory containing data to be copied.
        num_procs (int): Number of processes to use for msrsync
        msrsync_exec (str): Path to the executable msrsync binary
        dest_dir (str): Path to directory to be copied. Creates temporary directory
                        if None.

    Returns:
        dest_dir (str): Destination directory path
    """

    # Create destination directory
    if dest_dir is None:
        dest_dir = tempfile.mkdtemp()
    else:
        dest_dir = dest_dir.rstrip("/")
        os.makedirs(dest_dir,exist_ok=True)

    data_name = input_dir.rstrip("/").split("/")[-1]

    # Maintain flag if copying is complete
    destination_dir = f"{dest_dir}/{data_name}"
    complete_flag = f"{destination_dir}/copy_complete"

    if os.path.exists(complete_flag):
        print(f'Found data already copied to {destination_dir}')
        return destination_dir
    else:
        print(
            f"Copying {input_dir} to dir {destination_dir} using {num_procs} processes..."
        )
        # We have to do multi-threaded rsync to speed up copy.
        cmd = (
            f"{msrsync_exec} -p {num_procs} {input_dir.rstrip('/')} {dest_dir}"
        )
        os.system(cmd)
        open(complete_flag, "a").close()
        print(f"Copied to directory {destination_dir}")
        return destination_dir
    
wisp/app_utils.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import argparse
import logging
import os
import yaml
import inspect
import torch
from typing import Dict, List, Any

from wisp.models import nefs
from wisp.models import grids
from wisp import tracers
from wisp import datasets

# This file contains all the configuration and command-line parsing general to all app

__all__ = [
    'list_modules',
    'register_module',
    'get_module',
    'get_args_for_function',
    'get_grouped_args',
    'parse_args'
]

# str2mod ("str to module") are registered wisp blocks the config parser is aware of, and is able to dynamically load.
# You may register additional options by adding them to the dictionary here.
# ConfigParser expects the following default module categories:
str2mod = {
    'optim': {m.lower(): getattr(torch.optim, m) for m in dir(torch.optim) if m[0].isupper()},
    'dataset': {},
    'nef': {},
    'grid': {},
    'tracer': {}
}


def list_modules(module_type) -> List[str]:
    """ Returns a list of all available modules from a certain category.
    Args:
        type: a str from the following categories: ['nef', 'grid', 'tracer', 'dataset', 'optim']
    """
    return list(str2mod[module_type].keys())


def register_module(module_type, name, mod):
    """Register module to be used with config parser.
    Users should use this class to load their classes by name.
    """
    if module_type not in str2mod:
        raise ValueError(f"'{module_type}' is an unknown type")

    if name in str2mod[module_type]:
        raise KeyError(f"'{name}' already exist in type '{module_type}'")
    str2mod[module_type][name] = mod


def get_module(name, module_type=None):
    """Get module class by name, assuming it was registered with `register_module`.'"""
    types_to_check = []
    if module_type is None:
        types_to_check = str2mod
    else:
        if module_type not in str2mod:
            raise ValueError(f"'{module_type}' is an unknown type")
        types_to_check.append(module_type)

    for t in types_to_check:
        if name in str2mod[t]:
            return str2mod[t][name]

    raise ValueError(f"'{name}' is not a known module for any of the types '{types_to_check}'. "
                     f"registered modules are '{str2mod[module_type].keys()}'")


def get_args_for_function(args, func):
    """ Given a func (for example an __init__(..) function or from_X(..)), and also the parsed args,
    return the subset of args that func expects and args contains. """
    if isinstance(args, argparse.Namespace):
        args = vars(args)   # Namespace -> dict
    has_kwargs = inspect.getfullargspec(func).varkw != None
    if has_kwargs:
        collected_args = args
    else:
        parameters = dict(inspect.signature(func).parameters)
        collected_args = {a: args[a] for a in parameters if a in args}
    return collected_args


def get_grouped_args(parser, args) -> Dict[str, Any]:
    """Group args to a grouped hierarchy.

    Args:
        parser (argparse.parser): Parser object. Needed for the argument groups.
        args : The parsed arguments. Will compute from the parser if None.

    Returns:
        args    : The parsed arguments.
        arg_str : The string to be printed.
    """
    args_dict = {}
    for group in parser._action_groups:
        group_dict = {a.dest: getattr(args, a.dest, None) for a in group._group_actions}
        args_dict[group.title] = vars(argparse.Namespace(**group_dict))
    return args_dict


# -- Register all wisp library modules here -- this makes them loadable by specifying the class name to get_module() --
for name in dir(datasets):
    mod = getattr(datasets, name)
    if isinstance(mod, type) and \
            issubclass(mod, torch.utils.data.Dataset) and \
            mod != torch.utils.data.Dataset:
        register_module('dataset', name, mod)

for name in dir(nefs):
    mod = getattr(nefs, name)
    if isinstance(mod, type) and \
            issubclass(mod, nefs.BaseNeuralField) and \
            mod != nefs.BaseNeuralField:
        register_module('nef', name, mod)

for name in dir(grids):
    mod = getattr(grids, name)
    if isinstance(mod, type) and \
            issubclass(mod, grids.BLASGrid) and \
            mod != grids.BLASGrid:
        register_module('grid', name, mod)

for name in dir(tracers):
    mod = getattr(tracers, name)
    if isinstance(mod, type) and \
            issubclass(mod, tracers.BaseTracer) and \
            mod != tracers.BaseTracer:
        register_module('tracer', name, mod)

try:
    import apex

    for m in dir(apex.optimizers):
        if m[0].isupper():
            register_module('optim', m.lower(), getattr(apex.optimizers, m))
except ImportError:
    logger = logging.getLogger(__name__)
    logger.info("Cannot import apex for fused optimizers")


def parse_yaml_config(config_path, parser):
    """Parses and sets the parser defaults with a yaml config file.

    Args:
        config_path : path to the yaml config file.
        parser : The parser for which the defaults will be set.
        parent : True if parsing the parent yaml. Should never be set to True by the user.
    """
    with open(config_path) as f:
        config_dict = yaml.safe_load(f)

    list_of_valid_fields = []
    for group in parser._action_groups:
        group_dict = {list_of_valid_fields.append(a.dest) for a in group._group_actions}
    list_of_valid_fields = set(list_of_valid_fields)

    defaults_dict = {}
    defaults_dict_args = {}

    # Load the parent config if it exists
    parent_config_path = config_dict.pop("parent", None)

    if parent_config_path is not None:
        if not os.path.isabs(parent_config_path):
            parent_config_path = os.path.join(os.path.split(config_path)[0], parent_config_path)
        with open(parent_config_path) as f:
            parent_config_dict = yaml.safe_load(f)
        if "parent" in parent_config_dict.keys():
            raise Exception("Hierarchical configs of more than 1 level deep are not allowed.")
        for key in parent_config_dict:
            for field in parent_config_dict[key]:
                if field not in list_of_valid_fields:
                    raise ValueError(
                        f"ERROR: {field} is not a valid option. Check for typos in the config."
                    )
                if isinstance(parent_config_dict[key][field], dict):
                    defaults_dict_args[field] = parent_config_dict[key][field]
                    defaults_dict[field] = None
                else:
                    defaults_dict[field] = parent_config_dict[key][field]

    # Loads child parent and overwrite the parent configs
    # The yaml files assumes the argument groups, which aren't actually nested.
    for key in config_dict:
        for field in config_dict[key]:
            if field not in list_of_valid_fields:
                raise ValueError(
                    f"ERROR: {field} is not a valid option. Check for typos in the config."
                )
            if isinstance(config_dict[key][field], dict):
                defaults_dict_args[field] = config_dict[key][field]
                defaults_dict[field] = None
            else:
                defaults_dict[field] = config_dict[key][field]

    parser.set_defaults(**defaults_dict)
    return defaults_dict_args


def parse_args(parser) -> argparse.Namespace:
    """Parses args by priority into a flat configuration.
    The various options take the following precedence:
    1. CLI args, explicitly specified
    2. YAML configuration, defined with `--config <PATH>.yaml`
    3. argparse defaults

    Args:
        parser (argparse.parser): Parser object. Needed for the argument groups.

    Returns:
        args    : The parsed arguments, as a flat configuration.
    """

    args = parser.parse_args()
    defaults_dict = dict()

    if args.config is not None:
        config_defaults_dict = parse_yaml_config(args.config, parser)
        for key, val in config_defaults_dict.items():
            if key in defaults_dict:
                defaults_dict[key].update(val)
            else:
                defaults_dict[key] = val
        args = parser.parse_args()

    for key, val in defaults_dict.items():
        cmd_line_val = getattr(args, key)
        if cmd_line_val is not None:
            val.update(cmd_line_val)
        setattr(args, key, val)

    return args

wisp/config_parser.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .primitives import PrimitivesPack
from .rays import Rays
from .channels import *
from .channel_fn import *
from .colors import *
from .render_buffer import RenderBuffer
from .transforms import ObjectTransform
from .wisp_module import WispModule

wisp/core/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Callable, Any
import torch
import torch.nn.functional as F


BlendFunction = Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
"""
Blend functions receive: channel_a, channel_b, alpha_a, alpha_b, and return a blended channel value.
Channel A is assumed to be in front of channel B
"""

NormalizeFunction = Callable[[torch.Tensor, Any, Any], torch.Tensor]
"""
Normalization functions receive a channel, and possible min / max values acceptable for this channel type.
These functions output the normalized values for the channel in the range [0, 1]
"""

###############################################################################################################
#                                           How to use channels:                                              #
###############################################################################################################
#  Blending:                                                                                                  #
#  - Channel dimensions are flexible and assumed to be (..., C), where C is the number of channel dimensions. #
#  - Channel dtypes are assumed to be floating points (of arbitrary precision)                                #
#  - Alpha values are assumed to be in the range of [0, 1].                                                   #
#  - TODO (operel): add support for premultiplied alphas /w Porter Duff operators                             #
#                                                                                                             #
#  Normalization:                                                                                             #
#  - Channels may define a min value, a max value, none of them or both. It is up to the normalization        #
#    function to query for the existence of such boundaries. If any of them is None, a fallback should        #
#    exist.                                                                                                   #
#  - A normalized channel is constrains values to the range [0, 1].                                           #
###############################################################################################################


###############################################################################################################
#                                       Normalization Functions:                                              #
###############################################################################################################

def identity(c: torch.Tensor) -> torch.Tensor:
    """ A naive normalization function which assumes the value is already normalized and returned as is.

    Args:
        c (torch.Tensor): A single channel tensor of an arbitrary shape.

    Returns:
        (torch.Tensor): Input channel c is returned without a change.

    """
    return c


def normalize(c: torch.Tensor, min_val: Any = None, max_val: Any = None) -> torch.Tensor:
    """ A linear normalization function which maps the channel c to the range of [0, 1].
    If the min / max values bounds of the channel are not explicitly specified, they're determined by
    c's values.

    Args:
        c (torch.Tensor): A single channel tensor of an arbitrary shape.
        min_val (Any): Optional, the lower boundary for possible values in channel C
        max_val (Any): Optional, the upper boundary for possible values in channel C

    Returns:
        (torch.Tensor): Input channel c is linearly mapped to the range [0, 1].

    """
    min_val = torch.min(c) if min_val is None else min_val
    max_val = torch.max(c) if max_val is None else max_val
    return (c - min_val) / (max_val - min_val)


def normalize_linear_scale(c: torch.Tensor, min_val: Any = None, max_val: Any = None,
                           linear_scale: float = 1.0) -> torch.Tensor:
    """ A normalization function which linear scales the channel before normalizing it to the range of [0, 1].
    If the min / max values bounds of the channel are not explicitly specified, they're determined by
    c's values. If explicitly specified, the bounds are scaled as well.

    Args:
        c (torch.Tensor): A single channel tensor of an arbitrary shape.
        min_val (Any): Optional, the lower boundary for possible values in channel C
        max_val (Any): Optional, the upper boundary for possible values in channel C
        linear_scale (float): Channel scale value

    Returns:
        (torch.Tensor): Input channel c is linearly mapped to the range [0, 1].

    """
    c *= linear_scale
    min_val = linear_scale * min_val if min_val is not None else min_val
    max_val = linear_scale * max_val if max_val is not None else max_val
    return normalize(c=c, min_val=min_val, max_val=max_val)


def normalize_log_scale(c: torch.Tensor, min_val: Any = None, max_val: Any = None,
                        linear_scale: float = 1.0, log_scale: float = 1.0) -> torch.Tensor:
    """ A normalization function which applies log and linear scales to the channel before normalizing it
    to the range of [0, 1].
    If the min / max values bounds of the channel are not explicitly specified, they're determined by
    c's values. If explicitly specified, the bounds are scaled as well.

    Args:
        c (torch.Tensor): A single channel tensor of an arbitrary shape.
        min_val (Any): Optional, the lower boundary for possible values in channel C
        max_val (Any): Optional, the upper boundary for possible values in channel C
        linear_scale (float): Channel linear scale value
        log_scale (float: Channel log scale value.

    Returns:
        (torch.Tensor): Input channel c is mapped to the range [0, 1] by the formula:
            ``c' = a * log(b * c)`` where a is the linear scale, b is the log scale and c is the channel value.

    """
    c = linear_scale * torch.log(log_scale * c)
    min_val = linear_scale * torch.log(log_scale * min_val) if min_val is not None else min_val
    max_val = linear_scale * torch.log(log_scale * max_val) if max_val is not None else max_val
    return normalize(c=c, min_val=min_val, max_val=max_val)


def normalize_vector(c: torch.Tensor) -> torch.Tensor:
    """ A normalization function which applies a L2 normalization over a channel of vector data.

    Args:
        c (torch.Tensor): A single channel tensor of an arbitrary shape.

    Returns:
        (torch.Tensor): Input channel c is normalized by the L2 norm.
    """
    return F.normalize(c, dim=-1)


###############################################################################################################
#                                           Blending Functions:                                               #
###############################################################################################################


def blend_linear(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ A direct linear interpolation between c1 and c2.
    Useful for blending channels which do not consider the alpha value (i.e. the alpha channel itself).

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused
        alpha2 (torch.Tensor): Unused

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return c1 + c2 * (1.0 - c1)


def blend_alpha_composite_over(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ An alpha compositing op where a front pixel is alpha blended with the background pixel
    (in a usual painter's algorithm manner).
    Useful for blending channels such as RGB.
    See: https://en.wikipedia.org/wiki/Alpha_compositing

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): alpha channel tensor, corresponding to first channel, in the shape of c1.
        alpha2 (torch.Tensor): alpha channel tensor, corresponding to second channel, in the shape of c1.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    alpha_out = alpha1 + alpha2 * (1.0 - alpha1)
    c_out = torch.where(condition=alpha_out > 0,
                        input=(c1 * alpha1 + c2 * alpha2 * (1.0 - alpha1)) / alpha_out,
                        other=torch.zeros_like(c1))
    return c_out


def blend_alpha_lerp(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ A linear interpolation between c1 and c2, which uses the alpha channel as a weighting factor.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): alpha channel tensor, corresponding to first channel, in the shape of c1.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return c1 * alpha1 + c2 * (1.0 - alpha1)


def blend_alpha_slerp(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ A spherical linear interpolation, useful for interpolating rotations or blending directional vectors.
    c1 and c2 are normalized and interpolated over the unit hypersphere.
    alpha1 acts as the interpolation weight.
    See: https://en.wikipedia.org/wiki/Slerp

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): alpha channel tensor, corresponding to first channel, in the shape of c1.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    # To unit directions
    t = alpha1  # alpha1 is used as the interpolation weight
    c1 = F.normalize(c1, dim=1)
    c2 = F.normalize(c2, dim=1)
    dot = (c1*c2).sum(1)         # batched dot prod
    omega = torch.acos(dot)      # angle between directions
    sin_omega = torch.sin(omega)    # TODO (operel): Be careful of omega=0.0 case
    c2_weight = (torch.sin((1.0 - t) * omega) / sin_omega).unsqueeze(1)
    c1_weight = (torch.sin(t * omega) / sin_omega).unsqueeze(1)
    res = c2_weight * c2 + c1_weight * c1
    return res


def blend_normal(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ A standard blend mode which uses the front pixel value, without mixing.
        Useful when alpha blending is undesired, or the channel contains categorical info (i.e. semantic class ids).

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): Unused
        alpha1 (torch.Tensor): Unused
        alpha2 (torch.Tensor): Unused

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return c1


def blend_multiply(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ Commutative blend mode which preserves dark colors.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return c1 * c2


def blend_screen(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ Commutative blend mode which preserves light colors.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return 1.0 - (1.0 - c1)(1.0 - c2)


def blend_add(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ An additive blend mode, for aggregation of channel information.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return c1 + c2


def blend_sub(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ An subtractive blend mode, for removing channel information.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return c1 - c2


def blend_logical_and(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ For boolean channels, blends with a logical AND function.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return torch.logical_and(c1, c2).to(c1.dtype)


def blend_logical_or(c1: torch.Tensor, c2: torch.Tensor, alpha1: torch.Tensor, alpha2: torch.Tensor):
    """ For boolean channels, blends with a logical OR function.

    Args:
        c1 (torch.Tensor): first channel tensor of an arbitrary shape.
        c2 (torch.Tensor): second channel tensor, in the shape of c1.
        alpha1 (torch.Tensor): Unused.
        alpha2 (torch.Tensor): Unused.

    Returns:
        (torch.Tensor): Blended channel in the shape of c1
    """
    return torch.logical_or(c1, c2).to(c1.dtype)

wisp/core/channel_fn.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from wisp.core.channel_fn import *
from dataclasses import dataclass
from typing import Any, Optional, Dict
from functools import partial


@dataclass
class Channel:
    """ Defines how a Renderbuffer channel should behave in terms of functionalities like blending, normalization,
        and boundaries.
    """

    blend_fn: BlendFunction = None
    """" How to blend information from this channel between 2 RenderBuffers """

    normalize_fn: NormalizeFunction = normalize
    """ How to normalize the channel information to a scale of [0, 1] """

    min_val: Optional[Any] = None
    """ Minimal valid value supported by this channel type. None indicates the valid values range from -inf. """

    max_val: Optional[Any] = None
    """ Maximal valid value supported by this channel type . None indicates the valid values range to inf. """


def create_default_channel() -> Channel:
    """ A general channel template, to be used if no information about a channel have been recorded """
    return Channel(
        blend_fn=blend_alpha_composite_over,
        normalize_fn=normalize,
        min_val=None,
        max_val=None
    )


def channels_starter_kit() -> Dict[str, Channel]:
    """ Creates a predefined kit of channels commonly useful in the context of Wisp.
        Users may augment or replace this kit with additional custom channels.
    """
    return dict(
        rgb=Channel(
            blend_fn=blend_alpha_composite_over,
            normalize_fn=identity,
            min_val=0.0,
            max_val=1.0
        ),
        alpha=Channel(
            blend_fn=blend_linear,
            normalize_fn=normalize,
            min_val=0.0,
            max_val=1.0
        ),
        depth=Channel(
            blend_fn=blend_normal,
            normalize_fn=partial(normalize_linear_scale, linear_scale=1000.0),
            min_val=0.0
        ),
        normal=Channel(
            blend_fn=blend_alpha_slerp,
            normalize_fn=normalize_vector
        ),
        hit=Channel(
            blend_fn=blend_logical_or,
            normalize_fn=identity
        ),
        err=Channel(
            blend_fn=blend_add,
            normalize_fn=normalize
        ),
        gt=Channel(
            blend_fn=blend_alpha_composite_over,
            normalize_fn=identity,
            min_val=0.0,
            max_val=1.0
        )
    )

wisp/core/channels.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from typing import List, Tuple

white = (1.0, 1.0, 1.0)
black = (0.0, 0.0, 0.0)
dark_gray = (0.25, 0.25, 0.25)
light_purple = (0.788, 0.580, 1.0)
lime = (0.746, 1.0, 0.0)
red = (1.0, 0.0, 0.0)
green = (0.0, 1.0, 0.0)
blue = (0.0, 0.0, 1.0)
orange = (1.0, 0.5, 0.0)
light_cyan = (0.796, 1.0, 1.0)
light_pink = (1.0, 0.796, 1.0)
light_yellow = (1.0, 1.0, 0.796)
light_teal = (0.757, 1.0, 0.949)
gray = (0.5, 0.5, 0.5)
soft_blue = (0.721, 0.90, 1.0)
soft_red = (1.0, 0.0, 0.085)
lime_green = (0.519, 0.819, 0.0)
purple = (0.667, 0.0, 0.429)
gold = (1.0, 0.804, 0.0)


def color_wheel():
    """ Returns:
        (list) a list of all colors defined in the color module.
        Each entry is a tuple of 3 floats (RGB values).
    """
    return [
        white, black, dark_gray, light_purple, lime, red, green, blue, orange, light_cyan, light_pink,
        light_yellow, light_teal, gray, soft_blue, soft_red, lime_green, purple, gold
    ]


def colors_generator(skip_colors: List = None) -> Tuple[float, float, float]:
    """ Generates the next color in the color wheel on each invocation.
        This generator repeats the color wheel cyclically when exhausted.

    Args:
        skip_colors
    """
    if skip_colors is None:
        skip_colors = []
    while True:
        for color in color_wheel():
            if color in skip_colors:
                continue
            yield color

wisp/core/colors.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from dataclasses import dataclass, field, fields
from typing import List, Tuple, Optional
from wisp.core.transforms import ObjectTransform
import torch


@dataclass
class PrimitivesPack:
    """Packs of geometric primitives.
    Useful for visualizing vectorial data layers with visualizers such as wisp's interactive renderer,
    or external tools like Polyscope..
    """
    _lines_start: List[torch.Tensor] = field(default_factory=list)
    _lines_end: List[torch.Tensor] = field(default_factory=list)
    _lines_color: List[torch.Tensor] = field(default_factory=list)
    _points_pos: List[torch.Tensor] = field(default_factory=list)
    _points_color: List[torch.Tensor] = field(default_factory=list)
    transform: ObjectTransform = None

    line_width = 1.0
    point_size = 1.0

    def append(self, other: PrimitivesPack) -> None:
        """ Appends primitives from other to self, changes the dtype and device if needed.
            The transform is assumed to be coherent between the two packs and is not handled by this function.
        """
        def _append_field(field_name):
            # Get field attribute by name
            _self_field = getattr(self, field_name)
            _other_field = getattr(other, field_name)
            if len(_self_field) > 0:    # If self's field has any entry, we map other's field to this dtype / device
                _other_field = map(lambda t: t.to(_self_field[0].device, _self_field[0].dtype), _other_field)
            _self_field.extend(_other_field)    # Concat the lists of tensors

        for f in fields(self):
            if f.name == 'transform':
                continue
            _append_field(f.name)

    def add_lines(self, start: torch.Tensor, end: torch.Tensor, color: torch.Tensor) -> None:
        """ Adds a single or batch of line primitives to the pack.

        Args:
            start (torch.Tensor): A tensor of (B, 3) or (3,) marking the start point of the line(s).
            end (torch.Tensor): A tensor of (B, 3) or (3,) marking the end point of the line(s).
            color (torch.Tensor): A tensor of (B, 4) or (4,) marking the RGB color of the line(s).
        """
        if start.ndim == 1:
            start = start.unsqueeze(0)
        if end.ndim == 1:
            end = end.unsqueeze(0)
        if color.ndim == 1:
            color = color.unsqueeze(0)
        self._lines_start.append(start)
        self._lines_end.append(end)
        self._lines_color.append(color)

    @property
    def lines(self) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
        """
        Returns:
            (torch.Tensor, torch.Tensor, torch.Tensor) or None:
                - start, a tensor of (B, 3) marking the start vertex of each line.
                - end, a tensor of (B, 3) marking the end vertex of each line.
                - color, a tensor of (B, 4) marking the line color (the color of each start / end vertex).
        """
        # Squash the list of tensors into a single concatenated tensor in lazy load manner
        start = end = color = None
        if len(self._lines_start) > 0:
            start = torch.cat(self._lines_start, dim=0)
        if len(self._lines_end) > 0:
            end = torch.cat(self._lines_end, dim=0)
        if len(self._lines_color) > 0:
            color = torch.cat(self._lines_color, dim=0)
        if start is None or end is None:
            return None
        return start, end, color

    def __eq__(self, other):
        if not isinstance(other, type(self)):
            raise Exception("PrimitivePacks only support equality checks with other PrimitivePacks")
        bools = []
        for f in fields(self):
            self_attr = getattr(self, f.name)
            other_attr = getattr(other, f.name)
            if self_attr is None or other_attr is None:
                if self_attr != other_attr:
                    return False
                else:
                    continue
            if isinstance(self_attr, float) or isinstance(self_attr, ObjectTransform):
                bools.append(self_attr == other_attr)
            else:
                bools.append(all(torch.equal(s, o) for s, o in zip(self_attr, other_attr)))
        return all(bools)

wisp/core/primitives.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Union, Tuple, List
from dataclasses import dataclass
import torch


INFINITY = torch.finfo().max
""" Largest representable float in this system """


@dataclass
class Rays:
    """ A pack of rays represented as origin and direction.
    Ray packs are flexible and may use an arbitrary amount of spatial dimensions (e.g. flat, 2D, etc).
    """

    origins: torch.Tensor
    """ Ray's origin """

    dirs: torch.Tensor
    """ Ray's normalized direction """

    dist_min: Union[float, torch.Tensor] = 0.0
    """ Distance in which ray intersection test begins. Can be defined globally or per ray. """

    dist_max: Union[float, torch.Tensor] = INFINITY
    """ Distance in which ray intersection test ends. Can be defined globally or per ray. """

    # TODO (operel): Handle tensor case in functions below

    def __len__(self) -> int:
        """
        Returns:
            (int): Number of rays in the pack
        """
        if self.origins.shape != self.dirs.shape:
            raise Exception(
                f"Rays.origins shape should match Rays.dirs shape, but got {self.origins.shape} and {self.dirs.shape}.")
        return self.origins.shape[0]

    @property
    def shape(self) -> Tuple[...]:
        """
        Returns:
            (int): Shape of ray pack tensors, excluding the last dimension.
        """
        return self.origins.shape[:-1]

    @property
    def ndim(self) -> int:
        """
        Returns:
            (int): Number of spatial dimensions for ray pack
        """
        return self.origins.ndim - 1

    @classmethod
    def cat(cls, rays_list: List[Rays], dim: int = 0) -> Rays:
        """ Concatenates multiple ray packs into a single pack.

        Args:
            ray_list (List[Rays]): List of ray packs to concatenate, expected to have the same spatial dimensions,
                                  except of dimension dim.
            dim (int): Spatial dimension along which the concatenation should take place

        Returns:
            (Rays): A single ray pack with the concatenation of given ray packs
        """
        if dim < 0:
            dim -= 1
        if dim > rays_list[0].ndim - 1 or dim < -rays_list[0].ndim:
            raise IndexError("Dimension out of range (expected to be in range of "
                             f"[{-rays_list[0].ndim}, {rays_list[0].ndim-1}, but got {dim})")

        return Rays(
            origins=torch.cat([rays.origins for rays in rays_list], dim=dim),
            dirs=torch.cat([rays.dirs for rays in rays_list], dim=dim),
            dist_min=min([rays.dist_min for rays in rays_list]),
            dist_max=max([rays.dist_max for rays in rays_list])
        )

    @classmethod 
    def stack(cls, rays_list: List[Rays], dim: int = 0) -> Rays:
        """ Stacks multiple ray packs into a single pack.

        Args:
            ray_list (List[Rays): List of ray packs to stack, expected to have the same spatial dimensions,
                                  except of dimension dim.

        Returns:
            (Rays): A single ray pack with the stacked ray packs
        """
        return Rays(
            origins=torch.stack([rays.origins for rays in rays_list], dim=dim),
            dirs=torch.stack([rays.dirs for rays in rays_list], dim=dim),
            dist_min=min([rays.dist_min for rays in rays_list]),
            dist_max=max([rays.dist_max for rays in rays_list])
        )

    def __getitem__(self, idx) -> Rays:
        """ Returns a sliced view of the rays.

        Args:
            idx: Indices of rays to slice.

        Returns:
            (Rays): A sliced view of the rays struct
        """
        return Rays(
            origins=self.origins[idx],
            dirs=self.dirs[idx],
            dist_min=self.dist_min[idx] if isinstance(self.dist_min, torch.Tensor) else self.dist_min,
            dist_max=self.dist_max[idx] if isinstance(self.dist_max, torch.Tensor) else self.dist_max
        )

    def split(self, split_size) -> List[Rays]:
        """ Splits the rays pack to equally sized batches.
        The last chunk may be smaller if the number of rays is not divisible by batch_size.

        Args:
            batch_size: Size of each batch returned by this split.

        Returns:
            (List[Rays]): A list of smaller Rays batches, split from the current rays pack.
        """
        zipped = zip(torch.split(self.origins, split_size), torch.split(self.dirs, split_size))
        return [Rays(origins=origins, dirs=dirs, dist_min=self.dist_min, dist_max=self.dist_max)
                for origins, dirs in zipped]

    def reshape(self, *dims: Tuple) -> Rays:
        """ Reshapes the dimensions of the rays struct.

        Args:
            *dims: Tuple of new dimensions after reshape.

        Returns:
            (Rays): The reshaped Rays struct
        """
        return Rays(origins=self.origins.reshape(*dims),
                   dirs=self.dirs.reshape(*dims),
                   dist_min=self.dist_min,
                   dist_max=self.dist_max)

    def squeeze(self, dim: int) -> Rays:
        """ Squeezes the tensors of the rays struct.

        Args:
            dims: Dimension to squeeze

        Returns:
            (Rays): The squeezed Rays struct
        """
        return Rays(origins=self.origins.squeeze(dim),
                   dirs=self.dirs.squeeze(dim),
                   dist_min=self.dist_min,
                   dist_max=self.dist_max)

    def contiguous(self) -> Rays:
        """ Forces the rays tensors to use contiguous memory.
        If they already do, no modification will take place on the tensors.

        Returns:
            (Rays): The contiguous Rays struct
        """
        return Rays(origins=self.origins.contiguous(),
                   dirs=self.dirs.contiguous(),
                   dist_min=self.dist_min,
                   dist_max=self.dist_max)

    def to(self, *args, **kwargs) -> Rays:
        """ Shifts the rays struct to a different device / dtype.
        seealso::func:`torch.Tensor.to` for an elaborate explanation of using this method.

        Returns:
            The rays tensors will be ast to a different device / dtype
        """

        origins = self.origins.to(*args, **kwargs)
        dirs = self.dirs.to(*args, **kwargs)
        if origins is not self.origins or dirs is not self.dirs:
            return Rays(
                origins=origins,
                dirs=dirs,
                dist_min=self.dist_min,
                dist_max=self.dist_max
            )
        else:
            return self

wisp/core/rays.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from dataclasses import fields, dataclass, make_dataclass
from typing import Optional, List, Tuple, Set, Dict, Iterator
from functools import partial
import numpy as np
import torch
from wisp.core.channels import Channel, create_default_channel


__RB_VARIANTS__ = dict()


@dataclass
class RenderBuffer:
    """
    A torch based, multi-channel, pixel buffer object.
    RenderBuffers are "smart" data buffers, used for accumulating tracing results, blending buffers of information,
    and providing discretized images.

    The spatial dimensions of RenderBuffer channels are flexible, e.g: they can be multi-dimensional or flat
    (thus allowing gradual accumulation of traced pixels).

    All RenderBuffer objects support the default rgb, alpha and depth channels.
    Additional custom channels can be specified as **kwargs during construction.
    Access to these channels is identical to the default fields of the RenderBuffer class.
    Customizing how existing and new channels are blended and normalized can be defined via :class:`Channel`
    """

    rgb    : Optional[torch.Tensor] = None
    """ rgb is a shaded RGB color. """

    alpha  : Optional[torch.Tensor] = None
    """ alpha is the alpha component of RGB-A. """

    depth  : Optional[torch.Tensor] = None
    """ depth is usually a distance to the surface hit point."""

    # Renderbuffer supports additional custom channels passed to the Renderbuffer constructor.
    # Some example of custom channels used throughout wisp:
    #     xyz=None,         # xyz is usually the xyz position for the surface hit point.
    #     hit=None,         # hit is usually a segmentation mask of hit points.
    #     normal=None,      # normal is usually the surface normal at the hit point.
    #     shadow =None,     # shadow is usually some additional buffer for shadowing.
    #     ao=None,          # ao is usually some addition buffer for ambient occlusion.
    #     ray_o=None,       # ray_o is usually the ray origin.
    #     ray_d=None,       # ray_d is usually the ray direction.
    #     err=None,         # err is usually some error metric against the ground truth.
    #     gts=None,         # gts is usually the ground truth image.

    def __new__(cls, *args, **kwargs):
        """ If custom channels were specified, create a specialized Renderbuffer class containing them as dataclass
        fields. The returned object is a variant of Renderbuffer whose existance should be transparent to users.

        Args:
            *args: Variable arg of torch.Tensor values for default channels.
            **kwargs: Optional dict of default + optional custom channels, as pairs of [str, Optional[torch.Tensor]].
        """
        # kwargs contains all channels given to the constructor,
        # filter to keep only new channels which aren't listed as default fields under the Renderbuffer class
        class_fields = [f.name for f in fields(RenderBuffer)]
        new_fields = [k for k in kwargs.keys() if k not in class_fields]
        if len(new_fields) > 0:
            # If there are new custom channels, see if a specialized class for this combination of channels
            # have been created before. If so, reuse it.
            class_key = frozenset(new_fields)
            rb_class = __RB_VARIANTS__.get(class_key)
            if rb_class is None:
                # First time this combination of channels is encountered:
                # create a Renderbuffer dataclass variant with new additional channels as fields
                rb_class = make_dataclass(f'Renderbuffer_{len(__RB_VARIANTS__)}',
                                          fields=[(k, Optional[torch.Tensor], None) for k in kwargs.keys()],
                                          bases=(RenderBuffer,))
                __RB_VARIANTS__[class_key] = rb_class   # Cache for future __new__ calls
            return super(RenderBuffer, rb_class).__new__(rb_class)  # Construct the new Renderbuffer variant
        else:
            return super(RenderBuffer, cls).__new__(cls)    # No new fields, just build the default Renderbuffer

    def __iter__(self) -> Iterator[str, Optional[torch.Tensor]]:
        """ Creates an iterator on the Renderbuffer fields as {name: tensor}. """
        # A tensor safe version:
        # the dataclasse asdict function performs a deepcopy which does not respect tensors with gradients.
        return iter((f.name, getattr(self, f.name)) for f in fields(self))

    def __getattr__(self, item):
        """ Invoked when an attribute is not found: used to return a default None for unknown channels. """
        if item.startswith('__') and item.endswith('__'):
            raise AttributeError
        else:
            return None  # Renderbuffer silently returns None when unknown channels are accessed

    @property
    def rgba(self) -> Optional[torch.Tensor]:
        """
        Returns:
            (Optional[torch.Tensor]) A concatenated rgba. If rgb or alpha are none, this property will return None.
        """
        if self.alpha is None or self.rgb is None:
            return None
        else:
            return torch.cat((self.rgb, self.alpha), dim=-1)

    @rgba.setter
    def rgba(self, val: Optional[torch.Tensor]) -> None:
        """
        Args:
            val (Optional[torch.Tensor]) A concatenated rgba channel value, which sets values for the rgb and alpha
            internal channels simultaneously.
        """
        self.rgb = val[..., 0:-1]
        self.alpha = val[..., -1:]

    @property
    def channels(self) -> Set[str]:
        """ Returns a set of channels supported by this RenderBuffer """
        return set([f.name for f in fields(self)])

    def has_channel(self, name: str) -> bool:
        """ Returns whether the RenderBuffer supports the specified channel """
        return name in self.channels

    def get_channel(self, name: str) -> Optional[torch.Tensor]:
        """ Returns the pixels value of the specified channel,
        assuming this RenderBuffer supports the specified channel.
        """
        return getattr(self, name)

    @staticmethod
    def _join_fields(rb1, rb2):
        """ Creates a joint mapping of renderbuffer fields in a format of
            {
                channel1_name: (rb1.c1, rb2.c1),
                channel2_name: (rb1.c2, rb2.cb),
                channel3_name: (rb1.c1, None),  # rb2 doesn't define channel3
            }
            If a renderbuffer does not have define a specific channel, None is returned.
        """
        joint_fields = rb1.channels.union(rb2.channels)
        return {f: (getattr(rb1, f, None), getattr(rb2, f, None)) for f in joint_fields}

    def _apply(self, fn) -> RenderBuffer:
        """ Applies the function fn on each of the Renderbuffer channels, if not None.
            Returns a new instance with the processed channels.
        """
        data = {}
        for f in fields(self):
            attr = getattr(self, f.name)
            data[f.name] = None if attr is None else fn(attr)
        return RenderBuffer(**data)

    @staticmethod
    def _apply_on_pair(rb1, rb2, fn) -> RenderBuffer:
        """ Applies the function fn on each of the Renderbuffer channels, if not None.
            Returns a new instance with the processed channels.
        """
        joint_fields = RenderBuffer._join_fields(rb1, rb2)  # Union of field names and tuples of values
        combined_channels = map(fn, joint_fields.values())  # Invoke on pair per Renderbuffer field
        return RenderBuffer(**dict(zip(joint_fields.keys(), combined_channels)))    # Pack combined fields to a new rb

    def __add__(self, other: RenderBuffer) -> RenderBuffer:
        """ Renderbuffer support spatial acucmulation of pixels (that is, the RenderBuffer is not constrained to
        a 2D grid).
        By default, __add__ performs concatenation of values per channel over the first spatial dimension.
        """
        return self.cat(other)

    def cat(self, other: RenderBuffer, dim: int = 0) -> RenderBuffer:
        """ Concatenates the channels of self and other RenderBuffers.
        If a channel only exists in one of the RBs, that channel will be returned as is.
        For channels that exists in both RBs, the spatial dimensions are assumed to be identical except for the
        concatenated dimension.

        Args:
            other (RenderBuffer) A second buffer to concatenate to the current buffer.
            dim (int): The index of spatial dimension used to concat the channels

        Returns:
            A new RenderBuffer with the concatenated channels.
        """
        def _cat(pair):
            if None not in pair:
                # Concatenating tensors of different dims where one is unsqueezed with dimensionality 1
                if pair[0].ndim == (pair[1].ndim + 1) and pair[0].shape[-1] == 1:
                    pair = (pair[0], pair[1].unsqueeze(-1))
                elif pair[1].ndim == (pair[0].ndim + 1) and pair[1].shape[-1] == 1:
                    pair = (pair[0].unsqueeze(-1), pair[1])
                return torch.cat(pair, dim=dim)
            elif pair[0] is not None and pair[1] is None:   # Channel is None for other but not self
                return pair[0]
            elif pair[0] is None and pair[1] is not None:   # Channel is None for self but not other
                return pair[1]
            else:
                return None

        return RenderBuffer._apply_on_pair(self, other, _cat)

    def blend(self, other: RenderBuffer, channel_kit: Dict[str, Channel]) -> RenderBuffer:
        """ Blends the channels of self and other RenderBuffers.
        If a channel only exists in one of the RBs, that channel will be returned as is.
        For channels that exists in both RBs, the spatial dimensions are assumed to be identical.

        The exact blending process is applied per channel as follows:
            Let c1 be a channel of "self", and c2 a channel of the same type in "other".
            1. If other.depth is in front of self.depth, swap c1 and c2.
            2. If alpha channel exists in both RBs, alpha blending is enabled (step 3), otherwise c1 is returned.
            3. Invoke the blending function of this channel type as defined in channel_kit,
            using c1, c2, self.alpha and other.alpha

        Args:
            other (RenderBuffer) A second buffer to concatenate to the current buffer.
            channel_kit (Dict[str, Channel]): Describes which blending function should be applied to each channel.

        Returns:
            (RenderBuffer) A new RenderBuffer with the blended channels,
            according to depth and alpha and the defined blend func.
        """
        assert self.depth is not None and other.depth is not None, "Cannot blend renderbuffers without depth values."
        # TODO (operel): In the future depth front / back may depend on the choice of NDC space
        #   (currently objects in the front --> lower depth)
        mask: torch.Tensor = self.depth <= other.depth
        blended = dict()
        joint_fields = RenderBuffer._join_fields(self, other)  # Union of field names and tuples of values
        alpha_pair = (self.alpha, other.alpha)
        is_alpha_blending = None not in alpha_pair
        for field in joint_fields:
            pair = getattr(self, field), getattr(other, field)

            if None not in pair:    # Actual blending occurs when both RBs have information about this channel
                # Alpha channel available, invoke blending function
                if is_alpha_blending:
                    # Sort c1, c2 by depth:
                    # c1 is the channel from the RB with lower depth
                    # alpha1 is the corresponding alpha channel used for blending
                    c1 = torch.where(mask, pair[0], pair[1])
                    alpha1 = torch.where(mask, alpha_pair[0], alpha_pair[1])
                    c2 = torch.where(mask, pair[1], pair[0])
                    alpha2 = torch.where(mask, alpha_pair[1], alpha_pair[0])

                    # The channel kit stores information about how the channel should be blended.
                    # If no information have been registered for this channel type, resort to default settings.
                    channel_info = channel_kit.get(field, create_default_channel())
                    blend = channel_info.blend_fn
                    out = blend(c1, c2, alpha1, alpha2)
                else:   # Alpha channel is n/a, resort to painters algorithm and choose front facing pixel
                    out = torch.where(mask, pair[0], pair[1])
            elif pair[0] is not None and pair[1] is None:
                out = pair[0]   # Channel is n/a in other
            elif pair[0] is None and pair[1] is not None:
                out = pair[1]   # Channel is n/a in self
            else:
                out = None      # Channel is n/a in both self and other
            blended[field] = out
        return RenderBuffer(**blended)

    def transpose(self) -> RenderBuffer:
        """ Permutes dimensions 0 and 1 of each channel.
            The rest of the channel dimensions will remain in the same order.
        """
        fn = lambda x : x.permute(1, 0, *tuple(range(2, x.ndim)))
        return self._apply(fn)

    def scale(self, size: Tuple, interpolation='bilinear') -> RenderBuffer:
        """ Upsamples or downsamples the renderbuffer pixels using the specified interpolation.
        Scaling assumes renderbuffers with 2 spatial dimensions, e.g. (H, W, C) or (W, H, C).

        Warning: for non-floating point channels, this function will upcast to floating point dtype
        to perform interpolation, and will then re-cast back to the original dtype.
        Hence truncations due to rounding may occur.

        Args:
            size (Tuple): The new spatial dimensions of the renderbuffer.
            interpolation (str): Interpolation method applied to cope with missing or decimated pixels due to
            up / downsampling. The interpolation methods supported are aligned with
            :func:`torch.nn.functional.interpolate`.

        Returns:
            (RenderBuffer): A new RenderBuffer object with rescaled channels.
        """
        def _scale(x):
            assert x.ndim == 3, 'RenderBuffer scale() assumes channels have 2D spatial dimensions.'
            # Some versions of torch don't support direct interpolation of non-fp tensors
            dtype = x.dtype
            if not torch.is_floating_point(x):
                x = x.float()
            x = x.permute(2, 0, 1)[None]
            x = torch.nn.functional.interpolate(x, size=size, mode=interpolation)
            x = x[0].permute(1, 2, 0)
            if x.dtype != dtype:
                x = torch.round(x).to(dtype)
            return x
        return self._apply(_scale)

    def numpy_dict(self) -> Dict[str, np.array]:
        """This function returns a dictionary of numpy arrays containing the data of each channel.

        Returns:
            (Dict[str, numpy.Array])
                a dictionary with entries of (channel_name, channel_data)
        """
        _dict = dict(iter(self))
        _dict = {k: v.numpy() for k, v in _dict.items() if v is not None}
        return _dict

    def exr_dict(self) -> Dict[str, torch.Tensor]:
        """This function returns an EXR format compatible dictionary.

        Returns:
            (Dict[str, torch.Tensor])
                a dictionary suitable for use with `pyexr` to output multi-channel EXR images which can be
                viewed interactively with software like `tev`.
                This is suitable for debugging geometric quantities like ray origins and ray directions.
        """
        _dict = self.numpy_dict()
        if 'rgb' in _dict:
            _dict['default'] = _dict['rgb']
            del _dict['rgb']
        return _dict

    def image(self) -> RenderBuffer:
        """This function will return a copy of the RenderBuffer which will contain 8-bit [0,255] images.

        This function is used to output a RenderBuffer suitable for saving as a 8-bit RGB image (e.g. with
        Pillow). Since this quantization operation permanently loses information, this is not an inplace
        operation and will return a copy of the RenderBuffer. Currently this function will only return
        the hit segmentation mask, normalized depth, RGB, and the surface normals.

        If users want custom behaviour, users can implement their own conversion function which takes a
        RenderBuffer as input.
        """
        norm = lambda arr : ((arr + 1.0) / 2.0) if arr is not None else None
        bwrgb = lambda arr : torch.cat([arr]*3, dim=-1) if arr is not None else None
        rgb8 = lambda arr : (arr * 255.0) if arr is not None else None

        channels = dict()
        if self.rgb is not None:
            channels['rgb'] = rgb8(self.rgb)
        if self.alpha is not None:
            channels['alpha'] = rgb8(self.alpha)
        if self.depth is not None:
            # If the relative depth is respect to some camera clipping plane, the depth should
            # be clipped in advance.
            relative_depth = self.depth / (torch.max(self.depth) + 1e-8)
            channels['depth'] = rgb8(bwrgb(relative_depth))

        # TODO (operel): Write rgba channel
        # TODO (operel): Handle channels in a more general way

        if hasattr(self, 'hit') and self.hit is not None:
            channels['hit'] = rgb8(bwrgb(self.hit))
        else:
            channels['hit'] = None
        if hasattr(self, 'normal') and self.normal is not None:
            channels['normal'] = rgb8(norm(self.normal))
        else:
            channels['normal'] = None

        return RenderBuffer(**channels)

    @staticmethod
    def mean(*rblst) -> RenderBuffer:
        """This function will take a list of RenderBuffers and return the mean of the RenderBuffers.
        None channels count as zero towards the average (unless all channels are None, in which case the mean will be
        None).
        This is useful to implement things like antialiasing through aggregating multiple samples per ray.

        Args:
            *rblst: Variable list of RenderBuffers.

        Returns:
            The mean Renderbuffer
        """
        def _sum(pair):
            if None not in pair:
                return pair[0] + pair[1]
            elif pair[0] is not None and pair[1] is None:   # Channel is None for other but not self
                return pair[0]
            elif pair[0] is None and pair[1] is not None:   # Channel is None for self but not other
                return pair[1]
            else:
                return None

        rb = RenderBuffer()
        n = len(rblst)
        for item in rblst:
            rb = RenderBuffer._apply_on_pair(rb, item, _sum)
        div = partial(torch.div, other=float(n))
        return rb._apply(div)

    def reshape(self, *dims : List[int]) -> RenderBuffer:
        """ Reshapes the channels of the renderbuffer to the given dims """
        fn = lambda x : x.reshape(*dims)
        return self._apply(fn)

    def to(self, *args, **kwargs) -> RenderBuffer:
        """ Shifts the renderbuffer to a new dtype / device """
        fn = lambda x : x.to(*args, **kwargs)
        return self._apply(fn)

    def cuda(self) -> RenderBuffer:
        """ Shifts the renderbuffer to the default torch cuda device """
        fn = lambda x : x.cuda()
        return self._apply(fn)

    def cpu(self) -> RenderBuffer:
        """ Shifts the renderbuffer to the torch cpu device """
        fn = lambda x : x.cpu()
        return self._apply(fn)

    def detach(self) -> RenderBuffer:
        """ Detaches the gradients of all channel tensors of the renderbuffer """
        fn = lambda x : x.detach()
        return self._apply(fn)

    def byte(self) -> RenderBuffer:
        """ Returns a new RenderBuffer using the byte dtype for each channel """
        fn = lambda x : x.byte()
        return self._apply(fn)

    def half(self) -> RenderBuffer:
        """ Returns a new RenderBuffer using the half dtype for each channel """
        fn = lambda x : x.half()
        return self._apply(fn)

    def float(self) -> RenderBuffer:
        """ Returns a new RenderBuffer using the float dtype for each channel """
        fn = lambda x : x.float()
        return self._apply(fn)

    def double(self) -> RenderBuffer:
        """ Returns a new RenderBuffer using the double dtype for each channel """
        fn = lambda x : x.double()
        return self._apply(fn)

wisp/core/render_buffer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Union

import numpy as np
import torch


PI = torch.pi if hasattr(torch, 'pi') else np.pi  # Fallback for older torch versions


class ObjectTransform:
    """ObjectTransform represents the local transformations that convert the object from local space coordinates
    to world coordinates.
    In wisp, objects are usually optimized in some normalized space.
    The attached object transform allows to locate and orient them within the world.
    """

    def __init__(self, device=None, dtype=None):
        self.device: torch.device = device
        self.dtype: torch.dtype = dtype

        self._translation = torch.zeros(3, device=self.device, dtype=self.dtype)
        self._rotation = torch.zeros(3, device=self.device, dtype=self.dtype)
        self._scale = torch.ones(4, device=self.device, dtype=self.dtype)
        self._permutation = torch.eye(4, device=self.device, dtype=self.dtype)

    def reset(self):
        """ Restore object transform to unit scale, at the origin, with zero orientation. """
        self._translation = torch.zeros(3, device=self.device, dtype=self.dtype)
        self._rotation = torch.zeros(3, device=self.device, dtype=self.dtype)
        self._scale = torch.ones(4, device=self.device, dtype=self.dtype)
        self._permutation = torch.eye(4, device=self.device, dtype=self.dtype)

    def translate(self, translation: torch.Tensor):
        """
        Args:
            translation (torch.Tensor):
                Amount of translation in world space coordinates, as a (3,) shaped tensor
        """
        self._translation += translation

    def rotate(self, rotation: torch.Tensor):
        """
        Args:
            rotation (torch.Tensor):
                Amount of rotation as euler angles, as a (3,) shaped tensor
        """
        self._rotation += rotation

    def scale(self, scale: Union[float, torch.Tensor]):
        """
        Args:
            scale (float or torch.Tensor):
                Amount of scale in world space coordinates.
                For non uniform scale, pass a (3,) shaped tensor.
                For uniform scale, pass a float.
        """
        self._scale[:3] *= scale

    def permute(self, permutation):
        """
        Args:
            permutation (list[int]): The permutation of the axis.
                                     For example, [1, 0, 2] will swap the x and y axis.
        """
        permutation = torch.tensor(permutation, device=self.device, dtype=torch.long)
        if permutation.shape[0] != 3:
            raise Exception("Permutation only supports 3 axis.")
        if torch.any(permutation < 0):
            raise Exception("Permutation axis cannot be negative.")
        if torch.any(permutation > 2):
            raise Exception("Permutation axis out of bounds.")
        self._permutation.fill_(0)
        self._permutation[(0,1,2), permutation] = 1
        self._permutation[3, 3] = 1

    def _translation_mat(self, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            t (torch.Tensor): The translation component of the transform
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the translation component of the transform
        """
        mat = torch.eye(4, device=self.device, dtype=self.dtype)
        mat[:3, -1] = t
        return mat

    def _rotation_mat_x(self, angle: torch.Tensor) -> torch.Tensor:
        """
        Args:
            angle (torch.Tensor): The x axis euler angle component of the transform, in radians
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the rotation x component of the transform
        """
        mat = torch.eye(4, device=self.device, dtype=self.dtype)
        mat[1, 1] = torch.cos(angle)
        mat[1, 2] = torch.sin(angle)
        mat[2, 1] = -torch.sin(angle)
        mat[2, 2] = torch.cos(angle)
        return mat

    def _rotation_mat_y(self, angle: torch.Tensor) -> torch.Tensor:
        """
        Args:
            angle (torch.Tensor): The y axis euler angle component of the transform, in radians
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the rotation y component of the transform
        """
        mat = torch.eye(4, device=self.device, dtype=self.dtype)
        mat[0, 0] = torch.cos(angle)
        mat[0, 2] = -torch.sin(angle)
        mat[2, 0] = torch.sin(angle)
        mat[2, 2] = torch.cos(angle)
        return mat

    def _rotation_mat_z(self, angle: torch.Tensor) -> torch.Tensor:
        """
        Args:
            angle (torch.Tensor): The z axis euler angle component of the transform, in radians
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the rotation z component of the transform
        """
        mat = torch.eye(4, device=self.device, dtype=self.dtype)
        mat[0, 0] = torch.cos(angle)
        mat[0, 1] = -torch.sin(angle)
        mat[1, 0] = torch.sin(angle)
        mat[1, 1] = torch.cos(angle)
        return mat

    def _rotation_mat(self, rx, ry, rz) -> torch.Tensor:
        """
        Args:
            rx (torch.Tensor): The x axis euler angle component of the transform, in radians
            ry (torch.Tensor): The y axis euler angle component of the transform, in radians
            rz (torch.Tensor): The z axis euler angle component of the transform, in radians
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the complete rotation component of the transform
        """
        Rx = self._rotation_mat_x(rx)
        Ry = self._rotation_mat_y(ry)
        Rz = self._rotation_mat_z(rz)
        return Rz @ Ry @ Rx

    def _scale_mat(self, s: torch.Tensor) -> torch.Tensor:
        """
        Args:
            s (torch.Tensor): The scale component of the transform
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the scale component of the transform
        """
        return torch.diag(s).to(device=self.device, dtype=self.dtype)

    def _inv_translation_mat(self, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            t (torch.Tensor): The translation component of the transform
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the inverse translation component of the transform
        """
        mat = torch.eye(4, device=self.device, dtype=self.dtype)
        mat[:3, -1] = -t
        return mat

    def _inv_rotation_mat(self, rx, ry, rz) -> torch.Tensor:
        """
        Args:
            rx (torch.Tensor): The x axis euler angle component of the transform, in radians
            ry (torch.Tensor): The y axis euler angle component of the transform, in radians
            rz (torch.Tensor): The z axis euler angle component of the transform, in radians
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the complete inverse rotation component of the transform
        """
        Rx = self._rotation_mat_x(-rx)
        Ry = self._rotation_mat_y(-ry)
        Rz = self._rotation_mat_z(-rz)
        return Rx @ Ry @ Rz

    def _inv_scale_mat(self, s: torch.Tensor) -> torch.Tensor:
        """
        Args:
            s (torch.Tensor): The scale component of the transform
        Returns:
            (torch.Tensor): a 4x4 matrix which represents the inverse scale component of the transform
        """
        return torch.diag(1.0 / s).to(device=self.device, dtype=self.dtype)

    def model_matrix(self):
        """
        Returns:
            (torch.Tensor): a 4x4 model matrix which encodes the complete transform of the object from local
            coordinates to world coordinates.
        """
        # Transformations applied in order of
        # Translation @ Rotation @ Scale
        scale_mat = self._scale_mat(self._scale)
        rotation_rads = self._rotation.div(180.0).mul(PI)
        rotation_mat = self._rotation_mat(*rotation_rads)
        translation_mat = self._translation_mat(self._translation)
        model_mat = translation_mat @ rotation_mat @ scale_mat @ self._permutation
        return model_mat

    def inv_model_matrix(self):
        """
        Returns:
            (torch.Tensor): a 4x4 model matrix which encodes the complete transform of the object from world
            coordinates to local coordinates.
            This version can be useful, for example, for ray traced pipelines (where the rays are inverse transformed,
            rather than the object).
        """
        # Transformations applied in order of
        # Scale^(-1) @ Rotation^(-1) @ Translation^(-1)
        scale_mat = self._inv_scale_mat(self._scale)
        rotation_rads = self._rotation.div(180.0).mul(PI)
        rotation_mat = self._inv_rotation_mat(*rotation_rads)
        translation_mat = self._inv_translation_mat(self._translation)
        inv_mat = self._permutation @ scale_mat @ rotation_mat @ translation_mat
        return inv_mat

    def to(self, *args, **kwargs) -> ObjectTransform:
        """ Shifts the transform to a different device / dtype.
        seealso::func:`torch.Tensor.to` for an elaborate explanation of using this method.

        Returns:
            The object transform tensors will be ast to a different device / dtype
        """
        _translation = self._translation.to(*args, **kwargs)
        _rotation = self._rotation.to(*args, **kwargs)
        _scale = self._scale.to(*args, **kwargs)
        _permutation = self._permutation.to(*args, **kwargs)
        if _translation is not self._translation or \
           _rotation is not self._rotation or \
           _scale is not self._scale or \
           _permutation is not self._permutation:
            transform = ObjectTransform(device=_translation.device, dtype=_translation.dtype)
            transform._translation = _translation
            transform._rotation = _rotation
            transform._scale = _scale
            transform._permutation = _permutation
            return transform
        else:
            return self

    @property
    def tx(self):
        """
        Returns:
            The x-axis translation component tensor.
        """
        return self._translation[0]

    @tx.setter
    def tx(self, value):
        """ Sets the x-axis translation component tensor.
        Args:
            value: New translation-x value to set
        """
        self._translation[0] = value

    @property
    def ty(self):
        """
        Returns:
            The y-axis translation component tensor, shaped as (tx, ty, tz).
        """
        return self._translation[1]

    @ty.setter
    def ty(self, value):
        """ Sets the y-axis translation component tensor.
        Args:
            value: New translation-y value to set
        """
        self._translation[1] = value

    @property
    def tz(self):
        """
        Returns:
            The z-axis translation component tensor, shaped as (tx, ty, tz).
        """
        return self._translation[2]

    @tz.setter
    def tz(self, value):
        """ Sets the z-axis translation component tensor.
        Args:
            value: New translation-z value to set
        """
        self._translation[2] = value

    @property
    def rx(self):
        """
        Returns:
            The euler angle of the rotation component around the x axis, in degrees
        """
        return self._rotation[0]

    @rx.setter
    def rx(self, value):
        """ Sets the euler angle of the rotation component around the x axis
        Args:
            value: New angle-x value to set, in degrees
        """
        self._rotation[0] = value

    @property
    def ry(self):
        """
        Returns:
            The euler angle of the rotation component around the y axis, in degrees
        """
        return self._rotation[1]

    @ry.setter
    def ry(self, value):
        """ Sets the euler angle of the rotation component around the y axis
        Args:
            value: New angle-y value to set, in degrees
        """
        self._rotation[1] = value

    @property
    def rz(self):
        """
        Returns:
            The euler angle of the rotation component around the z axis, in degrees
        """
        return self._rotation[2]

    @rz.setter
    def rz(self, value):
        """ Sets the euler angle of the rotation component around the z axis
        Args:
            value: New angle-z value to set, in degrees
        """
        self._rotation[2] = value

    @property
    def sx(self):
        """
        Returns:
            The scale component for the x axis.
        """
        return self._scale[0]

    @sx.setter
    def sx(self, value):
        """ Sets the x-axis scale component tensor.
        Args:
            value: New scale-x value to set
        """
        self._scale[0] = value

    @property
    def sy(self):
        """
        Returns:
            The scale component for the y axis.
        """
        return self._scale[1]

    @sy.setter
    def sy(self, value):
        """ Sets the y-axis scale component tensor.
        Args:
            value: New scale-y value to set
        """
        self._scale[1] = value

    @property
    def sz(self):
        """
        Returns:
            The scale component for the z axis.
        """
        return self._scale[2]

    @sz.setter
    def sz(self, value):
        """ Sets the z-axis scale component tensor.
        Args:
            value: New scale-z value to set
        """
        self._scale[2] = value

wisp/core/transforms.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from abc import ABC, abstractmethod
from typing import Dict, Any
import torch.nn as nn


class WispModule(nn.Module, ABC):
    """ A general interface for all Wisp building blocks, such as neural fields, grids and tracers.
        WispModules should:
        1. Provide their name & dictionary of public properties. That makes them compatible with systems like
        logging & gui.
        2. WispModules extend torch's nn.Module out of convenience.
        Modules are not required however, to implement a forward() function.
    """

    def __init__(self):
        super().__init__()

    def name(self) -> str:
        """
        Returns:
            (str) A WispModule should be given a meaningful, human readable name.
            By default, the class name is used.
        """
        return type(self).__name__

    @abstractmethod
    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        raise NotImplementedError('Wisp modules should implement the `public_properties` method')

wisp/core/wisp_module.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .base_datasets import *
from .utils import *
from .formats import *
from .transforms import *

wisp/datasets/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Callable, List, Dict, Union, Tuple
import torch
from torch.utils.data import Dataset
from kaolin.render.camera import Camera
from wisp.datasets.batch import Batch, MultiviewBatch, SDFBatch, ImageBatch


class WispDataset(Dataset, ABC):
    """ A general abstract class all wisp dataset classes should implement to integrate with the datasets framework.
    WispDatasets subclass directly from torch.utils.data.Dataset, and make little assumptions on the dataset structure.

    At the bare minimum, load_singleprocess(), __getitem__ and __len__ should be implemented.
    Additional methods may i.e. extend the dataset to allow it to loader faster,
    load more conveniently, or support better management of dataset splits.

    The following assumptions hold:
        - Datasets can be loaded from some path, but this isn't mandatory. In the case of in memory datasets,
        dataset_path may default to None.
        - Datasets should always support loading from the main process. In addition, datasets may support faster loading
         using multiple workers. When `load()` is invoked, `dataset_num_workers` will be queried to decide
         if `load_singleprocess()` or `load_multiprocess()` should be called.
         By default, if `load_multiprocess` is left unimplemented, `load_singleprocess` will be invoked irrespective
         of the number of workers.
        - Datasets should support transformation functors applied before __getitem__() calls are returned.
        These transformations may modify the content of the returned batch.
        - Where applicable, datasets should support loading specific splits of train, validation and test.
        The exact partition is left up to the dataset logic (i.e. by external files, by some specified ratio, etc).
        - Datasets which rely on external files should designate a set of rules to specify if the contents of some
        directory is loadable by the class.

    Pertaining to the last point, where no ambiguity hold over the content of some input dataset files,
    Wisp subclasses can be created with the convenience functions of wisp.utils.load_X_dataset()
    """

    def __init__(self, dataset_path: str = None, dataset_num_workers: int = -1,
                 transform: Callable = None, split: str = None):
        """
        Args:
            dataset_path (str): The root directory of the dataset, where the dataset's external assets should reside.
                Default: None (in memory dataset).
            dataset_num_workers (int): The number of workers to spawn for multiprocessed loading.
                If dataset_num_workers < 1, processing will take place on the main process.
                Default: -1 (dataset always loads on main process, without using multiprocessing).
            transform (Callable): Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
                Default: None (no transformation is applied on batches).
            split (str): The dataset split to use. Options: 'train', 'val', 'test'.
                Default: None (the dataset does not distinguish between different splits).
        """
        self.dataset_path = dataset_path
        self.dataset_num_workers = dataset_num_workers
        self.transform = transform
        self.split = split

    def create_split(self, split: str, transform: Callable = None, *args, **kwargs) -> WispDataset:
        """ Creates a dataset with the same parameters and a different split.
        This is a convenient way of creating validation and test datasets, while making sure they're compatible
        with the train dataset.

        All settings except for split and transform will be copied from the current dataset.

        Example:
            ```
            train_dataset = Dataset(...)
            validation_dataset = train_dataset.create_split(split="val")
            test_dataset = train_dataset.create_split(split="test")
            ```

        Args:
            split (str): The dataset split to use, corresponding to the transform file to load.
                Options: 'train', 'val', 'test'.
            transform (Optional[Callable]):
                Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        """
        raise NotImplementedError("This dataset does not support generation of new splits from a given dataset.")

    def name(self) -> str:
        """ Returns a human readable name of this dataset class. By default, the class name is returned. """
        return type(self).__name__

    def load(self):
        """ Loads the contents of a dataset object.
        This operation should usually take place once when the dataset is first created.
        This function may return some data values to be cached, or update the dataset fields in place.

        load() is usually intended to be invoked by the subclass __init__.
        load() will route between `load_multiprocess()` and `load_singleprocess()`, depending on the number of workers.
        """
        if self.dataset_num_workers > 0:
            return self.load_multiprocess()
        else:
            return self.load_singleprocess()

    @abstractmethod
    def load_singleprocess(self):
        """Loads the contents of a dataset object on the main process.
        This operation is invoked by load() when the number of workers is 0 or negative.
        This function may return some data values to be cached, or update the dataset fields in place.
        """
        raise NotImplementedError('Wisp Datasets should override load_singleprocess()')

    def load_multiprocess(self):
        """Loads the contents of a dataset object using multiple workers.
        This operation is invoked by load() when the number of workers is 1 or higher.
        This function may return some data values to be cached, or update the dataset fields in place.
        """
        return self.load_singleprocess()  # By default, datasets are not obliged to support multiprocessed loading

    @classmethod
    def is_root_of_dataset(cls, root: str, files_list: List[str]) -> bool:
        """ Each dataset may implement a simple set of rules to distinguish it from other datasets.
        Rules should be unique for this dataset type, such that given a general root path, Wisp will know
        to associate it with this dataset class.

        Datasets which don't implement this function should be created explicitly.

        Args:
                root (str): A path to the root directory of the dataset.
                files_list (List[str]): List of files within the dataset root, without their prefix path.
        Returns:
                True if the root folder points to content loadable by this dataset.
        """
        return False

    def __len__(self):
        """ Returns the number of entries in the dataset. """
        raise NotImplementedError('Wisp Datasets should override __len__')

    def __getitem__(self, idx: int) -> Batch:
        """ Samples the dataset for a batch of information. The exact batch returned depends on idx. """
        raise NotImplementedError('Wisp Datasets should override __getitem__')


class MultiviewDataset(WispDataset):
    """ Extends the WispDataset with dataset behavior common to all multiview datasets. """

    @property
    @abstractmethod
    def img_shape(self) -> Union[torch.Size, Tuple[torch.Size]]:
        """ Returns the shape of ground truth images of each view used by this dataset.
        If all images have a common size, the return value is torch.Size.
        If the gt images have different sizes, a Tuple of sizes is returned per view.
        """
        pass

    @property
    @abstractmethod
    def num_images(self) -> int:
        """ Returns the number of images / views this dataset contain. """
        pass

    @property
    def cameras(self) -> Dict[str, Camera]:
        """ Returns the set of cameras this dataset uses to generate rays.
        Cameras are identifiable by unique ids / names
        """
        return dict()

    def as_pointcloud(self) -> torch.FloatTensor:
        """ If `supports_depth()` is True, the current dataset contains depth information.
        This function can be used to query the depth information in the form of a pointcloud tensor.
        """
        raise NotImplementedError('MultiviewDatasets that support depth information should '
                                  'return their data as a pointcloud. Otherwise, set supports_depth to return False.')

    def supports_depth(self) -> bool:
        """ Returns if this dataset have loaded depth information. """
        return False

    def __getitem__(self, idx: int) -> MultiviewBatch:
        """ Samples the dataset for a multiview batch of information.
        The exact ray channels the batch contains are up to the dataset implementation to determine.
        Callers may treat the batch as a dictionary, or query MultiviewBatch.ray_values() to view which
        supervision channels are available.
        """
        raise NotImplementedError('MultiviewDatasets should override __getitem__')

    def __len__(self):
        """Length of the dataset as number of views. """
        return self.num_images


class SDFDataset(WispDataset):
    """ Extends the WispDataset with dataset behavior common to all signed distance function based datasets. """

    @property
    @abstractmethod
    def coordinates(self) -> torch.Tensor:
        """ Returns a tensor of sample coordinates this dataset currently holds. """
        raise NotImplementedError('SDFDatasets should return a (N, d) tensor of sample coordinates.')

    def resample(self) -> None:
        """ Resamples to generate a new working set of coordinates values and their supervision values.
        When implemented, this function allows the dataset to refresh the available samples with a new set.
        Resampling is expected to be implemented as an in-place operation.
        """
        pass

    def __getitem__(self, idx: int) -> SDFBatch:
        """ Samples the dataset for a coordinate+sdf batch of information.
        The exact coordinate supervision channels the batch contains are up to the dataset implementation to determine.
        Callers may treat the batch as a dictionary, or query SDFBatch.coord_values() to view which
        supervision channels are available.
        """
        raise NotImplementedError('SDFDataset should override __getitem__')

    def __len__(self):
        """Length of the dataset as number of coordinate samples. """
        raise NotImplementedError('SDFDataset should override __len__')
    



class BaseImageDataset(WispDataset):
    """ Extends the WispDataset with dataset behavior common to all image based datasets. """

    @property
    @abstractmethod
    def coordinates(self) -> torch.Tensor:
        """ Returns a tensor of image coordinates. """
        raise NotImplementedError('ImageDatasets should return a (N, 2) tensor of sample coordinates.')

    def __getitem__(self, idx: int) -> ImageBatch:
        """ Samples the dataset for a coordinate+rgb batch of information.
        The exact coordinate supervision channels the batch contains are up to the dataset implementation to determine.
        Callers may treat the batch as a dictionary, or query ImageBatch.coord_values() to view which
        supervision channels are available.
        """
        raise NotImplementedError('ImageDataset should override __getitem__')

    def __len__(self):
        """Length of the dataset as number of coordinate samples. """
        raise NotImplementedError('ImageDataset should override __len__')
    

wisp/datasets/base_datasets.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from attrdict import AttrDict
from typing import Optional, List, Dict, Any
import torch
from kaolin.render.camera import Camera
from wisp.core import Rays


class Batch(AttrDict):
    """ Represents a single batch of information sampled and collated from a WispDataset.
    Batches in Wisp keep a general structure by subclassing python's dictionaries and using their semantics.
    The exact fields each batch contain depend on the dataset type.
    """
    def __init__(self, *args, **kwargs):
        super().__init__()
        for k, v in dict(*args, **kwargs).items():
            self[k] = v

    @property
    def fields(self) -> List[str]:
        """ Returns a list of all field names the batch currently contains.
        Note that the content of some field values may be None in practice.
        """
        return list(self.keys())


class MultiviewBatch(Batch):
    """ Represents a single batch of multiview information sampled and collated from a WispDataset.
    MultiviewBatch objects are always assumed to carry ray information, and may optionally return the camera information
    used to generate them.
    Batches in Wisp keep a general structure by subclassing python's dictionaries and using their semantics,
    so the exact channels of information each ray carries is flexible and are left up to datasets to initialize.
    Callers should be mindful of which fields are available by calling the `fields()` function.
    Out of convenience, the common rgb channels is explicitly listed here (though are left as optional).
    """

    def __init__(self,
                 rays: Rays,
                 cameras: Optional[List[Camera]] = None,
                 rgb: Optional[torch.Tensor] = None,    # RGB is default for convenience, but not mandatory
                 *args, **kwargs):
        """
        Creates a new MultiviewBatch of ray samples and optional supervision signals they carry.

        Args:
            rays (wisp.core.Rays): a pack of ray origins and directions, usually generated from the dataset cameras.
            cameras (kaolin.render.camera.Camera): Optional.
                The list of Camera objects used to generate the rays in this batch.
            rgb (torch.Tensor): Optional.
                A torch.Tensor of rgb color which corresponds the the gt image pixel rgb each ray intersects.

                "rays" - a wisp.core.Rays pack of ray origins and directions, pre-generated from the dataset camera.
                "rgb" - a torch.Tensor of rgb color which corresponds the the gt image pixel rgb
                    each ray intersects.
                *args, **kwargs - may specify any additional channels of information a ray or view can carry.
        """
        super().__init__(rays=rays, cameras=cameras, rgb=rgb)

    def ray_values(self) -> Dict[str, Any]:
        """ Specifies a dictionary of the ray specific supervision channels this MultiviewBatch carries. """
        out = dict()
        if self['rgb'] is not None:
            out['rgb'] = self['rgb']
        return out


class ImageBatch(Batch):
    """ Represents a single batch of coordinates + pixel information sampled and collated from a WispDataset.
    ImageBatch objects are always assumed to carry the coordinate information and corresponding pixel values.
    Batches in Wisp keep a general structure by subclassing python's dictionaries and using their semantics,
    so the exact channels of information each coordinate carries is flexible and are left up to datasets to initialize.
    Callers should be mindful of which fields are available by calling the `fields()` function.
    """

    def __init__(self,
                 coords: torch.Tensor,
                 rgb: Optional[torch.Tensor],
                 *args, **kwargs):
        """
        Creates a new ImageBatchFull of coordinate samples and supervision signals they carry.

        Args:
            coords (torch.Tensor): a torch Tensor of coordinate samples of shape (B, 2) where B is the number of
                samples in this batch.
            rgb (torch.Tensor):.
                A torch.Tensor of rgb color which corresponds the interpolated texture value of the coordinate.
            *args, **kwargs - may specify any additional channels of information a coordinate or batch can carry.
        """
        super().__init__(coords=coords, rgb=rgb)

    def coord_values(self) -> Dict[str, Any]:
        """ Specifies a dictionary of the coordinate specific supervision channels this ImageBatch carries. """
        out = dict(rgb=self['rgb'])
        return out
    

class SDFBatch(Batch):
    """ Represents a single batch of coordinates + sdf information sampled and collated from a WispDataset.
    SDFBatch objects are always assumed to carry the coordinate information and signed distance from some sampled
    surface, and may optionally return additional channels of information as well such as sampled materials of
    the surface.
    Batches in Wisp keep a general structure by subclassing python's dictionaries and using their semantics,
    so the exact channels of information each coordinate carries is flexible and are left up to datasets to initialize.
    Callers should be mindful of which fields are available by calling the `fields()` function.
    Out of convenience, the common rgb and normals channels are explicitly listed here (though are left as optional).
    """

    def __init__(self,
                 coords: torch.Tensor,
                 sdf: torch.Tensor,
                 rgb: Optional[torch.Tensor] = None,
                 normals: Optional[torch.Tensor] = None,
                 *args, **kwargs):
        """
        Creates a new SDFBatch of coordinate samples and optional supervision signals they carry.

        Args:
            coords (torch.Tensor): a torch Tensor of coordinate samples of shape (B, d) where B is the number of
                samples in this batch and d is their dimensionality (usually d is the number of spatiotemporal dims).
            sdf (torch.Tensor): A torch.Tensor of the signed distance of each sample coordinate from the surface it
                was generated against. The shape is (B, 1).
            rgb (torch.Tensor): Optional.
                A torch.Tensor of rgb color which corresponds the interpolated texture value of the coordinate.
            normals (torch.Tensor): Optional.
                A torch.Tensor of normal value which corresponds the estimated normal value of the coordinate against
                the nearest point on the surface.
                *args, **kwargs - may specify any additional channels of information a coordinate or batch can carry.
        """
        super().__init__(coords=coords, sdf=sdf, rgb=rgb, normals=normals)

    def coord_values(self) -> Dict[str, Any]:
        """ Specifies a dictionary of the coordinate specific supervision channels this SDFBatch carries. """
        out = dict(sdf=self['sdf'])
        if self['rgb'] is not None:
            out['rgb'] = self['rgb']
        if self['normals'] is not None:
            out['normals'] = self['normals']
        return out
    


wisp/datasets/batch.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .rtmv_dataset import *
from .nerf_standard_dataset import *
from .mesh_sdf_dataset import *
from .octree_sdf_dataset import *
from .random_view_dataset import *
from .multi_image_dataset import *

wisp/datasets/formats/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import os
from typing import Callable, List
import logging as log
import torch
from wisp.datasets.base_datasets import SDFDataset
from wisp.datasets.batch import SDFBatch
import wisp.ops.mesh as mesh_ops


_SUPPORTED_FORMATS = ['obj']
""" Supported mesh formats this dataset can load """


class MeshSampledSDFDataset(SDFDataset):
    """ A dataset class for signed-distance functions sampled over a single mesh shape's faces.
    The dataset supports online resampling during training. However, sampling time can take up to tens of seconds.
    Samples can also optionally include normals and rgb information when the mesh is textured with materials.
    """

    def __init__(self,
                 dataset_path: str,
                 split: str,
                 transform: Callable = None,
                 sample_mode: list = None,
                 num_samples: int = 100000,
                 get_normals: bool = False,
                 sample_tex: bool = False,
                 mode_norm: str = 'sphere'
                 ):
        """Construct by loading the mesh and generating an initial dataset of coords + sdf samples.

        Args:
            dataset_path (str): Path to file with mesh to initialize the dataset.
                Supported formats: .obj only.
            split (str): Any of 'test', 'train', 'val'.
                Currently used for keeping track of the dataset purpose and not used internally.
            transform (Optional[Callable]):
                A transform function applied per batch when data is accessed with __get_item__.
            sample_mode (list of str): List of different sample methods to apply over the mesh.
                Any sequential combo of:
                    'tracer' - samples generated on the mesh surface
                    'near' - samples generated near the mesh surface
                    'rand' - samples generated unifromly in space
                    See `mesh_ops.point_sample` for additional implementation details for the sampler.
                    Default: ['rand', 'rand', 'near', 'near', 'trace']
            num_samples (int): Number of data points to keep in the working set at any given time.
                A new working set can be regenerated with resample().
            get_normals (bool): If True, __get_item__ will also return normals (estimated by the SDF field).
            sample_tex (bool): If True, __get_item__  will also return RGB values from the nearest texture.
            mode_norm (str): The mode at which the mesh will be normalized in [-1, 1] space:
                'sphere' - the mesh will be normalized within the unit sphere.
                'aabb' - the mesh will be normalized within an axis aligned bounding box extending [-1, 1].
                'planar' - the mesh will be normalized using a non-uniform scale applied to the x,z axes separately.
                'none' - no normalization will take place.
        """
        super().__init__(dataset_path=dataset_path, transform=transform, split=split)

        # Sampling args
        self.sample_mode = sample_mode if sample_mode is not None else ['rand', 'rand', 'near', 'near', 'trace']
        self.num_samples = num_samples
        self.get_normals = get_normals
        self.sample_tex = sample_tex
        self.mode_norm = mode_norm

        # Mesh info
        self.verts = self.faces = self.texv = self.texf = self.mats = None

        # Validate & load mesh, then sample from it
        self.validate(dataset_path)
        self.data = None
        self.load()

    def validate(self, dataset_path) -> None:
        """ Validates the dataset path represents a valid mesh.
        If the mesh is valid, this function returns gracefully, otherwise it asserts with an informative error.

        Args:
            dataset_path (str): Path to a mesh file to initialize the dataset.
                Supported formats: .obj only.
        """
        if not os.path.exists(dataset_path):
            raise FileNotFoundError(f"MeshSampledSDFDataset requires a mesh path, "
                                    f"the dataset_path does not exist: {self.dataset_path}")
        if not any([dataset_path.endswith(ext) for ext in _SUPPORTED_FORMATS]):
            raise FileNotFoundError(f"MeshSampledSDFDataset does not support the mesh format of {self.dataset_path}. "
                                    f"Please use any of the supported formats: {_SUPPORTED_FORMATS}")

    def _load_mesh(self, dataset_path, mode_norm='sphere', sample_tex=False):
        """ Loads obj file, optionally with texture & materials information.

        Args:
            dataset_path (str): Path to a mesh file to initialize the dataset.
                Supported formats: .obj only.
            sample_tex (bool): If True, will load texture coordinates and materials.
            mode_norm (str): The mode at which the mesh will be normalized in [-1, 1] space:
                'sphere' - the mesh will be normalized within the unit sphere.
                'aabb' - the mesh will be normalized within an axis aligned bounding box extending [-1, 1].
                'planar' - the mesh will be normalized using a non-uniform scale applied to the x,z axes separately.
                'none' - no normalization will take place.
        Returns:
            (torch.FloatTensor) mesh: vertex coordinates of triangular faces, of shape [F, 3]
            (torch.FloatTensor) vertices: mesh vertices of shape [V, 3]
            (torch.LongTensor) faces: mesh face indices of shape [F, 3]
            (torch.LongTensor or None) texv
            (torch.LongTensor or None) texf
            (torch.FloatTensor or None) materials
        """
        texv = texf = materials = None
        if sample_tex:
            out = mesh_ops.load_obj(dataset_path, load_materials=True)
            verts, faces, texv, texf, materials = out
        else:
            verts, faces, = mesh_ops.load_obj(dataset_path)

        verts, faces, = mesh_ops.normalize(verts, faces, mode_norm)
        mesh = verts[faces]
        return mesh, verts, faces, texv, texf, materials

    def load_singleprocess(self) -> None:
        """Initializes the dataset by loading a mesh and sampling SDFs from it.
        This function uses the main process to load the dataset, without spawning any workers.
        """
        mesh_fields = self._load_mesh(self.dataset_path, self.mode_norm, self.sample_tex)
        _, self.verts, self.faces, self.texv, self.texf, self.mats = mesh_fields
        self.resample()

    @classmethod
    def is_root_of_dataset(cls, root: str, files_list: List[str]) -> bool:
        """ Each dataset may implement a simple set of rules to distinguish it from other datasets.
        Rules should be unique for this dataset type, such that given a general root path, Wisp will know
        to associate it with this dataset class.

        Datasets which don't implement this function should be created explicitly.

        Args:
                root (str): A path to the root directory of the dataset.
                files_list (List[str]): List of files within the dataset root, without their prefix path.
        Returns:
                True if the root folder points to content loadable by this dataset.
        """
        # Sampled meshes are currently always assumed to be given as mesh files, check for any of the compatible formats
        return any([root.endswith(ext) for ext in _SUPPORTED_FORMATS])

    def __len__(self):
        """Return length of dataset, as number of samples currently resampled."""
        return self.data["coords"].size()[0]

    def __getitem__(self, idx) -> SDFBatch:
        """Retrieve a batch of sample coordinates and their sdf values.
        If the dataset was loaded with rgb and normals information, these channels will be returned as well.

        Returns:
            (SDFBatch): A batch of coordinates and their SDF values. The fields can be accessed as a dictionary:
                "coords" - a torch.Tensor of the 3d coordinates of each sample
                "sdf" - a torch.Tensor of the signed distance of each sample from the mesh surface.
                "rgb" - optional, a torch.Tensor of the "textured color" the sample interpolates from the nearest point
                    on the mesh surface.
                    If the dataset with built with `sample_tex`=False, this entry is None.
                "normals" - optional,  a torch.Tensor of the normal to the mesh surface.
                    If the dataset with built with `get_normals`=False, this entry is None.
        """
        out = SDFBatch(
            coords=self.data["coords"][idx],
            sdf=self.data["sdf"][idx],
            rgb=self.data["rgb"][idx] if "rgb" in self.data else None,
            normals=self.data["normals"][idx] if "normals" in self.data else None
        )

        if self.transform is not None:
            out = self.transform(out)

        return out

    def resample(self) -> None:
        """Resamples a new working set of SDFs, by computing an entirely new set of SDF samples.
        resample() updates the self.data field in place.
        !! Warning: this operation can take up to a minute !!
        """
        log.info(f"Resampling mesh for new sdf samples...")
        self.data = dict()
        rgb = nrm = None
        if self.get_normals:
            pts, nrm = mesh_ops.sample_surface(self.verts, self.faces, self.num_samples * len(self.sample_mode))
        else:
            pts = mesh_ops.point_sample(self.verts, self.faces, self.sample_mode, self.num_samples)

        if self.sample_tex:
            rgb, _, d = mesh_ops.closest_tex(self.verts.cuda(), self.faces.cuda(), self.texv, self.texf, self.mats, pts)
        else:
            d = mesh_ops.compute_sdf(self.verts.cuda(), self.faces.cuda(), pts.cuda())

        self.data['coords'] = pts.cpu()
        self.data['sdf'] = d.cpu()
        if rgb is not None:
            self.data['rgb'] = rgb.cpu()
        if nrm is not None:
            self.data['normals'] = nrm.cpu()

    @property
    def coordinates(self) -> torch.Tensor:
        """ Returns the coordinates of samples stored in this dataset. """
        return self.data["coords"]

wisp/datasets/formats/mesh_sdf_dataset.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import os
from typing import Callable, List
import logging as log
import torch
from math import prod
from PIL import Image
from wisp.datasets.base_datasets import BaseImageDataset
from wisp.datasets.batch import ImageBatch
import wisp.ops.mesh as mesh_ops
from wisp.ops.image import load_rgb, load_rgb_tensor
from wisp.ops.raygen import generate_2d_grid

_SUPPORTED_FORMATS = (".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif", ".tiff", ".webp",
                      ".JPG", ".JPEG")
""" Supported image formats this dataset can load """
    
class MultiImageDataset(BaseImageDataset):
    """ 
    A dataset class for images and coordinate loading. The dataset loads images in the directory 
    with the supported formats.
    """

    def __init__(self,
                 dataset_path: str,
                 dataset_num_workers: int,
                 split: str,
                 transform: Callable = None,
                 num_samples: int = -1,
                 sample_mode: str='full',
                 ):
        """Construct by loading the image and generating an initial dataset of coords.

        Args:
            dataset_path (str): Path to images.
            dataset_num_workers (int): The number of workers to spawn for multiprocessed loading.
                Included for compatibility
            split (str): Any of 'train', 'val'.
                Currently used for keeping track of the dataset purpose and not used internally.
            transform (Optional[Callable]):
                A transform function applied per batch when data is accessed with __get_item__.
                Not used for image coordinate batch loading.
            num_samples (int): Number of samples/coords per batch.
                -1 for all coords in the image.
            sample_mode (str): Different sample methods to apply for each batch of coords:
            pregenerated coordinate grid methods: (Full coordinate grid stored in memory for sampling)
                'full' - all coords will be sampled from the coordinate grid.
                'woreplace' - num_samples coords will be sampled randomly from the coordinate grid
                            without replacement.
                'sequential' - num_samples coords will be sampled sequentially from the coordinate grid
            on-the-fly coordinate grid methods:
                'wreplace' - num_samples coords will be sampled randomly from the coordinate grid
                            with replacement, more memory efficient for large images.
                'eval' - num_samples coords sampled in sequential order
        """
        super().__init__(dataset_path=dataset_path, 
                         dataset_num_workers=dataset_num_workers,
                         split=split, 
                         transform=transform)

        self.dataset_path = dataset_path

        self.image_list = []
        for imagename in sorted(os.listdir(self.dataset_path)):
            if any([imagename.endswith(ext) for ext in _SUPPORTED_FORMATS]):
                self.image_list.append(os.path.join(self.dataset_path,imagename))

        self.num_images = len(self.image_list)
        # Sampling args
        self.sample_mode = sample_mode
        # If sampling all coords at once, set num_samples to -1
        if self.sample_mode == 'full': 
            num_samples = -1
        self.num_samples = num_samples

        # Maintain index of image in folder
        self.image_idx = 0

    def create_split(self, split: str):
        """ Creates a dataset with the same parameters and a different split.
        This is a convenient way of creating validation datasets, while making sure they're compatible
        with the train dataset.

        All settings except for split and transform will be copied from the current dataset.

        Args:
            split (str): The dataset split to use, corresponding to the transform file to load.
                Options: 'train', 'val'
            transform (Optional[Callable]):
                Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        """
        return MultiImageDataset(
            dataset_path=self.dataset_path,
            dataset_num_workers=self.dataset_num_workers,
            split=split,
            transform=self.transform,
            num_samples=self.num_samples,
            sample_mode='eval'
        )
    
    def load_next(self):
        """Load next image in multi image dataset"""
        self.data = dict()
        self.image_path = self.image_list[self.image_idx]

        # Calls load_singleprocess or load_multiprocess
        self.load()

        # Controls behavior of trainer if coordinates are same across iterations
        # For faster training, no need to iterate over dataloader if data is static
        self.static_coords = ((self.num_samples == -1) or 
                              (self.num_samples >= prod(self.image_size))) 
        # Increment image index
        self.image_idx += 1
        
    def resample(self) -> None:
        """Resample coordinate grid if using woreplace sampling mode. No resampling
        necessary for other modes.
        """
        if self.sample_mode == 'woreplace':
            self.shuffle_idx = torch.randperm(self.data['orig_coords'].shape[0])
            self.data['coords'] = self.data['orig_coords'][self.shuffle_idx]
            self.data['rgb'] = self.data['orig_rgb'][self.shuffle_idx]

    def load_singleprocess(self) -> None:
        """Initializes the dataset by loading an image and sampling coordinates from it.
        This function uses the main process to load the dataset, without spawning any workers.
        """

        image = load_rgb_tensor(self.image_path)
        self.data['orig_rgb'] = image.reshape(3,-1).permute(1,0)
        C, H, W = image.shape
        self.image_size = (H,W)
        # Create the coordinate grid if not using on-the-fly sampling methods such as 'eval' and 'wreplace'
        # Will be slow and take a lot of memory for large images
        if (self.sample_mode!='eval' and self.sample_mode != 'wreplace') or \
            (self.num_samples == -1 or self.num_samples > prod(self.image_size)):
            grid_y, grid_x = generate_2d_grid(W,H)
            # Normalize grid to [-1,1]
            # grid_y, grid_x = ((grid_y+0.5)/H-0.5)*2, ((grid_x+0.5)/W-0.5)*2 
            grid_y, grid_x = ((grid_y)/H-0.5)*2, ((grid_x)/W-0.5)*2 
            # grid_y, grid_x = (grid_y/H, grid_x/W)
            self.data['orig_coords'] = torch.stack([grid_y.reshape(-1,1), grid_x.reshape(-1,1)], dim=-1).reshape(-1,2)
            # If not sequential, shuffle the coordinates and corresponding rgb values
            if self.sample_mode != 'sequential':
                self.shuffle_idx = torch.randperm(self.data['orig_coords'].shape[0])
                self.data['coords'] = self.data['orig_coords'][self.shuffle_idx]
                self.data['rgb'] = self.data['orig_rgb'][self.shuffle_idx]
            else:
                self.data['coords'] = self.data['orig_coords']
                self.data['rgb'] = self.data['orig_rgb']

    @classmethod
    def is_root_of_dataset(cls, root: str, files_list: List[str]) -> bool:
        """ Each dataset may implement a simple set of rules to distinguish it from other datasets.
        Rules should be unique for this dataset type, such that given a general root path, Wisp will know
        to associate it with this dataset class.

        Datasets which don't implement this function should be created explicitly.

        Args:
                root (str): A path to the root directory of the dataset.
                files_list (List[str]): List of files within the dataset root, without their prefix path.
        Returns:
                True if the root dir contains atleast one image with supported format
        """
        return any([any([filename.endswith(ext) for ext in _SUPPORTED_FORMATS]) \
                    for filename in files_list])

    def __len__(self):
        """Return length of dataset, as number of batches of samples for current image.
           NOT number of images in dataset."""
        if self.num_samples == -1 or self.num_samples > prod(self.image_size):
            num_batches = 1
        else:
            num_batches = prod(self.image_size) // self.num_samples
            num_batches += 1 if prod(self.image_size) % self.num_samples > 0 else 0
        return num_batches

    def __getitem__(self, idx) -> ImageBatch:
        """Retrieve a batch of sample coordinates and their rgb values.

        Returns:
            (ImageBatch): A batch of coordinates and their RGB values. The fields can be accessed as a dictionary:
                "coords" - a torch.Tensor of the 2d coordinates of each sample
                "rgb" - a torch.Tensor of the RGB pixel values at each coordinate.
        """
        if self.num_samples == -1 or self.num_samples > prod(self.image_size):
            # Single batch consisting of all coordinates of the image
            out = ImageBatch(
                coords=self.data["coords"],
                rgb=self.data["rgb"]
            )
        elif self.sample_mode == 'woreplace' or self.sample_mode == 'sequential':
            # Sample current batch of coordinates from pregenerated coordinate grid
            start_idx = idx * self.num_samples
            end_idx = min(start_idx + self.num_samples, prod(self.image_size))
            out = ImageBatch(
                coords=self.data["coords"][start_idx:end_idx],
                rgb=self.data["rgb"][start_idx:end_idx]
            )
        else:
            if self.sample_mode == 'eval':
                start_idx = idx * self.num_samples
                end_idx = min(start_idx + self.num_samples, prod(self.image_size))
                idx = torch.arange(start_idx, end_idx, dtype=torch.long).reshape(-1,1)
            elif self.sample_mode == 'wreplace':
                idx = torch.randint(0, prod(self.image_size), size=(self.num_samples,1))
            # H,W = self.image_size
            # coord_idx = torch.div(idx, torch.Tensor([W,1]).reshape(1,-1), rounding_mode='floor').long()
            # coord_idx = torch.remainder(coord_idx, torch.Tensor([H,W]).reshape(1,-1).to(coord_idx)).long()
            # # coord_idx = ((coord_idx+0.5)/torch.Tensor([H,W]).reshape(1,-1)-0.5)*2
            # coord_idx = ((coord_idx)/torch.Tensor([H,W]).reshape(1,-1).to(coord_idx)-0.5)*2
            # # coord_idx = coord_idx/torch.Tensor([H,W]).reshape(1,-1)

            out = ImageBatch(
                coords=idx,
                rgb=self.data["orig_rgb"][idx[:,0]]
            )
        return out
    
    def transform_coords(self, coord_idx, device):
        """Transforms the coordinates to the range [-1,1]"""
        if self.num_samples == -1 or self.num_samples > prod(self.image_size) or \
           self.sample_mode == 'woreplace' or self.sample_mode == 'sequential':
            return coord_idx
        else:
            H,W = self.image_size
            coord_idx = torch.div(coord_idx.to(device), torch.Tensor([W,1]).reshape(1,-1).to(device), rounding_mode='floor').long()
            coord_idx = torch.remainder(coord_idx, torch.Tensor([H,W]).reshape(1,-1).to(coord_idx)).long()
            coord_idx = ((coord_idx)/torch.Tensor([H,W]).reshape(1,-1).to(coord_idx)-0.5)*2
            return coord_idx

    
    @property
    def coordinates(self) -> torch.Tensor:
        """ Returns the coordinates of samples stored in this dataset. """
        return self.data.get("coords")
    
    @property
    def pixels(self) -> torch.Tensor:
        """ Returns the coordinates of samples stored in this dataset. """
        return self.data.get("rgb")
wisp/datasets/formats/multi_image_dataset.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import os
import re
import glob
import cv2
import json
from tqdm import tqdm
import logging as log
from typing import Callable, List, Dict, Union
import numpy as np
import torch
from torch.multiprocessing import Pool
from kaolin.render.camera import Camera, blender_coords
from wisp.core import Rays
from wisp.ops.raygen import generate_pinhole_rays, generate_ortho_rays, generate_centered_pixel_coords
from wisp.ops.image import resize_mip, load_rgb
from wisp.datasets.base_datasets import MultiviewDataset
from wisp.datasets.batch import MultiviewBatch


class NeRFSyntheticDataset(MultiviewDataset):
    """ A dataset for files in the standard NeRF format, including extensions to the format
        supported by Instant Neural Graphics Primitives.
        See: https://github.com/NVlabs/instant-ngp
        NeRF-synthetic scenes include RGBA / RGB information.
    """

    def __init__(self, dataset_path: str, split: str, bg_color: str, mip: int = 0,
                 dataset_num_workers: int = -1, transform: Callable = None):
        """ Loads the NeRF-synthetic data and applies dataset specific transforms required for compatibility with the
        framework.
        The loaded data is cached inside the `data` field.

        Args:
            dataset_path (str): The root directory of the dataset, where images and json files of a single multiview
                scene reside.
            split (str): The dataset split to use, corresponding to the transform file to load, when splits are
                available. In case of a single transform file, it will be considered as a single split of
                'train' by default.
                Options: 'train', 'val', 'test'.
            bg_color (str): The background color to use for when alpha=0.
                Options: 'black', 'white'.
            mip (int): If provided, will rescale images by 2**mip. Useful when large images are loaded.
            dataset_num_workers (int): The number of workers to spawn for multiprocessed loading.
                If dataset_num_workers < 1, processing will take place on the main process.
            transform (Optional[Callable]):
                Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        """
        super().__init__(dataset_path=dataset_path, dataset_num_workers=dataset_num_workers,
                         transform=transform, split=split)
        self.mip = mip
        self.bg_color = bg_color

        self.coords = self.data = self.coords_center = self.coords_scale = None
        self._transform_file = self._validate_and_find_transform()
        self.data = self.load()

        self._img_shape = self.data["rgb"].shape[1:3]
        self.flatten_tensors()

    def create_split(self, split: str, transform: Callable = None) -> NeRFSyntheticDataset:
        """ Creates a dataset with the same parameters and a different split.
        This is a convenient way of creating validation and test datasets, while making sure they're compatible
        with the train dataset.

        All settings except for split and transform will be copied from the current dataset.

        Args:
            split (str): The dataset split to use, corresponding to the transform file to load.
                Options: 'train', 'val', 'test'.
            transform (Optional[Callable]):
                Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        """
        return NeRFSyntheticDataset(
            dataset_path=self.dataset_path,
            split=split,
            bg_color=self.bg_color,
            mip=self.mip,
            dataset_num_workers=self.dataset_num_workers,
            transform=transform
        )

    def __getitem__(self, idx) -> MultiviewBatch:
        """Retrieve a batch of rays and their corresponding values.
        Rays are precomputed from the dataset's cameras, and are cached within the dataset.
        By default, rays are assumed to have corresponding rgb values, sampled from the dataset's images.

        Returns:
            (MultiviewBatch): A batch of rays and their rgb values. The fields can be accessed as a dictionary:
                "rays" - a wisp.core.Rays pack of ray origins and directions, pre-generated from the dataset camera.
                "rgb" - a torch.Tensor of rgb color which corresponds the gt image's pixel each ray intersects.
                "masks" - a torch.BoolTensor specifying if the ray hits a dense area or not.
                 This is estimated from the alpha channel of the gt image, where mask=True if alpha > 0.5.
        """
        out = MultiviewBatch(
            rays=self.data["rays"][idx],
            rgb=self.data["rgb"][idx],
            masks=self.data["masks"][idx]
        )

        if self.transform is not None:
            out = self.transform(out)

        return out

    @classmethod
    def is_root_of_dataset(cls, root: str, files_list: List[str]) -> bool:
        """ Each dataset may implement a simple set of rules to distinguish it from other datasets.
        Rules should be unique for this dataset type, such that given a general root path, Wisp will know
        to associate it with this dataset class.

        Datasets which don't implement this function should be created explicitly.

        Args:
                root (str): A path to the root directory of the dataset.
                files_list (List[str]): List of files within the dataset root, without their prefix path.
        Returns:
                True if the root folder points to content loadable by this dataset.
        """
        # NeRF-synthetic data is distinguished by the transform jsons and an additional image folder.
        try:
            regex = re.compile(r"transform.+\.json")
            transform_files = list(filter(regex.match, files_list))
            if 'transforms.json' in transform_files and 'images' in files_list:
                # Single transform file
                images_path = os.path.join(root, 'images')
                return os.path.isdir(images_path)
            elif 'transforms_train.json' in transform_files and 'train' in files_list:
                # Three transform files, with splits. Check only for train split
                train_path = os.path.join(root, 'train')
                return os.path.isdir(train_path)
            else:
                return False
        except ValueError:
            return False

    def _validate_and_find_transform(self) -> str:
        """
        Validates the file structure and returns the filename for the dataset's split / transform.
        There are two pairs of standard file structures this dataset can parse:

        ```
        /path/to/dataset/transform.json
        /path/to/dataset/images/____.png
        ```

        or

        ```
        /path/to/dataset/transform_{split}.json
        /path/to/dataset/{split}/_____.png
        ```

        In the former case, the single transform file is assumed to be loaded as a train set,
        for the latter split is assumed to be any of: 'train', 'val', 'test'.
        """
        if not os.path.exists(self.dataset_path):
            raise FileNotFoundError(f"NeRF dataset path does not exist: {self.dataset_path}")

        transforms = sorted(glob.glob(os.path.join(self.dataset_path, "*.json")))
        if len(transforms) == 0:
            raise RuntimeError(f"NeRF dataset folder has no transform *.json files with camera data: {self.dataset_path}")
        elif len(transforms) > 3 or len(transforms) == 2:
            raise RuntimeError(f"NeRF dataset folder has an unsupported number of splits, "
                               f"there should be ['test', 'train', 'val'], but found: {transforms}.")
        transform_dict = {}
        if len(transforms) == 1:
            transform_dict['train'] = transforms[0]
        elif len(transforms) == 3:
            fnames = [os.path.basename(transform) for transform in transforms]

            # Create dictionary of split to file path, probably there is simpler way of doing this
            for _split in ['test', 'train', 'val']:
                for i, fname in enumerate(fnames):
                    if _split in fname:
                        transform_dict[_split] = transforms[i]

        if self.split not in transform_dict:
            log.warning(
                f"WARNING: Split type ['{self.split}'] does not exist in the dataset. Falling back to train data.")
            self.split = 'train'
        return transform_dict[self.split]

    @staticmethod
    def _load_single_entry(frame, root, mip=None):
        """ Loads a single image: takes a frame from the JSON to load image and associated poses from json.
        This is a helper function which also supports multiprocessing for the standard dataset.

        Args:
            root (str): The root of the dataset.
            frame (dict): The frame object from the transform.json. The frame contains the metadata.
            mip (int): Optional, If set, rescales the image by 2**mip.

        Returns:
            (dict): Dictionary of the image and pose.
        """
        fpath = os.path.join(root, frame['file_path'].replace("\\", "/"))

        basename = os.path.basename(os.path.splitext(fpath)[0])
        if os.path.splitext(fpath)[1] == "":
            # Assume PNG file if no extension exists... the NeRF synthetic data follows this convention.
            fpath += '.png'

        # For some reason instant-ngp allows missing images that exist in the transform but not in the data.
        # Handle this... also handles the above case well too.
        if os.path.exists(fpath):
            img = load_rgb(fpath)
            if mip is not None:
                img = resize_mip(img, mip, interpolation=cv2.INTER_AREA)
            return dict(basename=basename,
                        img=torch.FloatTensor(img), pose=torch.FloatTensor(np.array(frame['transform_matrix'])))
        else:
            # log.info(f"File name {fpath} doesn't exist. Ignoring.")
            return None

    def load_singleprocess(self):
        """Standard parsing function for loading nerf-synthetic files on the main process.
        This follows the conventions defined in https://github.com/NVlabs/instant-ngp.

        Returns:
            (dict of torch.FloatTensors): Channels of information from NeRF:
                - 'rays': a list of ray packs, each entry corresponds to a single camera view
                - 'rgb', 'masks': a list of torch.Tensors, each entry corresponds to a single gt image
                - 'cameras': a list of Camera objects, one camera per view
        """
        with open(self._transform_file, 'r') as f:
            metadata = json.load(f)

        imgs = []
        poses = []
        basenames = []

        for frame in tqdm(metadata['frames'], desc='loading data'):
            _data = self._load_single_entry(frame, self.dataset_path, mip=self.mip)
            if _data is not None:
                basenames.append(_data["basename"])
                imgs.append(_data["img"])
                poses.append(_data["pose"])

        return self._collect_data_entries(metadata=metadata, basenames=basenames, imgs=imgs, poses=poses)

    @staticmethod
    def _parallel_load_standard_imgs(args):
        """ Internal function used by the multiprocessing loader: allocates a single entry task for a worker.
        """
        torch.set_num_threads(1)
        result = NeRFSyntheticDataset._load_single_entry(args['frame'], args['root'], mip=args['mip'])
        if result is None:
            return dict(basename=None, img=None, pose=None)
        else:
            return dict(basename=result['basename'], img=result['img'], pose=result['pose'])

    def load_multiprocess(self):
        """Standard parsing function for loading nerf-synthetic files with multiple workers.
        This follows the conventions defined in https://github.com/NVlabs/instant-ngp.

        Returns:
            (dict of torch.FloatTensors): Channels of information from NeRF:
                - 'rays': a list of ray packs, each entry corresponds to a single camera view
                - 'rgb', 'masks': a list of torch.Tensors, each entry corresponds to a single gt image
                - 'cameras': a list of Camera objects, one camera per view
        """
        with open(self._transform_file, 'r') as f:
            metadata = json.load(f)

        imgs = []
        poses = []
        basenames = []

        p = Pool(self.dataset_num_workers)
        try:
            mp_entries = [dict(frame=frame, root=self.dataset_path, mip=self.mip)
                          for frame in metadata['frames']]
            iterator = p.imap(NeRFSyntheticDataset._parallel_load_standard_imgs, mp_entries)

            for _ in tqdm(range(len(metadata['frames']))):
                result = next(iterator)
                basename = result['basename']
                img = result['img']
                pose = result['pose']
                if basename is not None:
                    basenames.append(basename)
                if img is not None:
                    imgs.append(img)
                if pose is not None:
                    poses.append(pose)
        finally:
            p.close()
            p.join()

        return self._collect_data_entries(metadata=metadata, basenames=basenames, imgs=imgs, poses=poses)

    def _collect_data_entries(self, metadata, basenames, imgs, poses) -> Dict[str, Union[torch.Tensor, Rays, Camera]]:
        """ Internal function for aggregating the pre-loaded multi-views.
        This function will:
            1. Read the metadata & compute the intrinsic parameters of the camera view, (such as fov and focal length
                i.e., in the case of a pinhole camera).
            2. Apply various scaling and offsets transformations to the extrinsics,
                as specified by the metadata by parameters
               such as 'scale', 'offset' and 'aabb_scale'
            3. Create kaolin Camera objects out of the computed extrinsics and intrinsics.
            4. Invoke ray generation on each camera view.
            5. Stack the images pixel values and rays as per-view information entries.
        """
        imgs = torch.stack(imgs)
        poses = torch.stack(poses)

        # TODO(ttakikawa): Assumes all images are same shape and focal. Maybe breaks in general...
        h, w = imgs[0].shape[:2]

        if 'x_fov' in metadata:
            # Degrees
            x_fov = metadata['x_fov']
            fx = (0.5 * w) / np.tan(0.5 * float(x_fov) * (np.pi / 180.0))
            if 'y_fov' in metadata:
                y_fov = metadata['y_fov']
                fy = (0.5 * h) / np.tan(0.5 * float(y_fov) * (np.pi / 180.0))
            else:
                fy = fx
        elif 'fl_x' in metadata and False:
            fx = float(metadata['fl_x']) / float(2**self.mip)
            if 'fl_y' in metadata:
                fy = float(metadata['fl_y']) / float(2**self.mip)
            else:
                fy = fx
        elif 'camera_angle_x' in metadata:
            # Radians
            camera_angle_x = metadata['camera_angle_x']
            fx = (0.5 * w) / np.tan(0.5 * float(camera_angle_x))

            if 'camera_angle_y' in metadata:
                camera_angle_y = metadata['camera_angle_y']
                fy = (0.5 * h) / np.tan(0.5 * float(camera_angle_y))
            else:
                fy = fx

        else:
            fx = 0.0
            fy = 0.0

        if 'fix_premult' in metadata:
            log.info("WARNING: The dataset expects premultiplied alpha correction, "
                     "but the current implementation does not handle this.")

        if 'k1' in metadata:
            log.info \
                ("WARNING: The dataset expects distortion correction, but the current implementation does not handle this.")

        if 'rolling_shutter' in metadata:
            log.info("WARNING: The dataset expects rolling shutter correction,"
                     "but the current implementation does not handle this.")

        # The principal point in wisp are always a displacement in pixels from the center of the image.
        x0 = 0.0
        y0 = 0.0
        # The standard dataset generally stores the absolute location on the image to specify the principal point.
        # Thus, we need to scale and translate them such that they are offsets from the center.
        if 'cx' in metadata:
            x0 = (float(metadata['cx']) / (2**self.mip)) - (w//2)
        if 'cy' in metadata:
            y0 = (float(metadata['cy']) / (2**self.mip)) - (h//2)

        offset = metadata['offset'] if 'offset' in metadata else [0 ,0 ,0]
        scale = metadata['scale'] if 'scale' in metadata else 1.0
        aabb_scale = metadata['aabb_scale'] if 'aabb_scale' in metadata else 1.25

        # TODO(ttakikawa): Actually scale the AABB instead? Maybe
        poses[..., :3, 3] /= aabb_scale
        poses[..., :3, 3] *= scale
        poses[..., :3, 3] += torch.FloatTensor(offset)

        # nerf-synthetic uses a default far value of 6.0
        default_far = 6.0

        rays = []

        cameras = dict()
        for i in range(imgs.shape[0]):
            view_matrix = torch.zeros_like(poses[i])
            view_matrix[:3, :3] = poses[i][:3, :3].T
            view_matrix[:3, -1] = torch.matmul(-view_matrix[:3, :3], poses[i][:3, -1])
            view_matrix[3, 3] = 1.0
            camera = Camera.from_args(view_matrix=view_matrix,
                                      focal_x=fx,
                                      focal_y=fy,
                                      width=w,
                                      height=h,
                                      far=default_far,
                                      near=0.0,
                                      x0=x0,
                                      y0=y0,
                                      dtype=torch.float64)
            camera.change_coordinate_system(blender_coords())
            cameras[basenames[i]] = camera
            ray_grid = generate_centered_pixel_coords(camera.width, camera.height,
                                                      camera.width, camera.height, device='cuda')
            rays.append \
                (generate_pinhole_rays(camera.to(ray_grid[0].device), ray_grid).reshape(camera.height, camera.width, 3).to
                    ('cpu'))

        rays = Rays.stack(rays).to(dtype=torch.float)

        rgbs = imgs[... ,:3]
        alpha = imgs[... ,3:4]
        if alpha.numel() == 0:
            masks = torch.ones_like(rgbs[... ,0:1]).bool()
        else:
            masks = (alpha > 0.5).bool()

            if self.bg_color == 'black':
                rgbs[... ,:3] -= ( 1 -alpha)
                rgbs = np.clip(rgbs, 0.0, 1.0)
            else:
                rgbs[... ,:3] *= alpha
                rgbs[... ,:3] += ( 1 -alpha)
                rgbs = np.clip(rgbs, 0.0, 1.0)

        return {"rgb": rgbs, "masks": masks, "rays": rays, "cameras": cameras}

    def flatten_tensors(self) -> None:
        """ Flattens the cached data tensors to (NUM_VIEWS, NUM_RAYS, *).
        """
        num_imgs = len(self)
        self.data["rgb"] = self.data["rgb"].reshape(num_imgs, -1, 3)
        self.data["rays"] = self.data["rays"].reshape(num_imgs, -1, 3)
        if "masks" in self.data:
            self.data["masks"] = self.data["masks"].reshape(num_imgs, -1, 1)

    @property
    def img_shape(self) -> torch.Size:
        """ Returns the shape of the rescaled dataset images (cached values are flattened) """
        return self._img_shape

    @property
    def cameras(self) -> List[Camera]:
        """ Returns the list of camera views used to generate rays for this dataset. """
        return self.data["cameras"]

    @property
    def num_images(self) -> int:
        """ Returns the number of views this dataset stores. """
        return self.data["rgb"].shape[0]

wisp/datasets/formats/nerf_standard_dataset.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from typing import Callable, List
import torch
import logging as log
import kaolin.ops.spc as spc_ops
from wisp.datasets.base_datasets import SDFDataset
from wisp.datasets.batch import SDFBatch
from wisp.accelstructs import BaseAS, OctreeAS
import wisp.ops.mesh as mesh_ops
import wisp.ops.spc as wisp_spc_ops


class OctreeSampledSDFDataset(SDFDataset):
    """ A dataset class for signed-distance functions sampled from a sparse octree acceleration structure,
    which was constructed from single mesh shape. Sampling is executed on the mesh faces,
    and is limited to occupied octree cells. This dataset has the benefit of limiting the sampling region to areas
    which are actually occupied by the mesh. It also allows for equal distribution of samples per the octree cells.
    The dataset supports online resampling during training. Resampling is quick as a pool of narrowband values is
    cached for the duration of training, this pool is quickly subsampled each time resample() is invoked.
    Samples can also optionally include rgb information when the mesh is textured with materials.
    """

    def __init__(self,
                 occupancy_struct: OctreeAS,
                 split: str,
                 transform: Callable = None,
                 sample_mode: list = None,
                 num_samples: int = 100000,
                 sample_tex: bool = False,
                 samples_per_voxel: int = 32
                 ):
        """Construct a sample-set coords + sdf samples by sampling the mesh faces and dropping samples
        that fall into empty octree cells.

        Args:
            occupancy_struct (OctreeAS): An octree acceleration structure, initialized from a mesh.
            split (str): Any of 'test', 'train', 'val'.
                Currently used for keeping track of the dataset purpose and not used internally.
            transform (Optional[Callable]):
                A transform function applied per batch when data is accessed with __get_item__.
            sample_mode (list of str): List of different sample methods to apply over the mesh.
                Any sequential combo of:
                    'tracer' - samples generated on the mesh surface
                    'near' - samples generated near the mesh surface
                    'rand' - samples generated uniformly in the space of occupied grid cells
                    See `mesh_ops.point_sample` for additional implementation details for the sampler.
                    Default: ['rand', 'rand', 'near', 'near', 'trace']
            num_samples (int): Number of data points to keep in the working set at any given time.
                A new working set can be regenerated with resample().
            sample_tex (bool): If True, __get_item__  will also return RGB values from the nearest texture.
            samples_per_voxel (int): Number of samples to generate per octree grid cell to create a preliminary
                equally distributed pool of samples.
        """
        super().__init__(transform=transform, split=split)

        self.blas = occupancy_struct
        self.sample_mode = sample_mode if sample_mode is not None else ['rand', 'rand', 'near', 'near', 'trace']
        self.num_samples = num_samples
        self.sample_tex = sample_tex
        self.samples_per_voxel = samples_per_voxel

        self.validate()
        self.data_pool = None
        self.data = None
        self.load()

    @staticmethod
    def supports_blas(blas: BaseAS) -> bool:
        """ This dataset requires an octree to compute the narrowband.
        Given a bottom-level-acceleration-structure, this function checks if it contains the information needed
        to create this dataset from it (i.e. the blas was initialized from a mesh).
        """
        # This dataset requires the bottom level acceleration structure to keep the mesh info in its optional fields
        return isinstance(blas, OctreeAS) and 'vertices' in blas.extent and 'faces' in blas.extent

    def validate(self) -> None:
        """ Ensures the current dataset is valid. If it is valid, this function will terminate gracefully.
        Otherwise an informative error is prompted.
        """
        if not self.supports_blas(self.blas):
            raise RuntimeError("The Octree acceleration structure was not initialized from a mesh. To use "
                               "an OctreeAS with this dataset, make sure to construct it with a mesh.")

    def _sample_from_grid(self, blas: OctreeAS, samples_per_voxel=32):
        """Initializes the dataset by sampling SDFs from an OctreeGrid created from a mesh.

        Args:
            blas (OctreeAS): An OctreeAS class initialized from mesh.
            samples_per_voxel (int): Number of data points to uniformly sample per voxel,
                irrespective of samples generated on / near the mesh surface.
        """
        # TODO (operel): TBD when kaolin adds a mesh class:
        #  grid is only really needed for filtering out points and more efficient 'rand',
        #  better give the mesh as another input and not store the mesh contents in the extent field
        vertices = blas.extent['vertices']
        faces = blas.extent['faces']
        level = blas.max_level

        # Here, corners mean "the bottom left corner of the voxel to sample from"
        corners = spc_ops.unbatched_get_level_points(blas.points, blas.pyramid, level)

        # Two pass sampling to figure out sample size
        pts = []
        for mode in self.sample_mode:
            if mode == "rand":
                # Sample the points.
                pts.append(wisp_spc_ops.sample_spc(corners, level, samples_per_voxel).cpu())
        for mode in self.sample_mode:
            if mode == "rand":
                pass
            elif mode == "near":
                pts.append(mesh_ops.sample_near_surface(vertices.cuda(), faces.cuda(), pts[0].shape[0],
                                                        variance=1.0 / (2 ** level)).cpu())
            elif mode == "trace":
                pts.append(mesh_ops.sample_surface(vertices.cuda(), faces.cuda(), pts[0].shape[0])[0].cpu())
            else:
                raise Exception(f"Sampling mode {mode} not implemented")

        # Filter out points which do not belong to the narrowband
        pts = torch.cat(pts, dim=0).cuda()
        query_results = blas.query(pts, 0)
        pts = pts[query_results.pidx > -1]

        # Sample distances and textures.
        rgb = None
        if self.sample_tex:
            texv = blas.extent['texv']
            texf = blas.extent['texf']
            mats = blas.extent['mats']
            rgb, hit_pts, d = mesh_ops.closest_tex(vertices, faces, texv, texf, mats, pts)
        else:
            d = mesh_ops.compute_sdf(vertices, faces, pts)
            assert (d.shape[0] == pts.shape[0]), 'SDF validation logic failed: the number of returned sdf samples' \
                                                 'does not match the number of input coordinates.'

        data = dict(
            coords=pts.cpu(),
            sdf=d.cpu(),
        )
        if rgb is not None:
            data['rgb'] = rgb.cpu()
        return data

    def resample(self) -> None:
        """Resamples a working set of coordaintes + sdf from the existing pool of samples.
        This is essentially a quick subsampling operation which permutes and selects random indices.
        """
        log.info(f"Resampling {self.num_samples} samples..")
        idx = torch.randperm(self.pool_size - 1, device='cuda')
        if self.num_samples is not None and self.num_samples < self.pool_size:
            idx = idx[:self.num_samples]
        self.data = {k: v[idx.to(v.device)] for k, v in self.data_pool.items() if v is not None}

    def load_singleprocess(self):
        """Initializes the dataset by sampling SDF values from a grid and a mesh contained within it.
        This function uses the main process to load the dataset, without spawning any workers.
        """
        log.info(f"Computing SDFs for entire samples pool (may take a while)..")
        self.data_pool = self._sample_from_grid(blas=self.blas, samples_per_voxel=self.samples_per_voxel)
        log.info(f"Total Samples in Pool: {self.data_pool['coords'].shape[0]}")
        self.resample()

    @classmethod
    def is_root_of_dataset(cls, root: str, files_list: List[str]) -> bool:
        """ Each dataset may implement a simple set of rules to distinguish it from other datasets.
        Rules should be unique for this dataset type, such that given a general root path, Wisp will know
        to associate it with this dataset class.

        Datasets which don't implement this function should be created explicitly.
        This class does not load external files and therefore always returns False (e.g. should be created explicitly).
        """
        # This dataset is created from a preloaded grid and not from a file. Therefore it has to be created explicitly
        return False

    def __len__(self):
        """Return length of dataset, as number of samples currently resampled."""
        return self.data["coords"].size()[0]

    @property
    def pool_size(self) -> int:
        """ Total number of samples available in data pool, prior to subsampling the data from it with resample() """
        return self.data_pool['coords'].shape[0]

    def __getitem__(self, idx: int) -> SDFBatch:
        """Retrieve a batch of sample coordinates and their sdf values.
        If the dataset was loaded with rgb information, this channel will be returned as well.

        Returns:
            (SDFBatch): A batch of coordinates and their SDF values. The fields can be accessed as a dictionary:
                "coords" - a torch.Tensor of the 3d coordinates of each sample
                "sdf" - a torch.Tensor of the signed distance of each sample from the mesh surface.
                "rgb" - optional, a torch.Tensor of the "textured color" the sample interpolates from the nearest point
                    on the mesh surface.
                    If the dataset with built with `sample_tex`=False, this entry is None.
        """
        out = SDFBatch(
            coords=self.data["coords"][idx],
            sdf=self.data["sdf"][idx],
            rgb=self.data["rgb"][idx] if "rgb" in self.data else None,
        )

        if self.transform is not None:
            out = self.transform(out)

        return out

    @property
    def coordinates(self) -> torch.Tensor:
        """ Returns the coordinates of samples stored in this dataset. """
        return self.data["coords"]

wisp/datasets/formats/octree_sdf_dataset.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from typing import Callable, Tuple, Union
from copy import deepcopy
import unittest
import random
import numpy as np

import torch

from kaolin.render.camera import Camera
from kaolin.render.camera.extrinsics import CameraExtrinsics
from torch.utils.data import Dataset
from wisp.utils import DotDict
from wisp.ops.raygen import generate_pinhole_rays, generate_centered_pixel_coords

# TODO (operel): integrate with WispDataset interface, currently unused by any of the apps

def spherical_eye(
    radius=1,
    theta=np.pi / 3,
    phi=0,
):
    return torch.FloatTensor(
        [
            radius * np.sin(theta) * np.sin(phi),
            radius * np.cos(theta),
            radius * np.sin(theta) * np.cos(phi),
        ],
    )  # [3]


# TODO (cchoy): move the radius/theta initialization to kaolin CameraExtrinsics
def spherical_coord_to_pose(
    radius=1, theta=np.pi / 3, phi=0, up=torch.FloatTensor([0, 1, 0])
):
    """generate camera pose from a spherical coordinate

    Args:
        size: batch size of generated poses.
        device: where to allocate the output.
        radius: camera radius
        theta_range: [min, max], should be in [0, pi]
        phi_range: [min, max], should be in [0, 2 * pi]
    Return:
        poses: [size, 4, 4] in OpenGL convention
    """

    eye = spherical_eye(radius, theta, phi)

    # lookat
    def normalize(vec):
        return torch.nn.functional.normalize(vec, dim=-1)

    backward = -normalize(eye)
    right = normalize(torch.cross(backward, up, dim=-1))
    up = normalize(torch.cross(right, backward, dim=-1))

    world_rot = torch.stack((right, up, -backward), dim=1)
    world_tran = -world_rot @ eye.unsqueeze(-1)

    return CameraExtrinsics._from_world_in_cam_coords(
        rotation=world_rot, translation=world_tran, device="cpu", requires_grad=False
    )


class RandomViewDataset(Dataset):
    def __init__(
        self,
        # TODO(cchoy) add different random view types e.g. forward_facing, inward
        n_size=100,  # length of this dataset. Used to define number of iterations per epoch
        view_radius_range: Tuple = (2, 4),
        view_theta_range: Tuple = (np.pi / 4, np.pi / 2 - np.pi / 8),
        view_phi_range: Tuple = (0, 2 * np.pi),
        viewport_height: int = 320,
        viewport_width: int = 320,
        fov: float = 30 * np.pi / 180,
        ray_dist_range: Tuple = (0.01, 8),
        look_at: Tuple = (0, 0, 0),
        num_rays: int = -1,  # number of rays. If -1, return all rays
        transform: Callable = None,
        **kwargs,
    ):
        self.n_size = n_size
        self.cam = DotDict(
            dict(
                fov=fov,
                width=viewport_width,
                height=viewport_height,
            )
        )
        self.view_radius_range = view_radius_range
        self.view_theta_range = view_theta_range
        self.view_phi_range = view_phi_range
        self.ray_dist_range = ray_dist_range
        assert len(look_at) == 3
        self.look_at = look_at

        self.num_rays = num_rays
        self.transform = transform

    def __len__(self):
        """Length of the dataset in number of rays."""
        return self.n_size

    def __getitem__(self, idx: int):
        """Returns a ray."""
        # TODO (cchoy): uniform sphere sampling (http://corysimon.github.io/articles/uniformdistn-on-sphere/)
        radius = random.uniform(*self.view_radius_range)
        theta = random.uniform(*self.view_theta_range)
        phi = random.uniform(*self.view_phi_range)
        cam = Camera.from_args(
            eye=spherical_eye(radius, theta, phi),
            at=torch.tensor([0.0, 0.0, 0.0]),
            up=torch.tensor([0.0, 1.0, 0.0]),
            fov=self.cam.fov,
            width=self.cam.width,
            height=self.cam.height,
            device="cpu",
        )

        ray_grid = generate_centered_pixel_coords(
            cam.width, cam.height, cam.width, cam.height, device="cpu"
        )
        out = DotDict(dict(rays=generate_pinhole_rays(cam, ray_grid), cam=cam))
        if self.num_rays > 0:
            ray_idx = random.sample(range(len(out.rays)), self.num_rays)
            out.rays = out.rays[ray_idx]

        if self.transform is not None:
            out = self.transform(out)

        return out


class TestRandViewDataset(unittest.TestCase):
    def load(self):
        dataset = RandomViewDataset()
        print(dataset[0])

wisp/datasets/formats/random_view_dataset.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import copy
import os
import re
import glob
import json
import cv2
from tqdm import tqdm
import math
import numpy as np
import torch
from torch.multiprocessing import Pool
from typing import Callable, List, Dict
from kaolin.render.camera import Camera, blender_coords
from wisp.core import Rays
import wisp.ops.image as img_ops
from wisp.ops.raygen import generate_pinhole_rays, generate_centered_pixel_coords
from wisp.ops.pointcloud import create_pointcloud_from_images, normalize_pointcloud
from wisp.datasets.base_datasets import MultiviewDataset
from wisp.datasets.batch import MultiviewBatch


class RTMVDataset(MultiviewDataset):
    """ A dataset for loading data files in the RTMV format:
        RTMV: A Ray-Traced Multi-View Synthetic Dataset for Novel View Synthesis
        See: http://www.cs.umd.edu/~mmeshry/projects/rtmv/
        RTMV scenes include RGB + depth information.
    """

    def __init__(self, dataset_path: str, bg_color: str, mip: int = 0, split: str = 'train',
                 train_ratio: float = 2.0 / 3.0, val_ratio: float = 1.0 / 30.0,
                 coords_center: torch.Tensor = None, coords_scale: torch.Tensor = None,
                 dataset_num_workers: int = -1, transform: Callable = None):
        """Loads the RTMV data and applies dataset specific transforms required for compatibility with the framework.
        The loaded data is cached inside the `data` field.

        Args:
            dataset_path (str): The root directory of the dataset, where images and json files of a single multiview
                scene reside.
            bg_color (str): The background color to use for when alpha=0.
                Options: 'black', 'white'.
            mip (int): If provided, will rescale images by 2**mip. Useful when large images are loaded.
            split (str): The dataset split to use. The exact split's content
                are determined by train_ratio, val_ratio.
                Options: 'train', 'val', 'test'.
            train_ratio (float): Ratio of files allocated for the 'train' split.
            val_ratio (float): Ratio of files allocated for the 'validation' split.
                Note: test_ratio = 1 - train_ratio - val_ratio
            coords_center Optional[torch.Tensor]: optional, since rtmv contains depth information, if coords_center is
                specified, a point cloud will be computed from the depth rays, and centered using coords_center,
                and then used to normalize the ray origins.
                If coords center and scale are both not specified, they're calculated from the data.
                Usually this parameter is specified to ensure multiple splits are normalized in the same way
                (i.e. train set computes coords_center, and validation / test set normalize using that value).
            coords_scale Optional[torch.Tensor]: optional, since rtmv contains depth information, if coords_scale is
                specified, a point cloud will be computed from the depth rays, and scaled using coords_scale,
                and then used to normalize the ray origins.
                If coords center and scale are both not specified, they're calculated from the data.
                Usually this parameter is specified to ensure multiple splits are normalized in the same way
                (i.e. train set computes coords_scale, and validation / test set normalize using that value).
            dataset_num_workers (int): The number of workers to spawn for multiprocessed loading.
                If dataset_num_workers < 1, processing will take place on the main process.
            transform (Optional[Callable]):
                Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        """
        super().__init__(dataset_path=dataset_path, dataset_num_workers=dataset_num_workers,
                         transform=transform, split=split)
        self.mip = mip
        self.bg_color = bg_color

        self._train_ratio = train_ratio
        self._val_ratio = val_ratio
        self._json_files = self._list_jsons_for_split(dataset_path, split, train_ratio, val_ratio)

        images, alphas, depths, cameras = self.load()
        rays = self._raygen(cameras)
        self.coords = None

        # If coords center and scale were not specified, calculate them from data
        # This logic will also normalize the ray origins and calculate the depth values
        should_normalize = coords_center is None or coords_scale is None
        if should_normalize:
            depths, rays, coords_center, coords_scale = self._normalize(images, alphas, depths, cameras, rays)
            self.coords, _ = create_pointcloud_from_images(images, alphas, rays, depths)  # RGBD -> point cloud
        self.coords_center = coords_center
        self.coords_scale = coords_scale

        self.data = dict(
            rgb=torch.stack(images)[..., :3],
            alpha=torch.stack(alphas),
            rays=Rays.stack(rays),
            depth=torch.stack(depths),
            cameras=cameras
        )

        if not should_normalize:
            self.data["depth"] = self.data["depth"] * self.coords_scale
            self.data["rays"].origins = (self.data["rays"].origins - self.coords_center) * self.coords_scale

        self._img_shape = self.data["rgb"].shape[1:3]
        self.flatten_tensors()

    def create_split(self, split: str, transform: Callable = None) -> RTMVDataset:
        """ Creates a dataset with the same parameters and a different split.
        This is a convenient way of creating validation and test datasets, while making sure they're compatible
        with the train dataset.

        All settings except for split and transform will be copied from the current dataset.
        RTMV will also make sure that the created split uses the same normalization values of
        coords_center, coords_scale, such that both dataset splits are compatible.

        Args:
            split (str): The dataset split to use, corresponding to the transform file to load.
                Options: 'train', 'val', 'test'.
            transform (Optional[Callable]):
                Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        """
        return RTMVDataset(
            dataset_path=self.dataset_path,
            bg_color=self.bg_color,
            mip=self.mip,
            split=split,
            train_ratio=self._train_ratio,
            val_ratio=self._val_ratio,
            coords_center=self.coords_center,
            coords_scale=self.coords_scale,
            dataset_num_workers=self.dataset_num_workers,
            transform=transform
        )

    def __getitem__(self, idx: int) -> MultiviewBatch:
        """Retrieve a batch of rays and their corresponding values.
        Rays are precomputed from the dataset's cameras, and are cached within the dataset.
        By default, rays are assumed to have corresponding rgb and depth values, sampled from the dataset's images.

        Returns:
            (MultiviewBatch): A batch of rays and their rgbd values. The fields can be accessed as a dictionary:
                "rays" - a wisp.core.Rays pack of ray origins and directions, pre-generated from the dataset camera.
                "rgb" - a torch.Tensor of rgb color which corresponds the gt image pixel rgb
                    each ray intersects.
                "depth" - a torch.Tensor of depth which corresponds the gt image pixel depth
                    each ray intersects.
                "alpha" - a torch.Tensor of alpha value which corresponds to the gt image alpha.
        """
        out = MultiviewBatch(
            rays=self.data["rays"][idx],
            rgb=self.data["rgb"][idx],
            depth=self.data["depth"][idx],
            alpha=self.data["alpha"][idx]
        )

        if self.transform is not None:
            out = self.transform(out)

        return out

    @classmethod
    def is_root_of_dataset(cls, root: str, files_list: List[str]) -> bool:
        """ Each dataset may implement a simple set of rules to distinguish it from other datasets.
        Rules should be unique for this dataset type, such that given a general root path, Wisp will know
        to associate it with this dataset class.

        Datasets which don't implement this function should be created explicitly.

        Args:
                root (str): A path to the root directory of the dataset.
                files_list (List[str]): List of files within the dataset root, without their prefix path.
        Returns:
                True if the root folder points to content loadable by this dataset.
        """
        # Assume the existance of:
        # 00000.json, 00000.exr, 00001.json, 00001.exr, ...
        json_regex = re.compile(r"\d{5}\.json")
        json_files = list(filter(json_regex.match, files_list))
        exr_regex = re.compile(r"\d{5}\.exr")
        exr_files = list(filter(exr_regex.match, files_list))
        return len(json_files) == len(exr_files) and len(json_files) > 0

    @staticmethod
    def _list_jsons_for_split(dataset_path: str, split: str, train_ratio: float, val_ratio: float) -> List[str]:
        """
        Returns a list of dataset json files to use for the given split.

        Args:
            dataset_path (str): The root directory of the dataset, where images and json files of a single multiview
                scene reside.
            split (str): The dataset split to use. The exact split's content are determined by train_ratio, val_ratio.
                Options: 'train', 'val', 'test'.
            train_ratio (float): Ratio of files allocated for the 'train' split.
            val_ratio (float): Ratio of files allocated for the 'validation' split.
                Note: test_ratio = 1 - train_ratio - val_ratio
        """
        json_files = sorted(glob.glob(os.path.join(dataset_path, '*.json')))

        # Hard-coded train-val-test splits for now
        train_split_idx = math.floor(len(json_files) * train_ratio)
        eval_split_idx = train_split_idx + math.floor(len(json_files) * val_ratio)

        if split == 'train':
            subset_idxs = np.arange(0, train_split_idx)
        elif split in 'val':
            subset_idxs = np.arange(train_split_idx, eval_split_idx)
        elif split == 'test':
            subset_idxs = np.arange(eval_split_idx, len(json_files))
        else:
            raise RuntimeError(f"Unknown split type: {split}, split should be any of: ('train', 'val', 'test')")

        json_files = [json_files[i] for i in subset_idxs]
        if len(json_files) == 0:
            raise ValueError(f"RTMVDataset: No JSON files found for split {split} under {dataset_path}")
        return json_files

    def load_singleprocess(self):
        """Standard parsing function for loading rtmv files on the main process.
        This follows the conventions defined in http://www.cs.umd.edu/~mmeshry/projects/rtmv/

        Returns:
            (tuple of torch.FloatTensors): Channels of information from NeRF:
                - images: a list of torch.FloatTensor of size [H,W,3], each entry corresponds to a single rgb image.
                - alphas: a list of torch.FloatTensor of size [H,W,1], each entry corresponds to a single alpha image.
                - depths: a list of torch.FloatTensor of size [H,W,1], each entry corresponds to a single depth image.
                - cameras: a dictionary of filenames (view id) -> kaolin Camera object
        """
        images = []
        alphas = []
        depths = []
        cameras = dict()

        for img_index, json_file in tqdm(enumerate(self._json_files), desc='loading data'):
            with torch.no_grad():
                basename = os.path.splitext(os.path.basename(json_file))[0]
                image, alpha, depth = RTMVDataset._load_rtmv_images(self.dataset_path, basename,
                                                                    use_depth=True, mip=self.mip,
                                                                    srgb=True, bg_color=self.bg_color)
                json_path = os.path.join(self.dataset_path, basename + ".json")
                camera = RTMVDataset._load_rtmv_camera(path=json_path)
                transformed_camera = RTMVDataset._transform_rtmv_camera(copy.deepcopy(camera), mip=self.mip)

                images.append(image)
                alphas.append(alpha)
                depths.append(depth)
                cameras[basename] = transformed_camera
        return images, alphas, depths, cameras

    @staticmethod
    def _load_rtmv_images(root, basename, use_depth=False, mip=None, srgb=False, bg_color='white'):
        """Loads a set of RTMV images by path and basename.
        Args:
            root (str): Path to the root of the dataset.
            basename (str): Basename of the RTMV image set to load.
            use_depth (bool): if True, loads the depth data
                              by default, this assumes the depth is stored in the "depth" buffer
            mip (int): if not None, then each image will be resized by 2^mip
            srgb (bool): if True, convert to SRGB
        Returns:
            (dictionary of torch.Tensors)
            Keys:
                image : torch.FloatTensor of size [H,W,3]
                alpha : torch.FloatTensor of size [H,W,1]
                depth : torch.FloatTensor of size [H,W,1]
        """
        # TODO(ttakikawa): There is a lot that this function does... break this up
        from wisp.ops.image import resize_mip, linear_to_srgb

        image_exts = ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG']
        exr_exts = ['.exr', '.EXR']
        npz_exts = ['.npz', '.NPZ']

        img = None
        depth = None

        # Try to load RGB first
        found_image_path = None
        for ext in image_exts + exr_exts:
            image_path = os.path.join(root, basename + ext)
            if os.path.exists(image_path):
                found_image_path = image_path
                break
        if found_image_path is None:
            raise Exception("No images found! Check if your dataset path contains an actual RTMV dataset.")

        img_ext = os.path.splitext(found_image_path)[1]
        if img_ext in image_exts:
            img = img_ops.load_rgb(found_image_path)
        elif img_ext in exr_exts:
            try:
                import pyexr
            except:
                raise Exception(
                    "The RTMV dataset provided uses EXR, but module pyexr is not available. "
                    "To install, run `pip install pyexr`. "
                    "You will likely also need `libopenexr`, which through apt you can install with "
                    "`apt-get install libopenexr-dev` and on Windows you can install with "
                    "`pipwin install openexr`.")
            f = pyexr.open(found_image_path)
            img = f.get("default")
        else:
            raise Exception(f"Invalid image extension for the image path {found_image_path}")

        found_depth_path = None
        for ext in [".depth.npz", ".depth.exr"] + exr_exts:
            depth_path = os.path.join(root, basename + ext)
            if os.path.exists(depth_path):
                found_depth_path = depth_path
                break
        if found_depth_path is None:
            raise Exception("No depth found! Check if your dataset path contains an actual RTMV dataset.")

        depth_ext = os.path.splitext(found_depth_path)[1]
        # Load depth
        if depth_ext == ".npz":
            depth = np.load(found_depth_path)['arr_0'][..., 0]
        elif depth_ext == ".exr":
            try:
                import pyexr
            except:
                raise Exception(
                    "The RTMV dataset provided uses EXR, but module pyexr is not available. "
                    "To install, run `pip install pyexr`. "
                    "You will likely also need `libopenexr`, which through apt you can install with "
                    "`apt-get install libopenexr-dev` and on Windows you can install with "
                    "`pipwin install openexr`.")

            f = pyexr.open(found_depth_path)

            components = os.path.basename(found_depth_path).split('.')
            if len(components) > 2 and components[-1] == "exr" and components[-2] == "depth":
                depth = f.get('default')[:, :, 0]
            else:
                if len(f.channel_map['depth']) > 0:
                    depth = f.get("depth")
                else:
                    raise Exception("Depth channel not found in the EXR file provided!")
        else:
            raise Exception(f"Invalid depth extension for the depth path {found_depth_path}")

        alpha = img[..., 3:4]

        if bg_color == 'black':
            img[..., :3] -= (1 - alpha)
            img = np.clip(img, 0.0, 1.0)
        else:
            img[..., :3] *= alpha
            img[..., :3] += (1 - alpha)
            img = np.clip(img, 0.0, 1.0)

        if mip is not None:
            # TODO(ttakikawa): resize_mip causes the mask to be squuezed... why?
            img = resize_mip(img, mip, interpolation=cv2.INTER_AREA)
            if use_depth:
                depth = resize_mip(depth, mip, interpolation=cv2.INTER_NEAREST)
                # mask_depth = resize_mip(mask_depth[...,None].astype(np.float), mip, interpolation=cv2.INTER_NEAREST)

        img = torch.from_numpy(img)
        if use_depth:
            depth = torch.from_numpy(depth)

        if use_depth:
            mask_depth = torch.logical_and(depth > -1000, depth < 1000)
            depth[~mask_depth] = -1.0
            depth = depth[:, :, np.newaxis]

        if srgb:
            img = linear_to_srgb(img)

        alpha = mask_depth

        return img, alpha, depth

    @staticmethod
    def _parallel_load_rtmv_data(args):
        """ An internal wrapper function to allow rtmv load faster with multiprocessing.
        Represents a task performed by a single worker.
        Note: Internal logic should occur on the cpu, cuda behavior for multiprocessing is unstable.
        """
        torch.set_num_threads(1)
        with torch.no_grad():
            image, alpha, depth = RTMVDataset._load_rtmv_images(**args['exr_args'])
            camera = RTMVDataset._load_rtmv_camera(args['camera_args']['path'])
            transformed_camera = RTMVDataset._transform_rtmv_camera(copy.deepcopy(camera),
                                                                    mip=args['camera_args']['mip'])
            return dict(
                task_basename=args['task_basename'],
                image=image,
                alpha=alpha,
                depth=depth,
                camera=transformed_camera.cpu()
            )

    def load_multiprocess(self):
        """Standard parsing function for loading rtmv files with multiple workers.
        This follows the conventions defined in http://www.cs.umd.edu/~mmeshry/projects/rtmv/

        Returns:
            (tuple of torch.FloatTensors): Channels of information from NeRF:
                - images: a list of torch.FloatTensor of size [H,W,3], each entry corresponds to a single rgb image.
                - alphas: a list of torch.FloatTensor of size [H,W,1], each entry corresponds to a single alpha image.
                - depths: a list of torch.FloatTensor of size [H,W,1], each entry corresponds to a single depth image.
                - cameras: a dictionary of filenames (view id) -> kaolin Camera object
        """
        images = []
        alphas = []
        depths = []
        cameras = dict()

        p = Pool(self.dataset_num_workers)
        try:
            basenames = (os.path.splitext(os.path.basename(json_file))[0] for json_file in self._json_files)
            iterator = p.imap(RTMVDataset._parallel_load_rtmv_data, [
                dict(
                    task_basename=basename,
                    exr_args=dict(
                        root=self.dataset_path,
                        basename=basename,
                        use_depth=True,
                        mip=self.mip,
                        srgb=True,
                        bg_color=self.bg_color),
                    camera_args=dict(
                        path=os.path.join(self.dataset_path, basename + '.json'),
                        mip=self.mip)) for basename in basenames])
            for _ in tqdm(range(len(self._json_files)), desc='loading data'):
                result = next(iterator)
                images.append(result["image"])
                alphas.append(result["alpha"])
                depths.append(result["depth"])
                cameras[result['task_basename']] = result["camera"]
        finally:
            p.close()
            p.join()
        return images, alphas, depths, cameras

    @staticmethod
    def _rescale_rtmv_intrinsics(camera: Camera, target_size, original_width, original_height):
        """ Internal loading function: rescales the camera intrinsics according to the given ratios. """
        assert original_height == original_width, 'Current code assumes a square image.'
        resize_ratio = target_size * 1. / original_height
        camera.x0 *= resize_ratio
        camera.y0 *= resize_ratio
        camera.focal_x *= resize_ratio
        camera.focal_y *= resize_ratio

    @staticmethod
    def _load_rtmv_camera(path: str) -> Camera:
        """ Internal loading function: loads a single RTMV camera object from json metadata.
        """
        with open(path.replace("exr", "json"), 'r') as f:
            meta = json.load(f)

        # About RTMV conventions:
        # RTMV: meta['camera_data']['cam2world']), row major
        # kaolin cameras: cam2world is camera.inv_view_matrix(), column major
        # Therefore the following are equal up to a 1e-7 difference:
        # camera.inv_view_matrix() == torch.tensor(meta['camera_data']['cam2world']).T
        cam_data = meta['camera_data']
        camera = Camera.from_args(eye=torch.Tensor(cam_data['camera_look_at']['eye']),
                                  at=torch.Tensor(cam_data['camera_look_at']['at']),
                                  up=torch.Tensor(cam_data['camera_look_at']['up']),
                                  width=cam_data['width'],
                                  height=cam_data['height'],
                                  focal_x=cam_data['intrinsics']['fx'],
                                  focal_y=cam_data['intrinsics']['fy'],
                                  x0=0.0,
                                  y0=0.0,
                                  near=0.0,
                                  far=6.0,  # inheriting default for nerf-synthetic
                                  dtype=torch.float64,
                                  device='cpu')

        # RTMV cameras use Blender coordinates, which are right handed with Z axis pointing upwards instead of Y.
        camera.change_coordinate_system(blender_coords())

        return camera.cpu()

    @staticmethod
    def _transform_rtmv_camera(camera: Camera, mip: int) -> Camera:
        """ Transforms the RTMV camera's image plane according to the mip (essentially downsampling the image plane
        resolution).
        Returns the transformed camera with updated image plane and intrinsics.
        """
        original_width, original_height = camera.width, camera.height
        if mip is not None:
            camera.width = camera.width // (2 ** mip)
            camera.height = camera.height // (2 ** mip)

        # assume no resizing
        RTMVDataset._rescale_rtmv_intrinsics(camera, camera.width, original_width, original_height)
        return camera

    @staticmethod
    def _raygen(cameras: List[Camera]) -> List[Rays]:
        """ Generates a full ray pack for each camera view. """
        rays = []
        for idx in cameras:
            camera = cameras[idx]
            ray_grid = generate_centered_pixel_coords(camera.width, camera.height,
                                                      camera.width, camera.height,
                                                      device='cuda')
            _rays = generate_pinhole_rays(camera.to(ray_grid[0].device), ray_grid).reshape(
                camera.height, camera.width, 3).to('cpu')
            rays.append(_rays.to(dtype=torch.float32))
        return rays

    @staticmethod
    def _normalize(images: torch.Tensor, alphas: torch.Tensor, depths: torch.Tensor,
                   cameras: Dict[str, Camera], rays: List[Rays]):
        """ Normalizes the content of all views to fit within an axis aligned bounding box of [-1, 1]:
        1. RTMV scenes contain depth information, which is used to compute a point cloud from the depth rays.
        2. The pointcloud is normalized within the AABB of [-1, 1].
        3. The depth information, generated rays and cameras are rescaled according to normalization factors:
            coords_center, coords_scale.

        Returns:
            - (torch.Tensor) depths: the rescaled depth values of each ray
            - (wisp.core.Rays) rays: the rescaled rays
            - (torch.Tensor) coords_center: Value used to centeralize the point cloud around 0, 0, 0.
            - (torch.Tensor) coords_scale: Value used to scale the point cloud within [-1, 1].
        """
        coords, _ = create_pointcloud_from_images(images, alphas, rays, depths)
        normalized_coords, coords_center, coords_scale = normalize_pointcloud(coords, return_scale=True)

        depths = torch.stack(depths)
        rays = Rays.stack(rays)
        depths = depths * coords_scale
        rays.origins = (rays.origins - coords_center) * coords_scale
        depths = list(depths)
        rays = list(rays)

        for cam_id, cam in cameras.items():
            cam.translate(-coords_center.to(cam.dtype))
            cam.t = cam.t * coords_scale.to(cam.dtype)

        return depths, rays, coords_center, coords_scale

    def flatten_tensors(self) -> None:
        """ Flattens the cached data tensors to (NUM_VIEWS, NUM_RAYS, *).
        """
        num_imgs = len(self)
        self.data["rays"] = self.data["rays"].reshape(num_imgs, -1, 3)
        self.data["rgb"] = self.data["rgb"].reshape(num_imgs, -1, 3)
        self.data["depth"] = self.data["depth"].reshape(num_imgs, -1, 1)
        self.data["alpha"] = self.data["alpha"].reshape(num_imgs, -1, 1)

    @property
    def img_shape(self) -> torch.Size:
        """ Returns the shape of the rescaled dataset images (cached values are flattened) """
        return self._img_shape

    @property
    def cameras(self) -> List[Camera]:
        """ Returns the list of camera views used to generate rays for this dataset. """
        return self.data["cameras"]

    def as_pointcloud(self) -> torch.Tensor:
        """ If supports_depth=True, this function can be used to query the depth information in
        the form of a pointcloud tensor.
        """
        return self.coords

    def supports_depth(self) -> bool:
        """ Returns if this dataset have loaded depth information. """
        return self.coords is not None

    @property
    def num_images(self) -> int:
        """ Returns the number of views this dataset stores. """
        return self.data["rgb"].shape[0]

wisp/datasets/formats/rtmv_dataset.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from .ray_sampler import *

wisp/datasets/transforms/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from wisp.datasets.batch import MultiviewBatch


class SampleRays:
    """ A dataset transform for sub-sampling a fixed amount of rays. """
    def __init__(self, num_samples):
        self.num_samples = num_samples

    @torch.cuda.nvtx.range("SampleRays")
    def __call__(self, inputs: MultiviewBatch):
        device = inputs['rays'].origins.device
        ray_idx = torch.randint(0, inputs['rays'].shape[0], [self.num_samples], device=device)

        out = {}
        out['rays'] = inputs['rays'][ray_idx].contiguous()

        # Loop over ray values in this batch
        for channel_name, ray_value in inputs.ray_values().items():
            out[channel_name] = ray_value[ray_idx].contiguous()
        return out

wisp/datasets/transforms/ray_sampler.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import os
from typing import Callable, Optional, Type
import collections
import inspect
import torch
from torch._six import string_classes
from torch.utils.data._utils.collate import default_convert, default_collate_err_msg_format
from wisp.core import Rays
from wisp.datasets.base_datasets import WispDataset, MultiviewDataset, SDFDataset, BaseImageDataset
from wisp.datasets.batch import Batch


def load_multiview_dataset(dataset_path: str, dataset_num_workers: int = -1, transform: Callable = None,
                           split: str = None, **kwargs) -> MultiviewDataset:
    """
        A convenience method which loads the MultiviewDataset class which best matches the files under dataset_path.
        The implementation relies on the `WispDataset.is_root_of_dataset()` function being implemented by
        WispDataset implementations.
        Dataset classes are allowed to specify unique terms which set them apart from other datasets, i.e.:
        a dataset path with the following contents can be assumed to be attributed to the nerf-synethetic dataset:
            /path/to/dataset/transform.json
            /path/to/dataset/images/____.png

        Caveats:
        1. Where more than one dataset class matches the contents of dataset_path,
        this function will trigger a RuntimeError due to a non-resolved ambiguity.
        In such cases, it is recommended to construct the datasets explicitly rather than using this function.

        2. This function can only search for MultiviewDataset subclasses which were already loaded by python.
        When specifying custom MultiviewDataset classes which do not exist under this folder,
        make sure python imported them before invoking this function.

        Args:
            dataset_path (str): The root directory of the dataset, where dataset files should reside.
            dataset_num_workers (int): The number of workers to spawn for multiprocessed loading.
                If dataset_num_workers < 1, processing will take place on the main process.
            transform (Callable): Transform function applied per batch when data is accessed with __get_item__.
                For example: ray sampling, to filter the amount of rays returned per batch.
                When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
            split (str): The dataset split to use, may correspond to which files to load, or how to partion the data.
                Options: 'train', 'val', 'test'.
            **kwargs: Any other data specific arguments will be passed to the dataset class once it is matched.
                Excessive kwargs the matched dataset class does not use will be safely disregarded.
        """
    return _load_dataset(MultiviewDataset, dataset_path, dataset_num_workers, transform, split, **kwargs)

def load_image_dataset(dataset_path: str, dataset_num_workers: int = -1, transform: Callable = None,
                           split: str = None, **kwargs) -> MultiviewDataset:
    
    return _load_dataset(BaseImageDataset, dataset_path, dataset_num_workers, transform, split, **kwargs)

def _load_dataset(dataset_baseclass: Type[WispDataset],
                  dataset_path: str, dataset_num_workers: int = -1, transform: Callable = None,
                  split: str = None, **kwargs) -> WispDataset:
    """
    A convenience method which loads the dataset class which best matches the files under dataset_path.
    The implementation relies on the `WispDataset.is_root_of_dataset()` function being implemented by
    WispDataset implementations.
    Dataset classes are allowed to specify unique terms which set them apart from other datasets, i.e.:
    a dataset path with the following contents can be assumed to be attributed to the nerf-synethetic dataset:
        /path/to/dataset/transform.json
        /path/to/dataset/images/____.png

    Caveats:
    1. Where more than one dataset class matches the contents of dataset_path, this function will trigger a RuntimeError
    due to a non-resolved ambiguity. In such cases, it is recommended to construct the datasets explicitly rather
    than using this function.

    2. This function can only search for `dataset_baseclass` subclasses which were already loaded by python.
    When specifying custom WispDataset classes which do not exist under this folder,
    make sure python imported them before invoking this function.

    Args:
        dataset_baseclass (Type[WispDataset]): WispDataset or one of it's variants, to further scope which parts
            of the WispDataset class hierarchy are matched against the root directory.
        dataset_path (str): The root directory of the dataset, where dataset files should reside.
        dataset_num_workers (int): The number of workers to spawn for multiprocessed loading.
            If dataset_num_workers < 1, processing will take place on the main process.
        transform (Callable): Transform function applied per batch when data is accessed with __get_item__.
            For example: ray sampling, to filter the amount of rays returned per batch.
            When multiple transforms are needed, the transform callable may be a composition of multiple Callable.
        split (str): The dataset split to use, may correspond to which files to load, or how to partion the data.
            Options: 'train', 'val', 'test'.
        **kwargs: Any other data specific arguments will be passed to the dataset class once it is matched.
            Excessive kwargs the matched dataset class does not use will be safely disregarded.
    """
    from wisp.config_parser import get_args_for_function
    def __subclasses_hierarchy(cls): # Get hierarchy of subclasses of cls
        return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in __subclasses_hierarchy(c)])
    files_list = os.listdir(dataset_path)
    matching_dataset: Optional[Type[WispDataset]] = None
    for dataset in __subclasses_hierarchy(dataset_baseclass):
        if inspect.isabstract(dataset):
            continue
        is_match = dataset.is_root_of_dataset(root=dataset_path, files_list=files_list)
        if is_match:
            if matching_dataset is None:
                matching_dataset = dataset
            else:
                raise RuntimeError(f"load_dataset was given an ambiguous path which matches more than one dataset"
                                   f" class: {dataset} and {matching_dataset}. This is a result of "
                                   f"'is_root_of_dataset()' implementations which do not distinguish well enough"
                                   f"between the two dataset classes. A quick workaround is to load the desired "
                                   f"dataset class explicitly. Otherwise, update the dataset "
                                   f"'is_root_of_dataset()' logics. ")
    if matching_dataset is None:
        raise RuntimeError(f"load_dataset did match any dataset class which fits the contents of {dataset_path}. "
                           f"Is the dataset-path valid?")

    # Out of all given kwargs, select those which the dataset __init__ function uses
    ds_args = get_args_for_function(args=kwargs, func=matching_dataset.__init__)
    return matching_dataset(dataset_path=dataset_path, dataset_num_workers=dataset_num_workers,
                            transform=transform, split=split, **ds_args)


def default_collate(batch):
    r"""
    Function that extends torch.utils.data._utils.collate.default_collate
    to support custom wisp structures such as Rays and Batches.
    """
    elem = batch[0]
    elem_type = type(elem)
    if isinstance(elem, Rays):
        return Rays.cat(batch)
    elif isinstance(elem, Batch):
        # Assumes that if first element in batch has a certain None field, this field is None for all other entries too
        return elem_type(**{key: default_collate([d[key] for d in batch])
                            for key in elem if elem.get(key) is not None})
    if isinstance(elem, torch.Tensor):
        out = None
        if torch.utils.data.get_worker_info() is not None:
            # If we're in a background process, concatenate directly into a
            # shared memory tensor to avoid an extra copy
            numel = sum(x.numel() for x in batch)
            storage = elem.storage()._new_shared(numel, device=elem.device)
            out = elem.new(storage).resize_(len(batch), *list(elem.size()))
        return torch.stack(batch, 0, out=out)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':
            # array of string classes and object
            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                raise TypeError(default_collate_err_msg_format.format(elem.dtype))

            return default_collate([torch.as_tensor(b) for b in batch])
        elif elem.shape == ():  # scalars
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float64)
    elif isinstance(elem, int):
        return torch.tensor(batch)
    elif isinstance(elem, string_classes):
        return batch
    elif isinstance(elem, collections.abc.Mapping):
        try:
            return elem_type({key: default_collate([d[key] for d in batch]) for key in elem})
        except TypeError:
            # The mapping type may not support `__init__(iterable)`.
            return {key: default_collate([d[key] for d in batch]) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
        return elem_type(*(default_collate(samples) for samples in zip(*batch)))
    elif isinstance(elem, collections.abc.Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError('each element in list of batch should be of equal size')
        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.

        if isinstance(elem, tuple):
            return [default_collate(samples) for samples in transposed]  # Backwards compatibility.
        else:
            try:
                return elem_type([default_collate(samples) for samples in transposed])
            except TypeError:
                # The sequence type may not support `__init__(iterable)` (e.g., `range`).
                return [default_collate(samples) for samples in transposed]

    raise TypeError(default_collate_err_msg_format.format(elem_type))

wisp/datasets/utils.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .event import watch, watchedfields
from .state import *

wisp/framework/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from pydispatch import dispatcher

""" This module contains logic required to set up wisp's event handlers framework. """

# TODO(operel): Add "fire" event

##############################################
#            Public Handlers API             #
##############################################

def watch(watched_obj, field, status, handler):
    """ registers the handler for status updates on watched_obj.field.
        For example: watch(scene_status, "cam_controller", "changed", app.on_camera_controller_changed)
    """
    dispatcher.connect(handler, (status, field), sender=watched_obj)


def watchedfields(cls=None):
    """ Returns the class augmented with a custom __setattr__ implementation which notifies subscribers
        when class fields are updated.
    """
    def wrap(cls):
        return _register_func(cls)
    if cls is None:         # Called as @watchedfields()
        return wrap
    else:
        return wrap(cls)    # Called as @watchedfields

##############################################
#            watchedfields hooks             #
##############################################

def __setattr_notify__(obj, key, value):
    if hasattr(obj, key):
        prev_val = obj.__getattribute__(key) # First time this attribute is set
    else:
        prev_val = None
    obj._setattr(key, value)                 # Invoke internal setter here
    if prev_val != value:
        dispatcher.send(('changed', key), obj, value=value)

def _register_func(cls):
    # __setattr__ already explicitly defined, use it as internal setter implementation
    if '__setattr__' in cls.__dict__:
        setter_func = cls.__dict__['__setattr__']
    else:   # __setattr__ not defined, use the default implementation which simply sets the attribute
        def _setter_func(obj, key, value):
            obj.__dict__[key] = value
        setter_func = _setter_func
    setattr(cls, '_setattr', setter_func)
    setattr(cls, '__setattr__', __setattr_notify__)
    return cls


##############################################
#             watched iterables              #
##############################################

# TODO (operel): reconsider if we remove those, probably too slow but need to profile more carefully..

def wrap_iterable_fields(value):
    if isinstance(value, dict):
        return watcheddict(value)
    elif isinstance(value, list):
        return watchedlist(value)
    else:
        return value


class watcheddict(dict):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__class__ = type(dict.__name__, (self.__class__, dict), {})

    def __setitem__(self, item, value):
        super().__setitem__(item, value)
        dispatcher.send(('updated', self), self, value=item)


class watchedlist(list):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__class__ = type(list.__name__, (self.__class__, list), {})

    def __setitem__(self, item, value):
        super().__setitem__(item, value)
        dispatcher.send(('updated', self), self, value=item)

wisp/framework/event.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from __future__ import annotations
import torch
from typing import List, Dict, Type, DefaultDict, Tuple, TYPE_CHECKING
from collections import defaultdict
from dataclasses import dataclass, field
from kaolin.render.camera import Camera
from wisp.framework.event import watchedfields
from wisp.core import Channel, PrimitivesPack, ObjectTransform, channels_starter_kit
if TYPE_CHECKING:  # Prevent circular imports mess due to typed annotations
    from wisp.models import Pipeline
    from wisp.renderer.core.control import CameraControlMode
    from wisp.renderer.core import BottomLevelRenderer

"""
The WispState object is a shared resource between the various wisp components,
and is used to keep the various modules decoupled and synced on transient information, current configurations, and
state of the scene.

Wisp components share information and sync through the WispState subcomponents and their fields:
1. Wisp Components may write to certain fields directly, which are periodically read by other components.
   For example: The interactive renderer reads the InteractiveRendererState.canvas_dirty to determine when
   the canvas should be redrawn. Other components may request a redraw by setting this flag.
2. The @watchedfields decorator prompts an event when the state field values changes. 
   Components may register event handlers to immediately respond on such changes by using wisp.framework.event.watch
   For example: The interactive renderer listens on the OptimizationState.epoch field to redraw the canvas
   when an epoch ends.
   Note: Iterable fields will not respond to objects being added / removed.
   
Wisp applications may extend the state object via the WispState.extent field.
"""

@watchedfields
@dataclass
class InteractiveRendererState:
    """
    Holds settings directly used by the interactive renderer of wisp apps.
    """

    fps: float = 0.0
    """ FPS measured by renderer clocked """

    target_fps: float = None
    """ Ideal FPS the renderer should attempt to be on par with.
        If None, the ideal FPS is unbounded.
    """

    dt: float = 0.0
    """ Time elapsed since last renderer clock tick """

    canvas_width: int = 1600
    """ Default canvas width """

    canvas_height: int = 1200
    """ Default canvas height """

    clear_color_value: Tuple[float, float, float] = (0.0, 0.0, 0.0)
    """ Clear value for canvas background color channel, in RGB format """

    clear_depth_value: float = 1.0
    """ Clear value for canvas background depth channel """

    available_canvas_channels: List[str] = field(default_factory=list)
    """ List of available output channels to view on the canvas.
        This field should be set by the app, and should be compatible with the objects expected to be traced.
    """

    selected_canvas_channel: str = 'rgb'
    """ Currently selected output channel to view on the canvas.
        Choices: any of the channels listed in available_canvas_channels. 
    """

    canvas_dirty = False
    """ Setting this flag forces the renderer to redraw the canvas on the next frame. """

    cam_controller: Type[CameraControlMode] = None
    """ The current camera controller which manipulates the canvas camera.
        Choices: First Person, Trackball or Turntable. 
    """

    selected_camera: Camera = None
    """ The camera currently used by renderer to view the scene.
        If None, the renderer will create a default camera upon demand.
    """

    selected_camera_lens: str = 'perspective'
    """ Projection mode for selected camera used to view the canvas.
        Choices: 'perspective', 'orthographic'
    """

    interactive_mode: bool = False
    """ Whether the renderer is currently set in interactive mode: resolution may be adjusted to maintain FPS.
    """

    background_tasks_paused: bool = True
    """ When true, background tasks (i.e. optimizations) are paused. """

    antialiasing: str = "msaa_4x"
    """ Antialising method to use for the canvas.
        Due to the universal nature of the renderer, this setting may not affect object renderers, but
        mostly primitives drawn directly on the canvas with the graphics api
        (such as, i.e. vectorial layers drawn with OpenGL).
        Choices:
        'msaa_4x' - Framebuffer will be generated with 4x multisampling. 
                    More fragments will be generated to eliminate jagged borders, at the cost of additional computation. 
        'none' - No inherent antialising mode will be activated.
    """

    gl_version: str = "GL 3.3"
    """ Wisp applications rely on glumpy + OpenGL to render specific features and blit content to the window.
    This setting configures glumpy to load with a specific OpenGL backend.
    OpenGL 3.3 is widely supported and is therefore assumed to be the default.
    Users are free to bump this version to a higher number in case more advanced OpenGL capabilities are required.
    Format: (api major.minor. profile)
    For example: "GL 3.3 core"  
    """

    reference_grids: List[str] = field(default_factory=list)
    """ List of world grids to use as reference planes, for both rendering and some camera controllers.
        Choices: Any combination of the values: 'xy', 'xz', 'yz'. An empty list will turn off the grid.
    """

    device: torch.device = 'cpu'
    """ Default device for interactive renderer and bottom level renderers to use """

    enable_amp: bool = True
    """ Enables mixed precision with torch.cuda.amp.autocast on the interactive renderer's redraw() and render()
    functions. By default, this setting is enabled to allow for faster optimized rendering.
    Users with custom tracers or neural fields may opt to turn this off if for some reason their pipelines do not
    support mixed precision.  
    """

@watchedfields
@dataclass
class BottomLevelRendererState:
    """
    Holds the settings of a single Bottom-Level renderer.
    Wisp supports joint rendering of various pipelines (NeRF, SDFs, meshes, and so forth),
    where each pipeline is wrapped by a bottom level renderer configured by this state object.
    The state object exists throughout the lifecycle of the renderer / pipeline,
    and is used to determine how the bottom level renderer of an existing pipeline is constructed.
    """

    renderer: BottomLevelRenderer = None
    """ The instance of the bottom level renderer. Initialized only if the status field is 'loaded' """

    data_layers: Dict[str, PrimitivesPack] = field(default_factory=dict)  # layer id -> prims
    toggled_data_layers: Dict[str, bool] = field(default_factory=dict)    # layer id -> bool

    transform: ObjectTransform = ObjectTransform()
    """ The object transform maintains the 4x4 model matrix, which describes the object transformation from
        local object coordinates to world space.
        Manipulating the transform results in moving, scaling and orienting the object.
    """

    ### Lifecycle fields ###

    status: str = 'pending'
    """ The current status of this renderer instance:
        'pending' - the renderer is part of the scene graph but was not initialized yet.
        'loaded' - the renderer is part of the scene graph, and the instance was initialized (renderer field exists).
        'ignored' - this renderer should be excluded from the scene graph.  
    """

    setup_args: Dict[str, object] = field(default_factory=dict)
    """ Optional args, used to load or construct the renderer instance """


@watchedfields
@dataclass
class SceneGraphState:
    """
    Holds the entire scene graph of objects.
    """
    neural_pipelines: Dict[str, Pipeline] = field(default_factory=dict)
    """ Wisp objects are represented by neural pipelines which pair a BaseNeuralField and a BaseTracer.
        A BottomLevelRenderer knows how to invoke the pipeline to render the nef and obtain a 
        Renderbuffer of the object. 
    """

    bl_renderers: Dict[str, BottomLevelRendererState] = field(default_factory=dict)
    """ Renderable wisp objects are represented by BottomLevelRenderers which wrap neural pipelines.
        A BottomLevelRenderer knows how to invoke the rendering pipeline of some object type and obtain
        a Renderbuffer of the object. For example, a BottomLevelRenderer may wrap a NeRF object and its RFTracer.
    """

    cameras: Dict[str, Camera] = field(default_factory=dict)
    """ Set of cameras currently available for this scene. """

    visible_objects: Dict[str, bool] = field(default_factory=dict)
    """ Which objects are currently toggled on as visible in the scene.
        This mapping includes all object types: bl_renderers, cameras, etc.
    """

    channels: Dict[str, Channel] = field(default_factory=channels_starter_kit)
    """ The channels kit containing information per channel, such as blending function, normalization function, etc.
    """

@watchedfields
@dataclass
class OptimizationState:
    """
    Holds information about an optimization task wisp is currently running.
    This is used to cache information in trainers and share them with other components like the renderer.
    """

    running: bool = False
    """ Is the optimization currently running or paused / stopped. """

    validation_enabled: bool = False
    """ If true, validation will run every N epochs. """

    epoch: int = 0
    """ Current epoch of the optimization. An epoch is task / dataset dependent, but is normally
        equivalent to a single pass over a finite dataset.
    """

    iteration: int = 0
    """ Current iteration within an epoch. An iteration is task / dataset dependent, but usually refers
        to a step of the trainer on a single batch.
    """

    max_epochs: int = 0
    """ Total number of epochs set for this optimization task. """

    iterations_per_epoch: int = 0
    """ Total iterations used within each epoch / the current epoch. """

    elapsed_time: float = 0
    """ Total time elapsed for running the optimization. Measured in seconds. """

    lr: float = 0.0
    """ The learning rate currently set by the optimizaer. """

    losses: DefaultDict[str, List[float]] = field(default_factory=lambda: defaultdict(list))
    """ Losses currently reported by this optimization task (name to list of values per epoch) """

    metrics: DefaultDict[str, List[float]] = field(default_factory=lambda: defaultdict(float))
    """ Metrics currently reported by this optimization task (name to list of values per epoch) """

    train_data: List[torch.utils.data.Dataset] = field(default_factory=list)
    """ List of training sets employed by this optimization task """

    validation_data: List[torch.utils.data.Dataset] = field(default_factory=list)
    """ List of validation sets employed by this optimization task """


@dataclass
class WispState:
    """ Holds all information shared between components of a single wisp app. """

    renderer: InteractiveRendererState = InteractiveRendererState()
    """ Holds settings directly used by the interactive renderer of wisp apps. """

    graph: SceneGraphState = SceneGraphState()
    """ Holds the entire scene graph of objects. """

    optimization: OptimizationState = OptimizationState()
    """ Holds information about an optimization task wisp is currently running. """

    extent: Dict[str, object] = field(default_factory=dict)
    """ Extensible field: custom wisp applications information can be added here. """



@dataclass
class ImageOptimizationState:
    """
    Holds information about an optimization task wisp is currently running.
    This is used to cache information in trainers.
    """

    running: bool = False
    """ Is the optimization currently running or paused / stopped. """

    validation_enabled: bool = False
    """ If true, validation will run every N epochs. """

    epoch: int = 0
    """ Current epoch of the optimization. An epoch is task / dataset dependent, but is normally
        equivalent to a single pass over a finite dataset.
    """

    iteration: int = 0
    """ Current iteration within an epoch. An iteration is task / dataset dependent, but usually refers
        to a step of the trainer on a single batch.
    """

    max_epochs: int = 0
    """ Total number of epochs set for this optimization task. """

    iterations_per_epoch: int = 0
    """ Total iterations used within each epoch / the current epoch. """

    elapsed_time: float = 0
    """ Total time elapsed for running the optimization. Measured in seconds. """

    lr: float = 0.0
    """ The learning rate currently set by the optimizaer. """

    losses: DefaultDict[str, List[float]] = field(default_factory=lambda: defaultdict(list))
    """ Losses currently reported by this optimization task (name to list of values per epoch) """

    metrics: DefaultDict[str, List[float]] = field(default_factory=lambda: defaultdict(list))
    """ Metrics currently reported by this optimization task (name to list of values per epoch) """


@dataclass
class ImageDataState:
    """
    Holds information about the data, task wisp is currently running.
    """

    train_data: List[torch.utils.data.Dataset] = field(default_factory=list)
    """ List of training sets employed by this optimization task """

    validation_data: List[torch.utils.data.Dataset] = field(default_factory=list)
    """ List of validation sets employed by this optimization task """

@dataclass
class ImageState:

    """ Holds information about the image optimization task wisp is currently running. """
    optimization: ImageOptimizationState = field(default_factory=ImageOptimizationState)

    """ Holds information about the data wisp is currently running on. """
    data: ImageDataState = field(default_factory=ImageDataState)
wisp/framework/state.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .datalayers import *

wisp/gfx/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .datalayers import Datalayers
from .camera_datalayers import CameraDatalayers
from .octree_datalayers import OctreeDatalayers
from .aabb_datalayers import AABBDatalayers

wisp/gfx/datalayers/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from typing import Dict
import kaolin.ops.spc as spc_ops
from wisp.core import PrimitivesPack
from wisp.accelstructs import AxisAlignedBBoxAS
from wisp.gfx.datalayers import Datalayers
from wisp.core.colors import soft_blue, soft_red, lime_green, purple, gold


class AABBDatalayers(Datalayers):

    def __init__(self):
        self._last_state = dict()

    def needs_redraw(self, blas: AxisAlignedBBoxAS) -> True:
        return True

    def regenerate_data_layers(self, blas: AxisAlignedBBoxAS) -> Dict[str, PrimitivesPack]:
        data_layers = dict()
        color_tensor = torch.tensor((*soft_blue, 1.0))

        cells = PrimitivesPack()
        lod = 0
        level_points = spc_ops.unbatched_get_level_points(blas.points, blas.pyramid, 0)
        corners = spc_ops.points_to_corners(level_points) / (2 ** lod)
        corners = corners * 2.0 - 1.0
        grid_lines = corners[:, [(0, 1), (1, 3), (3, 2), (2, 0),
                                 (4, 5), (5, 7), (7, 6), (6, 4),
                                 (0, 4), (1, 5), (2, 6), (3, 7)]]

        grid_lines_start = grid_lines[:, :, 0].reshape(-1, 3)
        grid_lines_end = grid_lines[:, :, 1].reshape(-1, 3)
        grid_lines_color = color_tensor.repeat(grid_lines_start.shape[0], 1)
        cells.add_lines(grid_lines_start, grid_lines_end, grid_lines_color)

        data_layers[f'AABB'] = cells
        return data_layers

wisp/gfx/datalayers/aabb_datalayers.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from wisp.core.colors import white, black, red, green, blue
from typing import Dict
from kaolin.render.camera import Camera
from wisp.core.primitives import PrimitivesPack
from wisp.gfx.datalayers import Datalayers


class CameraDatalayers(Datalayers):
    def needs_redraw(self) -> bool:
        raise True

    def regenerate_data_layers(self, cameras: Dict[str, Camera],
                               bg_color: torch.Tensor = None) -> Dict[str, PrimitivesPack]:
        if bg_color == white:
            color = torch.tensor((*black, 1.0))
        else:
            color = torch.tensor((*white, 1.0))

        red_rgba = torch.tensor((*red, 1.0))
        green_rgba = torch.tensor((*green, 1.0))
        blue_rgba = torch.tensor((*blue, 1.0))

        layers = dict()

        for camera_id, camera in cameras.items():
            camera_outlines = PrimitivesPack()
            origin = camera.cam_pos().squeeze()
            forward = camera.cam_forward().squeeze()
            right = camera.cam_right().squeeze()
            up = camera.cam_up().squeeze()
            forward_length = 0.2
            right_length = 0.2
            up_length = 0.2
            corner1 = origin - forward * forward_length + right * right_length + up * up_length
            corner2 = origin - forward * forward_length - right * right_length + up * up_length
            corner3 = origin - forward * forward_length + right * right_length - up * up_length
            corner4 = origin - forward * forward_length - right * right_length - up * up_length
            camera_outlines.add_lines(start=origin, end=corner1, color=color)
            camera_outlines.add_lines(start=origin, end=corner2, color=color)
            camera_outlines.add_lines(start=origin, end=corner3, color=color)
            camera_outlines.add_lines(start=origin, end=corner4, color=color)
            camera_outlines.add_lines(start=corner1, end=corner2, color=color)
            camera_outlines.add_lines(start=corner3, end=corner4, color=color)
            camera_outlines.add_lines(start=corner2, end=corner4, color=color)
            camera_outlines.add_lines(start=corner1, end=corner3, color=color)

            # Draw axes above camera
            above = origin - forward * forward_length + up * up_length * 1.1
            cam_axes = camera.R[0]  # Returns the camera axes as rows
            permuted_axes = cam_axes.T @ camera.extrinsics.basis_change_matrix
            cam_right = permuted_axes[:, 0]
            cam_up = permuted_axes[:, 1]
            cam_forward = permuted_axes[:, 2]
            camera_outlines.add_lines(start=above, end=above + cam_right * 0.1, color=red_rgba)
            camera_outlines.add_lines(start=above, end=above + cam_up * 0.1, color=green_rgba)
            camera_outlines.add_lines(start=above, end=above + cam_forward * 0.1, color=blue_rgba)
            layers[camera_id] = camera_outlines
        return layers


wisp/gfx/datalayers/camera_datalayers.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from abc import ABC, abstractmethod
from typing import Dict
from wisp.core import PrimitivesPack

class Datalayers(ABC):

    @abstractmethod
    def needs_redraw(self, *args, **kwargs) -> bool:
        raise NotImplementedError

    @abstractmethod
    def regenerate_data_layers(self, *args, **kwargs) -> Dict[str, PrimitivesPack]:
        raise NotImplementedError

wisp/gfx/datalayers/datalayers.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from typing import Dict
import kaolin.ops.spc as spc_ops
from wisp.core import PrimitivesPack
from wisp.accelstructs import OctreeAS
from wisp.gfx.datalayers import Datalayers
from wisp.core.colors import soft_blue, soft_red, lime_green, purple, gold


class OctreeDatalayers(Datalayers):

    def __init__(self):
        self._last_state = dict()

    def needs_redraw(self, blas: OctreeAS) -> True:
        # Pyramids contain information about the number of cells per level,
        # it's a plausible heuristic to determine whether the frame should be redrawn
        return not ('pyramids' in self._last_state and
                    torch.equal(self._last_state['pyramids'], blas.pyramid))

    def regenerate_data_layers(self, blas: OctreeAS) -> Dict[str, PrimitivesPack]:
        data_layers = dict()
        lod_colors = [
            torch.tensor((*soft_blue, 1.0)),
            torch.tensor((*soft_red, 1.0)),
            torch.tensor((*lime_green, 1.0)),
            torch.tensor((*purple, 1.0)),
            torch.tensor((*gold, 1.0)),
        ]

        for lod in range(blas.max_level):
            cells = PrimitivesPack()
            
            level_points = spc_ops.unbatched_get_level_points(blas.points, blas.pyramid, lod)
            
            corners = spc_ops.points_to_corners(level_points) / (2 ** lod)
            
            corners = corners * 2.0 - 1.0

            grid_lines = corners[:, [(0, 1), (1, 3), (3, 2), (2, 0),
                                     (4, 5), (5, 7), (7, 6), (6, 4),
                                     (0, 4), (1, 5), (2, 6), (3, 7)]]

            grid_lines_start = grid_lines[:, :, 0].reshape(-1, 3)
            grid_lines_end = grid_lines[:, :, 1].reshape(-1, 3)
            color_tensor = lod_colors[lod % len(lod_colors)]
            grid_lines_color = color_tensor.repeat(grid_lines_start.shape[0], 1)
            cells.add_lines(grid_lines_start, grid_lines_end, grid_lines_color)

            data_layers[f'Octree LOD{lod}'] = cells

        self._last_state['pyramids'] = blas.pyramid
        return data_layers

wisp/gfx/datalayers/octree_datalayers.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .pipeline import Pipeline

wisp/models/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .basic_activations import *

wisp/models/activations/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import torch.nn as nn

class FullSort(nn.Module):
    """The "FullSort" activation function from https://arxiv.org/abs/1811.05381.
    """
    def forward(self, x):
        """Sorts the feature dimension.
        
        Args:
            x (torch.FloatTensor): Some tensor of shape [..., feature_size]
        
        Returns:
            (torch.FloatTensor): Activation of shape [..., feature_size]
        """
        return torch.sort(x, dim=-1)[0]

class MinMax(nn.Module):
    """The "MinMax" activation function from https://arxiv.org/abs/1811.05381.
    """
    def forward(self, x):
        """Partially sorts the feature dimension.
        
        The feature dimension needs to be a multiple of 2.

        Args:
            x (torch.FloatTensor): Some tensor of shape [batch, feature_size]
        
        Returns:
            (torch.FloatTensor): Activation of shape [batch, feature_size]
        """
        N, M = x.shape
        x = x.reshape(N, M//2, 2)
        return torch.cat([x.min(-1, keepdim=True)[0], x.max(-1, keepdim=True)[0]], dim=-1).reshape(N, M)

class Identity(nn.Module):
    """Identity function. Occasionally useful.
    """
    def forward(self, x):
        """Returns the input. :)

        Args:
            x (Any): Anything

        Returns:
            (Any): The input!
        """
        return x

class Clamp(torch.nn.Module):
    """Clamping module.
    Args:
        min (float): Minimum value.
        max (float): Maximum value.
    """
    def __init__(self, min=None, max=None):
        super().__init__()
        self.min = min
        self.max = max

    def forward(self, x):
        return torch.clamp(x,min=self.min,max=self.max)
    
class SineScaled(torch.nn.Module):
    """Sine activation with scaling.
    Args:
        w0 (float): Omega_0 parameter.
    """
    def __init__(self, w0=1.):
        super().__init__()
        self.w0 = w0

    def forward(self, x):
        return torch.sin(self.w0 * x)


def get_activation_class(activation_type):
    """Utility function to return an activation function class based on the string description.

    Args:
        activation_type (str): The name for the activation function.
    
    Returns:
        (Function): The activation function to be used. 
    """
    if activation_type == 'none' or activation_type == 'identity':
        return Identity()
    elif activation_type == 'fullsort':
        return FullSort()
    elif activation_type == 'minmax':
        return MinMax()
    elif activation_type == 'relu':
        return torch.relu
    elif activation_type == 'sin':
        return torch.sin
    elif activation_type == 'sinescaled':
        return SineScaled(30.0)
    elif activation_type == 'sigmoid':
        return torch.sigmoid
    else:
        assert False and "activation type does not exist"

wisp/models/activations/basic_activations.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .basic_conditioners import *

wisp/models/conditioners/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

#TODO(ttakikawa): Either use or remove this

def position(position, features, layers, activation):
    """Use the position as input (i.e. no conditioning)

    Args:
        position   : [N, ..., d] tensor of coordinates
        features   : [N, ..., f] tensor of features
        layers     : nn.ModuleList of layers
        activation : activation function
    """
    
    h = position

    for i, l in enumerate(layers):
        h = activation(l(h))
    
    return h

def feature(position, features, layers, activation):
    """Use the features as input.

    Args:
        position   : [N, ..., d] tensor of coordinates
        features   : [N, ..., f] tensor of features
        layers     : nn.ModuleList of layers
        activation : activation function
    """
    
    h = features

    for i, l in enumerate(layers):
        h = activation(l(h))
    
    return h

def concat(position, features, layers, activation):
    """Concatenates the input onto the features, and then feeds into the input of the neural network.
    
    Args:
        position   : [N, ..., d] tensor of coordinates
        features   : [N, ..., f] tensor of features
        layers     : nn.ModuleList of layers
        activation : activation function
    """
    
    h = torch.cat([position, features], dim=-1)

    for i, l in enumerate(layers):
        h = activation(l(h))
    return h

def film_linear(position, features, layers, activation):
    """Applies film conditioning (multiply only) on the network.

    Args:
        position   : [N, ..., d] tensor of coordinates
        features   : [N, ..., f] tensor of features
        layers     : nn.ModuleList of layers
        activation : activation function
    """
    feature_shape = features.shape[:-1]
    feature_dim = features.shape[-1]
    num_hidden = len(layers)
    # Maybe add assertion here... but if it errors, your feature_dim size is wrong
    features = features.reshape(features_shape, num_hidden, feature_dim // num_hidden)
    
    h = position

    for i, l in enumerate(layers):
        # Maybe also add another assertion here
        h = activation(l(h) * features[..., i, :])
    return h

def film_translate(position, features, layers, activation):
    """Applies film conditioning (add only) on the network.

    Args:
        position   : [N, ..., d] tensor of coordinates
        features   : [N, ..., f] tensor of features
        layers     : nn.ModuleList of layers
        activation : activation function
    """
    feature_shape = features.shape[:-1]
    feature_dim = features.shape[-1]
    num_hidden = len(layers)
    # Maybe add assertion here... but if it errors, your feature_dim size is wrong
    features = features.reshape(features_shape, num_hidden, feature_dim // num_hidden)
    
    h = position

    for i, l in enumerate(layers):
        # Maybe also add another assertion here
        h = activation(l(h) + features[..., i, :])
    return h

def film(position, features, layers, activation):
    """Applies film conditioning (add only) on the network.

    Args:
        position   : [N, ..., d] tensor of coordinates
        features   : [N, ..., f] tensor of features
        layers     : nn.ModuleList of layers
        activation : activation function
    """
    feature_shape = features.shape[:-1]
    feature_dim = features.shape[-1]
    num_hidden = len(layers)
    # Maybe add assertion here... but if it errors, your feature_dim size is wrong
    features = features.reshape(features_shape, 2, num_hidden, feature_dim // num_hidden)
    
    h = position

    for i, l in enumerate(layers):
        # Maybe also add another assertion here
        h = activation(l(h) * features[..., 0, i, :] + features[..., 1, i, :])
    return h




wisp/models/conditioners/basic_conditioners.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .basic_decoders import *

wisp/models/decoders/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from typing import Dict, Any
import torch
import numpy as np
import torch.nn as nn
from wisp.core import WispModule

from scipy.stats import ortho_group

class BasicDecoder(WispModule):
    """Super basic but super useful MLP class.
    """
    def __init__(self, 
        input_dim, 
        output_dim, 
        activation,
        bias,
        layer = nn.Linear,
        num_layers = 1, 
        hidden_dim = 128, 
        skip       = []
    ):
        """Initialize the BasicDecoder.

        Args:
            input_dim (int): Input dimension of the MLP.
            output_dim (int): Output dimension of the MLP.
            activation (function): The activation function to use.
            bias (bool): If True, use bias.
            layer (nn.Module): The MLP layer module to use.
            num_layers (int): The number of hidden layers in the MLP.
            hidden_dim (int): The hidden dimension of the MLP.
            skip (List[int]): List of layer indices where the input dimension is concatenated.

        Returns:
            (void): Initializes the class.
        """
        super().__init__()
        
        self.input_dim = input_dim
        self.output_dim = output_dim        
        self.activation = activation
        self.bias = bias
        self.layer = layer
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.skip = skip
        if self.skip is None:
            self.skip = []
        
        self.make()

    def make(self):
        """Builds the actual MLP.
        """
        layers = []
        for i in range(self.num_layers):
            if i == 0: 
                layers.append(self.layer(self.input_dim, self.hidden_dim, bias=self.bias))
            elif i in self.skip:
                layers.append(self.layer(self.hidden_dim+self.input_dim, self.hidden_dim, bias=self.bias))
            else:
                layers.append(self.layer(self.hidden_dim, self.hidden_dim, bias=self.bias))
        self.layers = nn.ModuleList(layers)
        self.lout = self.layer(self.hidden_dim, self.output_dim, bias=self.bias)

    def forward(self, x, return_h=False):
        """Run the MLP!

        Args:
            x (torch.FloatTensor): Some tensor of shape [batch, ..., input_dim]
            return_h (bool): If True, also returns the last hidden layer.

        Returns:
            (torch.FloatTensor, (optional) torch.FloatTensor):
                - The output tensor of shape [batch, ..., output_dim]
                - The last hidden layer of shape [batch, ..., hidden_dim]
        """
        N = x.shape[0]
        for i, l in enumerate(self.layers):
            if i == 0:
                h = self.activation(l(x))
            elif i in self.skip:
                h = self.activation(l(h))
                h = torch.cat([x, h], dim=-1)
            else:
                h = self.activation(l(h))
        
        out = self.lout(h)
        
        if return_h:
            return out, h
        else:
            return out

    def initialize(self, get_weight):
        """Initializes the MLP layers with some initialization functions.

        Args:
            get_weight (function): A function which returns a matrix given a matrix.

        Returns:
            (void): Initializes the layer weights.
        """
        ms = []
        for i, w in enumerate(self.layers):
            m = get_weight(w.weight)
            ms.append(m)
        for i in range(len(self.layers)):
            self.layers[i].weight = nn.Parameter(ms[i])
        m = get_weight(self.lout.weight)
        self.lout.weight = nn.Parameter(m)

    def name(self) -> str:
        """ A human readable name for the given wisp module. """
        return "BasicDecoder"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        return {
            "Input Dim": self.input_dim,
            "Hidden Dim": self.hidden_dim,
            "Outpt Dim": self.output_dim,
            "Num. Layers": self.num_layers,
            "Layer Type": self.layer.__name__,
            "Activation": self.activation.__name__,
            "Bias": self.bias,
            "Skip Connections": self.skip,
        }


def orthonormal(weight):
    """Initialize the layer as a random orthonormal matrix.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N]. Only used for the shape.

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    m = ortho_group.rvs(dim=max(weight.shape))
    #m = np.dot(m.T, m)
    m = m[:weight.shape[0],:weight.shape[1]]
    return torch.from_numpy(m).float()

def svd(weight):
    """Initialize the layer with the U,V of SVD.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    U,S,V = torch.svd(weight)
    return torch.matmul(U, V.T)

def spectral_normalization(weight):
    """Initialize the layer with spectral normalization.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    U,S,V = torch.svd(weight)
    return weight / S.max()

def identity(weight):
    """Initialize the layer with identity matrix.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    return torch.diag(torch.ones(weight.shape[0]))

def average(weight):
    """Initialize the layer by normalizing the weights.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    return weight / weight.sum()


wisp/models/decoders/basic_decoders.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .positional_embedder import *

wisp/models/embedders/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from typing import Dict, Any
import torch
import torch.nn as nn
from wisp.core import WispModule


class PositionalEmbedder(WispModule):
    """PyTorch implementation of regular positional embedding, as used in the original NeRF and Transformer papers.
    """
    def __init__(self, num_freq, max_freq_log2, log_sampling=True, include_input=True, input_dim=3):
        """Initialize the module.

        Args:
            num_freq (int): The number of frequency bands to sample. 
            max_freq_log2 (int): The maximum frequency.
                                 The bands will be sampled at regular intervals in [0, 2^max_freq_log2].
            log_sampling (bool): If true, will sample frequency bands in log space.
            include_input (bool): If true, will concatenate the input.
            input_dim (int): The dimension of the input coordinate space.

        Returns:
            (void): Initializes the encoding.
        """
        super().__init__()

        self.num_freq = num_freq
        self.max_freq_log2 = max_freq_log2
        self.log_sampling = log_sampling
        self.include_input = include_input
        self.out_dim = 0
        if include_input:
            self.out_dim += input_dim

        if self.log_sampling:
            self.bands = 2.0**torch.linspace(0.0, max_freq_log2, steps=num_freq)
        else:
            self.bands = torch.linspace(1, 2.0**max_freq_log2, steps=num_freq)

        # The out_dim is really just input_dim + num_freq * input_dim * 2 (for sin and cos)
        self.out_dim += self.bands.shape[0] * input_dim * 2
        self.bands = nn.Parameter(self.bands).requires_grad_(False)
    
    def forward(self, coords):
        """Embeds the coordinates.

        Args:
            coords (torch.FloatTensor): Coordinates of shape [N, input_dim]

        Returns:
            (torch.FloatTensor): Embeddings of shape [N, input_dim + out_dim] or [N, out_dim].
        """
        N = coords.shape[0]
        winded = (coords[:,None] * self.bands[None,:,None]).reshape(
            N, coords.shape[1] * self.num_freq)
        encoded = torch.cat([torch.sin(winded), torch.cos(winded)], dim=-1)
        if self.include_input:
            encoded = torch.cat([coords, encoded], dim=-1)
        return encoded

    def name(self) -> str:
        """ A human readable name for the given wisp module. """
        return "Positional Encoding"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        return {
            "Output Dim": self.out_dim,
            "Num. Frequencies": self.num_freq,
            "Max Frequency": f"2^{self.max_freq_log2}",
            "Include Input": self.include_input
        }


def get_positional_embedder(frequencies, input_dim=3, include_input=True):
    """Utility function to get a positional encoding embedding.

    Args:
        frequencies (int): The number of frequencies used to define the PE:
            [2^0, 2^1, 2^2, ... 2^(frequencies - 1)].
        input_dim (int): The input coordinate dimension.
        include_input (bool): If true, will concatenate the input coords.

    Returns:
        (nn.Module, int):
        - The embedding module
        - The output dimension of the embedding.
    """
    encoder = PositionalEmbedder(frequencies, frequencies-1, input_dim=input_dim, include_input=include_input)
    return encoder, encoder.out_dim

wisp/models/embedders/positional_embedder.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .blas_grid import *
from .octree_grid import *
from .codebook_grid import *
from .hash_grid import * 
from .triplanar_grid import *
from .latent_grid import *

wisp/models/grids/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from abc import ABC, abstractmethod
from typing import Dict, Any, Set, Type
from wisp.core import WispModule
from wisp.accelstructs import BaseAS, ASQueryResults, ASRaytraceResults, ASRaymarchResults


class BLASGrid(WispModule, ABC):
    """
    BLASGrids (commonly referred in documentation as simply "grids"), represent feature grids in Wisp.
    BLAS: "Bottom Level Acceleration Structure", to signify this structure is the backbone that captures
    a neural field's contents, in terms of both features and occupancy for speeding up queries.

    This is an abstract base class that uses some spatial acceleration structure under the hood, to speed up operations
    such as coordinate based queries or ray tracing.
    Classes which inherit the BLASGrid are generally compatible with BaseTracers to support such operations
    (see: raymarch(), raytrace(), query()).

    Grids are usually employed as building blocks within neural fields (see: BaseNeuralField),
    possibly paired with decoders to form a neural field.
    """

    def __init__(self, blas: BaseAS):
        """ BLASGrids are generally assumed to contain bottom level acceleration structures. """
        super().__init__()
        self.blas = blas

    def raymarch(self, *args, **kwargs) -> ASRaymarchResults:
        """By default, this function will use the equivalent BLAS function unless overridden for custom behaviour.
        """
        return self.blas.raymarch(*args, **kwargs)

    def raytrace(self, *args, **kwargs) -> ASRaytraceResults:
        """By default, this function will use the equivalent BLAS function unless overridden for custom behaviour.
        """
        return self.blas.raytrace(*args, **kwargs)

    def query(self, *args, **kwargs) -> ASQueryResults:
        """By default, this function will use the equivalent BLAS function unless overridden for custom behaviour.
        """
        return self.blas.query(*args, **kwargs)

    @abstractmethod
    def interpolate(self, coords, lod_idx):
        """ Interpolates a feature value for the given coords using the grid support, in the given lod_idx
        Args:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3] or [batch, 3]
            lod_idx  (int): int specifying the index to the desired level of detail, if supported.
        """
        raise NotImplementedError('A BLASGrid should implement the interpolation functionality according to '
                                  'the grid structure.')

    def supported_blas(self) -> Set[Type[BaseAS]]:
        """ Returns a set of bottom-level acceleration structures this grid type supports """
        return set()

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.

        BLASGrids are generally assumed to contain a bottom level acceleration structure.
        """
        return {
            "Acceleration Structure": self.blas
        }

wisp/models/grids/blas_grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Dict, Any
import logging as log
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import kaolin.ops.spc as spc_ops
import wisp.ops.spc as wisp_spc_ops
from wisp.models.grids.octree_grid import OctreeGrid
from wisp.accelstructs import OctreeAS


class CodebookOctreeGrid(OctreeGrid):
    """This is a multiresolution feature grid where the octree stores indices into a fixed size codebook.
    """
    def __init__(
        self,
        accelstruct,
        feature_dim         : int,
        base_lod            : int,
        num_lods            : int          = 1,
        interpolation_type  : str          = 'linear',
        multiscale_type     : str          = 'cat',
        feature_std         : float        = 0.0,
        feature_bias        : float        = 0.0,
        codebook_bitwidth   : int          = 8
    ):
        """
        Args:
            accelstruct: Spatial acceleration structure which tracks the occupancy state of this grid.
                         Used to speed up spatial queries and ray tracing operations.
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
        """
        self.bitwidth = codebook_bitwidth
        super().__init__(accelstruct=accelstruct,
                         feature_dim=feature_dim,
                         base_lod=base_lod,
                         num_lods=num_lods,
                         interpolation_type=interpolation_type,
                         multiscale_type=multiscale_type,
                         feature_std=feature_std,
                         feature_bias=feature_bias)

    @classmethod
    def make_dense(cls,
                   feature_dim        : int,
                   base_lod           : int,
                   num_lods           : int          = 1,
                   interpolation_type : str          = 'linear',
                   multiscale_type    : str          = 'cat',
                   feature_std        : float        = 0.0,
                   feature_bias       : float        = 0.0,
                   codebook_bitwidth   : int         = 8) -> CodebookOctreeGrid:
        """Builds a fully occupied CodebookOctreeGrid with `base_lod + num_lods - 1` levels.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth

        Returns:
            (CodebookOctreeGrid): A new instance of an OctreeGrid at full occupancy.
        """
        max_lod = OctreeGrid.max_octree_lod(base_lod, num_lods)
        blas = OctreeAS.make_dense(level=max_lod)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth)
        return grid

    @classmethod
    def from_mesh(cls,
                  mesh_path: str,
                  num_samples_on_mesh: int,
                  feature_dim: int,
                  base_lod: int,
                  num_lods: int = 1,
                  interpolation_type: str = 'linear',
                  multiscale_type: str = 'cat',
                  feature_std: float = 0.0,
                  feature_bias: float = 0.0,
                  codebook_bitwidth: int = 8) -> CodebookOctreeGrid:
        """Builds the CodebookOctreeGrid from point samples over mesh faces.
        Samples will be quantized to populate the octree cells.
        This method is not guaranteed to reproduce a perfect mesh structure without "holes",
        but achieves a very good approximation with high probability.

        Args:
            mesh_path (str): Path to an OBJ file with a mesh to initialize the octree (only supports OBJ for now).
            num_samples_on_mesh (int): The number of samples to be generated on the mesh surface.
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth

        Returns:
            (CodebookOctreeGrid): A new instance of an CodebookOctreeGrid with occupancy initialized from
            samples over the mesh faces.
        """
        max_lod = OctreeGrid.max_octree_lod(base_lod, num_lods)
        blas = OctreeAS.from_mesh(mesh_path, level=max_lod, sample_tex=False, num_samples=num_samples_on_mesh)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth)
        return grid

    @classmethod
    def from_pointcloud(cls,
                        pointcloud: torch.FloatTensor,
                        feature_dim: int,
                        base_lod: int,
                        num_lods: int = 1,
                        interpolation_type: str = 'linear',
                        multiscale_type: str = 'cat',
                        feature_std: float = 0.0,
                        feature_bias: float = 0.0,
                        codebook_bitwidth: int = 8) -> CodebookOctreeGrid:
        """Builds the CodebookOctreeGrid, initializing it with a pointcloud.
        The cells occupancy will be determined by points occupying the octree cells.

        Args:
            pointcloud (torch.FloatTensor): 3D coordinates of shape [num_coords, 3] in normalized space [-1, 1].
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth

        Returns:
            (CodebookOctreeGrid): A new instance of an CodebookOctreeGrid with occupancy initialized from the pointcloud.
        """
        max_lod = OctreeGrid.max_octree_lod(base_lod, num_lods)
        blas = OctreeAS.from_pointcloud(pointcloud, level=max_lod)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth)
        return grid

    def init_feature_structure(self):
        """ Initializes everything related to the features stored in the codebook octree structure. """
        # Assumes the occupancy structure have been initialized (the BLAS: Bottom Level Accelerated Structure).
        if self.interpolation_type == 'linear':
            self.points_dual, self.pyramid_dual, self.trinkets, self.parents = \
                    wisp_spc_ops.make_trilinear_spc(self.blas.points, self.blas.pyramid)
            log.info("Built dual octree and trinkets")
            
        # Create the pyramid of features.
        fpyramid = []
        for al in self.active_lods:
            if self.interpolation_type == 'linear':
                fpyramid.append(self.pyramid_dual[0,al]+1)
            elif self.interpolation_type == 'closest':
                fpyramid.append(self.blas.pyramid[0,al]+1)
            else:
                raise Exception(f"Interpolation mode {self.interpolation_type} is not supported.")
        self.num_feat = sum(fpyramid).long()
        log.info(f"# Feature Vectors: {self.num_feat}")

        self.dictionary_size = 2 ** self.bitwidth

        self.dictionary = nn.ParameterList([])
        for i in range(len(self.active_lods)):
            fts = torch.zeros(self.dictionary_size, self.feature_dim)
            fts += torch.randn_like(fts) * self.feature_std
            self.dictionary.append(nn.Parameter(fts))

        self.features = nn.ParameterList([]) 
        
        for i in range(len(self.active_lods)):
            fts = torch.zeros(fpyramid[i], self.dictionary_size)
            fts += torch.randn_like(fts) * self.feature_std
            self.features.append(nn.Parameter(fts))

    def bake(self):
        for i, f in enumerate(self.features):
            self.features[i] = nn.Parameter(f.max(dim=-1)[1].float())

    def size(self, use_torchac=False, use_prob_model=False):
        dict_size = sum([d.numel() * torch.finfo(d.dtype).bits for d in self.dictionary])
        index_bits = 0
        for dim in range(len(self.features)):
            weight = torch.argmax(self.features[dim], dim=-1)
            unique_vals, counts = torch.unique(weight, return_counts = True)
            probs = counts/torch.sum(counts)

            if not use_torchac:
                information_bits = torch.clamp(-1.0 * torch.log(probs + 1e-10) / np.log(2.0), 0, 1000)
                size_bits = torch.sum(information_bits*counts).item()
                index_bits += size_bits
            else:
                import torchac
                cdf = torch.cumsum(probs,dim=0)
                cdf = torch.cat((torch.Tensor([0.0]).to(cdf),cdf))
                cdf = cdf.unsqueeze(0).repeat(self.codebook.size(0),1)
                cdf = cdf/cdf[:,-1:] # Normalize the final cdf value just to keep torchac happy
                
                weight = weight - weight.min()
                unique_vals, counts = torch.unique(weight, return_counts = True)
                mapping = torch.zeros((weight.max().item()+1))
                mapping[unique_vals] = torch.arange(unique_vals.size(0)).to(mapping)
                weight = mapping[weight]
                cdf = torch.cumsum(counts/counts.sum(),dim=0)
                cdf = torch.cat((torch.Tensor([0.0]),cdf))
                cdf = cdf.unsqueeze(0).repeat(weight.size(0),1)
                cdf = cdf/cdf[:,-1:] # Normalize the final cdf value just to keep torchac happy
                byte_stream = torchac.encode_float_cdf(cdf.detach().cpu(), weight.detach().cpu().to(torch.int16), \
                                                        check_input_bounds=True)
                index_bits += len(byte_stream)*8

        return 0.0, index_bits + dict_size
    def _index_features(self, feats, idx, lod_idx):
        """Internal function. Returns the feats based on indices.

        Args:
            feats (torch.FloatTensor): tensor of feats of shape [num_feats, feat_dim]
            idx (torch.LongTensor): indices of shape [num_indices]
            lod_idx (int): index to `self.active_lods`.

        Returns:
            (torch.FloatTensor): tensor of feats of shape [num_indices, feat_dim]
        """
        # idx -> [N, 8]
        # [1, 1, 256, 32] * [N, 8, 256, 1] -> [N, 8, 32]
        
        if self.training:
            logits = feats[idx.long()]
            y_soft = F.softmax(logits, dim=-1)
            index = y_soft.max(-1, keepdim=True)[1]
            y_hard = torch.zeros_like(
                logits, memory_format=torch.legacy_contiguous_format
            ).scatter_(-1, index, 1.0)
            keys = y_hard - y_soft.detach() + y_soft
            return (self.dictionary[lod_idx][None, None] * keys[..., None]).sum(-2)
            
            # TODO(ttakikawa): Replace with a cleaner / faster softmax implementation
            #keys = F.softmax(feats[idx.long()], dim=-1)
            #return softmax_dictionary(keys, self.dictionary[lod_idx])
        else:
            # [N, 8, 256] -> [N, 8]
            #keys = feats[idx.long()].long()
            keys = torch.max(feats[idx.long()], dim=-1)[1]
            return self.dictionary[lod_idx][keys]
    
    def _interpolate(self, coords, feats, pidx, lod_idx):
        """Query multiscale features.

        Args:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3]
            lod_idx  (int): int specifying the index to ``active_lods`` 
            pidx (torch.LongTensor): point_hiearchy indices of shape [batch]
            features (torch.FloatTensor): features to interpolate. If ``None``, will use `self.features`.

        Returns:
            (torch.FloatTensor): interpolated features of shape [batch, num_samples, feature_dim]
        """
        batch, num_samples = coords.shape[:2]

        if self.interpolation_type == 'linear':
            
            fs = torch.zeros(batch, num_samples, self.feature_dim, device=coords.device)
            
            valid_mask = pidx > -1
            valid_pidx = pidx[valid_mask]
            if valid_pidx.shape[0] == 0:
                return fs

            corner_feats = self._index_features(feats, 
                    self.trinkets.index_select(0, valid_pidx).long(), lod_idx)[:, None]
            
            pts = self.blas.points.index_select(0, valid_pidx)[:,None].repeat(1, coords.shape[1], 1)

            coeffs = spc_ops.coords_to_trilinear_coeffs(coords[valid_mask], pts, self.active_lods[lod_idx])[..., None]
            fs[valid_mask] = (corner_feats * coeffs).sum(-2)

        elif self.interpolation_type == 'closest':
            raise NotImplementedError
        else:
            raise Exception(f"Interpolation mode {self.interpolation_type} is not supported.")
        
        return fs

    def name(self) -> str:
        return "Codebook Grid"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        parent_properties = super().public_properties()
        properties = {
            "Bitwidth": self.bitwidth
        }
        return {**parent_properties, **properties}

wisp/models/grids/codebook_grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Dict, Set, Any, Type, List
import os
import torch
import torch.nn as nn
import numpy as np
import kaolin.ops.spc as spc_ops
import wisp.ops.grid as grid_ops
from wisp.accelstructs import OctreeAS, BaseAS, ASRaymarchResults
from wisp.models.grids import BLASGrid


class HashGrid(BLASGrid):
    """A feature grid where hashed feature pointers are stored as multi-LOD grid nodes,
    and actual feature contents are stored in a hash table.
    (see: Muller et al. 2022, Instant-NGP: https://nvlabs.github.io/instant-ngp/)

    The occupancy state (e.g. BLAS, Bottom Level Acceleration Structure) is tracked separately from the feature
    volume, and relies on heuristics such as pruning for keeping it aligned with the feature structure.
    """
    def __init__(self,
        feature_dim        : int,
        resolutions        : List[int],
        multiscale_type    : str   = 'sum',
        resolution_dim     : int = 3,
        feature_std        : float = 0.0,
        feature_bias       : float = 0.0,
        codebook_bitwidth  : int   = 8,
        blas_level         : int = 7
    ):
        """Builds a HashGrid instance, including the feature structure and an underlying BLAS for fast queries.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
             resolutions (List[int]): A list of resolutions to be used for each feature grid lod of the hash structure.
                i.e. resolutions=[562, 777, 1483, 2048] means that at LOD0, a grid of 562x562x562 nodes will be used,
                where each node is a hashed pointer to the feature table
                (note that feature table size at level L >= resolution of level L).
            multiscale_type (str): The type of multiscale aggregation.
                'sum' - aggregates features from different LODs with summation.
                'cat' - aggregates features from different LODs with concatenation.
                Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure).
        """
        # Occupancy Structure
        self.blas_level = blas_level
        blas = OctreeAS.make_dense(level=blas_level)
        super().__init__(blas)
        self.dense_points = spc_ops.unbatched_get_level_points(self.blas.points,
                                                               self.blas.pyramid,
                                                               self.blas_level).clone()
        self.num_cells = self.dense_points.shape[0]
        self.occupancy = torch.zeros(self.num_cells)

        # Feature Structure - dims
        self.feature_dim = feature_dim
        self.multiscale_type = multiscale_type
        self.feature_std = feature_std
        self.feature_bias = feature_bias
        self.codebook_bitwidth = codebook_bitwidth

        # Feature Structure - setup grid LODs
        self.resolutions = resolutions
        self.num_lods = len(resolutions)
        self.active_lods = [x for x in range(self.num_lods)]
        self.max_lod = self.num_lods - 1

        self.codebook_size = 2 ** self.codebook_bitwidth
        self.register_buffer("codebook_lod_sizes", torch.zeros(self.num_lods, dtype=torch.int32))
        self.register_buffer("codebook_lod_first_idx", torch.zeros(self.num_lods, dtype=torch.int32))
        

        self.codebook = []

        offset = 0
        for lod, res in enumerate(resolutions):
            num_pts = res ** resolution_dim
            fts = torch.zeros(min(self.codebook_size, num_pts), self.feature_dim)
            # fts += torch.randn_like(fts) * self.feature_std
            fts += torch.randn_like(fts) * self.feature_std
            self.codebook.append(fts)
            self.codebook_lod_sizes[lod] = fts.shape[0]
            self.codebook_lod_first_idx[lod] = offset

            offset += fts.shape[0]
        self.codebook = nn.Parameter(torch.cat(self.codebook, dim=0))

    def size(self, use_torchac=False, use_prob_model=False):
        return 0.0, self.codebook.numel() * torch.finfo(self.codebook.dtype).bits
    
    @classmethod
    def from_octree(cls,
                    feature_dim        : int,
                    base_lod           : int   = 2,
                    num_lods           : int   = 1,
                    multiscale_type    : str   = 'sum',
                    resolution_dim     : int = 3,
                    feature_std        : float = 0.0,
                    feature_bias       : float = 0.0,
                    codebook_bitwidth  : int   = 8,
                    blas_level         : int   = 7) -> HashGrid:
        """
        Builds a hash grid using an octree sampling pattern.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure).
                The HashGrid is backed
        """
        octree_lods = [base_lod + x for x in range(num_lods)]
        resolutions = [2 ** lod for lod in octree_lods]
        return cls(feature_dim=feature_dim, resolutions=resolutions, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth,
                   blas_level=blas_level, resolution_dim=resolution_dim)

    @classmethod
    def from_geometric(cls,
                       feature_dim        : int,
                       num_lods           : int,
                       multiscale_type    : str   = 'sum',
                       resolution_dim     : int = 3,
                       feature_std        : float = 0.0,
                       feature_bias       : float = 0.0,
                       codebook_bitwidth  : int   = 8,
                       min_grid_res       : int   = 16,
                       max_grid_res       : int   = None,
                       blas_level: int = 7) -> HashGrid:
        """
        Builds a hash grid using the geometric sequence initialization pattern from Muller et al. 2022 (Instant-NGP).
        This is an implementation of the geometric multiscale grid from
        instant-ngp (https://nvlabs.github.io/instant-ngp/).
        See Section 3 Equations 2 and 3 for more details.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            num_lods (int): The number of LODs for which features are defined. Starts at lod=0.
                            i.e.  num_lods=16 means features are kept for levels 0, 1, 2, .., 14, 15.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            min_grid_res (int): min resolution of the feature grid.
            max_grid_res (int): max resolution of the feature grid.
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure).
        """
        b = np.exp((np.log(max_grid_res) - np.log(min_grid_res)) / (num_lods-1))
        resolutions = [int(1 + np.floor(min_grid_res*(b**l))) for l in range(num_lods)]
        return cls(feature_dim=feature_dim, resolutions=resolutions, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth,
                   blas_level=blas_level, resolution_dim=resolution_dim)

    @classmethod
    def from_resolutions(cls,
                         feature_dim: int,
                         resolutions: List[int],
                         multiscale_type: str = 'sum',
                         resolution_dim: int = 3,
                         feature_std: float = 0.0,
                         feature_bias: float = 0.0,
                         codebook_bitwidth: int = 8,
                         blas_level: int = 7) -> HashGrid:
        """
        Builds a hash grid from a list of resolution sizes (each entry contains a RES for the RES x RES x RES
        lod of nodes pointing at the actual hash table) .

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            resolutions (List[int]): A list of resolutions to be used for each feature grid lod of the hash structure.
                i.e. resolutions=[562, 777, 1483, 2048] means that at LOD0, a grid of 562x562x562 nodes will be used,
                where each node is a hashed pointer to the feature table
                (note that feature table size at level L >= resolution of level L).
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure)
        """
        return cls(feature_dim=feature_dim, resolutions=resolutions, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth,
                   blas_level=blas_level, resolution_dim=resolution_dim)

    def freeze(self):
        """Freezes the feature grid.
        """
        self.codebook.requires_grad_(False)

    def interpolate(self, coords, lod_idx):
        """Query multiscale features.

        Args:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3] or [batch, 3]
                For some grid implementations, specifying num_samples may allow for slightly faster trilinear
                interpolation. HashGrid doesn't use this optimization, but allows this input type for compatability.
            lod_idx  (int): int specifying the index to ``active_lods``

        Returns:
            (torch.FloatTensor): interpolated features of shape
             [batch, num_samples, feature_dim] or [batch, feature_dim]
        """
        # Remember desired output shape
        output_shape = coords.shape[:-1]
        if coords.ndim == 3:    # flatten num_samples dim with batch for cuda call
            batch, num_samples, coords_dim = coords.shape  # batch x num_samples
            coords = coords.reshape(batch * num_samples, coords_dim)

        coords_dim = coords.shape[-1]
        hashgrid_fn = grid_ops.hashgrid2d if coords_dim == 2 else grid_ops.hashgrid
        feats = hashgrid_fn(coords, self.resolutions, self.codebook_bitwidth, lod_idx, self.codebook, self.codebook_lod_sizes, self.codebook_lod_first_idx)

        if "RENDERING_FINAL" in os.environ:
            mask = torch.zeros_like(feats)
            mask[:,:lod_idx*self.feature_dim] = 1
            feats = feats * mask

        if self.multiscale_type == 'cat':
            return feats.reshape(*output_shape, feats.shape[-1])
        elif self.multiscale_type == 'sum':
            return feats.reshape(*output_shape, len(self.resolutions), feats.shape[-1] // len(self.resolutions)).sum(-2)
        else:
            raise NotImplementedError

    def raymarch(self, rays, raymarch_type, num_samples, level=None) -> ASRaymarchResults:
        """Mostly a wrapper over OctreeAS.raymarch. See corresponding function for more details.

        Important detail: the OctreeGrid raymarch samples over the coarsest LOD where features are available.
        """
        return self.blas.raymarch(rays, raymarch_type=raymarch_type, num_samples=num_samples, level=self.blas_level)

    def supported_blas(self) -> Set[Type[BaseAS]]:
        """ Returns a set of bottom-level acceleration structures this grid type supports """
        return {OctreeAS}

    def name(self) -> str:
        return "Hash Grid"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        parent_properties = super().public_properties()
        active_lods = None if self.active_lods is None or len(self.active_lods) == 0 else \
            f'{min(self.active_lods)} - {max(self.active_lods)}'
        properties = {
            "Feature Dims": self.feature_dim,
            "Total LODs": self.max_lod,
            "Active feature LODs": active_lods,
            "Interpolation": 'linear',
            "Multiscale aggregation": self.multiscale_type,
            "HashTable Size": f"2^{self.codebook_bitwidth}"
        }
        return {**parent_properties, **properties}

wisp/models/grids/hash_grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Dict, Set, Any, Type, List
import os
import torch
import torch.nn as nn
import numpy as np
import kaolin.ops.spc as spc_ops
import wisp.ops.grid as grid_ops
from wisp.accelstructs import OctreeAS, BaseAS, ASRaymarchResults
from wisp.models.grids import BLASGrid
from wisp.models.latent_decoders import *
from wisp.models.prob_models import *


class LatentGrid(BLASGrid):
    """A feature grid where hashed feature pointers are stored as multi-LOD grid nodes,
    and actual feature contents are stored in a latent hash table.
    The latents are quantized and decoded into a hashtable as in Instant-NGP
    (see: Muller et al. 2022, Instant-NGP: https://nvlabs.github.io/instant-ngp/)

    The occupancy state (e.g. BLAS, Bottom Level Acceleration Structure) is tracked separately from the feature
    volume, and relies on heuristics such as pruning for keeping it aligned with the feature structure.
    """
    def __init__(self,
        feature_dim        : int,
        latent_dim         : int,
        resolutions        : List[int],
        multiscale_type    : str   = 'sum',
        resolution_dim     : int = 3,
        feature_std        : float = 0.0,
        feature_bias       : float = 0.0,
        codebook_bitwidth  : int   = 8,
        blas_level         : int = 7,
        init_grid          : str = 'normal',
        conf_latent_decoder: Dict[str, Any] = {},
        conf_entropy_reg   : Dict[str, Any] = {},
    ):
        """Builds a LatentGrid instance, including the feature structure and an underlying BLAS for fast queries.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
             resolutions (List[int]): A list of resolutions to be used for each feature grid lod of the hash structure.
                i.e. resolutions=[562, 777, 1483, 2048] means that at LOD0, a grid of 562x562x562 nodes will be used,
                where each node is a hashed pointer to the feature table
                (note that feature table size at level L >= resolution of level L).
            multiscale_type (str): The type of multiscale aggregation.
                'sum' - aggregates features from different LODs with summation.
                'cat' - aggregates features from different LODs with concatenation.
                Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            init_grid (str): The type of initialization for the feature grid. Choices are 'normal' and 'uniform'.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure).
            conf_latent_decoder (dict): Hyperparameters for the latent decoder if enabled
            conf_entropy_reg (dict): Hyperparameters for the probability models for entropy regularization
        """
        # Occupancy Structure
        self.blas_level = blas_level
        blas = OctreeAS.make_dense(level=blas_level)
        super().__init__(blas)
        self.dense_points = spc_ops.unbatched_get_level_points(self.blas.points,
                                                               self.blas.pyramid,
                                                               self.blas_level).clone()
        self.num_cells = self.dense_points.shape[0]
        self.occupancy = torch.zeros(self.num_cells)

        # Feature Structure - dims
        self.feature_dim = feature_dim
        # Latent dimension; By default is set to 0, which means that the latent dimension is the same as the feature dimension.
        self.latent_dim = feature_dim if latent_dim == 0 else latent_dim
        self.multiscale_type = multiscale_type
        self.feature_std = feature_std
        self.feature_bias = feature_bias
        self.codebook_bitwidth = codebook_bitwidth

        # Feature Structure - setup grid LODs
        self.resolutions = resolutions
        self.num_lods = len(resolutions)
        self.active_lods = [x for x in range(self.num_lods)]
        self.max_lod = self.num_lods - 1

        self.codebook_size = 2 ** self.codebook_bitwidth
        self.register_buffer("codebook_lod_sizes", torch.zeros(self.num_lods, dtype=torch.int32))
        self.register_buffer("codebook_lod_first_idx", torch.zeros(self.num_lods, dtype=torch.int32))
        

        self.codebook = []

        offset = 0
        for lod, res in enumerate(resolutions):
            num_pts = res ** resolution_dim
            fts = torch.zeros(min(self.codebook_size, num_pts), self.latent_dim)
            if init_grid == 'uniform':
                fts += (torch.rand_like(fts)-0.5) * 2 * self.feature_std
            elif init_grid == 'normal':
                fts += torch.randn_like(fts) * self.feature_std
            self.codebook.append(fts)
            self.codebook_lod_sizes[lod] = fts.shape[0]
            self.codebook_lod_first_idx[lod] = offset
            offset += fts.shape[0]
        self.codebook = nn.Parameter(torch.cat(self.codebook, dim=0))
        
        self.latent_dec = self.setup_decoders(conf_latent_decoder)
        self.prob_model = None
        self.noise = None
        if conf_latent_decoder['ldecode_enabled'] and (conf_entropy_reg['entropy_reg']>0.0
            or conf_entropy_reg['entropy_reg_end']>0.0):
            self.prob_model = BitEstimator(self.latent_dim, num_layers=conf_entropy_reg['num_prob_layers'])
            self.noise_freq = conf_entropy_reg['noise_freq']

    def ent_loss(self, idx, is_val=False):
        if self.prob_model is None:
            return 0.0, 0.0
        else:
            noise = self.noise
            if self.noise_freq == 1:
               noise = torch.rand(self.codebook.shape).to(self.codebook)-0.5
            elif idx % self.noise_freq == 0:
               self.noise = torch.rand(self.codebook.shape).to(self.codebook)-0.5
               noise = self.noise
            weight = (self.codebook + noise) if not is_val else torch.round(self.codebook) 
            weight_p, weight_n = weight + 0.5, weight - 0.5
            prob = self.prob_model(weight_p) - self.prob_model(weight_n)
            total_bits = torch.sum(torch.clamp(-1.0 * torch.log(prob + 1e-10) / np.log(2.0), 0, 50))
            return total_bits / self.codebook.shape[0], total_bits

    def size(self, use_torchac=False, use_prob_model=False):
        ldec_size = self.latent_dec.size(use_torchac)
        codebook_bits = 0
        for dim in range(self.codebook.size(1)):
            weight = torch.round(self.codebook[:,dim]).long()
            unique_vals, counts = torch.unique(weight, return_counts = True)
            if not use_prob_model:
                probs = counts/torch.sum(counts)
            else:
                assert self.prob_model is not None
                probs = self.prob_model(unique_vals+0.5,single_channel=dim) - self.prob_model(unique_vals-0.5,single_channel=dim)

            if not use_torchac:
                information_bits = torch.clamp(-1.0 * torch.log(probs + 1e-10) / np.log(2.0), 0, 1000)
                size_bits = torch.sum(information_bits*counts).item()
                codebook_bits += size_bits
            else:
                import torchac
                cdf = torch.cumsum(probs,dim=0)
                cdf = torch.cat((torch.Tensor([0.0]).to(cdf),cdf))
                cdf = cdf.unsqueeze(0).repeat(self.codebook.size(0),1)
                cdf = cdf/cdf[:,-1:] # Normalize the final cdf value just to keep torchac happy
                
                weight = weight - weight.min()
                unique_vals, counts = torch.unique(weight, return_counts = True)
                mapping = torch.zeros((weight.max().item()+1))
                mapping[unique_vals] = torch.arange(unique_vals.size(0)).to(mapping)
                weight = mapping[weight]
                cdf = torch.cumsum(counts/counts.sum(),dim=0)
                cdf = torch.cat((torch.Tensor([0.0]).to(cdf),cdf))
                cdf = cdf.unsqueeze(0).repeat(weight.size(0),1)
                cdf = cdf/cdf[:,-1:] # Normalize the final cdf value just to keep torchac happy
                byte_stream = torchac.encode_float_cdf(cdf.detach().cpu(), weight.detach().cpu().to(torch.int16), \
                                                        check_input_bounds=True)
                codebook_bits += len(byte_stream)*8

        return ldec_size, codebook_bits
        
    def setup_decoders(self, decoder_cfg):
        if not decoder_cfg['ldecode_enabled']:
            return DecoderIdentity()
        decoder_cfg['feature_dim'] = self.feature_dim
        decoder_cfg['latent_dim'] = self.latent_dim
        if decoder_cfg['ldecode_type'] == 'hierarchical':
            offsets = torch.cat((self.codebook_lod_first_idx, self.codebook_lod_sizes[-1:]))
            decoder = HierarchicalLatentDecoder(self.num_lods, offsets, decoder_cfg)
        elif decoder_cfg['ldecode_type'] == 'multi':
            decoder_cfg['num_entries'] = self.codebook.size(0)
            decoder = MultiLatentDecoder(**decoder_cfg)
            del decoder_cfg['num_entries']
        elif decoder_cfg['ldecode_type'] == 'single':
            decoder = LatentDecoder(**decoder_cfg)
        return decoder
    
    @classmethod
    def from_octree(cls,
                    feature_dim        : int,
                    latent_dim         : int   = 0,
                    base_lod           : int   = 2,
                    num_lods           : int   = 1,
                    multiscale_type    : str   = 'sum',
                    resolution_dim     : int = 3,
                    feature_std        : float = 0.0,
                    feature_bias       : float = 0.0,
                    codebook_bitwidth  : int   = 8,
                    blas_level         : int   = 7,
                    init_grid          : str = 'normal',
                    conf_latent_decoder: dict  = {},
                    conf_entropy_reg   : dict  = {}) -> LatentGrid:
        """
        Builds a hash grid using an octree sampling pattern.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure).
                The LatentGrid is backed
            init_grid (str): The type of initialization for the feature grid. Choices are 'normal' and 'uniform'.
            conf_latent_decoder (dict): Hyperparameters for the latent decoder if enabled
            conf_entropy_reg (dict): Hyperparameters for the probability models for entropy regularization
        """
        octree_lods = [base_lod + x for x in range(num_lods)]
        resolutions = [2 ** lod for lod in octree_lods]
        return cls(feature_dim=feature_dim, resolutions=resolutions, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth,
                   blas_level=blas_level, latent_dim=latent_dim, conf_latent_decoder=conf_latent_decoder,
                   conf_entropy_reg=conf_entropy_reg, resolution_dim=resolution_dim, init_grid=init_grid)

    @classmethod
    def from_geometric(cls,
                       feature_dim        : int,
                       num_lods           : int,
                       latent_dim         : int   = 0,
                       multiscale_type    : str   = 'sum',
                       resolution_dim     : int = 3,
                       feature_std        : float = 0.0,
                       feature_bias       : float = 0.0,
                       codebook_bitwidth  : int   = 8,
                       min_grid_res       : int   = 16,
                       max_grid_res       : int   = None,
                       blas_level         : int   = 7,
                       init_grid          : str = 'normal',
                       conf_latent_decoder: dict  = {},
                       conf_entropy_reg   : dict  = {}) -> LatentGrid:
        """
        Builds a hash grid using the geometric sequence initialization pattern from Muller et al. 2022 (Instant-NGP).
        This is an implementation of the geometric multiscale grid from
        instant-ngp (https://nvlabs.github.io/instant-ngp/).
        See Section 3 Equations 2 and 3 for more details.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            num_lods (int): The number of LODs for which features are defined. Starts at lod=0.
                            i.e.  num_lods=16 means features are kept for levels 0, 1, 2, .., 14, 15.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            min_grid_res (int): min resolution of the feature grid.
            max_grid_res (int): max resolution of the feature grid.
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure).
            init_grid (str): The type of initialization for the feature grid. Choices are 'normal' and 'uniform'.
            conf_latent_decoder (dict): Hyperparameters for the latent decoder if enabled
            conf_entropy_reg (dict): Hyperparameters for the probability models for entropy regularization
        """
        b = np.exp((np.log(max_grid_res) - np.log(min_grid_res)) / (num_lods-1))
        resolutions = [int(1 + np.floor(min_grid_res*(b**l))) for l in range(num_lods)]
        return cls(feature_dim=feature_dim, resolutions=resolutions, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth,
                   blas_level=blas_level, latent_dim=latent_dim, conf_latent_decoder=conf_latent_decoder,
                   conf_entropy_reg=conf_entropy_reg, resolution_dim=resolution_dim, init_grid=init_grid)

    @classmethod
    def from_resolutions(cls,
                         feature_dim: int,
                         resolutions: List[int],
                         latent_dim: int = 0,
                         multiscale_type: str = 'sum',
                         resolution_dim: int = 3,
                         feature_std: float = 0.0,
                         feature_bias: float = 0.0,
                         codebook_bitwidth: int = 8,
                         blas_level: int = 7,
                         init_grid: str = 'normal',
                         conf_latent_decoder: dict = {},
                         conf_entropy_reg: dict = {}) -> LatentGrid:
        """
        Builds a hash grid from a list of resolution sizes (each entry contains a RES for the RES x RES x RES
        lod of nodes pointing at the actual hash table) .

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            resolutions (List[int]): A list of resolutions to be used for each feature grid lod of the hash structure.
                i.e. resolutions=[562, 777, 1483, 2048] means that at LOD0, a grid of 562x562x562 nodes will be used,
                where each node is a hashed pointer to the feature table
                (note that feature table size at level L >= resolution of level L).
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            resolution_dim (int): Grid resolution dimension; 3(for videos/NeRFs) or 2 (for images)
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.
            codebook_bitwidth (int): Codebook dictionary_size is set as 2**bitwidth
            blas_level (int): The level of the octree to be used as the BLAS (bottom level acceleration structure)
            init_grid (str): The type of initialization for the feature grid. Choices are 'normal' and 'uniform'.
            conf_latent_decoder (dict): Hyperparameters for the latent decoder if enabled
            conf_entropy_reg (dict): Hyperparameters for the probability models for entropy regularization
        """
        return cls(feature_dim=feature_dim, resolutions=resolutions, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias, codebook_bitwidth=codebook_bitwidth,
                   blas_level=blas_level, latent_dim=latent_dim, conf_latent_decoder=conf_latent_decoder,
                   conf_entropy_reg=conf_entropy_reg, resolution_dim=resolution_dim, init_grid=init_grid)

    def freeze(self):
        """Freezes the feature grid.
        """
        self.codebook.requires_grad_(False)
        for p in self.latent_dec.parameters():
            p.requires_grad_(False)
        if self.prob_model is not None:
            for p in self.prob_model.parameters():
                p.requires_grad_(False)

    def interpolate(self, coords, lod_idx):
        """Query multiscale features.

        Args:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3] or [batch, 3]
                For some grid implementations, specifying num_samples may allow for slightly faster trilinear
                interpolation. LatentGrid doesn't use this optimization, but allows this input type for compatability.
            lod_idx  (int): int specifying the index to ``active_lods``

        Returns:
            (torch.FloatTensor): interpolated features of shape
             [batch, num_samples, feature_dim] or [batch, feature_dim]
        """
        # Remember desired output shape
        output_shape = coords.shape[:-1]
        if coords.ndim == 3:    # flatten num_samples dim with batch for cuda call
            batch, num_samples, coords_dim = coords.shape  # batch x num_samples
            coords = coords.reshape(batch * num_samples, coords_dim)

        codebook = self.latent_dec(self.codebook)
        rep = False
        if codebook.size(1)==1:
            codebook = codebook.repeat(1,2)
            rep = True

        coords_dim = coords.shape[-1]
        hashgrid_fn = grid_ops.hashgrid2d if coords_dim == 2 else grid_ops.hashgrid
        
        feats = hashgrid_fn(coords, self.resolutions, self.codebook_bitwidth, lod_idx, codebook, self.codebook_lod_sizes, self.codebook_lod_first_idx)

        feats = feats[:,::2] if rep else feats

        if "RENDERING_FINAL" in os.environ:
            mask = torch.zeros_like(feats)
            mask[:,:lod_idx*self.feature_dim] = 1
            feats = feats * mask

        if self.multiscale_type == 'cat':
            return feats.reshape(*output_shape, feats.shape[-1])
        elif self.multiscale_type == 'sum':
            return feats.reshape(*output_shape, len(self.resolutions), feats.shape[-1] // len(self.resolutions)).sum(-2)
        else:
            raise NotImplementedError

    def raymarch(self, rays, raymarch_type, num_samples, level=None) -> ASRaymarchResults:
        """Mostly a wrapper over OctreeAS.raymarch. See corresponding function for more details.

        Important detail: the OctreeGrid raymarch samples over the coarsest LOD where features are available.
        """
        return self.blas.raymarch(rays, raymarch_type=raymarch_type, num_samples=num_samples, level=self.blas_level)

    def supported_blas(self) -> Set[Type[BaseAS]]:
        """ Returns a set of bottom-level acceleration structures this grid type supports """
        return {OctreeAS}

    def name(self) -> str:
        return "Latent Grid"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        parent_properties = super().public_properties()
        active_lods = None if self.active_lods is None or len(self.active_lods) == 0 else \
            f'{min(self.active_lods)} - {max(self.active_lods)}'
        properties = {
            "Feature Dims": self.feature_dim,
            "Latent Dims": self.latent_dim,
            "Total LODs": self.max_lod,
            "Active feature LODs": active_lods,
            "Interpolation": 'linear',
            "Multiscale aggregation": self.multiscale_type,
            "HashTable Size": f"2^{self.codebook_bitwidth}"
        }
        return {**parent_properties, **properties}

wisp/models/grids/latent_grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import logging as log
from typing import Dict, Set, Any, Type
import torch
import torch.nn as nn
import wisp.ops.spc as wisp_spc_ops
from wisp.models.grids import BLASGrid
import kaolin.ops.spc as spc_ops
from wisp.accelstructs import BaseAS, OctreeAS, ASRaymarchResults


class OctreeGrid(BLASGrid):
    """This is a multiscale feature grid where the features are defined on the BLAS, the octree.
    """

    def __init__(
        self,
        accelstruct: OctreeAS,
        feature_dim         : int,
        base_lod            : int,
        num_lods            : int          = 1,
        interpolation_type  : str          = 'linear',
        multiscale_type     : str          = 'cat',
        feature_std         : float        = 0.0,
        feature_bias        : float        = 0.0
    ):
        """Initialize the octree grid class.

        Args:
            accelstruct: Spatial acceleration structure which tracks the occupancy state of this grid.
                         Used to speed up spatial queries and ray tracing operations.
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid. This is the lowest LOD of the SPC octree
                            for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
            interpolation_type (str): The type of interpolation function.
            multiscale_type (str): The type of multiscale aggregation. Usually 'sum' or 'cat'.
                                   Note that 'cat' will change the decoder input dimension.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The mean of the Gaussian distribution.
            sample_tex (bool): If True, will also sample textures and store it in the accelstruct.
            num_samples (int): The number of samples to be generated on the mesh surface.
        Returns:
            (void): Initializes the class.
        """
        super().__init__(accelstruct)
        self.feature_dim = feature_dim
        self.base_lod = base_lod
        self.num_lods = num_lods
        self.interpolation_type = interpolation_type
        self.multiscale_type = multiscale_type

        self.feature_std = feature_std
        self.feature_bias = feature_bias

        # List of octree levels which are optimized.
        self.active_lods = [self.base_lod + x for x in range(self.num_lods)]
        self.max_lod = OctreeGrid.max_octree_lod(self.base_lod, self.num_lods)

        log.info(f"Active LODs: {self.active_lods}")    # TODO(operel): move into trainer

        if self.num_lods > 0:
            self.init_feature_structure()

    @staticmethod
    def max_octree_lod(base_lod, num_lods) -> int:
        """
        Returns:
            (int): The highest level-of-detail maintaining features in this Octree grid.
        """
        return base_lod + num_lods - 1

    @classmethod
    def make_dense(cls,
                   feature_dim        : int,
                   base_lod           : int,
                   num_lods           : int          = 1,
                   interpolation_type : str          = 'linear',
                   multiscale_type    : str          = 'cat',
                   feature_std        : float        = 0.0,
                   feature_bias       : float        = 0.0) -> OctreeGrid:
        """Builds a fully occupied OctreeGrid with `base_lod + num_lods - 1` levels.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.

        Returns:
            (OctreeGrid): A new instance of an OctreeGrid at full occupancy.
        """
        max_lod = OctreeGrid.max_octree_lod(base_lod, num_lods)
        blas = OctreeAS.make_dense(level=max_lod)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias)
        return grid

    @classmethod
    def from_mesh(cls,
                  mesh_path: str,
                  num_samples_on_mesh: int,
                  feature_dim: int,
                  base_lod: int,
                  num_lods: int = 1,
                  interpolation_type: str = 'linear',
                  multiscale_type: str = 'cat',
                  feature_std: float = 0.0,
                  feature_bias: float = 0.0) -> OctreeGrid:
        """Builds the OctreeGrid from point samples over mesh faces.
        Samples will be quantized to populate the octree cells.
        This method is not guaranteed to reproduce a perfect mesh structure without "holes",
        but achieves a very good approximation with high probability.

        Args:
            mesh_path (str): Path to an OBJ file with a mesh to initialize the octree (only supports OBJ for now).
            num_samples (int): The number of samples to be generated on the mesh surface.
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.

        Returns:
            (OctreeGrid): A new instance of an OctreeGrid with occupancy initialized from samples over the mesh faces.
        """
        max_lod = OctreeGrid.max_octree_lod(base_lod, num_lods)
        blas = OctreeAS.from_mesh(mesh_path, level=max_lod, sample_tex=False, num_samples=num_samples_on_mesh)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias)
        return grid

    @classmethod
    def from_pointcloud(cls,
                        pointcloud: torch.FloatTensor,
                        feature_dim: int,
                        base_lod: int,
                        num_lods: int = 1,
                        interpolation_type: str = 'linear',
                        multiscale_type: str = 'cat',
                        feature_std: float = 0.0,
                        feature_bias: float = 0.0) -> OctreeGrid:
        """Builds the OctreeGrid, initializing it with a pointcloud.
        The cells occupancy will be determined by points occupying the octree cells.

        Args:
            pointcloud (torch.FloatTensor): 3D coordinates of shape [num_coords, 3] in normalized space [-1, 1].
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.

        Returns:
            (OctreeGrid): A new instance of an OctreeGrid with occupancy initialized from the pointcloud.
        """
        max_lod = OctreeGrid.max_octree_lod(base_lod, num_lods)
        blas = OctreeAS.from_pointcloud(pointcloud, level=max_lod)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias)
        return grid

    @classmethod
    def from_spc(cls,
                 spc_octree: torch.FloatTensor,
                 feature_dim: int,
                 base_lod: int,
                 num_lods: int = 1,
                 interpolation_type: str = 'linear',
                 multiscale_type: str = 'cat',
                 feature_std: float = 0.0,
                 feature_bias: float = 0.0) -> OctreeGrid:
        """Builds the OctreeGrid, initializing it's occupancy state with a Structured Point Cloud (SPC).
        The octree cells occupancy will be determined by the SPC cells.
        SPCs are compressed, sparse volumetric data structures used for accelerated point queries and ray tracing ops.

        Args:
            spc_octree (torch.ByteTensor):
                A tensor which holds the topology of a Structured Point Cloud (SPC).
                Each byte represents a single octree cell's occupancy (that is, each bit of that byte represents
                the occupancy status of a child octree cell), yielding 8 bits for 8 cells.
                See also https://kaolin.readthedocs.io/en/latest/notes/spc_summary.html
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid.
                            This is the lowest LOD of the  octree for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
                            i.e. base_lod=4 and num_lods=5 means features are kept for levels 5, 6, 7, 8.
            interpolation_type (str): The type of interpolation function used when querying features on the grid.
                                      'linear' - uses trilinear interpolation from nearest 8 nodes.
                                      'closest' - uses feature from nearest grid node.
            multiscale_type (str): The type of multiscale aggregation.
                                   'sum' - aggregates features from different LODs with summation.
                                   'cat' - aggregates features from different LODs with concatenation.
                                   Note that 'cat' will change the decoder input dimension to num_lods * feature_dim.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The features are initialized with a Gaussian distribution with the given mean.

        Returns:
            (OctreeGrid): A new instance of an OctreeGrid with occupancy initialized from the SPC topology.
        """
        blas = OctreeAS(spc_octree)
        grid = cls(accelstruct=blas, feature_dim=feature_dim, base_lod=base_lod, num_lods=num_lods,
                   interpolation_type=interpolation_type, multiscale_type=multiscale_type,
                   feature_std=feature_std, feature_bias=feature_bias)
        return grid

    def init_feature_structure(self):
        """ Initializes everything related to the features stored in the codebook octree structure. """

        # Assumes the occupancy structure have been initialized (the BLAS: Bottom Level Accelerated Structure).
        # Build the trinket structure
        if self.interpolation_type in ['linear']:
            self.points_dual, self.pyramid_dual, self.trinkets, self.parents = \
                wisp_spc_ops.make_trilinear_spc(self.blas.points, self.blas.pyramid)
            log.info("Built dual octree and trinkets")
        
        # Build the pyramid of features
        fpyramid = []
        for al in self.active_lods:
            if self.interpolation_type == 'linear':
                fpyramid.append(self.pyramid_dual[0,al]+1)
            elif self.interpolation_type == 'closest':
                fpyramid.append(self.blas.pyramid[0,al]+1)
            else:
                raise Exception(f"Interpolation mode {self.interpolation_type} is not supported.")
        self.num_feat = sum(fpyramid).long()
        log.info(f"# Feature Vectors: {self.num_feat}")

        self.features = nn.ParameterList([])
        for i in range(len(self.active_lods)):
            fts = torch.zeros(fpyramid[i], self.feature_dim) + self.feature_bias
            fts += torch.randn_like(fts) * self.feature_std
            self.features.append(nn.Parameter(fts))
        
        log.info(f"Pyramid:{self.blas.pyramid[0]}")
        log.info(f"Pyramid Dual: {self.pyramid_dual[0]}")

    def freeze(self):
        """Freezes the feature grid.
        """
        for lod_idx in range(self.num_lods):
            self.features[lod_idx].requires_grad_(False)

    def _index_features(self, feats, idx):
        """Internal function. Returns the feats based on indices.

        This function exists to override in case you want to implement a different method of indexing,
        i.e. a differentiable one as in Variable Bitrate Neural Fields (VQAD).

        Args:
            feats (torch.FloatTensor): tensor of feats of shape [num_feats, feat_dim]
            idx (torch.LongTensor): indices of shape [num_indices]

        Returns:
            (torch.FloatTensor): tensor of feats of shape [num_indices, feat_dim]
        """
        return feats[idx.long()]

    def _interpolate(self, coords, feats, pidx, lod_idx):
        """Interpolates the given feature using the coordinates x. 

        This is a more low level interface for optimization.

        Inputs:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3]
            feats (torch.FloatTensor): feats of shape [num_feats, feat_dim]
            pidx (torch.LongTensor) : point_hiearchy indices of shape [batch]
            lod_idx (int) : int specifying the index fo ``active_lods``
        Returns:
            (torch.FloatTensor): acquired features of shape [batch, num_samples, feat_dim]
        """
        batch, num_samples = coords.shape[:2]
        lod = self.active_lods[lod_idx]

        if self.interpolation_type == 'linear':
            fs = spc_ops.unbatched_interpolate_trilinear(
                coords, pidx.int(), self.blas.points, self.trinkets.int(),
                feats.half(), lod).float()
        elif self.interpolation_type == 'closest':
            fs = self._index_features(feats, pidx.long()-self.blas.pyramid[1, lod])[...,None,:]
            fs = fs.expand(batch, num_samples, feats.shape[-1])
       
        # Keep as backup
        elif self.interpolation_type == 'trilinear_old':
            corner_feats = feats[self.trinkets.index_select(0, pidx).long()]
            coeffs = spc_ops.coords_to_trilinear_coeffs(coords, self.points.index_select(0, pidx)[:,None].repeat(1, coords.shape[1], 1), lod)
            fs = (corner_feats[:, None] * coeffs[..., None]).sum(-2)

        else:
            raise Exception(f"Interpolation mode {self.interpolation_type} is not supported.")
        
        return fs

    def interpolate(self, coords, lod_idx):
        """Query multiscale features.

        Args:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3] or [batch, 3]
            lod_idx  (int): int specifying the index to ``active_lods``
            features (torch.FloatTensor): features to interpolate. If ``None``, will use `self.features`.

        Returns:
            (torch.FloatTensor): interpolated features of shape
            [batch, num_samples, feature_dim] or [batch, feature_dim]
        """
        # Remember desired output shape, and inflate to (batch, num_samples, 3) format
        output_shape = coords.shape[:-1]
        if coords.ndim < 3:
            coords = coords[:, None]    # (batch, 3) -> (batch, num_samples, 3)

        if lod_idx == 0:
            query_results = self.blas.query(coords[:,0], self.active_lods[lod_idx], with_parents=False)
            pidx = query_results.pidx
            feat = self._interpolate(coords, self.features[0], pidx, 0)
            return feat.reshape(*output_shape, feat.shape[-1])
        else:
            feats = []
            
            # In the multiscale case, the raytrace _currently_  does not return multiscale indices.
            # As such, no matter what the pidx will be recalculated to get the multiscale indices.
            num_feats = lod_idx + 1
            
            # This might look unoptimal since it assumes that samples are _not_ in the same voxel.
            # This is the correct assumption here, because the point samples are from the base_lod,
            # not the highest LOD.
            query_results = self.blas.query(coords.reshape(-1, 3), self.active_lods[lod_idx], with_parents=True)
            pidx = query_results.pidx[...,self.base_lod:]
            pidx = pidx.reshape(-1, coords.shape[1], num_feats)
            pidx = torch.split(pidx, 1, dim=-1)
            
            # list of [batch, num_samples, 1]

            for i in range(num_feats):
                feat = self._interpolate(
                    coords.reshape(-1, 1, 3), self.features[i], pidx[i].reshape(-1), i)[:,0]
                feats.append(feat)
            
            feats = torch.cat(feats, dim=-1)

            if self.multiscale_type == 'sum':
                feats = feats.reshape(*feats.shape[:-1], num_feats, self.feature_dim)
                if self.training:
                    feats = feats.sum(-2)
                else:
                    feats = feats.sum(-2)
            
            return feats.reshape(*output_shape, self.feature_dim)

    def raymarch(self, rays, raymarch_type, num_samples, level=None) -> ASRaymarchResults:
        """Mostly a wrapper over OctreeAS.raymarch. See corresponding function for more details.

        Important detail: the OctreeGrid raymarch samples over the coarsest LOD where features are available.
        """
        return self.blas.raymarch(rays, raymarch_type=raymarch_type, num_samples=num_samples, level=self.base_lod)

    def supported_blas(self) -> Set[Type[BaseAS]]:
        """ Returns a set of bottom-level acceleration structures this grid type supports """
        return {OctreeAS}

    def name(self) -> str:
        return "Octree Grid"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        parent_properties = super().public_properties()
        properties = {
            "Feature Dims": self.feature_dim,
            "Total LODs": self.max_lod,
            "Active feature LODs": [str(x) for x in self.active_lods],
            "Interpolation": self.interpolation_type,
            "Multiscale aggregation": self.multiscale_type
        }
        return {**parent_properties, **properties}

wisp/models/grids/octree_grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import logging as log
from typing import Dict, Set, Any, Type
import torch
import torch.nn as nn
import torch.nn.functional as F
from wisp.core import WispModule
from wisp.models.grids import BLASGrid
from wisp.accelstructs import AxisAlignedBBoxAS, BaseAS, ASRaytraceResults, ASRaymarchResults


class TriplanarGrid(BLASGrid):
    """A feature grid where the features are stored on a multiresolution pyramid of triplanes.
    Each LOD consists of a triplane, e.g. a triplet of orthogonal planes.

    The shape of the triplanar feature grid means the support region is bounded by an AABB,
    therefore spatial queries / ray tracing ops can use an AABB as an acceleration structure.
    Hence the class is compatible with BaseTracer implementations.
    """

    def __init__(self,
                 feature_dim: int,
                 base_lod: int,
                 num_lods: int = 1,
                 interpolation_type: str = 'linear',
                 multiscale_type: str = 'sum',
                 feature_std: float = 0.0,
                 feature_bias: float = 0.0
                 ):
        """Constructs an instance of a TriplanarGrid.

        Args:
            feature_dim (int): The dimension of the features stored on the grid.
            base_lod (int): The base LOD of the feature grid. This is the lowest LOD of the SPC octree
                            for which features are defined.
            num_lods (int): The number of LODs for which features are defined. Starts at base_lod.
            interpolation_type (str): The type of interpolation function.
            multiscale_type (str): The type of multiscale aggregation. Usually 'sum' or 'cat'.
                                   Note that 'cat' will change the decoder input dimension.
            feature_std (float): The features are initialized with a Gaussian distribution with the given
                                 standard deviation.
            feature_bias (float): The mean of the Gaussian distribution.

        Returns:
            (void): Initializes the class.
        """
        # The bottom level acceleration structure is an axis aligned bounding box
        super().__init__(blas=AxisAlignedBBoxAS())
        # The actual feature_dim is multiplied by 3 because of the triplanar maps.
        self.feature_dim = feature_dim * 3
        self.base_lod = base_lod
        self.num_lods = num_lods
        self.interpolation_type = interpolation_type
        self.multiscale_type = multiscale_type
        self.feature_std = feature_std
        self.feature_bias = feature_bias

        # TODO(ttakikawa) The Triplanar API should look more like the ImagePyramid class with
        #   similar initialization mechanisms since it's not limited to octrees.
        self.active_lods = [self.base_lod + x for x in range(self.num_lods)]
        self.max_lod = self.num_lods + self.base_lod - 1

        log.info(f"Active LODs: {self.active_lods}")

        self.num_feat = 0
        self.init_feature_structure()

    def init_feature_structure(self):
        """ Initializes everything related to the features stored in the triplanar grid structure. """
        self.features = nn.ModuleList([])
        self.num_feat = 0
        for i in self.active_lods:
            self.features.append(
                TriplanarFeatureVolume(self.feature_dim // 3, 2 ** i, self.feature_std, self.feature_bias))
            self.num_feat += ((2 ** i + 1) ** 2) * self.feature_dim * 3

        log.info(f"# Feature Vectors: {self.num_feat}")

    def freeze(self):
        """Freezes the feature grid.
        """
        self.features.requires_grad_(False)

    def interpolate(self, coords, lod_idx):
        """Query multiscale features.

        Args:
            coords (torch.FloatTensor): coords of shape [batch, num_samples, 3] or [batch, 3]
            lod_idx  (int): int specifying the index to ``active_lods``

        Returns:
            (torch.FloatTensor): interpolated features of
            shape [batch, num_samples, feature_dim] or [batch, feature_dim]
        """
        output_shape = coords.shape[:-1]
        if coords.ndim < 3:
            coords = coords[:, None]  # (batch, 3) -> (batch, num_samples, 3)

        feats = []
        for i in range(lod_idx + 1):
            feats.append(self._interpolate(coords, self.features[i], i))

        feats = torch.cat(feats, dim=-1)

        if self.multiscale_type == 'sum':
            feats = feats.reshape(*output_shape, lod_idx + 1, feats.shape[-1] // (lod_idx + 1)).sum(-2)

        return feats

    def _interpolate(self, coords, feats, lod_idx):
        """Interpolates the given feature using the coordinates x.

        This is a more low level interface for optimization.

        Inputs:
            coords     : float tensor of shape [batch, num_samples, 3]
            feats : float tensor of shape [num_feats, feat_dim]
            lod_idx   : int specifying the lod
        Returns:
            float tensor of shape [batch, num_samples, feat_dim]
        """
        batch, num_samples = coords.shape[:2]

        if self.interpolation_type == 'linear':
            fs = feats(coords).reshape(batch, num_samples, 3 * feats.fdim)
        else:
            raise ValueError(f"Interpolation mode '{self.interpolation_type}' is not supported")

        return fs

    def raymarch(self, rays, raymarch_type, num_samples, level=None) -> ASRaymarchResults:
        """Mostly a wrapper over OctreeAS.raymarch. See corresponding function for more details.

        Important detail: this is just used as an AABB tracer.
        """
        return self.blas.raymarch(rays, raymarch_type=raymarch_type, num_samples=num_samples, level=0)

    def raytrace(self, rays, level=None, with_exit=False) -> ASRaytraceResults:
        """By default, this function will use the equivalent BLAS function unless overridden for custom behaviour.

        Important detail: this is just used as an AABB tracer.
        """
        return self.blas.raytrace(rays, level=0, with_exit=with_exit)

    def supported_blas(self) -> Set[Type[BaseAS]]:
        """ Returns a set of bottom-level acceleration structures this grid type supports """
        return {AxisAlignedBBoxAS}

    def name(self) -> str:
        return "Triplanar Grid"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        parent_properties = super().public_properties()
        properties = {
            "Feature Dims": self.feature_dim,
            "Total LODs": self.max_lod,
            "Active feature LODs": [str(x) for x in self.active_lods],
            "Interpolation": self.interpolation_type,
            "Multiscale aggregation": self.multiscale_type,
        }
        for idx, module in enumerate(self.features):
            properties[f"Pyramid Layer #{idx + 1}"] = module

        return {**parent_properties, **properties}


class TriplanarFeatureVolume(WispModule):
    """Triplanar feature volume represents a single triplane, e.g. a single LOD in a TriplanarGrid. """

    def __init__(self, fdim, fsize, std, bias):
        """Initializes the feature triplane.

        Args:
            fdim (int): The feature dimension.
            fsize (int): The height and width of the texture map.
            std (float): The standard deviation for the Gaussian initialization.
            bias (float): The mean for the Gaussian initialization.
        """
        super().__init__()
        self.fsize = fsize
        self.fdim = fdim
        self.fmx = nn.Parameter(torch.randn(1, fdim, fsize+1, fsize+1) * std + bias)
        self.fmy = nn.Parameter(torch.randn(1, fdim, fsize+1, fsize+1) * std + bias)
        self.fmz = nn.Parameter(torch.randn(1, fdim, fsize+1, fsize+1) * std + bias)
        self.padding_mode = 'reflection'

    def forward(self, x):
        """Interpolates from the feature volume.

        Args:
            x (torch.FloatTensor): Coordinates of shape [batch, num_samples, 3] or [batch, 3].

        Returns:
            (torch.FloatTensor): Features of shape [batch, num_samples, fdim] or [batch, fdim].
        """
        N = x.shape[0]
        if len(x.shape) == 3:
            sample_coords = x.reshape(1, N, x.shape[1], 3)  # [N, 1, 1, 3]
            samplex = F.grid_sample(self.fmx, sample_coords[...,[1,2]], 
                                    align_corners=True, padding_mode=self.padding_mode)[0,:,:,:].transpose(0,1)
            sampley = F.grid_sample(self.fmy, sample_coords[...,[0,2]], 
                                    align_corners=True, padding_mode=self.padding_mode)[0,:,:,:].transpose(0,1)
            samplez = F.grid_sample(self.fmz, sample_coords[...,[0,1]], 
                                    align_corners=True, padding_mode=self.padding_mode)[0,:,:,:].transpose(0,1)
            sample = torch.stack([samplex, sampley, samplez], dim=1).permute(0,3,1,2)
        else:
            sample_coords = x.reshape(1, N, 1, 3)  # [N, 1, 1, 3]
            samplex = F.grid_sample(self.fmx, sample_coords[...,[1,2]], 
                                    align_corners=True, padding_mode=self.padding_mode)[0,:,:,0].transpose(0,1)
            sampley = F.grid_sample(self.fmy, sample_coords[...,[0,2]], 
                                    align_corners=True, padding_modes=self.padding_mode)[0,:,:,0].transpose(0,1)
            samplez = F.grid_sample(self.fmz, sample_coords[...,[0,1]], 
                                    align_corners=True, padding_mode=self.padding_mode)[0,:,:,0].transpose(0,1)
            sample = torch.stack([samplex, sampley, samplez], dim=1)
        return sample

    def name(self) -> str:
        """ A human readable name for the given wisp module. """
        return "TriplanarFeatureVolume"

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        return {
            'Resolution': f'3x{self.fsize}x{self.fsize}'
        }

wisp/models/grids/triplanar_grid.py
Sharath-girish/Shacira
from .basic_latent_decoder import *
from .hierarchical_latent_decoder import *
from .multi_latent_decoder import *

wisp/models/latent_decoders/__init__.py
Sharath-girish/Shacira
import math
import torch
import numpy as np
import torch.nn as nn
from torch import Tensor
from torch.nn import Module, Parameter, init
from typing import Optional, List, Tuple, Union
from wisp.models.activations import SineScaled
from torch.nn.modules.utils import _single, _pair, _triple, _reverse_repeat_tuple, _ntuple

epsilon = 1e-6
def get_dft_matrix(conv_dim, channels):
    dft = torch.zeros(conv_dim,channels)
    for i in range(conv_dim):
        for j in range(channels):
            # Each row of dft is a bias vector
            dft[i,j] = math.cos(torch.pi/channels*(i+0.5)*j)/math.sqrt(channels) 
            dft[i,j] = dft[i,j]*(math.sqrt(2) if j>0 else 1)
    return dft

class BNGlobal(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return (x - torch.mean(x,dim=0,keepdim=True))/(torch.std(x,dim=0,keepdim=True)+1e-8)

class StraightThrough(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output

class StraightThroughFloor(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x):
        return torch.floor(x)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output
    
class DecoderLayer(Module):

    def __init__(self, in_features: int, out_features: int, ldecode_matrix: str, bias: bool = False) -> None:
        super(DecoderLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        if 'dft' in ldecode_matrix:
            self.dft = Parameter(get_dft_matrix(in_features, out_features), requires_grad=False)
        if 'dft' in ldecode_matrix:
            self.scale = Parameter(torch.empty((1,out_features)))
        else:
            self.scale = Parameter(torch.empty((in_features,out_features)))
        if bias:
            self.shift = Parameter(torch.empty(1,out_features))
        else:
            self.register_parameter('shift', None)

        self.ldecode_matrix = ldecode_matrix
        if ldecode_matrix == 'dft_fixed':
            self.scale.requires_grad_(False)
            if not bias:
                self.shift.requires_grad_(False)

    def reset_parameters(self, param=1.0, init_type = 'normal') -> None:
        if init_type == 'normal':
            init.normal_(self.scale, std=param)
        elif init_type == 'uniform':
            init.uniform_(self.scale, -param, param)
        elif init_type == 'constant':
            init.constant_(self.scale, val=param)
        if self.shift is not None:
            init.zeros_(self.shift)

    def clamp(self, val: float = 0.5) -> None:
        with torch.no_grad():
            self.scale.clamp_(-val, val)

    def forward(self, input: Tensor) -> Tensor:
        if 'dft' in self.ldecode_matrix:
            w_out = torch.matmul(input,self.dft)*self.scale+(self.shift if self.shift is not None else 0)
        else:
            w_out = torch.matmul(input,self.scale)+(self.shift if self.shift is not None else 0)
        return w_out

    def extra_repr(self) -> str:
        return 'in_features={}, out_features={}, bias={}'.format(
            self.in_features, self.out_features, self.shift is not None
        )

class LatentDecoder(Module):

    def __init__(
        self,
        latent_dim: int,
        feature_dim: int,
        norm: str,
        ldecode_matrix:str,
        use_shift: bool,
        num_layers_dec:int = 0,
        hidden_dim_dec:int = 0,
        activation:str = 'none',
        final_activation:str = 'none',
        clamp_weights:float = 0.0,
        ldec_std:float = 1.0,
        use_sga:bool = False,
        diff_sampling:bool = False,
        **kwargs,
    ) -> None:
        super(LatentDecoder, self).__init__()
        latent_dim = feature_dim if latent_dim == 0 else latent_dim
        self.ldecode_matrix = ldecode_matrix
        self.channels = feature_dim
        self.latent_dim = latent_dim
        self.norm = norm
        self.div = 1 if norm == 'none' else None
        self.num_layers_dec =  num_layers_dec
        if num_layers_dec>0:
            if hidden_dim_dec == 0:
                hidden_dim_dec = feature_dim
            self.hidden_dim_dec = _ntuple(num_layers_dec)(hidden_dim_dec)
        self.use_shift = use_shift
        act_dict = {
                    'none':torch.nn.Identity(), 'sigmoid':torch.nn.Sigmoid(), 'tanh':torch.nn.Tanh(),
                    'relu':torch.nn.ReLU(), 'sine':SineScaled(30.0)
                    }
        self.act = act_dict[activation]
        self.final_activation = act_dict[final_activation]
        self.clamp_weights = clamp_weights
        
        layers = []
        for l in range(num_layers_dec):
            feature_dim = self.hidden_dim_dec[l]
            feature_dim = latent_dim if feature_dim == 0 else feature_dim
            layers.append(DecoderLayer(latent_dim, feature_dim, ldecode_matrix, bias=self.use_shift))
            layers.append(self.act)
            latent_dim = feature_dim
        feature_dim = self.channels
        layers.append(DecoderLayer(latent_dim,feature_dim,ldecode_matrix,bias=self.use_shift))

        self.use_sga = use_sga
        self.temperature = 1.0
        self.layers = nn.Sequential(*layers)
        self.reset_parameters('normal', ldec_std)
        self.diff_sampling = diff_sampling
        
    def reset_parameters(self, init_type, param=0.5) -> None:
        for layer in list(self.layers.children()):
            if isinstance(layer, DecoderLayer):
                layer.reset_parameters(param,init_type)

    def get_scale(self):
        assert self.num_layers_dec == 0, "Can only get scale for 0 hidden layers decoder!"
        return list(self.layers.children())[0].scale

    def clamp(self, val: float = 0.2) -> None:
        for layer in list(self.layers.children()):
            if isinstance(layer, DecoderLayer):
                layer.clamp(val)

    def size(self, use_torchac=False):
        return sum([p.numel()*torch.finfo(p.dtype).bits for p in self.parameters()])

    def scale_norm(self):
        if self.num_layers_dec>0:
            print("Warning: norm is not implemented for multiple layer decoder>0, returning default value 1")
            return 1
        return list(self.layers.children())[0].scale.norm()

    def scale_grad_norm(self):
        if self.num_layers_dec>0:
            print("Warning: norm is not implemented for multiple layer decoder>0, returning default value 1")
            return 1
        return list(self.layers.children())[0].scale.grad.norm()
    
    def forward(self, weight: Tensor) -> Tensor:
        if self.use_sga:
            weightf = torch.floor(weight) if self.diff_sampling else StraightThroughFloor.apply(weight)
            weightc = weightf+1
            logits_f = -torch.tanh(torch.clamp(weight-weightf, min=-1+epsilon, max=1-epsilon)).unsqueeze(-1)/self.temperature
            logits_c = -torch.tanh(torch.clamp(weightc-weight, min=-1+epsilon, max=1-epsilon)).unsqueeze(-1)/self.temperature
            logits = torch.cat((logits_f,logits_c),dim=-1)
            dist = torch.distributions.relaxed_categorical.RelaxedOneHotCategorical(self.temperature, logits=logits)
            sample = dist.rsample() if self.diff_sampling else dist.sample()
            weight = weightf*sample[...,0]+weightc*sample[...,1]
        else:
            weight = StraightThrough.apply(weight)
        w_out = self.layers(weight/self.div)
        w_out = self.final_activation(w_out)
        if self.clamp_weights>0.0:
            w_out = torch.clamp(w_out, min=-self.clamp_weights, max=self.clamp_weights)
        return w_out


# Class for identity decoder with placeholder variables/functions
class DecoderIdentity(Module):

    def __init__(
        self,
    ) -> None:
        super(DecoderIdentity, self).__init__()
        self.latent_dim = 1
        self.num_layers_dec = 0
        self.shift = False
        self.norm = 'none'
        
    # For compatibility with Decoder
    def reset_parameters(self, init_type, param=1.0) -> None:
        return

    def forward(self, input: Tensor) -> Tensor:
        # print(input.min(),input.max())
        return input

    def scale_norm(self):
        return 1
    
    def scale_grad_norm(self):
        return 1
    
    def size(self, use_torchac=False) -> int:
        return 0


wisp/models/latent_decoders/basic_latent_decoder.py
Sharath-girish/Shacira
from wisp.models.latent_decoders.basic_latent_decoder import *

class HierarchicalLatentDecoder(nn.Module):

    def __init__(self, num_decoders, offsets, conf_decoder):
        super().__init__()
        self.num_decoders = num_decoders
        self.decoders = nn.ModuleList([LatentDecoder(**conf_decoder) for _ in range(self.num_decoders)])
        self.offsets = offsets

    def forward(self, input):
        output = torch.empty((input.size(0),self.decoders[0].channels)).to(input)
        for l in range(self.num_decoders):
            output[self.offsets[l]:self.offsets[l+1]] = self.decoders[l](input[self.offsets[l]:self.offsets[l+1]])
        return output

    @property
    def temperature(self):
        return self.decoders[0].temperature
    
    @property
    def use_sga(self):
        return self.decoders[0].use_sga
    
    @temperature.setter
    def temperature(self, value):
        for l in range(self.num_decoders):
            self.decoders[l].temperature = value

    @use_sga.setter
    def use_sga(self, value):
        for l in range(self.num_decoders):
            self.decoders[l].use_sga = value
            
    def size(self, use_torchac=False):
        return sum([p.numel()*torch.finfo(p.dtype).bits for p in self.parameters()])
wisp/models/latent_decoders/hierarchical_latent_decoder.py
Sharath-girish/Shacira
from wisp.models.latent_decoders.basic_latent_decoder import *
from torch.nn.modules.utils import _ntuple


class MultiSequential(nn.Sequential):
    def forward(self, input, alpha):
        for module in self._modules.values():
            if isinstance(module,MultiLatentDecoderLayer):
                input = module(input, alpha)
                # print(input.min(),input.max())
            else:
                input = module(input)
        return input

class StraightThroughOneHot(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x):
        out = torch.argmax(x,dim=0)
        out = torch.nn.functional.one_hot(out,num_classes=x.size(0)).permute(1,0).to(x)
        return out

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output

class MultiLatentDecoderLayer(Module):

    def __init__(self, in_features: int, out_features: int, ldecode_matrix: str, num_decoders: int = 1, bias: bool = False) -> None:
        super(MultiLatentDecoderLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        if 'dft' in ldecode_matrix:
            self.dft = Parameter(get_dft_matrix(in_features, out_features), requires_grad=False)
        if 'dft' in ldecode_matrix:
            self.scale = Parameter(torch.empty((num_decoders,1,out_features)))
        else:
            self.scale = Parameter(torch.empty((num_decoders,in_features,out_features)))
        if bias:
            self.use_shift = Parameter(torch.empty(num_decoders,1,out_features))
        else:
            self.register_parameter('use_shift', None)

        self.ldecode_matrix = ldecode_matrix
        if ldecode_matrix == 'dft_fixed':
            self.scale.requires_grad_(False)
            if not bias:
                self.use_shift.requires_grad_(False)

    def reset_parameters(self, param=1.0, init_type = 'normal') -> None:
        if init_type == 'normal':
            init.normal_(self.scale, std=param)
        elif init_type == 'uniform':
            init.uniform_(self.scale, -param, param)
        elif init_type == 'constant':
            init.constant_(self.scale, val=param)
        if self.use_shift is not None:
            init.zeros_(self.use_shift)

    def clamp(self, val: float = 0.5) -> None:
        with torch.no_grad():
            self.scale.clamp_(-val, val)

    def forward(self, input: Tensor, alpha: Tensor) -> Tensor:

        if 'dft' in self.ldecode_matrix:
            out = torch.matmul(input,self.dft).unsqueeze(0)
            w_out = out*self.scale+(self.use_shift if self.use_shift is not None else 0)
        else:
            # out = torch.matmul(input.unsqueeze(1).unsqueeze(0),self.scale.unsqueeze(1)).squeeze(2)
            out = torch.empty((self.scale.size(0),*input.size()),dtype=input.dtype,device=input.device)
            for i in range(self.scale.size(0)):
                out[i] = torch.matmul(input,self.scale[i])
            out = torch.sum(out*alpha.unsqueeze(-1),dim=0)
            w_out = out+(self.use_shift if self.use_shift is not None else 0)
        w_out = torch.sum(w_out*alpha.unsqueeze(-1),dim=0)
        return w_out

    def extra_repr(self) -> str:
        return 'in_features={}, out_features={}, bias={}'.format(
            self.in_features, self.out_features, self.use_shift is not None
        )

class MultiLatentDecoder(Module):

    def __init__(
        self,
        latent_dim: int,
        feature_dim: int,
        norm: str,
        ldecode_matrix:str,
        use_shift: bool,
        num_entries:int,
        num_layers_dec:int = 0,
        hidden_dim_dec:int = 0,
        activation:str = 'none',
        final_activation:str = 'none',
        clamp_weights:float = 0.0,
        ldec_std:float = 1.0,
        num_decoders:int = 1,
        alpha_std: float = 1.0,
        use_sga: bool = False,
        **kwargs,
    ) -> None:
        super(MultiLatentDecoder, self).__init__()
        latent_dim = feature_dim if latent_dim == 0 else latent_dim
        self.ldecode_matrix = ldecode_matrix
        self.channels = feature_dim
        self.latent_dim = latent_dim
        self.norm = norm
        self.div = 1 if norm == 'none' else None
        self.num_layers_dec =  num_layers_dec
        if num_layers_dec>0:
            if hidden_dim_dec == 0:
                hidden_dim_dec = feature_dim
            self.hidden_dim_dec = _ntuple(num_layers_dec)(hidden_dim_dec)
        self.use_shift = use_shift
        
        act_dict = {
                    'none':nn.Identity(), 
                    'sigmoid':nn.Sigmoid(), 
                    'tanh':nn.Tanh(),
                    'relu':nn.ReLU(), 
                    'sine':SineScaled(30.0)
                    }
        self.act = act_dict[activation]
        self.final_activation = act_dict[final_activation]
        self.clamp_weights = clamp_weights
        self.num_decoders = num_decoders

        layers = []
        for l in range(num_layers_dec):
            feature_dim = self.hidden_dim_dec[l]
            feature_dim = latent_dim if feature_dim == 0 else feature_dim
            layers.append(MultiLatentDecoderLayer(latent_dim, feature_dim, ldecode_matrix, num_decoders=num_decoders, bias=self.use_shift))
            layers.append(self.act)
            latent_dim = feature_dim
        feature_dim = self.channels
        layers.append(MultiLatentDecoderLayer(latent_dim,feature_dim,ldecode_matrix, num_decoders=num_decoders,bias=self.use_shift))

        self.alpha = nn.Parameter(torch.randn(num_decoders,num_entries)*alpha_std, requires_grad=True)

        self.temperature = 1.0
        self.layers = MultiSequential(*layers)
        self.reset_parameters('normal', ldec_std)
        self.straight_through = True
        self.use_sga = use_sga
        self.diff_sampling = False
        
    def reset_parameters(self, init_type, param=0.5) -> None:
        for layer in list(self.layers.children()):
            if isinstance(layer, MultiLatentDecoderLayer):
                layer.reset_parameters(param,init_type)

    def get_scale(self):
        assert self.num_layers_dec == 0, "Can only get scale for 0 hidden layers decoder!"
        return list(self.layers.children())[0].scale

    def clamp(self, val: float = 0.2) -> None:
        for layer in list(self.layers.children()):
            if isinstance(layer, MultiLatentDecoderLayer):
                layer.clamp(val)

    def size(self, use_torchac=False):
        fp_size = sum([p.numel()*torch.finfo(p.dtype).bits for n,p in self.named_parameters() if 'alpha' not in n])
        alpha = torch.argmax(self.alpha,dim=0)
        if not use_torchac:
            unique_vals, counts = torch.unique(alpha, return_counts = True)
            probs = counts/torch.sum(counts)
            information_bits = torch.clamp(-1.0 * torch.log(probs + 1e-10) / np.log(2.0), 0, 1000)
            size_bits = torch.sum(information_bits*counts).item()
            return size_bits+fp_size
        else:
            import torchac
            alpha = alpha - alpha.min()
            unique_vals, counts = torch.unique(alpha, return_counts = True)
            mapping = torch.zeros((alpha.max().item()+1))
            mapping[unique_vals] = torch.arange(unique_vals.size(0)).to(mapping)
            alpha = mapping[alpha]
            cdf = torch.cumsum(counts/counts.sum(),dim=0)
            cdf = torch.cat((torch.Tensor([0.0]),cdf))
            cdf = cdf.unsqueeze(0).repeat(alpha.size(0),1)
            cdf = cdf/cdf[:,-1:] # Normalize the final cdf value just to keep torchac happy
            byte_stream = torchac.encode_float_cdf(cdf.detach().cpu(), alpha.detach().cpu().to(torch.int16), \
                                                    check_input_bounds=True)
            return len(byte_stream)*8+fp_size

    def forward(self, weight: Tensor) -> Tensor:

        alpha = nn.functional.softmax(self.alpha/self.temperature,dim=0)
        if self.straight_through:
            alpha = StraightThroughOneHot.apply(alpha)

        if self.use_sga:
            weightf = torch.floor(weight) if self.diff_sampling else StraightThroughFloor.apply(weight)
            weightc = weightf+1
            logits_f = -torch.tanh(torch.clamp(weight-weightf, min=-1+epsilon, max=1-epsilon)).unsqueeze(-1)/self.temperature
            logits_c = -torch.tanh(torch.clamp(weightc-weight, min=-1+epsilon, max=1-epsilon)).unsqueeze(-1)/self.temperature
            logits = torch.cat((logits_f,logits_c),dim=-1)
            dist = torch.distributions.relaxed_categorical.RelaxedOneHotCategorical(self.temperature, logits=logits)
            sample = dist.rsample() if self.diff_sampling else dist.sample()
            weight = weightf*sample[...,0]+weightc*sample[...,1]
        else:
            weight = StraightThrough.apply(weight)
        w_out = self.layers(weight/self.div, alpha)
        w_out = self.final_activation(w_out)
        if self.clamp_weights>0.0:
            w_out = torch.clamp(w_out, min=-self.clamp_weights, max=self.clamp_weights)

        return w_out

wisp/models/latent_decoders/multi_latent_decoder.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import torch.nn as nn
import torch.nn.functional as F

def normalize_frobenius(x):
    """Normalizes the matrix according to the Frobenius norm.

    Args:
        x (torch.FloatTensor): A matrix.

    Returns:
        (torch.FloatTensor): A normalized matrix.
    """
    norm = torch.sqrt((torch.abs(x)**2).sum())
    return x / norm

def normalize_L_1(x):
    """Normalizes the matrix according to the L1 norm.

    Args:
        x (torch.FloatTensor): A matrix.

    Returns:
        (torch.FloatTensor): A normalized matrix.
    """
    abscolsum = torch.sum(torch.abs(x), dim=0)
    abscolsum = torch.min(torch.stack([1.0/abscolsum, torch.ones_like(abscolsum)], dim=0), dim=0)[0]
    return x * abscolsum[None,:]

def normalize_L_inf(x):    
    """Normalizes the matrix according to the Linf norm.

    Args:
        x (torch.FloatTensor): A matrix.

    Returns:
        (torch.FloatTensor): A normalized matrix.
    """
    absrowsum = torch.sum(torch.abs(x), axis=1)
    absrowsum = torch.min(torch.stack([1.0/absrowsum, torch.ones_like(absrowsum)], dim=0), dim=0)[0]
    return x * absrowsum[:,None]

class FrobeniusLinear(nn.Module):
    """A standard Linear layer which applies a Frobenius normalization in the forward pass.
    """
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.linear = nn.Linear(*args, **kwargs)

    def forward(self, x):
        weight = normalize_frobenius(self.linear.weight)
        return F.linear(x, weight, self.linear.bias)

class L_1_Linear(nn.Module):
    """A standard Linear layer which applies a L1 normalization in the forward pass.
    """
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.linear = nn.Linear(*args, **kwargs)

    def forward(self, x):
        weight = normalize_L_1(self.linear.weight)
        return F.linear(x, weight, self.linear.bias)

class L_inf_Linear(nn.Module):
    """A standard Linear layer which applies a Linf normalization in the forward pass.
    """
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.linear = nn.Linear(*args, **kwargs)

    def forward(self, x):
        weight = normalize_L_inf(self.linear.weight)
        return F.linear(x, weight, self.linear.bias)
        
def spectral_norm_(*args, **kwargs):
    """Initializes a spectral norm layer.
    """
    return nn.utils.spectral_norm(nn.Linear(*args, **kwargs))

def get_layer_class(layer_type):
    """Convenience function to return the layer class name from text.

    Args:
        layer_type (str): Text name for the layer.

    Retunrs:
        (nn.Module): The layer to be used for the decoder.
    """
    if layer_type == 'none' or layer_type == 'linear':
        return nn.Linear
    elif layer_type == 'spectral_norm':
        return spectral_norm_
    elif layer_type == 'frobenius_norm':
        return FrobeniusLinear
    elif layer_type == "l_1_norm":
        return L_1_Linear
    elif layer_type == "l_inf_norm":
        return L_inf_Linear
    else:
        assert(False and "layer type does not exist")

wisp/models/layers.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .base_nef import *
from .neural_sdf import *
from .nerf import *
from .image import *
from .spc_field import *

wisp/models/nefs/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import inspect
from abc import abstractmethod
from typing import Dict, Any
import torch
from wisp.core import WispModule


class BaseNeuralField(WispModule):
    """The base class for all Neural Fields within Wisp.
    Neural Fields are defined as modules which take coordinates as input and output signals of some form.
    The term "Neural" is loosely used here to imply these modules are generally subject for optimization.

    The domain of neural fields in Wisp is flexible, and left up for the user to decide when implementing the subclass.
    Popular neural fields from the literature, such as Neural Radiance Fields (Mildenhall et al. 2020),
    and Neural Signed Distance Functions (SDFs) can be implemented by creating and registering
    the required forward functions (for i.e. rgb, density, sdf values).

    BaseNeuralField subclasses  usually consist of several optional components:
    - A feature grid (BLASGrid), sometimes also known as 'hybrid representations'.
      These are responsible for querying and interpolating features, often in the context of some 3D volume
      (but not limited to).
      Feature grids often employ some acceleration structure (i.e. OctreeAS),
      which can be used to accelerate spatial queries or raytracing ops,
      hence the term "BLAS" (Bottom Level Acceleration Structure).
    - A decoder (i.e. BasicDecoder) which can feeds on features (or coordinates / pos embeddings) and coverts
      them to output signals.
    - Other components such as positional embedders may be employed.

    BaseNeuralFields are generally meant to be compatible with BaseTracers, thus forming a complete pipeline of
    render-able neural primitives.
    """
    def __init__(self):
        super().__init__()
        self._forward_functions = {}
        self.register_forward_functions()
        self.supported_channels = set([channel for channels in self._forward_functions.values() for channel in channels])

    @property
    def device(self):
        """ Returns the device used to process inputs in this neural field.
        By default, the device is queried from the first registered torch nn.parameter.
        Override this property to explicitly specify the device.

        Returns:
            (torch.device): The expected device for inputs to this neural field.
        """
        return next(self.parameters()).device

    def _register_forward_function(self, fn, channels):
        """Registers a forward function.

        Args:
            fn (function): Function to register.
            channels (list of str): Channel output names.
        """
        if isinstance(channels, str):
            channels = [channels]
        self._forward_functions[fn] = set(channels)

    @abstractmethod
    def register_forward_functions(self):
        """Register forward functions with the channels that they output.
        
        This function should be overrided and call `self._register_forward_function` to 
        tell the class which functions output what output channels. The function can be called
        multiple times to register multiple functions.

        Example:

        ```
        self._register_forward_function(self.rgba, ["density", "rgb"])
        self._register_forward_function(self.sdf, ["sdf"])
        ```
        """
        pass

    def get_forward_function(self, channel):
        """Will return the function that will return the channel.
        
        Args: 
            channel (str): The name of the channel to return.

        Returns:
            (function): Function that will return the function. Will return None if the channel is not supported.
        """
        if channel not in self.get_supported_channels():
            raise Exception(f"Channel {channel} is not supported in {self.__class__.__name__}")
        for fn in self._forward_functions:
            output_channels = self._forward_functions[fn]
            if channel in output_channels:
                return lambda *args, **kwargs: fn(*args, **kwargs)[channel]

    def get_supported_channels(self):
        """Returns the channels that are supported by this class.

        Returns:
            (set): Set of channel strings.
        """
        return self.supported_channels

    def forward(self, channels=None, **kwargs):
        """Queries the neural field with channels.

        Args:
            channels (str or list of str or set of str): Requested channels. See return value for details.
            kwargs: Any keyword argument passed in will be passed into the respective forward functions.

        Returns:
            (list or dict or torch.Tensor): 
                If channels is a string, will return a tensor of the request channel. 
                If channels is a list, will return a list of channels.
                If channels is a set, will return a dictionary of channels.
                If channels is None, will return a dictionary of all channels.
        """
        if not (isinstance(channels, str) or isinstance(channels, list) or isinstance(channels, set) or channels is None):
            raise Exception(f"Channels type invalid, got {type(channels)}." \
                             "Make sure your arguments for the nef are provided as keyword arguments.")
        if channels is None:
            requested_channels = self.get_supported_channels()
        elif isinstance(channels, str):
            requested_channels = set([channels])
        else:
            requested_channels = set(channels)

        unsupported_channels = requested_channels - self.get_supported_channels()
        if unsupported_channels:
            raise Exception(f"Channels {unsupported_channels} are not supported in {self.__class__.__name__}")
        
        return_dict = {}
        for fn in self._forward_functions:
            torch.cuda.nvtx.range_push(f"{fn.__name__}")
            output_channels = self._forward_functions[fn]
            # Filter the set of channels supported by the current forward function
            supported_channels = output_channels & requested_channels

            # Check that the function needs to be executed
            if len(supported_channels) != 0:

                # Filter args to the forward function and execute
                argspec = inspect.getfullargspec(fn)
                required_args = argspec.args[:-len(argspec.defaults)][1:] # Skip first element, self
                optional_args = argspec.args[-len(argspec.defaults):]
                
                input_args = {}
                for _arg in required_args:
                    # TODO(ttakiakwa): This doesn't actually format the string, fix :) 
                    if _arg not in kwargs:
                        raise Exception(f"Argument {_arg} not found as input to in {self.__class__.__name__}.{fn.__name__}()")
                    input_args[_arg] = kwargs[_arg]
                for _arg in optional_args:
                    if _arg in kwargs:
                        input_args[_arg] = kwargs[_arg]
                output = fn(**input_args)

                for channel in supported_channels:
                    return_dict[channel] = output[channel]
            torch.cuda.nvtx.range_pop()
        
        if isinstance(channels, str):
            if channels in return_dict:
                return return_dict[channels]
            else:
                return None
        elif isinstance(channels, list):
            return [return_dict[channel] for channel in channels]
        else:
            return return_dict

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        return dict()

wisp/models/nefs/base_nef.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from typing import Dict, Any
from wisp.ops.geometric import sample_unif_sphere
from wisp.models.nefs import BaseNeuralField
from wisp.models.embedders import get_positional_embedder
from wisp.models.layers import get_layer_class
from wisp.models.activations import get_activation_class
from wisp.models.decoders import BasicDecoder
from wisp.models.grids import BLASGrid, HashGrid, LatentGrid

class NeuralImage(BaseNeuralField):
    """Model for encoding images as neural fields (RGB values).
    This implementation uses feature grids for a
    higher quality and more efficient implementation, following later trends in the literature,
    such as Neural Sparse Voxel Fields (Liu et al. 2020), Instant Neural Graphics Primitives (Muller et al. 2022)
    and Variable Bitrate Neural Fields (Takikawa et al. 2022).
    """

    def __init__(self,
                 grid: BLASGrid = None,
                 # embedder args
                 pos_embedder: str = 'none',
                 pos_multires: int = 10,
                 position_input: bool = False,
                 # decoder args
                 activation_type: str = 'relu',
                 final_activation: str = 'none',
                 layer_type: str = 'none',
                 hidden_dim: int = 128,
                 num_layers: int = 1,
                 ):
        """
        Creates a new NeuralImage instance, which maps 2D input coordinates to RGB

        This neural field consists of:
         * A feature grid (backed by an acceleration structure)
         * Color decoder
         * Optional: positional embedders for input position coords, concatenated to grid features.

         This neural field also supports:
          * Aggregation of multi-resolution features (more than one LOD) via summation or concatenation

        Args:
            grid: (BLASGrid): represents feature grids in Wisp. BLAS: "Bottom Level Acceleration Structure",
                to signify this structure is the backbone that captures
                a neural field's contents, in terms of both features and occupancy for speeding up queries.
                Notable examples: OctreeGrid, HashGrid, LatentGrid, TriplanarGrid, CodebookGrid.

            pos_embedder (str): Type of positional embedder to use for input coordinates.
                Options:
                 - 'none': No positional input is fed into the color decoder.
                 - 'identity': The sample coordinates are fed as is into the color decoder.
                 - 'positional': The sample coordinates are embedded with the Positional Encoding from
                    Mildenhall et al. 2020, before passing them into the color decoder.
            pos_multires (int): Number of frequencies used for 'positional' embedding of pos_embedder.
                 Used only if pos_embedder is 'positional'.
            position_input (bool): If True, the input coordinates will be passed into the decoder.
                 For 'positional': the input coordinates will be concatenated to the embedded coords.
                 For 'none' and 'identity': the embedder will behave like 'identity'.
            activation_type (str): Type of activation function to use in BasicDecoder:
                 'none', 'relu', 'sin', 'fullsort', 'minmax'.
            final_activation (str): Type of activation function to use after BasicDecoder:
                 'none', 'sigmoid', 'relu'.
            layer_type (str): Type of MLP layer to use in BasicDecoder:
                 'none' / 'linear', 'spectral_norm', 'frobenius_norm', 'l_1_norm', 'l_inf_norm'.
            hidden_dim (int): Number of neurons in hidden layers of both decoders.
            num_layers (int): Number of hidden layers in both decoders.
        """
        super().__init__()
        self.grid = grid

        # Init Embedders
        self.pos_embedder, self.pos_embed_dim = self.init_embedder(pos_embedder, pos_multires,
                                                                   include_input=position_input)

        # Init Decoder
        self.activation_type = activation_type
        self.final_activation = get_activation_class(final_activation)
        self.layer_type = layer_type
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.decoder_color = self.init_decoders(activation_type, layer_type, num_layers, hidden_dim)

        torch.cuda.empty_cache()

    def init_embedder(self, embedder_type, frequencies=None, include_input=False):
        """Creates positional embedding functions for the position.
        """
        if embedder_type == 'none' and not include_input:
            embedder, embed_dim = None, 0
        elif embedder_type == 'identity' or (embedder_type == 'none' and include_input):
            embedder, embed_dim = torch.nn.Identity(), 2    # Assumes pos input is always 2D
        elif embedder_type == 'positional':
            embedder, embed_dim = get_positional_embedder(frequencies=frequencies, 
                                                          input_dim=2,
                                                          include_input=include_input)
        else:
            raise NotImplementedError(f'Unsupported embedder type for NeuralImage: {embedder_type}')
        return embedder, embed_dim

    def init_decoders(self, activation_type, layer_type, num_layers, hidden_dim):
        """Initializes the decoder object.
        """
        decoder_color = BasicDecoder(input_dim=self.color_net_input_dim(),
                                     output_dim=3,
                                     activation=get_activation_class(activation_type),
                                     bias=True,
                                     layer=get_layer_class(layer_type),
                                     num_layers=num_layers + 1,
                                     hidden_dim=hidden_dim,
                                     skip=[])
        return decoder_color

    def register_forward_functions(self):
        """Register the forward functions.
        """
        self._register_forward_function(self.rgb, ["rgb"])

    def rgb(self, coords, lod_idx=None):
        """Compute color for the provided coordinates.

        Args:
            coords (torch.FloatTensor): tensor of shape [batch, 2]
            lod_idx (int): index into active_lods. If None, will use the maximum LOD.
        
        Returns:
            {"rgb": torch.FloatTensor}:
                - RGB tensor of shape [batch, 3]
        """
        if lod_idx is None:
            lod_idx = len(self.grid.active_lods) - 1
        batch, _ = coords.shape

        # Embed coordinates into high-dimensional vectors with the grid.
        feats = self.grid.interpolate(coords, lod_idx).reshape(batch, self.effective_feature_dim())

        # Optionally concat the positions to the embedding
        if self.pos_embedder is not None:
            embedded_pos = self.pos_embedder(coords).view(batch, self.pos_embed_dim)
            feats = torch.cat([feats, embedded_pos], dim=-1)

        # Colors are values [0, 1] floats (if final activation clips)
        # colors ~ (batch, 3)
        colors = self.final_activation(self.decoder_color(feats))

        return dict(rgb=colors)

    def effective_feature_dim(self):
        if self.grid.multiscale_type == 'cat':
            effective_feature_dim = self.grid.feature_dim * self.grid.num_lods
        else:
            effective_feature_dim = self.grid.feature_dim
        return effective_feature_dim

    def color_net_input_dim(self):
        return self.effective_feature_dim() + self.pos_embed_dim

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        properties = {
            "Grid": self.grid,
            "Pos. Embedding": self.pos_embedder,
            "Decoder (color)": self.decoder_color,
            "Final Activation": self.final_activation
        }
        return properties

wisp/models/nefs/image.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from typing import Dict, Any
from wisp.ops.geometric import sample_unif_sphere
from wisp.models.nefs import BaseNeuralField
from wisp.models.embedders import get_positional_embedder
from wisp.models.layers import get_layer_class
from wisp.models.activations import get_activation_class
from wisp.models.decoders import BasicDecoder
from wisp.models.grids import BLASGrid, HashGrid, LatentGrid

class NeuralRadianceField(BaseNeuralField):
    """Model for encoding Neural Radiance Fields (Mildenhall et al. 2020), e.g., density and view dependent color.
    Different to the original NeRF paper, this implementation uses feature grids for a
    higher quality and more efficient implementation, following later trends in the literature,
    such as Neural Sparse Voxel Fields (Liu et al. 2020), Instant Neural Graphics Primitives (Muller et al. 2022)
    and Variable Bitrate Neural Fields (Takikawa et al. 2022).
    """

    def __init__(self,
                 grid: BLASGrid = None,
                 # embedder args
                 pos_embedder: str = 'none',
                 view_embedder: str = 'none',
                 pos_multires: int = 10,
                 view_multires: int = 4,
                 position_input: bool = False,
                 # decoder args
                 activation_type: str = 'relu',
                 layer_type: str = 'none',
                 hidden_dim: int = 128,
                 num_layers: int = 1,
                 # pruning args
                 prune_density_decay: float = None,
                 prune_min_density: float = None,
                 ):
        """
        Creates a new NeRF instance, which maps 3D input coordinates + view directions to RGB + density.

        This neural field consists of:
         * A feature grid (backed by an acceleration structure to boost raymarching speed)
         * Color & density decoders
         * Optional: positional embedders for input position coords & view directions, concatenated to grid features.

         This neural field also supports:
          * Aggregation of multi-resolution features (more than one LOD) via summation or concatenation
          * Pruning scheme for HashGrids

        Args:
            grid: (BLASGrid): represents feature grids in Wisp. BLAS: "Bottom Level Acceleration Structure",
                to signify this structure is the backbone that captures
                a neural field's contents, in terms of both features and occupancy for speeding up queries.
                Notable examples: OctreeGrid, HashGrid, TriplanarGrid, CodebookGrid.

            pos_embedder (str): Type of positional embedder to use for input coordinates.
                Options:
                 - 'none': No positional input is fed into the density decoder.
                 - 'identity': The sample coordinates are fed as is into the density decoder.
                 - 'positional': The sample coordinates are embedded with the Positional Encoding from
                    Mildenhall et al. 2020, before passing them into the density decoder.
            view_embedder (str): Type of positional embedder to use for view directions.
                Options:
                 - 'none': No positional input is fed into the color decoder.
                 - 'identity': The view directions are fed as is into the color decoder.
                 - 'positional': The view directions are embedded with the Positional Encoding from
                    Mildenhall et al. 2020, before passing them into the color decoder.
            pos_multires (int): Number of frequencies used for 'positional' embedding of pos_embedder.
                 Used only if pos_embedder is 'positional'.
            view_multires (int): Number of frequencies used for 'positional' embedding of view_embedder.
                 Used only if view_embedder is 'positional'.
            position_input (bool): If True, the input coordinates will be passed into the decoder.
                 For 'positional': the input coordinates will be concatenated to the embedded coords.
                 For 'none' and 'identity': the embedder will behave like 'identity'.
            activation_type (str): Type of activation function to use in BasicDecoder:
                 'none', 'relu', 'sin', 'fullsort', 'minmax'.
            layer_type (str): Type of MLP layer to use in BasicDecoder:
                 'none' / 'linear', 'spectral_norm', 'frobenius_norm', 'l_1_norm', 'l_inf_norm'.
            hidden_dim (int): Number of neurons in hidden layers of both decoders.
            num_layers (int): Number of hidden layers in both decoders.
            prune_density_decay (float): Decay rate of density per "prune step",
                 using the pruning scheme from Muller et al. 2022. Used only for grids which support pruning.
            prune_min_density (float): Minimal density allowed for "cells" before they get pruned during a "prune step".
                 Used within the pruning scheme from Muller et al. 2022. Used only for grids which support pruning.
        """
        super().__init__()
        self.grid = grid

        # Init Embedders
        self.pos_embedder, self.pos_embed_dim = self.init_embedder(pos_embedder, pos_multires,
                                                                   include_input=position_input)
        self.view_embedder, self.view_embed_dim = self.init_embedder(view_embedder, view_multires,
                                                                     include_input=True)

        # Init Decoder
        self.activation_type = activation_type
        self.layer_type = layer_type
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.decoder_density, self.decoder_color = \
            self.init_decoders(activation_type, layer_type, num_layers, hidden_dim)

        self.prune_density_decay = prune_density_decay
        self.prune_min_density = prune_min_density

        torch.cuda.empty_cache()

    def init_embedder(self, embedder_type, frequencies=None, include_input=False):
        """Creates positional embedding functions for the position and view direction.
        """
        if embedder_type == 'none' and not include_input:
            embedder, embed_dim = None, 0
        elif embedder_type == 'identity' or (embedder_type == 'none' and include_input):
            embedder, embed_dim = torch.nn.Identity(), 3    # Assumes pos / view input is always 3D
        elif embedder_type == 'positional':
            embedder, embed_dim = get_positional_embedder(frequencies=frequencies, include_input=include_input)
        else:
            raise NotImplementedError(f'Unsupported embedder type for NeuralRadianceField: {embedder_type}')
        return embedder, embed_dim

    def init_decoders(self, activation_type, layer_type, num_layers, hidden_dim):
        """Initializes the decoder object.
        """
        decoder_density = BasicDecoder(input_dim=self.density_net_input_dim(),
                                       output_dim=16,
                                       activation=get_activation_class(activation_type),
                                       bias=True,
                                       layer=get_layer_class(layer_type),
                                       num_layers=num_layers,
                                       hidden_dim=hidden_dim,
                                       skip=[])
        decoder_density.lout.bias.data[0] = 1.0

        decoder_color = BasicDecoder(input_dim=self.color_net_input_dim(),
                                     output_dim=3,
                                     activation=get_activation_class(activation_type),
                                     bias=True,
                                     layer=get_layer_class(layer_type),
                                     num_layers=num_layers + 1,
                                     hidden_dim=hidden_dim,
                                     skip=[])
        return decoder_density, decoder_color

    def prune(self):
        """Prunes the blas based on current state.
        """
        if self.grid is not None:
            if isinstance(self.grid, HashGrid) or isinstance(self.grid, LatentGrid):
                density_decay = self.prune_density_decay
                min_density = self.prune_min_density

                self.grid.occupancy = self.grid.occupancy.cuda()
                self.grid.occupancy = self.grid.occupancy * density_decay
                points = self.grid.dense_points.cuda()
                res = 2.0**self.grid.blas_level
                samples = torch.rand(points.shape[0], 3, device=points.device)
                samples = points.float() + samples
                samples = samples / res
                samples = samples * 2.0 - 1.0
                sample_views = torch.FloatTensor(sample_unif_sphere(samples.shape[0])).to(points.device)
                with torch.no_grad():
                    density = self.forward(coords=samples, ray_d=sample_views, channels="density")
                self.grid.occupancy = torch.stack([density[:, 0], self.grid.occupancy], -1).max(dim=-1)[0]

                mask = self.grid.occupancy > min_density

                _points = points[mask]

                if _points.shape[0] == 0:
                    return

                if hasattr(self.grid.blas.__class__, "from_quantized_points"):
                    self.grid.blas = self.grid.blas.__class__.from_quantized_points(_points, self.grid.blas_level)
                else:
                    raise Exception(f"The BLAS {self.grid.blas.__class__.__name__} does not support initialization " 
                                     "from_quantized_points, which is required for pruning.")

            else:
                raise NotImplementedError(f'Pruning not implemented for grid type {self.grid}')

    def register_forward_functions(self):
        """Register the forward functions.
        """
        self._register_forward_function(self.rgba, ["density", "rgb"])

    def rgba(self, coords, ray_d, lod_idx=None):
        """Compute color and density [particles / vol] for the provided coordinates.

        Args:
            coords (torch.FloatTensor): tensor of shape [batch, 3]
            ray_d (torch.FloatTensor): tensor of shape [batch, 3]
            lod_idx (int): index into active_lods. If None, will use the maximum LOD.
        
        Returns:
            {"rgb": torch.FloatTensor, "density": torch.FloatTensor}:
                - RGB tensor of shape [batch, 3]
                - Density tensor of shape [batch, 1]
        """
        if lod_idx is None:
            lod_idx = len(self.grid.active_lods) - 1
        batch, _ = coords.shape

        # Embed coordinates into high-dimensional vectors with the grid.
        feats = self.grid.interpolate(coords, lod_idx).reshape(batch, self.effective_feature_dim())

        # Optionally concat the positions to the embedding
        if self.pos_embedder is not None:
            embedded_pos = self.pos_embedder(coords).view(batch, self.pos_embed_dim)
            feats = torch.cat([feats, embedded_pos], dim=-1)

        # Decode high-dimensional vectors to density features.
        density_feats = self.decoder_density(feats)

        # Concatenate embedded view directions.
        if self.view_embedder is not None:
            embedded_dir = self.view_embedder(-ray_d).view(batch, self.view_embed_dim)
            fdir = torch.cat([density_feats, embedded_dir], dim=-1)
        else:
            fdir = density_feats

        # Colors are values [0, 1] floats
        # colors ~ (batch, 3)
        colors = torch.sigmoid(self.decoder_color(fdir))

        # Density is [particles / meter], so need to be multiplied by distance
        # density ~ (batch, 1)
        density = torch.relu(density_feats[...,0:1])
        return dict(rgb=colors, density=density)

    def effective_feature_dim(self):
        if self.grid.multiscale_type == 'cat':
            effective_feature_dim = self.grid.feature_dim * self.grid.num_lods
        else:
            effective_feature_dim = self.grid.feature_dim
        return effective_feature_dim

    def density_net_input_dim(self):
        return self.effective_feature_dim() + self.pos_embed_dim

    def color_net_input_dim(self):
        return 16 + self.view_embed_dim

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        properties = {
            "Grid": self.grid,
            "Pos. Embedding": self.pos_embedder,
            "View Embedding": self.view_embedder,
            "Decoder (density)": self.decoder_density,
            "Decoder (color)": self.decoder_color
        }
        if self.prune_density_decay is not None:
            properties['Pruning Density Decay'] = self.prune_density_decay
        if self.prune_min_density is not None:
            properties['Pruning Min Density'] = self.prune_min_density
        return properties

wisp/models/nefs/nerf.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from typing import Dict, Any

from wisp.models.nefs import BaseNeuralField
from wisp.models.embedders import get_positional_embedder
from wisp.models.layers import get_layer_class
from wisp.models.activations import get_activation_class
from wisp.models.decoders import BasicDecoder
from wisp.models.grids import BLASGrid


class NeuralSDF(BaseNeuralField):
    """Model for encoding neural signed distance functions (implicit surfaces).
    This field implementation uses feature grids for faster and more efficient queries.
    For example, the usage of Octree follows the idea from Takikawa et al. 2021 (Neural Geometric Level of Detail).
    """
    def __init__(self,
                 grid: BLASGrid = None,
                 # embedder args
                 pos_embedder: str = 'none',
                 pos_multires: int = 10,
                 position_input: bool = True,
                 # decoder args
                 activation_type: str = 'relu',
                 layer_type: str = 'none',
                 hidden_dim: int = 128,
                 num_layers: int = 1
                 ):
        """
        Creates a new neural field of an implicit furface, which maps 3D input coordinates to SDF values.

        This neural field consists of:
         * A feature grid (backed by an acceleration structure to boost raymarching speed)
         * SDF decoder
         * Optional: positional embedders for input position coords, concatenated to grid features.

         This neural field also supports:
          * Aggregation of multi-resolution features (more than one LOD) via summation or concatenation

        Args:
            grid: (BLASGrid): represents feature grids in Wisp. BLAS: "Bottom Level Acceleration Structure",
                to signify this structure is the backbone that captures
                a neural field's contents, in terms of both features and occupancy for speeding up queries.
                Notable examples: OctreeGrid, HashGrid, TriplanarGrid.

            pos_embedder (str): Type of positional embedder to use for input coordinates.
                Options:
                 - 'none': No positional input is fed into the SDF decoder.
                 - 'identity': The sample coordinates are fed as is into the SDF decoder.
                 - 'positional': The sample coordinates are embedded with the Positional Encoding from
                    Mildenhall et al. 2020, before passing them into the SDF decoder.
            pos_multires (int): Number of frequencies used for 'positional' embedding of pos_embedder.
                 Used only if pos_embedder is 'positional'.
            position_input (bool): If True, the input coordinates will be passed into the decoder.
                 For 'positional': the input coordinates will be concatenated to the embedded coords.
                 For 'none' and 'identity': the embedder will behave like 'identity'.
            activation_type (str): Type of activation function to use in BasicDecoder:
                 'none', 'relu', 'sin', 'fullsort', 'minmax'.
            layer_type (str): Type of MLP layer to use in BasicDecoder:
                 'none' / 'linear', 'spectral_norm', 'frobenius_norm', 'l_1_norm', 'l_inf_norm'.
            hidden_dim (int): Number of neurons in hidden layers of SDF decoder.
            num_layers (int): Number of hidden layers in SDF decoder.
        """
        super().__init__()
        self.grid = grid

        # Init Embedders
        self.pos_multires = pos_multires
        self.position_input = position_input
        self.pos_embedder, self.pos_embed_dim = self.init_embedder(pos_embedder, pos_multires, position_input)

        # Init Decoder
        self.activation_type = activation_type
        self.layer_type = layer_type
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.decoder = self.init_decoder(activation_type, layer_type, num_layers, hidden_dim)

        torch.cuda.empty_cache()

    def init_embedder(self, embedder_type, frequencies=None, position_input=True):
        """Creates positional embedding functions for the position and view direction.
        """
        if embedder_type == 'none' and not position_input:
            embedder, embed_dim = None, 0
        elif embedder_type == 'identity' or (embedder_type == 'none' and position_input):
            embedder, embed_dim = torch.nn.Identity(), 3
        elif embedder_type == 'positional':
            embedder, embed_dim = get_positional_embedder(frequencies=frequencies, position_input=position_input)
        else:
            raise NotImplementedError(f'Unsupported embedder type for NeuralSDF: {embedder_type}')
        return embedder, embed_dim

    def init_decoder(self, activation_type, layer_type, num_layers, hidden_dim):
        """Initializes the decoder object.
        """
        decoder = BasicDecoder(input_dim=self.decoder_input_dim(),
                               output_dim=1,
                               activation=get_activation_class(activation_type),
                               bias=True,
                               layer=get_layer_class(layer_type),
                               num_layers=num_layers,
                               hidden_dim=hidden_dim,
                               skip=[])
        return decoder

    def register_forward_functions(self):
        """Register the forward functions.
        """
        self._register_forward_function(self.sdf, ["sdf"])

    def sdf(self, coords, lod_idx=None):
        """Computes the Signed Distance Function for input samples.

        Args:
            coords (torch.FloatTensor): tensor of shape [batch, num_samples, 3]
            lod_idx (int): index into active_lods. If None, will use the maximum LOD.
        
        Outputs:
            (torch.FloatTensor):
            - SDF of shape [batch, num_samples, 1]
        """
        shape = coords.shape
        
        if shape[0] == 0:
            return dict(sdf=torch.zeros_like(coords)[...,0:1])

        if lod_idx is None:
            lod_idx = self.grid.num_lods - 1

        if len(shape) == 2:
            coords = coords[:, None]
        num_samples = coords.shape[1]

        feats = self.grid.interpolate(coords, lod_idx)

        # Optionally concat the positions to the embedding
        if self.pos_embedder is not None:
            feats = torch.cat([self.pos_embedder(coords.view(-1, 3)).view(-1, num_samples, self.pos_embed_dim), 
                               feats], dim=-1)

        sdf = self.decoder(feats)

        if len(shape) == 2:
            sdf = sdf[:,0]
            
        return dict(sdf=sdf)

    def effective_feature_dim(self):
        if self.grid.multiscale_type == 'cat':
            effective_feature_dim = self.grid.feature_dim * self.grid.num_lods
        else:
            effective_feature_dim = self.grid.feature_dim
        return effective_feature_dim

    def decoder_input_dim(self):
        input_dim = self.effective_feature_dim()
        if self.position_input:
            input_dim += self.pos_embed_dim
        return input_dim

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        properties = {
            "Grid": self.grid,
            "Pos. Embedding": self.pos_embedder,
            "Decoder (sdf)": self.decoder,
        }
        return properties

wisp/models/nefs/neural_sdf.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import logging as log

from wisp.models.nefs import BaseNeuralField
from wisp.models.embedders import get_positional_embedder
from wisp.models.layers import get_layer_class
from wisp.models.activations import get_activation_class
from wisp.models.decoders import BasicDecoder
from wisp.models.grids import *


class NeuralSDFTex(BaseNeuralField):
    """Model for encoding neural signed distance functions + plenoptic color, e.g., implicit surfaces with albedo.
    """

    def __init__(self,
                 grid: BLASGrid = None,
                 # embedder args
                 embedder_type: str = 'none',
                 pos_multires: int = 10,
                 # decoder args
                 activation_type: str = 'relu',
                 layer_type: str = 'none',
                 hidden_dim: int = 128,
                 num_layers: int = 1
                 ):
        super().__init__()
        self.grid = grid

        # Init Embedders
        self.embedder_type = embedder_type
        self.pos_multires = pos_multires
        self.pos_embedder, self.pos_embed_dim = self.init_embedder(embedder_type, pos_multires)

        # Init Decoder
        self.activation_type = activation_type
        self.layer_type = layer_type
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.position_input = embedder_type != 'none'
        self.decoder, self.effective_feature_dim, self.input_dim = \
            self.init_decoder(activation_type, layer_type, num_layers, hidden_dim,
                              self.position_input, self.pos_embed_dim)

        torch.cuda.empty_cache()

    def init_embedder(self, embedder_type, pos_multires):
        """Creates positional embedding functions for the position and view direction.
        """
        is_active = embedder_type == "positional"
        pos_embedder, pos_embed_dim = get_positional_embedder(frequencies=pos_multires, active=is_active)
        log.info(f"Position Embed Dim: {pos_embed_dim}")
        return pos_embedder, pos_embed_dim

    def init_decoder(self, activation_type, layer_type, num_layers, hidden_dim, position_input, pos_embed_dim):
        """Initializes the decoder object.
        """
        if self.grid.multiscale_type == 'cat':
            effective_feature_dim = self.grid.feature_dim * self.num_lods
        else:
            effective_feature_dim = self.grid.feature_dim
        
        input_dim = effective_feature_dim
        if position_input:
            input_dim += pos_embed_dim

        decoder = BasicDecoder(input_dim=input_dim,
                               output_dim=4,
                               activation=get_activation_class(activation_type),
                               bias=True,
                               layer=get_layer_class(layer_type),
                               num_layers=num_layers,
                               hidden_dim=hidden_dim,
                               skip=[])
        return decoder, effective_feature_dim, input_dim

    def register_forward_functions(self):
        """Register the forward functions.
        """
        self._register_forward_function(self.rgbsdf, ["rgb", "sdf"])

    def rgbsdf(self, coords, lod_idx=None):
        """Computes the RGB + SDF for some samples.

        Args:
            coords (torch.FloatTensor): packed tensor of shape [batch, num_samples, 3]
            lod_idx (int): index into active_lods. If None, will use the maximum LOD.
        
        Outputs:
            {"rgb": torch.FloatTensor, "sdf": torch.FloatTensor}:
            - RGB of shape [batch, num_samples, 3]
            - SDF of shape [batch, num_samples, 1]
        """
        shape = coords.shape
        
        if shape[0] == 0:
            return dict(rgb=torch.zeros_like(coords)[...,:3], sdf=torch.zeros_like(coords)[...,0:1])

        if lod_idx is None:
            lod_idx = self.num_lods - 1
        
        if len(shape) == 2:
            coords = coords[:, None]

        feats = self.grid.interpolate(coords, lod_idx)

        if self.position_input:
            feats = torch.cat([self.pos_embedder(coords), feats], dim=-1)

        rgbsdf = self.decoder(feats)

        if len(shape) == 2:
            rgbsdf = rgbsdf[:,0]
            
        return dict(rgb=torch.sigmoid(rgbsdf[...,:3]), sdf=rgbsdf[...,3:4])

wisp/models/nefs/neural_sdf_tex.py
Sharath-girish/Shacira
from typing import Dict, Any
import numpy as np
import torch
from wisp.models.grids import OctreeGrid
from wisp.models.nefs.nerf import NeuralRadianceField, BaseNeuralField
import wisp.ops.spc as wisp_spc_ops
import kaolin.ops.spc as kaolin_ops_spc


class SPCField(BaseNeuralField):
    """ A field based on Structured Point Clouds (SPC) from kaolin.
    SPC is a hierarchical compressed data structure, which can be interpreted in various ways:
    * Quantized point cloud, where each sparse point is quantized to some (possibly very dense) grid.
      Each point is associated with some feature(s).
    * An Octree, where each cell center is represented by a quantized point.
    Throughout wisp, SPCs are used to implement efficient octrees or grid structures.
    This field class allows wisp to render SPCs directly with their feature content (hence no embedders or decoders
    are assumed).

    When rendered, SPCs behave like octrees which allow for efficient tracing.
    Feature samples per ray may be collected from each intersected "cell" of the structured point cloud.
    """

    def __init__(self, spc_octree, features_dict=None, device='cuda'):
        r"""Creates a new Structured Point Cloud (SPC), represented as a Wisp Field.

        In wisp, SPCs are considered neural fields, since their features may be optimized.
        See `examples/spc_browser` for an elaborate description of SPCs.

        Args:
            spc_octree (torch.ByteTensor):
                A tensor which holds the topology of the SPC.
                Each byte represents a single octree cell's occupancy (that is, each bit of that byte represents
                the occupancy status of a child octree cell), yielding 8 bits for 8 cells.
                See also https://kaolin.readthedocs.io/en/latest/notes/spc_summary.html
            features_dict (dict):
                A dictionary holding the features information of the SPC.
                Keys are assumed to be a subset of ('colors', 'normals').
                Values are torch feature tensors containing information per point, of shape
                :math:`(\text{num_points}, \text{feat_dim})`.
                Where `num_points` is the number of occupied cells in the SPC.
                See `kaolin.ops.conversions.pointcloud.unbatched_pointcloud_to_spc` for conversion of point
                cloud information to such features.
            device (torch.device):
                Torch device on which the features and topology of the SPC field will be stored.
        """
        super().__init__()
        self.spc_octree = spc_octree
        self.features_dict = features_dict if features_dict is not None else dict()
        self.spc_device = device
        self.grid = None
        self.colors = None
        self.normals = None
        self.init_grid(spc_octree)

    def init_grid(self, spc_octree):
        """ Uses the OctreeAS / OctreeGrid mechanism to quickly parse the SPC object into a Wisp Neural Field.

        Args:
            spc_octree (torch.ByteTensor):
                A tensor which holds the topology of the SPC.
                Each byte represents a single octree cell's occupancy (that is, each bit of that byte represents
                the occupancy status of a child octree cell), yielding 8 bits for 8 cells.
                See also https://kaolin.readthedocs.io/en/latest/notes/spc_summary.html
        """
        # Use features, either colors or normals
        spc_features = self.features_dict
        if "colors" in self.features_dict:
            colors = spc_features["colors"]
            colors = colors.reshape(-1, 4) / 255.0
            self.colors = colors
        if "normals" in self.features_dict:
            normals = spc_features["normals"]
            normals = normals.reshape(-1, 3)
            self.normals = normals

        if self.colors is None:  # If no color features exist, use normals information
            if self.normals is not None:
                colors = 0.5 * (normals + 1.0)
            else:
                # manufacture colors from point coordinates
                lengths = torch.tensor([len(spc_octree)], dtype=torch.int32)
                level, pyramids, exsum = kaolin_ops_spc.scan_octrees(spc_octree, lengths)
                point_hierarchies = kaolin_ops_spc.generate_points(spc_octree, pyramids, exsum)
                # get coordinate of highest level
                colors = point_hierarchies[pyramids[0,1,level]:]
                # normalize
                colors = colors/np.power(2, level)

            self.colors = colors

        # By default assume the SPC keeps features only at the highest LOD.
        # Compute the highest LOD:
        _, pyramid, _ = wisp_spc_ops.octree_to_spc(spc_octree)
        max_level = pyramid.shape[-1] - 2

        self.grid = OctreeGrid.from_spc(
            spc_octree=spc_octree,
            feature_dim=3,
            base_lod=max_level,
            num_lods=0  # SPCFields track features internally, avoiding full OctreeGrid initialization is faster
        )

    @property
    def device(self):
        """ Returns the device used to process inputs in this neural field.

        Returns:
            (torch.device): The expected device used for this Structured Point Cloud.
        """
        return self.spc_device

    def register_forward_functions(self):
        """Register the forward functions.
        """
        # TODO (operel): support normals channel explicitly
        self._register_forward_function(self.rgba, ["rgb"])

    def rgba(self, ridx_hit=None):
        """Compute color for the provided ray hits.

        Args:
            ridx_hit (torch.LongTensor): SPC point_hierarchy indices of shape [batch].
                                     used to indicate index of first hit voxel.

        Returns:
            {"rgb": torch.FloatTensor}:
                - RGB tensor of shape [batch, 1, 3]
        """
        # find offset to final level to make indices relative to final level
        level = self.grid.blas.max_level
        offset = self.grid.blas.pyramid[1, level]
        ridx_hit = ridx_hit - offset

        colors = self.colors[ridx_hit, :3].unsqueeze(1)
        return dict(rgb=colors)

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        properties = {
            "Grid": self.grid,
        }
        return properties

wisp/models/nefs/spc_field.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch.nn as nn
from wisp.models.nefs import BaseNeuralField
from wisp.tracers.base_tracer import BaseTracer


class Pipeline(nn.Module):
    """Base class for implementing neural field pipelines.

    Pipelines consist of several components:

        - Neural fields (``self.nef``) which take coordinates as input and outputs signals.
          These usually consist of several optional components:

            - A feature grid (``self.nef.grid``)
              Sometimes also known as 'hybrid representations'.
            - An acceleration structure (``self.nef.grid.blas``) which can be used to accelerate spatial queries.
            - A decoder (``self.net.decoder``) which can take the features (or coordinates, or embeddings) and covert it to signals.

        - A forward map (``self.tracer``) which is a function which will invoke the pipeline in
          some outer loop. Usually this consists of renderers which will output a RenderBuffer object.
    
    The 'Pipeline' classes are responsible for holding and orchestrating these components.
    """
    
    def __init__(self, nef: BaseNeuralField, tracer: BaseTracer = None):
        """Initialize the Pipeline.

        Args:
            nef (nn.Module): Neural fields module.
            tracer (nn.Module or None): Forward map module.
        """
        super().__init__()
    
        self.nef: BaseNeuralField = nef
        self.tracer: BaseTracer = tracer

    def forward(self, *args, **kwargs):
        """The forward function will use the tracer (the forward model) if one is available. 
        
        Otherwise, it'll execute the neural field.
        """
        if self.tracer is not None:
            return self.tracer(self.nef, *args, **kwargs)
        else:
            return self.nef(*args, **kwargs)

wisp/models/pipeline.py
Sharath-girish/Shacira
from .bit_estimator import *
from .bit_estimatorN import *
wisp/models/prob_models/__init__.py
Sharath-girish/Shacira
from numpy import sin
import torch.nn as nn
import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn.modules.conv import Conv1d
from wisp.models.prob_models import *

class Bitparm(nn.Module):
    '''
    save params
    '''
    def __init__(self, channel, is_symmetric=False, is_unimodal=False, final=False):
        super(Bitparm, self).__init__()
        self.final = final
        self.is_unimodal = is_unimodal
        self.h = nn.Parameter(torch.nn.init.normal_(torch.empty(channel).view(1,-1), 0, 0.01))
        if is_symmetric:
            self.b = nn.Parameter(torch.nn.init.zeros_(torch.empty(channel).view(1,-1)), requires_grad=False)
        else:
            self.b = nn.Parameter(torch.nn.init.normal_(torch.empty(channel).view(1,-1), 0, 0.01))
        if not final:
            self.a = nn.Parameter(torch.nn.init.normal_(torch.empty(channel).view(1,-1), 0, 0.01))
        else:
            self.a = None

    def forward(self, x, single_channel=None):
        if single_channel is not None:
            h = self.h[:,single_channel]
            b = self.b[:,single_channel]
            if not self.final:
                a = self.a[:,single_channel]
        else:
            h = self.h
            b = self.b
            if not self.final:
                a = self.a
        if self.is_unimodal:
            a = torch.abs(a)
        if self.final:
            return torch.sigmoid(x * F.softplus(h) + b)
        else:
            x = x * F.softplus(h) + b
            return x + torch.tanh(x) * torch.tanh(a)

class BitEstimator(nn.Module):
    '''
    Estimate bit
    '''
    def __init__(self, channel, is_symmetric=False, is_unimodal=False, num_layers = 4):
        super(BitEstimator, self).__init__()
        self.num_layers = num_layers
        self.f1 = Bitparm(channel, is_symmetric=is_symmetric, is_unimodal=is_unimodal)
        self.f2 = Bitparm(channel, is_symmetric=is_symmetric, is_unimodal=is_unimodal)
        self.f3 = Bitparm(channel, is_symmetric=is_symmetric, is_unimodal=is_unimodal)
        self.f4 = Bitparm(channel, is_symmetric=is_symmetric, is_unimodal=is_unimodal, final=True)
        
    def forward(self, x, single_channel=None):
        if self.num_layers>1:
            x = self.f1(x, single_channel)
        if self.num_layers>2:
            x = self.f2(x, single_channel)
        if self.num_layers>3:
            x = self.f3(x, single_channel)
        return self.f4(x, single_channel)
        

wisp/models/prob_models/bit_estimator.py
Sharath-girish/Shacira
from wisp.models.prob_models.bit_estimator import *
from wisp.models.prob_models import *

class Conv1DBit(nn.Conv1d):
    def __init__(self, act, width, **kwargs):
        super(Conv1DBit, self).__init__(**kwargs)
        self.act = act
        self.width = width

    def forward(self, input: Tensor, dims:tuple = None) -> Tensor:
        if dims is None:
            return self._conv_forward(input, self.act(self.weight), self.bias)
        else:
            return F.conv1d(input, self.act(self.weight[dims[0]:dims[1],:,:]), \
                            self.bias[dims[0]:dims[1]] if self.bias is not None else None,
                            groups=dims[2])

class BitparmN(nn.Module):
    '''
    save params
    '''
    def __init__(self, channel, width=1, final=False, initial=False):
        super(BitparmN, self).__init__()
        assert not (initial and final)
        self.final = final
        self.width = width
        in_channel = channel if initial else channel*width
        out_channel = channel if final else channel*width
        self.conv1 = Conv1DBit(F.sigmoid, width, in_channels=in_channel, out_channels=out_channel, \
                               kernel_size=1, groups=channel, bias=True)
        if not final:
            self.conv2 = Conv1DBit(torch.tanh, width, in_channels=out_channel,out_channels=out_channel,\
                                   kernel_size=1,groups=out_channel,bias=False)

    def forward(self, x, single_channel=None):
        if single_channel is not None:
            mult = 1 if self.final else self.width
            dims1 = (single_channel * mult, (single_channel+1) * mult,1)
            dims2 = (single_channel * self.width, (single_channel+1)*self.width,self.width)
        else:
            dims1 = None
            dims2 = None
        x = self.conv1(x, dims1)

        if self.final:
            return torch.sigmoid(x)
        else:
            return x + self.conv2(torch.tanh(x), dims2)

class BitEstimatorN(nn.Module):
    '''
    Estimate bit
    '''
    def __init__(self, channel, width):
        super(BitEstimatorN, self).__init__()
        self.f1 = BitparmN(channel, width, False, True)
        self.f2 = BitparmN(channel, width, False, False)
        self.f3 = BitparmN(channel, width, False, False)
        self.f4 = BitparmN(channel, width, True, False)
        
    def forward(self, x, single_channel=None):
        x = x.unsqueeze(-1)
        x = self.f1(x, single_channel)
        x = self.f2(x, single_channel)
        x = self.f3(x, single_channel)
        x = self.f4(x, single_channel)
        return x.squeeze(-1)

wisp/models/prob_models/bit_estimatorN.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import time
import numpy as np
import torch
import torch.nn.functional as F
from wisp.core import RenderBuffer, Rays
from wisp.ops.shaders import matcap_shader, pointlight_shadow_shader
from wisp.ops.differential import finitediff_gradient
from wisp.ops.geometric import normalized_grid, normalized_slice
from wisp.tracers import *


# --  This module will be deprecated, public usage of this functionality is discouraged -- """


def _look_at(f, t, height, width, mode='persp', fov=90.0, device='cuda'):
    """Vectorized look-at function, returns an array of ray origins and directions

    This function is mostly just a wrapper on top of generate_rays, but will calculate for you
    the view, right, and up vectors based on the from and to.

    Args:
        f (list of floats): [3] size list or tensor specifying the camera origin
        t (list of floats): [3] size list or tensor specifying the camera look at point
        height (int): height of the image
        width (int): width of the image
        mode (str): string that specifies the camera mode.
        fov (float): field of view of the camera.
        device (str): device the tensor will be allocated on.

    Returns:
        (torch.FloatTensor, torch.FloatTensor):
        - [height, width, 3] tensor of ray origins
        - [height, width, 3] tensor of ray directions
    """

    camera_origin = torch.FloatTensor(f).to(device)
    camera_view = F.normalize(torch.FloatTensor(t).to(device) - camera_origin, dim=0)
    camera_right = F.normalize(torch.cross(camera_view, torch.FloatTensor([0, 1, 0]).to(device)), dim=0)
    camera_up = F.normalize(torch.cross(camera_right, camera_view), dim=0)

    return _generate_rays(camera_origin, camera_view, camera_right, camera_up,
                          height, width, mode=mode, fov=fov, device=device)


def _generate_rays(camera_origin, camera_view, camera_right, camera_up, height, width,
                  mode='persp', fov=90.0, device='cuda'):
    """Generates rays from camera parameters.

    Args:
        camera_origin (torch.FloatTensor): [3] size tensor specifying the camera origin
        camera_view (torch.FloatTensor): [3] size tensor specifying the camera view direction
        camera_right (torch.FloatTensor): [3] size tensor specifying the camera right direction
        camera_up (torch.FloatTensor): [3] size tensor specifying the camera up direction
        height (int): height of the image
        width (int): width of the image
        mode (str): string that specifies the camera mode.
        fov (float): field of view of the camera.
        device (str): device the tensor will be allocated on.

    Returns:
        (torch.FloatTensor, torch.FloatTensor):
        - [height, width, 3] tensor of ray origins
        - [height, width, 3] tensor of ray directions
    """
    coord = normalized_grid(height, width, device=device)

    ray_origin = camera_right * coord[..., 0, np.newaxis] * np.tan(np.radians(fov / 2)) + \
                 camera_up * coord[..., 1, np.newaxis] * np.tan(np.radians(fov / 2)) + \
                 camera_origin + camera_view
    ray_origin = ray_origin.reshape(-1, 3)
    ray_offset = camera_view.unsqueeze(0).repeat(ray_origin.shape[0], 1)

    if mode == 'ortho':  # Orthographic camera
        ray_dir = F.normalize(ray_offset, dim=-1)
    elif mode == 'persp':  # Perspective camera
        ray_dir = F.normalize(ray_origin - camera_origin, dim=-1)
        ray_origin = camera_origin.repeat(ray_dir.shape[0], 1)
    else:
        raise ValueError('Invalid camera mode!')

    return ray_origin, ray_dir


class OfflineRenderer():
    """Renderer class to do simple offline, non interactive rendering.

    TODO(ttakikawa): This class will be deprecated soon when the main renderer class supports offline renders.
    """
    def __init__(self, 
        render_res   : list = [1024, 720], # [w, h]
        camera_proj  : str  = 'persp', # one of ['persp', 'ortho']
        render_batch : int  = -1, # -1 for no batching
        shading_mode : str  = 'rb', # one of ['matcap', 'rb']
        matcap_path  : str  = './data/matcap/Pearl.png', # set if shadming mode = matcap
        shadow       : bool = False, 
        ao           : bool = False, 
        perf         : bool = False,
        device              = 'cuda',
        **kwargs
    ):
        self.render_res = render_res
        self.render_batch = render_batch
        self.shading_mode = shading_mode
        self.matcap_path = matcap_path
        self.shadow = shadow
        self.ao = ao
        self.perf = perf
        self.device = device
        
        self.width, self.height = self.render_res

        self.kwargs = kwargs

    def render_lookat(self, 
            pipeline, 
            f            = [0,0,1], 
            t            = [0,0,0], 
            fov          = 30.0, 
            camera_proj  = 'persp',  
            device       = 'cuda', 
            mm           = None,
            lod_idx      = None,
            camera_clamp = [0, 5]
        ):
        """Render images from a lookat.

        Args:
            pipeline (wisp.core.Pipeline): the pipeline to render
            f (list[f,f,f]): camera from
            t (list[f,f,f]): camera to
            fov (float): field of view
            aa (int): Number of samples per pixel to render
            mm (torch.FloatTensor): model transformation matrix
            device: Device to initialize rays on.
            lod_idx (int): LOD to renderer
            camera_clamp (tuple of int): The near and far clipping planes.

        Returns: 
            (wisp.core.RenderBuffer): The rendered image buffers.
        """
        # Generate the ray origins and directions, from camera parameters
        ray_o, ray_d = _look_at(f, t, self.height, self.width, fov=fov, mode=camera_proj, device=device)
        # Rotate the camera into model space
        if mm is not None:
            mm = mm.to('cuda')
            ray_o = torch.mm(ray_o, mm)
            ray_d = torch.mm(ray_d, mm)
        
        rays = Rays(origins=ray_o, dirs=ray_d, dist_min=camera_clamp[0], dist_max=camera_clamp[1])

        rb = self.render(pipeline, rays, lod_idx=lod_idx)
        rb = rb.reshape(self.height, self.width, -1) 
        return rb


    def render(self, pipeline, rays, lod_idx=None):
        """Render images from a lookat.

        Args:
            pipeline (wisp.core.Pipeline): the pipeline to render
            rays (wisp.core.Rays): the rays to render
            lod_idx (int): LOD to renderer

        Returns: 
            (wisp.core.RenderBuffer): The renderer image.
        """
        # Differentiable Renderer
        if self.perf:
            _time = time.time()

        with torch.no_grad():
            if self.render_batch > 0:
                rb = RenderBuffer(xyz=None, hit=None, normal=None, shadow=None, ao=None, dirs=None)
                for ray_pack in rays.split(self.render_batch):
                    rb  += pipeline.tracer(pipeline.nef, rays=ray_pack, lod_idx=lod_idx, **self.kwargs)
            else:
                rb = pipeline.tracer(pipeline.nef, rays=rays, lod_idx=lod_idx, **self.kwargs)

        ######################
        # Shading Rendering
        ######################
        
        # This is executed if the tracer does not handle RGB
        # TODO(ttakikawa): Brittle
        if self.shading_mode == 'rb' and rb.rgb is None: 
            rb.rgb = pipeline.nef.rgb(rb.xyz, rays.dirs)

        if self.perf:
            print("Time Elapsed:{:.4f}".format(time.time() - _time))
        
        # Shade the image
        if self.shading_mode == 'matcap':
            # TODO(ttakikawa): Should use the mm
            rb = matcap_shader(rb, rays, self.matcap_path, mm=None)
        elif self.shading_mode == 'rb':
            assert rb.rgb is not None and "No rgb in buffer; change shading-mode"
            pass
        elif self.shading_mode == 'normal':    
            rb.rgb = (rb.normal + 1.0) / 2.0
        else:
            raise NotImplementedError
        
        # Use segmentation
        if rb.normal is not None:
            rb.normal[~rb.hit] = 1.0
        if rb.rgb is not None:
            rb.rgb[~rb.hit] = 1.0

        # Add secondary effects
        if self.shadow:
            rb = pointlight_shadow_shader(rb, rays, pipeline)
        
        if self.ao:
            # TODO(ttakikawa): Mystery function... how did I write this?
            acc = torch.zeros_like(rb.depth).to(rays.origins.device)
            r = torch.zeros_like(rb.depth).to(rays.origins.device)
            with torch.no_grad():
                weight = 3.5
                for i in range(40):
                    
                    # Visual constants
                    ao_width = 0.1
                    _d = ao_width * 0.25 * (float(i+1) / float(40+1)) ** 1.6
                    q = rb.xyz + rb.normal * _d

                    # AO for surface
                    with torch.no_grad():
                        r[rb.hit] = pipeline.nef(q[rb.hit])

                    if self.shadow:
                        net_d = pipeline.nef(q[plane_hit])
                        plane_d = torch.zeros_like(net_d) + _d
                        r[plane_hit] = torch.min(torch.cat([net_d, plane_d], dim=-1), dim=-1, keepdim=True)[0]
                        acc[plane_hit] += 3.5 * F.relu(_d - r[plane_hit] - 0.0015)
                    acc[rb.hit] += 3.5 * F.relu(_d - r[rb.hit] - 0.0015)
                    weight *= 0.84
        
            rb.ao = torch.clamp(1.0 - acc,  0.1, 1.0)
            rb.ao = rb.ao * rb.ao

        if self.ao:
            rb.rgb[...,:3] *= rb.ao        
        
        return rb
    
    # TODO(ttakikawa): These are useful functions but probably does not need to live in the renderer. Migrate.
    def normal_slice(self, fn, dim=0, depth=0.0):
        """Get a slice of the normals.
        """
        pts = normalized_slice(self.width, self.height, dim=dim, depth=depth, device=self.device).reshape(-1,3)
        normal = (F.normalize(finitediff_gradient(pts, fn).detach()) + 1.0) / 2.0
        normal = normal.reshape(self.width, self.height, 3).cpu().numpy()
        return normal

    def sdf_slice(self, fn, dim=0, depth=0):
        """Get a slice of SDFs.
        """
        pts = normalized_slice(self.width, self.height, dim=dim, depth=depth, device=self.device)

        d = torch.zeros(self.width * self.height, 1, device=pts.device)
        with torch.no_grad():
            d = fn(pts.reshape(-1,3))

        d = d.reshape(self.width, self.height, 1)

        d = d.squeeze().cpu().numpy()
        dpred = d
        d = np.clip((d + 1.0) / 2.0, 0.0, 1.0)
        blue = np.clip((d - 0.5)*2.0, 0.0, 1.0)
        yellow = 1.0 - blue
        vis = np.zeros([*d.shape, 3])
        vis[...,2] = blue
        vis += yellow[...,np.newaxis] * np.array([0.4, 0.3, 0.0])
        vis += 0.2
        vis[d - 0.5 < 0] = np.array([1.0, 0.38, 0.0])
        for i in range(50):
            vis[np.abs(d - 0.02*i) < 0.0015] = 0.8
        vis[np.abs(d - 0.5) < 0.004] = 0.0
        return vis
    
    def shade_images(self, pipeline, f=[0,0,1], t=[0,0,0], fov=30.0, aa=1, mm=None, 
                     lod_idx=None, camera_clamp=[0,10]):
        """
        Invokes the renderer and outputs images.

        Args:
            pipeline (wisp.core.Pipeline): the pipeline to render
            f (list[f,f,f]): camera from
            t (list[f,f,f]): camera to
            fov (float): field of view
            aa (int): Number of samples per pixel to render
            mm (torch.FloatTensor): model transformation matrix
            lod_idx (int): LOD to renderer
            camera_clamp (tuple of int): The near and far clipping planes.
        
        Returns: 
            (wisp.core.RenderBuffer): The rendered image buffers.
        """
        if mm is None:
            mm = torch.eye(3)
        
        if aa > 1:
            rblst = [] 
            for _ in range(aa):
                rblst.append(self.render_lookat(pipeline, f=f, t=t, fov=fov, mm=mm, lod_idx=lod_idx,
                    camera_clamp=camera_clamp))
            rb = RenderBuffer.mean(*rblst)
        else:
            rb = self.render_lookat(pipeline, f=f, t=t, fov=fov, mm=mm, lod_idx=lod_idx, 
                    camera_clamp=camera_clamp)
        rb = rb.cpu().transpose()
        return rb


wisp/offline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from . import geometric
from . import grid

wisp/ops/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .gradients import *

wisp/ops/differential/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch


# Differential operators

def autodiff_gradient(x, f):
    """Compute gradient using the PyTorch autodiff.

    Args:
        x (torch.FloatTensor): Coordinate tensor
        f (nn.Module): The function to perform autodiff on.
    """
    with torch.enable_grad():
        x = x.requires_grad_(True)
        y = f(x)
        grad = torch.autograd.grad(y, x,
                                   grad_outputs=torch.ones_like(y), create_graph=True)[0]
    return grad


def finitediff_gradient(x, f, eps=0.005):
    """Compute 3D gradient using finite difference.

    Args:
        x (torch.FloatTensor): Coordinate tensor of shape [..., 3]
        f (nn.Module): The function to perform autodiff on.
    """
    eps_x = torch.tensor([eps, 0.0, 0.0], device=x.device)
    eps_y = torch.tensor([0.0, eps, 0.0], device=x.device)
    eps_z = torch.tensor([0.0, 0.0, eps], device=x.device)

    grad = torch.cat([f(x + eps_x) - f(x - eps_x),
                      f(x + eps_y) - f(x - eps_y),
                      f(x + eps_z) - f(x - eps_z)], dim=-1)
    grad = grad / (eps * 2.0)

    return grad


def tetrahedron_gradient(x, f, eps=0.005):
    """Compute 3D gradient using finite difference (using tetrahedron method).

    Args:
        x (torch.FloatTensor): Coordinate tensor of shape [..., 3]
        f (nn.Module): The function to perform autodiff on.
    """
    h = eps
    k0 = torch.tensor([1.0, -1.0, -1.0], device=x.device, requires_grad=False)
    k1 = torch.tensor([-1.0, -1.0, 1.0], device=x.device, requires_grad=False)
    k2 = torch.tensor([-1.0, 1.0, -1.0], device=x.device, requires_grad=False)
    k3 = torch.tensor([1.0, 1.0, 1.0], device=x.device, requires_grad=False)
    h0 = torch.tensor([h, -h, -h], device=x.device, requires_grad=False)
    h1 = torch.tensor([-h, -h, h], device=x.device, requires_grad=False)
    h2 = torch.tensor([-h, h, -h], device=x.device, requires_grad=False)
    h3 = torch.tensor([h, h, h], device=x.device, requires_grad=False)
    h0 = x + h0
    h1 = x + h1
    h2 = x + h2
    h3 = x + h3
    h0 = h0.detach()
    h1 = h1.detach()
    h2 = h2.detach()
    h3 = h3.detach()
    h0 = k0 * f(h0)
    h1 = k1 * f(h1)
    h2 = k2 * f(h2)
    h3 = k3 * f(h3)
    grad = (h0 + h1 + h2 + h3) / (h * 4.0)
    return grad

wisp/ops/differential/gradients.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import numpy as np
import torch
import wisp._C as _C

# This is a collection of utility functions that broadly work with different 3D constructs.

def find_depth_bound(query, nug_depth, info, curr_idxes=None):
    r"""Associate query points to the closest depth bound in-order.
    
    TODO: Document the input.
    """
    if curr_idxes is None:
        curr_idxes = torch.nonzero(info).contiguous()
    return _C.render.find_depth_bound_cuda(query.contiguous(), curr_idxes.contiguous(), nug_depth.contiguous())


def sample_unif_sphere(n):
    """Sample uniformly random points on a sphere.
    
    Args:
        n (int): Number of samples.

    Returns:
        (np.array): Positions of shape [n, 3]
    """
    u = np.random.rand(2, n)
    z = 1 - 2*u[0,:]
    r = np.sqrt(1. - z * z)
    phi = 2 * np.pi * u[1,:]
    xyz = np.array([r * np.cos(phi), r * np.sin(phi), z]).transpose()
    return xyz


def sample_fib_sphere(n):
    """
    Evenly distributed points on sphere using Fibonnaci sequence.
    From <http://extremelearning.com.au/evenly-distributing-points-on-a-sphere>
    WARNING: Order is not randomized.
    
    Args:
        n (int): Number of samples.

    Returns:
        (np.array): Positions of shape [n, 3]
    """

    i = np.arange(0, n, dtype=float) + 0.5
    phi = np.arccos(1 - 2*i/n)
    golden_ratio = (1 + 5**0.5)/2
    theta = 2. * np.pi * i / golden_ratio
    xyz = np.array([np.cos(theta) * np.sin(phi),
                    np.sin(theta) * np.sin(phi),
                    np.cos(phi)]).transpose()
    return xyz


def normalized_grid(height, width, jitter=False, device='cuda', use_aspect=True):
    """Returns grid[x,y] -> coordinates for a normalized window.

    This is generally confusing and terrible, but in the [XYZ] 3D space, the width generally corresponds to
    the XZ axis and the height corresponds to the Y axis. So in the normalized [XY] space, width also
    corresponds to the X and the height corresponds to Y.

    However, PyTorch follows HW ordering. So that means, given a [HW] image, the [H+1, W] coordinate will
    see an increase in the Y axis (the 2nd axis) of the actual global coordinates.

    Args:
        height (int): grid height
        width (int): grid width
        jitter (bool): If True, will jitter the coordinates.
        device (str): Device to allocate the grid on.
        use_aspect (bool): If True, will scale the coords by the aspect ratio.

    Returns:
        (torch.FloatTensor): Coords tensor of shape [H, W, 2]
    """
    window_x = torch.linspace(-1, 1, steps=width, device=device)
    window_y = torch.linspace(1, -1, steps=height, device=device)

    if jitter:
        window_x += (2.0 * torch.rand(*window_x.shape, device=device) - 1.0) * (1. / width)
        window_y += (2.0 * torch.rand(*window_y.shape, device=device) - 1.0) * (1. / height)

    if use_aspect:
        if width > height:
            window_x = window_x * (width / height)
        elif height > width:
            window_y = window_y * (height / width)

    coord = torch.stack(torch.meshgrid(window_x, window_y)).permute(2, 1, 0)  # torch>=1.10: indexing='ij' (default)
    return coord


def normalized_slice(height, width, dim=0, depth=0.0, device='cuda'):
    """Returns a set of 3D coordinates for a slicing plane.

    Args:
        height (int): Grid height.
        width (int): Grid width.
        dim (int): Dimension to slice along.
        depth (float): The depth (from the 0 on the axis) for which the slicing will happen.
        device (str): Device to allocate the grid on.

    Returns:
        (torch.FloatTensor): Coords tensor of shape [height, width, 3].
    """
    window = normalized_grid(height, width, device)
    depth_pts = torch.ones(height, width, 1, device=device) * depth

    if dim == 0:
        pts = torch.cat([depth_pts, window[..., 0:1], window[..., 1:2]], dim=-1)
    elif dim == 1:
        pts = torch.cat([window[..., 0:1], depth_pts, window[..., 1:2]], dim=-1)
    elif dim == 2:
        pts = torch.cat([window[..., 0:1], window[..., 1:2], depth_pts], dim=-1)
    else:
        assert (False, "dim is invalid!")
    pts[..., 1] *= -1
    return pts


def spherical_envmap(ray_dir, normal):
    """Computes matcap UV-coordinates from the ray direction and normal.

    Args:
        ray_dir (torch.Tensor): incoming ray direction of shape [...., 3]
        normal (torch.Tensor): surface normal of shape [..., 3]

    Returns:
        (torch.FloatTensor): UV coordinates of shape [..., 2]
    """
    # Input should be size [...,3]
    # Returns [N,2] # Might want to make this [...,2]

    # TODO(ttakikawa): Probably should implement all this on GPU
    ray_dir_screen = ray_dir.clone()
    ray_dir_screen[..., 2] *= -1
    ray_dir_normal_dot = torch.sum(normal * ray_dir_screen, dim=-1, keepdim=True)
    r = ray_dir_screen - 2.0 * ray_dir_normal_dot * normal
    r[..., 2] -= 1.0
    m = 2.0 * torch.sqrt(torch.sum(r ** 2, dim=-1, keepdim=True))
    vN = (r[..., :2] / m) + 0.5
    vN = 1.0 - vN
    vN = vN[..., :2].reshape(-1, 2)
    vN = torch.clip(vN, 0.0, 1.0)
    vN[torch.isnan(vN)] = 0
    return vN


def spherical_envmap_numpy(ray_dir, normal):
    """Computes matcap UV-coordinates from the ray direction and normal.

    Args:
        ray_dir (torch.Tensor): incoming ray direction of shape [...., 3]
        normal (torch.Tensor): surface normal of shape [..., 3]

    Returns:
        (torch.FloatTensor): UV coordinates of shape [..., 2]
    """
    ray_dir_screen = ray_dir * np.array([1, 1, -1])
    # Calculate reflection
    ray_dir_normal_dot = np.sum(normal * ray_dir_screen, axis=-1)[..., np.newaxis]
    r = ray_dir_screen - 2.0 * ray_dir_normal_dot * normal
    m = 2.0 * np.sqrt(r[..., 0] ** 2 + r[..., 1] ** 2 + (r[..., 2] - 1) ** 2)
    vN = (r[..., :2] / m[..., np.newaxis]) + 0.5
    vN = 1.0 - vN
    vN = vN[..., :2].reshape(-1, 2)
    vN = np.clip(vN, 0, 1)
    vN[np.isnan(vN)] = 0
    return vN

wisp/ops/geometric.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from kaolin import _C
import wisp._C as wisp_C
import kaolin.ops.spc as spc_ops

# Alternative set of primes
#PRIMES = [2654436881, 5915587277, 1500450271, 3267000013, 5754853343, 
#          4093082899, 9576890767, 3628273133, 2860486313, 5463458053, 
#          3367900313, 5654500741, 5654500763, 5654500771, 5654500783, 
#          5654500801, 5654500811, 5654500861, 5654500879, 5654500889, 
#          5654500897, 5654500927, 5654500961, 5654500981, 5654500993,
#          9999999967, 7654487179, 7654489553, 7654495087, 7654486423,
#          7654488209, 8654487029, 8654489771, 8654494517, 8654495341]

PRIMES = [1, 265443567, 805459861]

def hashgrid_naive(coords, resolutions, codebook_bitwidth, lod_idx, codebook):
    """
    -- This function is deprecated and unused --
    This code exists here mostly as a reference:
    Do NOT expect a 1-to-1 numerical correspondence to the CUDA accelerated version.

    A naive PyTorch implementation of the hashgrid.
    This code is therefore very slow. :)
    
    Args:
        coords (torch.FloatTensor): 3D coordinates of shape [batch, 3]
        resolutions (torch.LongTensor): the resolution of the grid per level of shape [num_lods]
        codebook_bitwidth (int): The bitwidth of the codebook. The codebook will have 2^bw entries.
        lod_idx (int): The LOD to aggregate to.
        codebook (torch.ModuleList[torch.FloatTensor]): A list of codebooks of shapes [codebook_size, feature_dim].

    Returns:
        (torch.FloatTensor): Features of shape [batch, feature_dim]
    """
    _, feature_dim = codebook[0].shape
    batch, _ = coords.shape
    codebook_size = 2**codebook_bitwidth
    feats = []
    for i, res in enumerate(resolutions[:lod_idx+1]):
        tf_coords = torch.clip(((coords + 1.0) / 2.0) * res, 0, res-1-1e-5).reshape(-1, 3)
        cc000 = torch.floor(tf_coords).short()
        cc = spc_ops.points_to_corners(cc000).long()

        num_pts = res**3
        if num_pts > codebook_size:
            cidx = ((cc[...,0] * PRIMES[(i*3+0)%len(PRIMES)]) ^ \
                        (cc[...,1] * PRIMES[(i*3+1)%len(PRIMES)]) ^ \
                        (cc[...,2] * PRIMES[(i*3+2)%len(PRIMES)])) % codebook_size
        else:
            cidx = cc[...,0] + cc[...,1] * res + cc[...,2] * res * res
        fs = codebook[i][cidx]
        
        feats.append(cidx[...,0:1])
        feats.append(cidx[...,0:1])
        
        coeffs = _C.ops.spc.coords_to_trilinear_cuda(tf_coords.contiguous(), cc000.contiguous())[...,None]
        feats.append((fs * coeffs).sum(-2))
    return torch.cat(feats, -1)

class HashGridInterpolate(torch.autograd.Function):
    # TODO(ttakikawa): This class should also support the 2D case... which also means I have to write another kernel!

    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
    def forward(ctx, coords, resolutions, codebook_bitwidth, lod_idx, codebook, codebook_sizes, codebook_first_idx):
        if codebook[0].shape[-1] % 2 == 1:
            raise Exception("The codebook feature dimension needs to be a multiple of 2.")


        # TODO(ttakikawa): Make the kernel use the LOD
        feats_out = wisp_C.ops.hashgrid_interpolate_cuda(coords.float().contiguous(), 
                                                         codebook,
                                                         codebook_first_idx,
                                                         resolutions,
                                                         codebook_bitwidth).contiguous()
    
        ctx.save_for_backward(coords, codebook, codebook_first_idx)
        ctx.resolutions = resolutions
        ctx.num_lods = len(resolutions)
        ctx.codebook_size = 2**codebook_bitwidth
        ctx.codebook_bitwidth = codebook_bitwidth
        ctx.feature_dim = codebook.shape[-1]
        return feats_out
    
    @staticmethod
    @torch.cuda.amp.custom_bwd
    def backward(ctx, grad_output):
    
        coords = ctx.saved_tensors[0]
        codebook = ctx.saved_tensors[1]
        codebook_first_idx = ctx.saved_tensors[2]
        resolutions = ctx.resolutions
        feature_dim = ctx.feature_dim
        codebook_bitwidth = ctx.codebook_bitwidth
        

        grad_codebook = wisp_C.ops.hashgrid_interpolate_backward_cuda(
                coords.float().contiguous(), grad_output.contiguous(), codebook,
                codebook_first_idx,
                resolutions,  
                codebook_bitwidth, feature_dim, ctx.needs_input_grad[0])
        return (None, None, None, None, grad_codebook, None, None)
        
def hashgrid(coords, resolutions, codebook_bitwidth, lod_idx, codebook, codebook_sizes, codebook_first_idx):
    """A hash-grid query + interpolation function, accelerated with CUDA.
    
    Args:
        coords (torch.FloatTensor): 3D coordinates of shape [batch, 3]
        resolutions (torch.LongTensor): the resolution of the grid per level of shape [num_lods]
        codebook_bitwidth (int): The bitwidth of the codebook. The codebook will have 2^bw entries.
        lod_idx (int): The LOD to aggregate to.
        codebook (torch.ModuleList[torch.FloatTensor]): A list of codebooks of shapes [codebook_size, feature_dim].

    Returns:
        (torch.FloatTensor): Features of shape [batch, feature_dim]
    """
    batch, dim = coords.shape
    feats = HashGridInterpolate.apply(coords.contiguous(), resolutions,
                                      codebook_bitwidth, lod_idx, codebook,
                                      codebook_sizes, codebook_first_idx)
    feature_dim = codebook.shape[1] * len(resolutions)
    return feats.reshape(batch, feature_dim)



class HashGridInterpolate2D(torch.autograd.Function):

    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
    def forward(ctx, coords, resolutions, codebook_bitwidth, lod_idx, codebook, codebook_sizes, codebook_first_idx):
        if codebook[0].shape[-1] % 2 == 1:
            raise Exception("The codebook feature dimension needs to be a multiple of 2.")


        # TODO(ttakikawa): Make the kernel use the LOD
        feats_out = wisp_C.ops.hashgrid_interpolate2d_cuda(coords.float().contiguous(), 
                                                         codebook,
                                                         codebook_first_idx,
                                                         resolutions,
                                                         codebook_bitwidth).contiguous()
    
        ctx.save_for_backward(coords, codebook, codebook_first_idx)
        ctx.resolutions = resolutions
        ctx.num_lods = len(resolutions)
        ctx.codebook_size = 2**codebook_bitwidth
        ctx.codebook_bitwidth = codebook_bitwidth
        ctx.feature_dim = codebook.shape[-1]
        return feats_out
    
    @staticmethod
    @torch.cuda.amp.custom_bwd
    def backward(ctx, grad_output):
    
        coords = ctx.saved_tensors[0]
        codebook = ctx.saved_tensors[1]
        codebook_first_idx = ctx.saved_tensors[2]
        resolutions = ctx.resolutions
        feature_dim = ctx.feature_dim
        codebook_bitwidth = ctx.codebook_bitwidth
        

        grad_codebook = wisp_C.ops.hashgrid_interpolate2d_backward_cuda(
                coords.float().contiguous(), grad_output.contiguous(), codebook,
                codebook_first_idx,
                resolutions,  
                codebook_bitwidth, feature_dim, ctx.needs_input_grad[0])
        return (None, None, None, None, grad_codebook, None, None)
        
def hashgrid2d(coords, resolutions, codebook_bitwidth, lod_idx, codebook, codebook_sizes, codebook_first_idx):
    """A hash-grid query + interpolation function, accelerated with CUDA.
    
    Args:
        coords (torch.FloatTensor): 2D coordinates of shape [batch, 2]
        resolutions (torch.LongTensor): the resolution of the grid per level of shape [num_lods]
        codebook_bitwidth (int): The bitwidth of the codebook. The codebook will have 2^bw entries.
        lod_idx (int): The LOD to aggregate to.
        codebook (torch.ModuleList[torch.FloatTensor]): A list of codebooks of shapes [codebook_size, feature_dim].

    Returns:
        (torch.FloatTensor): Features of shape [batch, feature_dim]
    """
    batch, dim = coords.shape
    feats = HashGridInterpolate2D.apply(coords.contiguous(), resolutions,
                                      codebook_bitwidth, lod_idx, codebook,
                                      codebook_sizes, codebook_first_idx)
    feature_dim = codebook.shape[1] * len(resolutions)
    return feats.reshape(batch, feature_dim)

wisp/ops/grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .io import *
from .metrics import *
from .processing import *

wisp/ops/image/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import os
import glob
import PIL
import numpy as np
import torch
import torchvision

""" A module for reading / writing various image formats. """


def write_exr(path, data):
    """Writes an EXR image to some path.

    Data is a dict of form { "default" = rgb_array, "depth" = depth_array }

    Args:
        path (str): Path to save the EXR
        data (dict): Dictionary of EXR buffers.

    Returns:
        (void): Writes to path.
    """
    try:
        import pyexr
    except:
        raise Exception(
            "Module pyexr is not available. To install, run `pip install pyexr`. "
            "You will likely also need `libopenexr`, which through apt you can install with "
            "`apt-get install libopenexr-dev` and on Windows you can install with "
            "`pipwin install openexr`.")
    pyexr.write(path, data,
                channel_names={'normal': ['X', 'Y', 'Z'],
                               'x': ['X', 'Y', 'Z'],
                               'view': ['X', 'Y', 'Z']},
                precision=pyexr.HALF)

def write_png(path, data):
    """Writes an PNG image to some path.

    Args:
        path (str): Path to save the PNG.
        data (np.array): HWC image.

    Returns:
        (void): Writes to path.
    """
    torchvision.io.write_png(hwc_to_chw(data), path)

def glob_imgs(path, exts=['*.png', '*.PNG', '*.jpg', '*.jpeg', '*.JPG', '*.JPEG']):
    """Utility to find images in some path.

    Args:
        path (str): Path to search images in.
        exts (list of str): List of extensions to try.

    Returns:
        (list of str): List of paths that were found.
    """
    imgs = []
    for ext in exts:
        imgs.extend(glob.glob(os.path.join(path, ext)))
    return imgs

def load_rgb_tensor(path, normalize=True):
    """Loads an image.

    Args:
        path (str): Path to the image.
        noramlize (bool): If True, will return [0,1] floating point values. Otherwise returns [0,255] ints.

    Returns:
        (np.array): Image as an array of shape [H,W,C]
    """
    try:
        image = PIL.Image.open(path).convert("RGB")
    except Exception as e:
        PIL.Image.MAX_IMAGE_PIXELS = 1216601101
        image = PIL.Image.open(path).convert("RGB")

    transform = torchvision.transforms.ToTensor()
    tensor_image = transform(image)

    return tensor_image
    # return chw_to_hwc(tensor_image)

def load_rgb(path, normalize=True):
    """Loads an image.

    Args:
        path (str): Path to the image.
        noramlize (bool): If True, will return [0,1] floating point values. Otherwise returns [0,255] ints.

    Returns:
        (np.array): Image as an array of shape [H,W,C]
    """
    img = torchvision.io.read_image(path)
    if normalize:
        img = img.float() / 255.0
    return np.array(chw_to_hwc(img))

def hwc_to_chw(img):
    """Converts [H,W,C] to [C,H,W] for TensorBoard output.

    Args:
        img (torch.Tensor): [H,W,C] image.

    Returns:
        (torch.Tensor): [C,H,W] image.
    """
    return img.permute(2, 0, 1)

def chw_to_hwc(img):
    """Converts [C,H,W] to [H,W,C].

    Args:
        img (torch.Tensor): [C,H,W] image.

    Returns:
        (torch.Tensor): [H,W,C] image.
    """
    return img.permute(1, 2, 0)

wisp/ops/image/io.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import skimage
import skimage.metrics
import numpy as np
import torch


""" A module for image based metrics """


def psnr(rgb, gts):
    """Calculate the PSNR metric.

    Assumes the RGB image is in [0,1]

    Args:
        rgb (torch.FloatTensor): Image tensor of shape [H,W,3]
        gts (torch.FloatTensor): Image tensor of shape [H,W,3]

    Returns:
        (float): The PSNR score
    """
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    assert (rgb.shape[-1] == 3)
    assert (gts.shape[-1] == 3)

    mse = torch.mean((rgb[..., :3] - gts[..., :3]) ** 2).item()
    return 10 * np.log10(1.0 / mse)

def clamped_psnr(rgb, gts):
    """Calculate the image PSNR metric after clamping and quantizing to [0,255].

    Assumes the GT image is in [0,1]

    Args:
        rgb (torch.FloatTensor): Image tensor of shape [H,W,3]
        gts (torch.FloatTensor): Image tensor of shape [H,W,3]

    Returns:
        (float): The PSNR score
    """
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    assert (rgb.shape[-1] == 3)
    assert (gts.shape[-1] == 3)

    rgb = (torch.clamp(rgb, 0, 1)*255).to(torch.uint8)
    gts = (torch.clamp(gts, 0, 1)*255).to(torch.uint8)
    mse = torch.mean((rgb[..., :3].float() - gts[..., :3].float()) ** 2).item()
    return 20*np.log10(255.) - 10 * np.log10(mse)

def clamped_mse(rgb, gts):
    """Calculate the image MSE metric after clamping and quantizing to [0,255].

    Assumes the GT image is in [0,1]

    Args:
        rgb (torch.FloatTensor): Image tensor of shape [H,W,3]
        gts (torch.FloatTensor): Image tensor of shape [H,W,3]

    Returns:
        (float): The PSNR score
    """
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    assert (rgb.shape[-1] == 3)
    assert (gts.shape[-1] == 3)

    rgb = (torch.clamp(rgb, 0, 1)*255).to(torch.uint8)
    gts = (torch.clamp(gts, 0, 1)*255).to(torch.uint8)
    mse = torch.mean((rgb[..., :3].float() - gts[..., :3].float()) ** 2).item()
    return mse

def lpips(rgb, gts, lpips_model=None):
    """Calculate the LPIPS metric.

    Assumes the RGB image is in [0,1]

    Args:
        rgb (torch.FloatTensor): Image tensor of shape [H,W,3]
        gts (torch.FloatTensor): Image tensor of shape [H,W,3]

    Returns:
        (float): The LPIPS score
    """
    try:
        from lpips import LPIPS
    except:
        raise Exception(
            "Module lpips not available. To install, run `pip install lpips`")
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    assert (rgb.shape[-1] == 3)
    assert (gts.shape[-1] == 3)

    if lpips_model is None:
        lpips_model = LPIPS(net='vgg').cuda()

    return lpips_model(
        (2.0 * rgb[..., :3] - 1.0).permute(2, 0, 1),
        (2.0 * gts[..., :3] - 1.0).permute(2, 0, 1)).mean().item()


def ssim(rgb, gts):
    """Calculate the SSIM metric.

    Assumes the RGB image is in [0,1]

    Args:
        rgb (torch.FloatTensor): Image tensor of shape [H,W,3]
        gts (torch.FloatTensor): Image tensor of shape [H,W,3]

    Returns:
        (float): The SSIM score
    """
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    return skimage.metrics.structural_similarity(
        rgb[..., :3].cpu().numpy(),
        gts[..., :3].cpu().numpy(),
        multichannel=True,
        data_range=1,
        gaussian_weights=True,
        sigma=1.5,
        channel_axis=2)

wisp/ops/image/metrics.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import cv2
import torch


def srgb_to_linear(img):
    """Converts from SRGB to Linear colorspace.

    Args:
        img (torch.FloatTensor): SRGB image.

    Returns:
        (torch.FloatTensor): Linear image.
    """
    limit = 0.04045
    return torch.where(img > limit, torch.power((img + 0.055) / 1.055, 2.4), img / 12.92)


def linear_to_srgb(img):
    """Converts from Linear to SRGB colorspace.

    Args:
        img (torch.FloatTensor): Linear image.

    Returns:
        (torch.FloatTensor): SRGB image.
    """
    limit = 0.0031308
    img = torch.where(img > limit, 1.055 * (img ** (1.0 / 2.4)) - 0.055, 12.92 * img)
    img[img > 1] = 1
    return img


def resize_mip(img, mip, interpolation=cv2.INTER_LINEAR):
    """Resize image with cv2.

    Args:
        img (torch.FloatTensor): Image of shape [H, W, 3]
        mip (int): Rescaling factor. Will rescale by 2**mip.
        interpolation: Interpolation modes used by `cv2.resize`.

    Returns:
        (torch.FloatTensor): Rescaled image of shape [H/(2**mip), W/(2**mip), 3]
    """
    resize_factor = 2**mip
    # WARNING: cv2 expects (w,h) for the shape. God knows why :)
    shape = (int(img.shape[1] // resize_factor), int(img.shape[0] // resize_factor))
    img = cv2.resize(img, dsize=shape, interpolation=interpolation)
    return img

wisp/ops/image/processing.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .area_weighted_distribution import area_weighted_distribution
from .random_face import random_face
from .point_sample import point_sample
from .sample_surface import sample_surface
from .sample_near_surface import sample_near_surface
from .sample_uniform import sample_uniform
from .load_obj import load_obj
from .compute_sdf import compute_sdf
from .normalize import normalize
from .closest_point import closest_point
from .closest_tex import closest_tex
from .barycentric_coordinates import barycentric_coordinates
from .sample_tex import sample_tex
from .per_face_normals import per_face_normals

wisp/ops/mesh/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from .per_face_normals import per_face_normals

def area_weighted_distribution(
    V : torch.Tensor,
    F : torch.Tensor, 
    normals : torch.Tensor = None):
    """Construct discrete area weighted distribution over triangle mesh.

    Args:
        V (torch.Tensor): #V, 3 array of vertices
        F (torch.Tensor): #F, 3 array of indices
        normals (torch.Tensor): normals (if precomputed)
        eps (float): epsilon
    
    Returns:
        (torch.distributions): Distribution to be used
    """

    if normals is None:
        normals = per_face_normals(V, F)
    areas = torch.norm(normals, p=2, dim=1) * 0.5
    areas /= torch.sum(areas) + 1e-10
    
    # Discrete PDF over triangles
    return torch.distributions.Categorical(areas.view(-1))


wisp/ops/mesh/area_weighted_distribution.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import numpy as np

# Same API as https://github.com/libigl/libigl/blob/main/include/igl/barycentric_coordinates.cpp

def barycentric_coordinates(
    points : torch.Tensor, 
    A : torch.Tensor,
    B : torch.Tensor,
    C : torch.Tensor):
    """
    Return barycentric coordinates for a given set of points and triangle vertices

    Args:
        points (torch.FloatTensor): [N, 3]
        A (torch.FloatTensor): [N, 3] vertex0
        B (torch.FloatTensor): [N, 3] vertex1
        C (torch.FloatTensor): [N, 3] vertex2
    
    Returns:
        (torch.FloatTensor): barycentric coordinates of [N, 2] 
    """

    v0 = B-A
    v1 = C-A
    v2 = points-A
    d00 = (v0*v0).sum(dim=-1)
    d01 = (v0*v1).sum(dim=-1)
    d11 = (v1*v1).sum(dim=-1)
    d20 = (v2*v0).sum(dim=-1)
    d21 = (v2*v1).sum(dim=-1)
    denom = d00*d11 - d01*d01
    L = torch.zeros(points.shape[0], 3, device=points.device)
    # Warning: This clipping may cause undesired behaviour
    L[...,1] = torch.clip((d11*d20 - d01*d21)/denom, 0.0, 1.0)
    L[...,2] = torch.clip((d00*d21 - d01*d20)/denom, 0.0, 1.0)
    L[...,0] = torch.clip(1.0 - (L[...,1] + L[...,2]), 0.0, 1.0)
    return L

wisp/ops/mesh/barycentric_coordinates.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

# Closest point function + texture sampling
# https://en.wikipedia.org/wiki/Closest_point_method

import torch
import numpy as np

def closest_point(
    V : torch.Tensor, 
    F : torch.Tensor,
    points : torch.Tensor):

    assert False and "This function is not supported in the main branch. Please use ttakikawa/pysdf for this function."



wisp/ops/mesh/closest_point.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import numpy as np
from .barycentric_coordinates import barycentric_coordinates
from .closest_point import closest_point
from .sample_tex import sample_tex

def closest_tex(
    V : torch.Tensor, 
    F : torch.Tensor,
    TV : torch.Tensor,
    TF : torch.Tensor,
    materials,
    points : torch.Tensor):
    """Returns the closest texture for a set of points.

        V (torch.FloatTensor): mesh vertices of shape [V, 3] 
        F (torch.LongTensor): mesh face indices of shape [F, 3]
        TV (torch.FloatTensor): 
        TF (torch.FloatTensor):
        materials:
        points (torch.FloatTensor): sample locations of shape [N, 3]

    Returns:
        (torch.FloatTensor): texture samples of shape [N, 3]
    """

    TV = TV.to(V.device)
    TF = TF.to(V.device)
    points = points.to(V.device)

    dist, hit_pts, hit_tidx = closest_point(V, F, points)
    
    hit_F = F[hit_tidx]
    hit_V = V[hit_F]
    BC = barycentric_coordinates(hit_pts.cuda(), hit_V[:,0], hit_V[:,1], hit_V[:,2])

    hit_TF = TF[hit_tidx]
    hit_TM = hit_TF[...,3]
    hit_TF = hit_TF[...,:3]

    if TV.shape[0] > 0:
        hit_TV = TV[hit_TF]
        hit_Tp = (hit_TV * BC.unsqueeze(-1)).sum(1)
    else:
        hit_Tp = BC
    
    rgb = sample_tex(hit_Tp, hit_TM, materials)

    return rgb, hit_pts, dist

wisp/ops/mesh/closest_tex.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import math
import contextlib
import os
import sys

import torch
import numpy as np
import wisp._C as _C

def compute_sdf(
    V : torch.Tensor,
    F : torch.Tensor,
    points : torch.Tensor,
    split_size : int = 10**6):
    """Computes SDF given point samples and a mesh.

    Args:
        V (torch.FloatTensor): #V, 3 array of vertices
        F (torch.LongTensor): #F, 3 array of indices
        points (torch.FloatTensor): [N, 3] array of points to sample
        split_size (int): The batch at which the SDF will be computed. The kernel will break for too large
                          batches; when in doubt use the default.

    Returns:
        (torch.FloatTensor): [N, 1] array of computed SDF values.
    """
    mesh = V[F]

    _points = torch.split(points, split_size)
    sdfs = []
    for _p in _points:
        sdfs.append(_C.external.mesh_to_sdf_cuda(_p.cuda().contiguous(), mesh.cuda().contiguous())[0])
    return torch.cat(sdfs)[...,None]

wisp/ops/mesh/compute_sdf.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import os
import sys

import numpy as np
import tinyobjloader
import torch

from PIL import Image

import logging as log

Image.MAX_IMAGE_PIXELS = None

# Refer to 
# https://github.com/tinyobjloader/tinyobjloader/blob/master/tiny_obj_loader.h
# for conventions for tinyobjloader data structures.

texopts = [
    'ambient_texname',
    'diffuse_texname',
    'specular_texname',
    'specular_highlight_texname',
    'bump_texname',
    'displacement_texname',
    'alpha_texname',
    'reflection_texname',
    'roughness_texname',
    'metallic_texname',
    'sheen_texname',
    'emissive_texname',
    'normal_texname'
]

def load_mat(fname : str):
    """Loads material.
    """

    img = torch.FloatTensor(np.array(Image.open(fname)))
    img = img / 255.0

    return img


def load_obj(
    fname : str, 
    load_materials : bool = False):
    """Load .obj file using TinyOBJ and extract info.
    This is more robust since it can triangulate polygon meshes 
    with up to 255 sides per face.
    
    Args:
        fname (str): path to Wavefront .obj file
    """

    assert os.path.exists(fname), \
        'Invalid file path and/or format, must be an existing Wavefront .obj'

    reader = tinyobjloader.ObjReader()
    config = tinyobjloader.ObjReaderConfig()
    config.triangulate = True # Ensure we don't have any polygons

    reader.ParseFromFile(fname, config)

    # Get vertices
    attrib = reader.GetAttrib()
    vertices = torch.FloatTensor(attrib.vertices).reshape(-1, 3)

    # Get triangle face indices
    shapes = reader.GetShapes()
    faces = []
    for shape in shapes:
        faces += [idx.vertex_index for idx in shape.mesh.indices]
    faces = torch.LongTensor(faces).reshape(-1, 3)
    
    mats = {}

    if load_materials:
        # Load per-faced texture coordinate indices
        texf = []
        matf = []
        for shape in shapes:
            texf += [idx.texcoord_index for idx in shape.mesh.indices]
            matf.extend(shape.mesh.material_ids)
        # texf stores [tex_idx0, tex_idx1, tex_idx2, mat_idx]
        texf = torch.LongTensor(texf).reshape(-1, 3)
        matf = torch.LongTensor(matf).reshape(-1, 1)
        texf = torch.cat([texf, matf], dim=-1)

        # Load texcoords
        texv = torch.FloatTensor(attrib.texcoords).reshape(-1, 2)
        
        # Load texture maps
        parent_path = os.path.dirname(fname) 
        materials = reader.GetMaterials()
        for i, material in enumerate(materials):
            mats[i] = {}
            diffuse = getattr(material, 'diffuse')
            if diffuse != '':
                mats[i]['diffuse'] = torch.FloatTensor(diffuse)

            for texopt in texopts:
                mat_path = getattr(material, texopt)
                if mat_path != '':
                    img = load_mat(os.path.join(parent_path, mat_path))
                    mats[i][texopt] = img
                    #mats[i][texopt.split('_')[0]] = img
        return vertices, faces, texv, texf, mats

    return vertices, faces


wisp/ops/mesh/load_obj.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch

def normalize(
    V : torch.Tensor,
    F : torch.Tensor,
    mode : str):
    """Normalizes a mesh.

    Args:
        V (torch.FloatTensor): Vertices of shape [V, 3]
        F (torch.LongTensor): Faces of shape [F, 3]
        mode (str): Different methods of normalization.

    Returns:
        (torch.FloatTensor, torch.LongTensor):
        - Normalized Vertices
        - Faces
    """

    if mode == 'sphere':

        V_max, _ = torch.max(V, dim=0)
        V_min, _ = torch.min(V, dim=0)
        V_center = (V_max + V_min) / 2.
        V = V - V_center

        # Find the max distance to origin
        max_dist = torch.sqrt(torch.max(torch.sum(V**2, dim=-1)))
        V_scale = 1. / max_dist
        V *= V_scale
        return V, F

    elif mode == 'aabb':
        
        V_min, _ = torch.min(V, dim=0)
        V = V - V_min

        max_dist = torch.max(V)
        V *= 1.0 / max_dist

        V = V * 2.0 - 1.0

        return V, F

    elif mode == 'planar':
        
        V_min, _ = torch.min(V, dim=0)
        V = V - V_min

        x_max = torch.max(V[...,0])
        z_max = torch.max(V[...,2])

        V[...,0] *= 1.0 / x_max
        V[...,2] *= 1.0 / z_max

        max_dist = torch.max(V)
        V[...,1] *= 1.0 / max_dist
        #V *= 1.0 / max_dist

        V = V * 2.0 - 1.0

        y_min = torch.min(V[...,1])

        V[...,1] -= y_min

        return V, F

    elif mode == 'none':

        return V, F





wisp/ops/mesh/normalize.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch

def per_face_normals(
    V : torch.Tensor,
    F : torch.Tensor):
    """Compute normals per face.
    
    Args:
        V (torch.FloatTensor): Vertices of shape [V, 3]
        F (torch.LongTensor): Faces of shape [F, 3]
    
    Returns:
        (torch.FloatTensor): Normals of shape [F, 3]
    """
    mesh = V[F]

    vec_a = mesh[:, 0] - mesh[:, 1]
    vec_b = mesh[:, 1] - mesh[:, 2]
    normals = torch.cross(vec_a, vec_b)
    return normals


wisp/ops/mesh/per_face_normals.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from .sample_near_surface import sample_near_surface
from .sample_surface import sample_surface
from .sample_uniform import sample_uniform
from .area_weighted_distribution import area_weighted_distribution

def point_sample(
    V : torch.Tensor, 
    F : torch.Tensor, 
    techniques : list, 
    num_samples : int):
    """Sample points from a mesh.

    Args:
        V (torch.Tensor): #V, 3 array of vertices
        F (torch.Tensor): #F, 3 array of indices
        techniques (list[str]): list of techniques to sample with
        num_samples (int): points to sample per technique
    
    Returns:
        (torch.FloatTensor): Samples of shape [len(techniques)*num_samples, 3]
    """
    if 'trace' in techniques or 'near' in techniques:
        # Precompute face distribution
        distrib = area_weighted_distribution(V, F)

    samples = []
    for technique in techniques:
        if technique =='trace':
            samples.append(sample_surface(V, F, num_samples, distrib=distrib)[0])
        elif technique == 'near':
            samples.append(sample_near_surface(V, F, num_samples, distrib=distrib))
        elif technique == 'rand':
            samples.append(sample_uniform(num_samples).to(V.device))
    samples = torch.cat(samples, dim=0)
    return samples


wisp/ops/mesh/point_sample.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from .area_weighted_distribution import area_weighted_distribution
from .per_face_normals import per_face_normals

def random_face(
    V : torch.Tensor, 
    F : torch.Tensor, 
    num_samples : int, 
    distrib=None):
    """Return an area weighted random sample of faces and their normals from the mesh.

    Args:
        V (torch.Tensor): #V, 3 array of vertices
        F (torch.Tensor): #F, 3 array of indices
        num_samples (int): num of samples to return
        distrib: distribution to use. By default, area-weighted distribution is used.
    
    Returns:
        (torch.LongTensor, torch.FloatTensor):
        - Faces
        - Normals
    """
    if distrib is None:
        distrib = area_weighted_distribution(V, F)

    normals = per_face_normals(V, F)

    idx = distrib.sample([num_samples])

    return F[idx], normals[idx]


wisp/ops/mesh/random_face.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from .sample_surface import sample_surface
from .area_weighted_distribution import area_weighted_distribution

def sample_near_surface(
    V : torch.Tensor,
    F : torch.Tensor, 
    num_samples: int, 
    variance : float = 0.01,
    distrib=None):
    """Sample points near the mesh surface.

    Args:
        V (torch.Tensor): #V, 3 array of vertices
        F (torch.Tensor): #F, 3 array of indices
        num_samples (int): number of surface samples
        distrib: distribution to use. By default, area-weighted distribution is used
    
    Returns:
        (torch.FloatTensor): samples of shape [num_samples, 3]
    """
    if distrib is None:
        distrib = area_weighted_distribution(V, F)
    samples = sample_surface(V, F, num_samples, distrib)[0]
    samples += torch.randn_like(samples) * variance
    return samples

wisp/ops/mesh/sample_near_surface.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
from .random_face import random_face
from .area_weighted_distribution import area_weighted_distribution

def sample_surface(
    V : torch.Tensor,
    F : torch.Tensor,
    num_samples : int,
    distrib = None):
    """Sample points and their normals on mesh surface.

    Args:
        V (torch.Tensor): #V, 3 array of vertices
        F (torch.Tensor): #F, 3 array of indices
        num_samples (int): number of surface samples
        distrib: distribution to use. By default, area-weighted distribution is used
    
    Returns:
        (torch.FloatTensor): samples of shape [num_samples, 3]
    """
    if distrib is None:
        distrib = area_weighted_distribution(V, F)

    # Select faces & sample their surface
    fidx, normals = random_face(V, F, num_samples, distrib)
    f = V[fidx]

    u = torch.sqrt(torch.rand(num_samples)).to(V.device).unsqueeze(-1)
    v = torch.rand(num_samples).to(V.device).unsqueeze(-1)

    samples = (1 - u) * f[:,0,:] + (u * (1 - v)) * f[:,1,:] + u * v * f[:,2,:]
    
    return samples, normals


wisp/ops/mesh/sample_surface.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import torch.nn.functional as F

def sample_tex(
    Tp : torch.Tensor, # points [N ,2] 
    TM : torch.Tensor, # material indices [N]
    materials):
    """Sample from a texture.

    Args:
        Tp (torch.FloatTensor): 2D coordinates to sample of shape [N, 2]
        TM (torch.LongTensor): Indices of the material to sample of shape [N]
        materials (list of material): Materials

    Returns:
        (torch.FloatTensor): RGB samples of shape [N, 3]
    """
    max_idx = TM.max()
    assert(max_idx > -1 and "No materials detected! Check the material definiton on your mesh.")

    rgb = torch.zeros(Tp.shape[0], 3, device=Tp.device)

    Tp = (Tp * 2.0) - 1.0
    # The y axis is flipped from what UV maps generally expects vs in PyTorch
    Tp[...,1] *= -1

    for i in range(max_idx+1):
        mask = (TM == i)
        if mask.sum() == 0:
            continue
        if 'diffuse_texname' not in materials[i]:
            if 'diffuse' in materials[i]:
                rgb[mask] = materials[i]['diffuse'].to(Tp.device)
            continue

        map = materials[i]['diffuse_texname'][...,:3].permute(2, 0, 1)[None].to(Tp.device)
        grid = Tp[mask]
        grid = grid.reshape(1, grid.shape[0], 1, grid.shape[1])
        _rgb = F.grid_sample(map, grid, mode='bilinear', padding_mode='reflection', align_corners=True)
        _rgb = _rgb[0,:,:,0].permute(1,0)
        rgb[mask] = _rgb


    return rgb



wisp/ops/mesh/sample_tex.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch

def sample_uniform(num_samples : int):
    """Sample uniformly in [-1,1] bounding volume.
    
    Args:
        num_samples(int) : number of points to sample
    
    Returns:
        (torch.FloatTensor): samples of shape [num_samples, 3]
    """
    return torch.rand(num_samples, 3) * 2.0 - 1.0


wisp/ops/mesh/sample_uniform.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .conversions import *
from .processing import *

wisp/ops/pointcloud/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch


def create_pointcloud_from_images(rgbs, masks, rays, depths):
    """Given depth images, will create a RGB pointcloud.

    TODO (ttakikawa): Probably make the input a tensor not a list...

    Args:
        rgbs (list of torch.FloatTensor): List of RGB tensors of shape [H, W, 3].
        masks (list of torch.FloatTensor): List of mask tensors of shape [H, W, 1].
        rays (list of wisp.core.Rays): List of rays.origins and rays.dirs of shape [H, W, 3].
        depths (list of torch.FloatTensor): List of depth tensors of shape [H, W, 1].

    Returns:
        (torch.FloatTensor, torch.FloatTensor):
        - 3D coordinates of shape [N*H*W, 3]
        - colors of shape [N*H*W, 3]
    """
    cloud_coords = []
    cloud_colors = []

    for i in range(len(rgbs)):
        mask = masks[i].bool()
        h, w = mask.shape[:2]
        mask = mask.reshape(h, w)
        depth = depths[i].reshape(h, w, 1)
        assert(len(mask.shape) == 2 and "Mask shape is not correct... it should be [H,W], check size here")
        coords = rays[i].origins[mask] + rays[i].dirs[mask] * depth[mask]
        colors = rgbs[i][mask]
        cloud_coords.append(coords.reshape(-1, 3))
        cloud_colors.append(colors[...,:3].reshape(-1, 3))

    return torch.cat(cloud_coords, dim=0), torch.cat(cloud_colors, dim=0)

wisp/ops/pointcloud/conversions.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch


def normalize_pointcloud(coords, return_scale=False):
    """Normalizes pointcloud to an AABB within [-1, 1].

    Args:
        coords (torch.FloatTensor): 3D coordinates of shape [N, 3]
        return_scale (bool): If True, will return the center of the cloud and the scaling factor.

    Returns:
        (torch.FloatTensor) or (torch.FloatTensor, torch.FloatTensor, float):
        - Normalized 3D coordinates of shape [N, 3]
        - Center of the cloud of shape [3]
        - Scaling factor (floating point value)
    """
    coords_max, _ = torch.max(coords, dim=0)
    coords_min, _ = torch.min(coords, dim=0)
    coords_center = (coords_max + coords_min) / 2.0

    # AABB normalize
    coords = coords - coords_center
    max_dist = torch.max(coords)
    coords_scale = 1.0 / max_dist
    coords *= coords_scale

    if return_scale:
        return coords, coords_center, coords_scale
    else:
        return coords

wisp/ops/pointcloud/processing.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .raygen import *

wisp/ops/raygen/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import torch
from kaolin.render.camera import Camera
from kaolin.render.camera.intrinsics import CameraFOV
from wisp.core import Rays

# -- Supersample / Ray jitter --

def generate_3d_grid(width, height, time, device=None):
    h_coords = torch.arange(height, device=device)
    w_coords = torch.arange(width, device=device)
    t_coords = torch.arange(time, device=device)
    return torch.meshgrid(t_coords,h_coords, w_coords)  # return pixel_y, pixel_x, pixel_t

def generate_2d_grid(width, height, device=None):
    h_coords = torch.arange(height, device=device)
    w_coords = torch.arange(width, device=device)
    return torch.meshgrid(h_coords, w_coords)  # return pixel_y, pixel_x


def generate_centered_pixel_coords(img_width, img_height, res_x=None, res_y=None, device=None):
    pixel_y, pixel_x = generate_2d_grid(res_x, res_y, device)
    scale_x = 1.0 if res_x is None else float(img_width) / res_x
    scale_y = 1.0 if res_y is None else float(img_height) / res_y
    pixel_x = pixel_x * scale_x + 0.5   # scale and add bias to pixel center
    pixel_y = pixel_y * scale_y + 0.5   # scale and add bias to pixel center
    return pixel_y, pixel_x


# -- Ray gen --

def _to_ndc_coords(pixel_x, pixel_y, camera):
    pixel_x = 2 * (pixel_x / camera.width) - 1.0
    pixel_y = 2 * (pixel_y / camera.height) - 1.0
    return pixel_x, pixel_y


def generate_pinhole_rays(camera: Camera, coords_grid: torch.Tensor):
    """Default ray generation function for pinhole cameras.

    This function assumes that the principal point (the pinhole location) is specified by a 
    displacement (camera.x0, camera.y0) in pixel coordinates from the center of the image. 

    The Kaolin camera class does not enforce a coordinate space for how the principal point is specified,
    so users will need to make sure that the correct principal point conventions are followed for 
    the cameras passed into this function.

    Args:
        camera (kaolin.render.camera): The camera class. 
        coords_grid (torch.FloatTensor): Grid of coordinates of shape [H, W, 2].

    Returns:
        (wisp.core.Rays): The generated pinhole rays for the camera.
    """
    if camera.device != coords_grid[0].device:
        raise Exception(f"Expected camera and coords_grid[0] to be on the same device, but found {camera.device} and {coords_grid[0].device}.")
    if camera.device != coords_grid[1].device:
        raise Exception(f"Expected camera and coords_grid[1] to be on the same device, but found {camera.device} and {coords_grid[1].device}.")
    # coords_grid should remain immutable (a new tensor is implicitly created here)
    pixel_y, pixel_x = coords_grid
    pixel_x = pixel_x.to(camera.device, camera.dtype)
    pixel_y = pixel_y.to(camera.device, camera.dtype)

    # Account for principal point (offsets from the center)
    pixel_x = pixel_x - camera.x0
    pixel_y = pixel_y + camera.y0

    # pixel values are now in range [-1, 1], both tensors are of shape res_y x res_x
    pixel_x, pixel_y = _to_ndc_coords(pixel_x, pixel_y, camera)

    ray_dir = torch.stack((pixel_x * camera.tan_half_fov(CameraFOV.HORIZONTAL),
                           -pixel_y * camera.tan_half_fov(CameraFOV.VERTICAL),
                           -torch.ones_like(pixel_x)), dim=-1)

    ray_dir = ray_dir.reshape(-1, 3)    # Flatten grid rays to 1D array
    ray_orig = torch.zeros_like(ray_dir)

    # Transform from camera to world coordinates
    ray_orig, ray_dir = camera.extrinsics.inv_transform_rays(ray_orig, ray_dir)
    ray_dir /= torch.linalg.norm(ray_dir, dim=-1, keepdim=True)
    ray_orig, ray_dir = ray_orig[0], ray_dir[0]  # Assume a single camera

    return Rays(origins=ray_orig, dirs=ray_dir, dist_min=camera.near, dist_max=camera.far)


def generate_ortho_rays(camera: Camera, coords_grid: torch.Tensor):
    # coords_grid should remain immutable (a new tensor is implicitly created here)
    pixel_y, pixel_x = coords_grid
    pixel_y = pixel_y.to(camera.device, camera.dtype)
    pixel_x = pixel_x.to(camera.device, camera.dtype)

    # pixel values are now in range [-1, 1], both tensors are of shape res_y x res_x
    pixel_x, pixel_y = _to_ndc_coords(pixel_x, pixel_y, camera)

    # Rescale according to distance from camera
    aspect_ratio = camera.width / camera.height
    pixel_x *= camera.fov_distance * aspect_ratio
    pixel_y *= camera.fov_distance

    zeros = torch.zeros_like(pixel_x)
    ray_dir = torch.stack((zeros, zeros, -torch.ones_like(pixel_x)), dim=-1)    # Ortho rays are parallel
    ray_orig = torch.stack((pixel_x, -pixel_y, zeros), dim=-1)
    ray_dir = ray_dir.reshape(-1, 3)    # Flatten grid rays to 1D array
    ray_orig = ray_orig.reshape(-1, 3)  # Flatten grid rays to 1D array

    # Transform from camera to world coordinates
    ray_orig, ray_dir = camera.extrinsics.inv_transform_rays(ray_orig, ray_dir)
    ray_dir /= torch.linalg.norm(ray_dir, dim=-1, keepdim=True)
    ray_orig, ray_dir = ray_orig[0], ray_dir[0]  # Assume a single camera

    return Rays(origins=ray_orig, dirs=ray_dir, dist_min=camera.near, dist_max=camera.far)

wisp/ops/raygen/raygen.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .metrics import *

wisp/ops/sdf/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch

def compute_sdf_iou(pred, gts):
    """Compute intersection over union for SDFs.

    Args:
        pred (torch.FloatTensor): Predicted signed distances
        gts (torch.FloatTensor): Groundtruth signed distances

    Returns:
        (float): The IOU score between 0 and 100.
    """
    inside_pred = (pred < 0).byte()
    inside_gts = (gts < 0).byte()

    area_union = torch.sum((inside_pred | inside_gts).float()).item()
    area_intersect = torch.sum((inside_pred & inside_gts).float()).item()

    iou = area_intersect / area_union
    return 100.0 * iou


def compute_sparse_sdf_iou(nef, coords, gts, lod_idx=0):
    """Given a sparse SDF neural field, coordinates, and ground truth SDF, will calculate the narrowband IOU.

    In the case where the point does not exist in the bounds of the octree, will simply calculate
    those as intersections.

    Inputs:
        nef (wisp.models.NeuralFields) : The neural field. Assumed to be sparse.
        coords (torch.FloatTensor) : 3D coordinates of shape [N, 3]
        gts (torch.FloatTensor) : Ground truth SDF of shape [N, 1]
        lod_idx : level of detail (if specified)

    Returns:
        (float): The calculated IOU
    """
    if nef.grid is None:
        raise Exception(f"{nef.__class__.__name__} is incompatible with this function.")
    pred = torch.zeros([coords.shape[0], 1])
    query_results = nef.grid.query(gts, lod_idx=lod_idx)
    pidx = query_results.pidx
    mask = (pidx != -1)

    pred[mask] = nef(coords=gts[mask], pidx=pidx[mask], lod_idx=lod_idx, channels="sdf").cpu()
    pred[~mask] = gts[~mask]

    return compute_sdf_iou(pred, gts)

wisp/ops/sdf/metrics.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .matcap import *
from .shadow_rays import *

wisp/ops/shaders/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import os
import numpy as np
import torch
from scipy.interpolate import RegularGridInterpolator
from PIL import Image
from wisp.core import RenderBuffer, Rays
from wisp.ops.geometric import spherical_envmap

""" A collection of shader functions for matcaps """


def matcap_shader(rb: RenderBuffer, rays: Rays, matcap_path, mm=None) -> RenderBuffer:
    """Apply matcap shading.

    Args:
        rb (wisp.core.RenderBuffer): The RenderBuffer.
        rays (wisp.core.Rays): The rays object.
        matcap_path (str): Path to a matcap.
        mm (torch.FloatTensor): A 3x3 rotation matrix.

    Returns:
        (wisp.core.RenderBuffer): The output RenderBuffer.
    """
    if not os.path.exists(matcap_path):
        raise Exception(f"The path [{matcap_path}] does not exist. Check your working directory or use an absolute path to the matcap with --matcap-path")


    # TODO: Write a GPU version of the sampler...
    matcap = matcap_sampler(matcap_path)
    matcap_normal = rb.normal.clone()
    matcap_view = rays.dirs.clone()
    if mm is not None:
        mm = mm.to(matcap_normal.device)
        #matcap_normal = torch.mm(matcap_normal.reshape(-1, 3), mm.transpose(1,0))
        #matcap_normal = matcap_normal.reshape(self.width, self.height, 3)
        shape = matcap_view.shape
        matcap_view = torch.mm(matcap_view.reshape(-1, 3), mm.transpose(1,0))
        matcap_view = matcap_view.reshape(*shape)
    vN = spherical_envmap(matcap_view, matcap_normal).cpu().numpy()
    rb.rgb = torch.FloatTensor(matcap(vN)[...,:3].reshape(*matcap_view.shape)).to(matcap_normal.device) / 255.0
    return rb


def matcap_sampler(path, interpolate=True):
    """Fetches MatCap texture & converts to a interpolation function (if needed).

    TODO(ttakikawa): Replace this with something GPU compatible.

    Args:
        path (str): path to MatCap texture
        interpolate (bool): perform interpolation (default: True)

    Returns:
        (np.array) or (scipy.interpolate.Interpolator)
        - The matcap texture
        - A SciPy interpolator function to be used for CPU texture fetch.
    """

    matcap = np.array(Image.open(path)).transpose(1, 0, 2)
    if interpolate:
        return RegularGridInterpolator((np.linspace(0, 1, matcap.shape[0]),
                                        np.linspace(0, 1, matcap.shape[1])), matcap)
    else:
        return matcap

wisp/ops/shaders/matcap.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import torch.nn.functional as F
from scipy.ndimage import gaussian_filter
from wisp.core import RenderBuffer, Rays

""" A collection of shader functions for shadow rays """


def pointlight_shadow_shader(rb: RenderBuffer, rays: Rays, pipeline,
                             point_light=[1.5, 4.5, 1.5], min_y=-2.0) -> RenderBuffer:
    """Apply shadow rays with one secondary ray towards the pointlight.

    Args:
        rb (wisp.core.RenderBuffer): The RenderBuffer.
        rays (wisp.core.Rays): The rays object.
        pipeline (wisp.core.Pipeline): The neural field.
        point_light (list[3] of float): Position of the point light.
        min_y (float): The location of the xz plane.

    Returns:
        (wisp.core.RenderBuffer): The output RenderBuffer.
    """
    rb.shadow = torch.zeros_like(rb.depth)[:, 0].bool().to(rays.origins.device)

    with torch.no_grad():
        plane_hit = torch.zeros_like(rb.depth)[:, 0].bool().to(rays.origins.device)
        rate = -rays.dirs[:, 1]  # check negative sign probably lol
        plane_hit[torch.abs(rate) < 0.00001] = False
        delta = rays.origins[:, 1] - min_y
        plane_t = delta / rate
        plane_hit[(plane_t > 0) & (plane_t < 500)] = True
        plane_hit = plane_hit & (plane_t < rb.depth[..., 0])

        rb.hit = rb.hit & ~plane_hit

        rb.depth[plane_hit] = plane_t[plane_hit].unsqueeze(1)
        rb.xyz[plane_hit] = rays.origins[plane_hit] + rays.dirs[plane_hit] * plane_t[plane_hit].unsqueeze(1)
        rb.normal[plane_hit] = 0
        rb.normal[plane_hit, 1] = 1

        # x is shadow ray origin
        light_o = torch.FloatTensor([[point_light]]).to(rays.origins.device)

        shadow_ray_o = rb.xyz + 0.01 * rb.normal
        shadow_ray_d = torch.zeros_like(rb.xyz).normal_(0.0, 0.01) + \
                       light_o - shadow_ray_o
        shadow_ray_d = F.normalize(shadow_ray_d, dim=1)[0]

        light_hit = ((shadow_ray_d * rb.normal).sum(-1) > 0.0)

        shadow_rays = Rays(origins=shadow_ray_o, dirs=shadow_ray_d, dist_min=0, dist_max=rays.dist_max)
        rb.shadow = pipeline.tracer(pipeline.nef, rays=shadow_rays).hit

        # rb.shadow[~plane_hit] = 0.0
        rb.shadow[~light_hit] = 0.0
        # rb.hit = rb.hit | plane_hit
        shadow_map = torch.clamp((1.0 - rb.shadow.float()) + 0.7, 0.0, 1.0).cpu().numpy()[..., 0]
        shadow_map = torch.from_numpy(gaussian_filter(shadow_map, sigma=2)).unsqueeze(-1)
        rb.rgb[..., :3] *= shadow_map.cuda()
    return rb
wisp/ops/shaders/shadow_rays.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .constructors import *
from .conversions import *
from .metrics import *
from .processing import *
from .sampling import *

wisp/ops/spc/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import numpy as np
import kaolin.ops.spc as spc_ops


def create_dense_octree(level):
    """Creates a dense SPC model

    Args:
        level (int): The level at which the octree will be initialized to.

    Returns:
        (torch.ByteTensor): the octree tensor
    """
    coords = np.arange(2**level)
    points = np.array(np.meshgrid(coords, coords, coords, indexing='xy'))
    points = points.transpose(3,2,1,0).reshape(-1, 3)
    points = torch.from_numpy(points).short().cuda()
    octree = spc_ops.unbatched_points_to_octree(points, level)
    return octree


def make_trilinear_spc(points, pyramid):
    """Builds a trilinear spc from a regular spc.

    Args:
        points (torch.ShortTensor): The point_hierarchy.
        pyramid (torch.LongTensor): The pyramid tensor.

    Returns:
        (torch.ShortTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor)
        - The dual point_hierarchy.
        - The dual pyramid.
        - The trinkets.
        - The parent pointers.
    """
    points_dual, pyramid_dual = spc_ops.unbatched_make_dual(points, pyramid)
    trinkets, parents = spc_ops.unbatched_make_trinkets(points, pyramid, points_dual, pyramid_dual)
    return points_dual, pyramid_dual, trinkets, parents

wisp/ops/spc/constructors.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import kaolin.ops.spc as spc_ops
import wisp.ops.mesh as mesh_ops
from wisp.ops.spc.processing import dilate_points


def pointcloud_to_octree(pointcloud, level, attributes=None, dilate=0):
    """Converts floating point coordinates to an octree.

    Args:
        pointcloud (torch.FloatTensor): 3D coordinates in [-1, 1] of shape [N, 3]
        level (int): Depth of the octreee
        attributes (torch.FloatTensor): Attributes of shape [N, F]. Will be averaged within voxels.
        dilate (int): Dilates the octree if specified.

    Returns:
        (torch.ByteTensor) or (torch.ByteTensor, torch.FloatTensor)
        - octree tensor
        - the averaged attributes
    """
    points = spc_ops.quantize_points(pointcloud.contiguous().cuda(), level)

    for i in range(dilate):
        points = dilate_points(points, level)

    unique, unique_keys, unique_counts = torch.unique(points.contiguous(), dim=0,
                                                      return_inverse=True, return_counts=True)

    morton, keys = torch.sort(spc_ops.points_to_morton(unique.contiguous()).contiguous())

    points = spc_ops.morton_to_points(morton.contiguous())
    octree = spc_ops.unbatched_points_to_octree(points, level, sorted=True)

    if attributes is not None:
        att = torch.zeros_like(unique).float()
        att = att.index_add_(0, unique_keys, attributes) / unique_counts[... ,None].float()
        att = att[keys]
        return octree, att

    return octree


def mesh_to_spc(vertices, faces, level, num_samples=100000000):
    """Construct SPC from a mesh.

    Args:
        vertices (torch.FloatTensor): Vertices of shape [V, 3]
        faces (torch.LongTensor): Face indices of shape [F, 3]
        level (int): The level of the octree
        num_samples (int): The number of samples to be generated on the mesh surface.

    Returns:
        (torch.ByteTensor, torch.ShortTensor, torch.LongTensor, torch.BoolTensor):
        - the octree tensor
        - point hierarchy
        - pyramid
        - prefix
    """
    octree = mesh_to_octree(vertices, faces, level, num_samples)
    points, pyramid, prefix = octree_to_spc(octree)
    return octree, points, pyramid, prefix


def octree_to_spc(octree):
    """Convenience function to generate the usual SPC data.

    Args:
        octree (torch.ByteTensor): SPC octree

    Returns:
        (torch.LongTensor, torch.LongTensor, torch.BoolTensor)
        - point hierarcy
        - pyramid
        - prefix
    """
    lengths = torch.tensor([len(octree)], dtype=torch.int32)
    _, pyramid, prefix = spc_ops.scan_octrees(octree, lengths)
    points = spc_ops.generate_points(octree, pyramid, prefix)
    pyramid = pyramid[0]
    return points, pyramid, prefix


def mesh_to_octree(vertices, faces, level, num_samples=100000000):
    """Construct an octree from a mesh.

    Args:
        vertices (torch.FloatTensor): Vertices of shape [V, 3]
        faces (torch.LongTensor): Face indices of shape [F, 3]
        level (int): The level of the octree
        num_samples (int): The number of samples to be generated on the mesh surface.

    Returns:
        (torch.ByteTensor): the octree tensor
    """
    samples = mesh_ops.sample_surface(vertices.cuda(), faces.cuda(), num_samples)[0]
    # Augment samples... may be a hack that isn't actually needed
    samples = torch.cat([samples,
        samples + (torch.rand_like(samples) * 2.0 - 1.0) * (1.0/(2**(level+1)))], dim=0)
    samples = spc_ops.quantize_points(samples, level)
    octree = spc_ops.unbatched_points_to_octree(samples, level)
    return octree



def mesh_to_trilinear_spc(vertices, faces, level):
    """Builds a trilinear spc from a regular spc.

    Args:
        vertices (torch.FloatTensor): Vertices of shape [V, 3]
        faces (torch.LongTensor): Face indices of shape [F, 3]
        level (int): The level of the octree

    Returns:
        (torch.ByteTensor, torch.ShortTensor, torch.LongTensor, torch.BoolTensor,
         torch.ShortTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor)
        - the octree tensor
        - point hierarchy
        - pyramid
        - prefix
        - The dual point_hierarchy.
        - The dual pyramid.
        - The trinkets.
        - The parent pointers.
    """
    """Builds a trilinear spc from a mesh"""
    octree, points, pyramid, prefix = mesh_to_spc(vertices, faces, level)
    points_dual, pyramid_dual, trinkets, parents = build_trilinear_spc(points, pyramid)
    return octree, points, pyramid, prefix, points_dual, pyramid_dual, trinkets, parents

wisp/ops/spc/conversions.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch


def total_variation(pidx, trinkets, features, level):
    """Calculates total variation for the voxels specified by the pidx.

    Args:
        pidx : int tensor of size [N] specifying the point indices to calculate TV on.
        trinkets : the trinkets.
        features : the features for the given level. (assumes the correct level is given)
        level : int specifying the level of spc

    Returns:
        (torch.FloatTensor): Total variation of shape [1]
    """
    resolution = 2**level
    # N, 8, F tensor of features
    feats = features[trinkets.index_select(0, pidx).long()]

    # N, F
    diff_x = ((torch.abs(feats[:, [4,5,6,7]] - feats[:, [0,1,2,3]]) / resolution)**2).sum((1,2))
    diff_y = ((torch.abs(feats[:, [2,3,6,7]] - feats[:, [0,1,4,5]]) / resolution)**2).sum((1,2))
    diff_z = ((torch.abs(feats[:, [1,3,5,7]] - feats[:, [0,2,4,6]]) / resolution)**2).sum((1,2)) 
    return diff_x + diff_y + diff_z

wisp/ops/spc/metrics.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import kaolin.ops.spc as spc_ops


def dilate_points(points, level):
    """Dilates the octree points.

    Args:
        points (torch.ShortTensor): The SPC points from some level
        level (int): The level from which the points come from

    Returns:
        (torch.ShortTensor): The dilated points
    """
    _x = torch.ShortTensor([[1 ,0 ,0]]).to(points.device)
    _y = torch.ShortTensor([[0 ,1 ,0]]).to(points.device)
    _z = torch.ShortTensor([[0 ,0 ,1]]).to(points.device)
    points = torch.cat([
        points + _x, points - _x,
        points + _y, points - _y,
        points + _z, points - _z,
        points + _x + _y, points + _x - _y, points + _x + _z, points + _x - _z,
        points + _y + _x, points + _y - _x, points + _y + _z, points + _y - _z,
        points + _z + _x, points + _z - _x, points + _z + _y, points + _z - _y,
        points + _x + _y + _z, points + _x + _y - _z,
        points + _x - _y + _z, points + _x - _y - _z,
        points - _x + _y + _z, points - _x + _y - _z,
        points - _x - _y + _z, points - _x - _y - _z,
        ], dim=0)
    points = torch.clip(points, 0, 2** level - 1)

    unique, unique_keys, unique_counts = torch.unique(points.contiguous(), dim=0,
                                                      return_inverse=True, return_counts=True)

    morton, keys = torch.sort(spc_ops.points_to_morton(unique.contiguous()).contiguous())

    points = spc_ops.morton_to_points(morton.contiguous())

    return points
wisp/ops/spc/processing.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch


def sample_spc(
    corners     : torch.Tensor,
    level       : int,
    num_samples : int):
    """Sample uniformly in [-1,1] bounding volume within SPC voxels

    Args:
        corners (tensor)  : set of corners to sample from
        level (int)         : level to sample from
        num_samples (int) : number of points to sample

    Returns:
        (torch.FloatTensor): samples of shape [num_samples, 3]
    """

    res = 2.0**level
    samples = torch.rand(corners.shape[0], num_samples, 3, device=corners.device)
    samples = corners[...,None,:3] + samples
    samples = samples.reshape(-1, 3)
    samples /= res
    return samples * 2.0 - 1.0


def sample_from_depth_intervals(depth_intervals, num_samples):
    """Convert depth intervals to samples.

    SPC raytrace will return a [num_nuggets, 2] array where the first element is the entry depth
    and the second element is the exit depth. This function will convert them into a
    [num_nuggets, num_samples, 3] array of samples.

    Args:
        depth_intervals (torch.FloatTensor): intervals of shape [num_nuggets, 2]
        num_samples (int): sample size

    Returns:
        (torch.FloatTensor): Samples of shape [num_nuggets, num_samples, 3]
    """
    device = depth_intervals.device
    steps = torch.arange(num_samples, device=device)[None].float().repeat([depth_intervals.shape[0], 1])
    steps += torch.rand_like(steps)
    steps *= (1.0 / num_samples)
    samples = depth_intervals[..., 0:1] + (depth_intervals[..., 1:2] - depth_intervals[..., 0:1]) * steps

    return samples


def expand_pack_boundary(pack_boundary, num_samples):
    """Expands the pack boundaries according to the number of samples.

    Args:
        pack_boundary (torch.BoolTensor): pack boundaries [N]
        num_samples (int): Number of samples

    Returns:
        (torch.BoolTensor): pack boundaries of shape [N*num_samples]
    """
    bigpack_boundary = torch.zeros(pack_boundary.shape[0]*num_samples, device=pack_boundary.device).bool()
    bigpack_boundary[pack_boundary.nonzero().long() * num_samples] = True
    bigpack_boundary = bigpack_boundary.int()
    return bigpack_boundary

wisp/ops/spc/sampling.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

wisp/renderer/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from .wisp_app import WispApp
from .optimization_app import OptimizationApp

wisp/renderer/app/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Callable, Dict, List
from wisp.renderer.gui import WidgetImgui
from wisp.renderer.gui import WidgetInteractiveVisualizerProperties, WidgetGPUStats, WidgetSceneGraph, WidgetOptimization
from wisp.renderer.gizmos.gizmo import Gizmo
from wisp.renderer.app.wisp_app import WispApp
from wisp.renderer.core.api import request_redraw
from wisp.framework import WispState, watch


class OptimizationApp(WispApp):
    """ An app for running an optimization and visualizing it's progress interactively in real time. """

    def __init__(self, wisp_state: WispState, trainer_step_func: Callable[[], None], experiment_name: str):
        super().__init__(wisp_state, experiment_name)

        # Tell the renderer to invoke the optimization step() function on every background iteration.
        # The background tasks are constantly invoked by glumpy within the on_idle() event.
        # The actual rendering will occur in-between these calls, invoked by the on_draw() event (which checks if
        # it's time to render the scene again).
        self.register_background_task(trainer_step_func)

    def init_wisp_state(self, wisp_state: WispState) -> None:
        """ A hook for applications to initialize specific fields inside the wisp state object.
        This function is called at the very beginning of WispApp initialization, hence the initialized fields can
        be customized to affect the behaviour of the renderer.
        """
        # Channels available to view over the canvas
        wisp_state.renderer.available_canvas_channels = ["rgb", "depth"]
        wisp_state.renderer.selected_canvas_channel = "rgb"

        # For this app, we'll use only a world grid which resides in the 'xz' plane.
        # This attribute will signal the renderer that some camera controllers should align with this plane,
        # as well as what gizmos (e.g: reference grids, axes) to draw on the canvas.
        wisp_state.renderer.reference_grids = ['xz']

    def create_widgets(self) -> List[WidgetImgui]:
        """ Returns the list of widgets the gui will display, in order. """
        return [WidgetGPUStats(),            # Current FPS, memory occupancy, GPU Model
                WidgetOptimization(),        # Live status of optimization, epochs / iterations count, loss curve
                WidgetInteractiveVisualizerProperties(),  # Canvas dims, user camera controller & definitions
                WidgetSceneGraph()]          # A scene graph tree of the entire hierarchy of objects in the scene
                                             # and their properties

    def create_gizmos(self) -> Dict[str, Gizmo]:
        """ Override to control which gizmos appear on the canvas.
            For example:
            gizmos = dict(
                world_grid_xy=WorldGrid(squares_per_axis=20, grid_size=10,
                                        line_color=(128, 128, 128), line_size=1, plane="xy")
                grid_axis_painter=AxisPainter(axes_length=10, line_width=1, axes=('x', 'y'), is_bidirectional=True)
            )
        """
        return super().create_gizmos()

    def default_user_mode(self) -> str:
        """ Set the default camera controller mode.
            Possible choices: 'First Person View', 'Turntable', 'Trackball' """
        return "Turntable"

    def register_event_handlers(self) -> None:
        """ Register event handlers for various events that occur in a wisp app.
            For example, the renderer is able to listen on updates on fields of WispState objects.
            (note: events will not prompt when iterables like lists, dicts and tensors are internally updated!)
        """
        # Register default events, such as updating the renderer camera controller when the wisp state field changes
        super().register_event_handlers()

        # For this app, we define a custom event that prompts when an optimization epoch is done
        watch(watched_obj=self.wisp_state.optimization, field="epoch", status="changed", handler=self.on_epoch_ended)
        watch(watched_obj=self.wisp_state.optimization, field="running", status="changed",
              handler=self.on_optimization_running_changed)

    def update_renderer_state(self, wisp_state, dt) -> None:
        """
        Populate the scene state object with the most recent information about the interactive renderer.
        The scene state, for example, may be used by the GUI widgets to display up to date information.
        This function is invoked in the beginning of the render() function, before the gui and the canvas are drawn.
        :param wisp_state The WispState object holding shared information between components about the wisp app.
        :param dt Amount of time elapsed since the last update.
        """
        # Update the wisp state with new information about this frame.
        # i.e.: Current FPS, time elapsed.
        super().update_renderer_state(wisp_state, dt)

    def on_epoch_ended(self):
        """ A custom event used by the optimization renderer.
            When an epoch ends, this handler is invoked to force a redraw() and render() of the canvas if needed.
        """
        request_redraw(self.wisp_state)

        # Force render if target FPS is 0 (renderer only responds to specific events) or too much time have elapsed
        if self.is_time_to_render() or self.wisp_state.renderer.target_fps == 0:
            self.render()

    def on_optimization_running_changed(self, value: bool):
        # When training starts / resumes, invoke a redraw() to refresh the renderer core with newly
        # added objects to the scene graph (i.e. the optimized object, or some objects from the dataset).
        if value:
            self.redraw()

wisp/renderer/app/optimization_app.py
Sharath-girish/Shacira
# -*- coding: utf-8 -*-

# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from abc import ABC
import numpy as np
import torch
from glumpy import app, gloo, gl, ext
import imgui
from typing import Optional, Type, Callable, Dict, List, Tuple
from kaolin.render.camera import Camera
from wisp.framework import WispState, watch
from wisp.renderer.core import RendererCore
from wisp.renderer.core import cuda_register_gl_image, cuda_map_resource, cuda_2d_memcpy, cuda_unregister_resource
from wisp.renderer.core.control import CameraControlMode, WispKey, WispMouseButton
from wisp.renderer.core.control import FirstPersonCameraMode, TrackballCameraMode, TurntableCameraMode
from wisp.renderer.core.api import add_pipeline_to_scene_graph
from wisp.renderer.gizmos import Gizmo, WorldGrid, AxisPainter, PrimitivesPainter
from wisp.renderer.gui import WidgetInteractiveVisualizerProperties, WidgetGPUStats, WidgetSceneGraph, WidgetImgui


def enable_amp(func):
    """ An extension to @torch.cuda.amp.autocast which queries WispState to check if
    mixed precision should be enabled.
    """
    def _enable_amp(self: WispApp, *args, **kwargs):
        with torch.cuda.amp.autocast(enabled=self.wisp_state.renderer.enable_amp):
            return func(self, *args, **kwargs)
    return _enable_amp


class WispApp(ABC):
    """ WispApp is a base app implementation which takes care of the entire lifecycle of the rendering loop:
    this is the infinite queue of events which includes: handling of IO and OS events, rendering frames and running
    backgrounds tasks, i.e. optimizations.

    The app is initiated by calling the following functions:
    - register_background_task(): Registers a task to run alternately with the render() function per
      frame. Background tasks can be, i.e, functions that run a single optimization step for some neural object.
    - run(): Initiates the rendering loop. This method blocks the calling thread until the window is closed.

    Future custom interactive apps can subclass this base app, and inherit methods to customize the behaviour
    of the app:
    - init_wisp_state(): A hook for initializing the fields of the shared state object.
        The interactive renderer configuration, scene graph and user custom fields can be initialized here.
    - update_render_state(): A hook for updating the fields of the shared state object at the beginning of each
      frame.
    - create_widgets(): Controls which gui components the app uses
    - create_gizmos(): Controls which transient canvas drawable components will be used (OpenGL based).
    - default_user_mode(): The default camera controller mode (first person, trackball, turntable).
    - register_event_handlers(): Registers which methods are invoked in response to app events / wisp state changes.

    The rendering loop alternates between the following calls:
    - on_idle, which invokes user background tasks, registered via register_background_task before calling run().
    - on_draw, which invokes render() when it's time to draw a new frame

    Internally, the app uses the RendererCore object to manage the drawing of all objects in the scene graph.
    The app may request the RendererCore to switch into "interactive mode" to ensure the FPS remains interactive
    (this is done at the expense of rendering quality).
    Interactive mode is automatically initiated, i.e, during user interactions.

    The render() logic is composed by the following sub-functions:
    - update_render_state() - Updates the fields of the shared state object at the beginning of each frame.
    - render_gui() - Renders the imgui components over the canvas, fromm scratch (imgui uses immediate mode gui).
    - redraw() - A "heavier" function which forces the render-core to refresh its internal state.
      Newly created objects may get added to the scene graph, and obsolete objects may get removed.
      Vector-primitives data layers may regenerate here.
    - render_canvas() - Invokes the render-core to obtain a RenderBuffer of the rendered scene objects.
      The lion share of draw logic occurs within this call, in particular the drawing of neural objects.
    - _blit_to_gl_renderbuffer - Copies the RenderBuffer results to the screen buffer.
    - Gizmos are finally drawn directly to the screen framebuffer (as common OpenGL draw calls), these objects
      are considered transient in the sense that they don't belong to the scene graph.
    - Timer tick events may also get taken care of during the rendering loop (i.e: adjust velocity of user camera).
    Users should rarely override these functions, unless they're sure about what they're doing.
    Customizing the app behaviour should always be preferred via the initialization hooks.
    """

    # Period of time between user interactions before resetting back to full resolution mode
    COOLDOWN_BETWEEN_RESOLUTION_CHANGES = 0.35  # In seconds

    def __init__(self, wisp_state, window_name="wisp app"):
        # Initialize app state instance
        self.wisp_state: WispState = wisp_state
        self.init_wisp_state(wisp_state)

        # Create main app window & initialize GL context
        # glumpy with a specialized glfw backend takes care of that (backend is imgui integration aware)
        window = self._create_window(self.width, self.height, window_name, gl_version=wisp_state.renderer.gl_version)
        self.register_io_mappings()

        # Initialize gui, assumes the window is managed by glumpy with glfw
        imgui.create_context()
        self._is_imgui_focused = False
        self._is_imgui_hovered = False
        self._is_reposition_imgui_menu = True
        self.canvas_dirty = False
        self.redraw_every_frame = False

        # Tell torch to initialize the CUDA context
        torch.cuda.init()

        # Initialize applicative renderer, which independently paints images for the main canvas
        render_core = RendererCore(self.wisp_state)

        self.window = window                    # App window with a GL context & oversees event callbacks
        self.render_core = render_core          # Internal renderer, responsible for painting over canvas
        self.render_clock = app.clock.Clock()
        self.render_clock.tick()
        self.interactions_clock = app.clock.Clock()
        self.interactions_clock.tick()
        self._was_interacting_prev_frame = False

        # The initialization of these fields is deferred util "on_resize" is first prompted.
        # There we generate a simple billboard GL program (normally with a shared CUDA resource)
        # Canvas content will be blitted onto it
        self.canvas_program: Optional[gloo.Program] = None   # GL program used to paint a single billboard
        self.cugl_rgb_handle = None                              # CUDA buffer, as a shared resource with OpenGL
        self.cugl_depth_handle = None

        try:
            # WSL does not support CUDA-OpenGL interoperability, fallback to device2host2device copy instead
            from platform import uname
            is_wsl = 'microsoft-standard' in uname().release
            self.blitdevice2device = not is_wsl
        except Exception:
            # By default rendering results copy directly from torch/cuda mem to OpenGL Texture
            self.blitdevice2device = True

        self.user_mode: CameraControlMode = None    # Camera controller object (first person, trackball or turntable)

        self.widgets = self.create_widgets()        # Create gui widgets for this app
        self.gizmos = self.create_gizmos()          # Create canvas widgets for this app
        self.prim_painter = PrimitivesPainter()

        self.register_event_handlers()
        self.change_user_mode(self.default_user_mode())

        self.redraw()   # Refresh RendererCore
    
    def add_pipeline(self, name, pipeline, transform=None):
        """Register a neural fields pipeline into the scene graph.

        Args:
            name (str): The name of the pipeline.
            pipeline (wisp.models.Pipeline): The pipeline holding a tracer and nef.
            transform (wisp.core.ObjectTransform): The transform for the pipeline.
        """
        add_pipeline_to_scene_graph(self.wisp_state, name, pipeline, transform=transform)
    
    def add_widget(self, widget):
        """ Adds a widget to the list of widgets.

        Args:
            widget (wisp.renderer.gui.imgui.WidgetImgui): The widget to add.
        """
        self.widgets.append(widget)

    def add_gizmo(self, name, gizmo):
        """Adds a gizmo to the list of gizmos.

        Args:
            name (str): The name of the gizmo.
            gizmo (wisp.renderer.gizmos.Gizmo): The gizmo to add.
        """
        self.gizmos[name] = gizmo

    def init_wisp_state(self, wisp_state: WispState) -> None:
        """ A hook for applications to initialize specific fields inside the wisp state object.
        This function is called at the very beginning of WispApp initialization, hence the initialized fields can
        be customized to affect the behaviour of the renderer.
        """
        # Channels available to view over the canvas
        wisp_state.renderer.available_canvas_channels = ["rgb", "depth"]
        wisp_state.renderer.selected_canvas_channel = "rgb"

    def create_widgets(self) -> List[WidgetImgui]:
        """ Returns which widgets the gui will display, in order.
        Override to define which gui widgets are used by the wisp app.
        """
        return [WidgetGPUStats(), WidgetInteractiveVisualizerProperties(), WidgetSceneGraph()]

    def create_gizmos(self) -> Dict[str, Gizmo]:
        """ Override to define which gizmos are painted on the canvas by the wisp app.
        Gizmos are transient rasterized objects rendered by OpenGL on top of the canvas.
        For example: world grid, axes painter.
        """
        gizmos = dict()
        grid_size = 10.0
        planes = self.wisp_state.renderer.reference_grids
        axes = set(''.join(planes))
        for plane in planes:
            gizmos[f'world_grid_{plane}'] = WorldGrid(squares_per_axis=20, grid_size=grid_size,
                                                      line_color=(128, 128, 128), line_size=1, plane=plane)
            gizmos[f'world_grid_fine_{plane}'] = WorldGrid(squares_per_axis=200, grid_size=10.0,
                                                           line_color=(128, 128, 128), line_size=2, plane=plane)
        # Axes on top of the reference grid
        gizmos['grid_axis_painter'] = AxisPainter(axes_length=grid_size, line_width=1,
                                                  axes=axes, is_bidirectional=False)
        return gizmos

    def default_user_mode(self) -> str:
        """ Override to determine the default camera control mode.
        Possible choices: 'First Person View', 'Turntable', 'Trackball'
        """
        return "Turntable"

    def register_event_handlers(self) -> None:
        """ Override (and call super) to register additional event handlers """
        watch(watched_obj=self.wisp_state.renderer, field="cam_controller", status="changed",
              handler=self.on_cam_controller_changed)
        watch(watched_obj=self.wisp_state.renderer, field="selected_camera", status="changed",
              handler=self.on_selected_camera_changed)
        watch(watched_obj=self.wisp_state.renderer, field="selected_canvas_channel", status="changed",
              handler=self.on_selected_canvas_channel_changed)
        watch(watched_obj=self.wisp_state.renderer, field="selected_camera_lens", status="changed",
              handler=self.on_selected_camera_lens_changed)
        watch(watched_obj=self.wisp_state.renderer, field="clear_color_value", status="changed",
              handler=self.on_clear_color_value_changed)

    def on_cam_controller_changed(self, value: Type[CameraControlMode]):
        # Stay synced with state change: generate new instance of mode type
        mode_type = value
        self.user_mode = mode_type(render_core=self.render_core, wisp_state=self.wisp_state)

    def on_selected_camera_changed(self, value: Camera):
        # Rebuild camera controller to free any cached info from previous camera
        cam_controller_cls = type(self.user_mode)
        self.user_mode = cam_controller_cls(render_core=self.render_core, wisp_state=self.wisp_state)

        # Adjust the width / height according to current state of the renderer window
        self.render_core.resize_canvas(height=self.height, width=self.width)

    def on_selected_camera_lens_changed(self, value: str):
        self.render_core.change_camera_projection_mode(value)

    def on_selected_canvas_channel_changed(self, value: str):
        self.canvas_dirty = True    # Request canvas redraw

    def on_clear_color_value_changed(self, value: Tuple[float, float, float]):
        self.canvas_dirty = True    # Request canvas redraw

    def run(self):
        """ Initiate events message queue, which triggers the rendering loop.
        This call will block the thread until the app window is closed.
        """
        app.run()   # App clock should always run as frequently as possible (background tasks should not be limited)

    def _create_window(self, width, height, window_name, gl_version):
        # Currently assume glfw backend due to integration with imgui
        app.use(f"glfw_imgui ({gl_version})")
        win_config = app.configuration.Configuration()
        if self.wisp_state.renderer.antialiasing == 'msaa_4x':
            win_config.samples = 4

        # glumpy implicitly sets the GL context as current
        window = app.Window(width=width, height=height, title=window_name, config=win_config)
        window.on_draw = self.on_draw
        window.on_resize = self.on_resize
        window.on_key_press = self.on_key_press
        window.on_key_release = self.on_key_release
        window.on_mouse_press = self.on_mouse_press
        window.on_mouse_drag = self.on_mouse_drag
        window.on_mouse_release = self.on_mouse_release
        window.on_mouse_scroll = self.on_mouse_scroll
        window.on_mouse_motion = self.on_mouse_motion

        if self.wisp_state.renderer.antialiasing == 'msaa_4x':
            gl.glEnable(gl.GL_MULTISAMPLE)

        return window

    @staticmethod
    def _create_gl_depth_billboard_program(texture: np.ndarray, depth_texture: np.ndarray):
        vertex = """
                    uniform float scale;
                    attribute vec2 position;
                    attribute vec2 texcoord;
                    varying vec2 v_texcoord;
                    void main()
                    {
                        v_texcoord = texcoord;
                        gl_Position = vec4(scale*position, 0.0, 1.0);
                    } """

        fragment = """
                    uniform sampler2D tex;
                    uniform sampler2D depth_tex;
                    varying vec2 v_texcoord;
                    void main()
                    {
                        gl_FragColor = texture2D(tex, v_texcoord);
                        gl_FragDepth = texture2D(depth_tex, v_texcoord).r;
                    } """
        # TODO (operel): r component is a waste?

        # Compile GL program
        canvas = gloo.Program(vertex, fragment, count=4)

        # Upload fixed values to GPU
        canvas['position'] = [(-1, -1), (-1, +1), (+1, -1), (+1, +1)]
        canvas['texcoord'] = [(0, 0), (0, 1), (1, 0), (1, 1)]
        canvas['scale'] = 1.0
        canvas['tex'] = texture
        canvas['depth_tex'] = depth_texture
        return canvas

    @staticmethod
    def _create_screen_texture(res_h, res_w, channel_depth, dtype=np.uint8):
        """ Create and return a Texture2D with gloo and a cuda handle. """
        if issubclass(dtype, np.integer):
            tex = np.zeros((res_h, res_w, channel_depth), dtype).view(gloo.Texture2D)
        elif issubclass(dtype, np.floating):
            tex = np.zeros((res_h, res_w, channel_depth), dtype).view(gloo.TextureFloat2D)
        else:
            raise ValueError(f'_register_cugl_shared_texture invoked with unsupported texture dtype: {dtype}')
        # Force gloo to create GL object on GPU
        tex.activate()
        tex.deactivate()
        return tex

    def _register_cugl_shared_texture(self, tex):
        if self.blitdevice2device:
            # Create shared GL / CUDA resource
            handle = cuda_register_gl_image(image=int(tex.handle), target=tex.target)
        else:
            # No shared resource required, as we copy from cuda buffer -> cpu -> GL texture
            handle = None
        return handle

    def _reposition_gui_menu(self, menu_width, main_menu_height):
        window_height = self.window.height
        window_width = self.window.width
        imgui.set_next_window_size(width=menu_width, height=window_height-main_menu_height, condition=imgui.ALWAYS)
        imgui.set_next_window_position(x=window_width-menu_width, y=main_menu_height, condition=imgui.ALWAYS)
        self._is_reposition_imgui_menu = False

    def render_gui(self, state):
        """ Render the entire gui window per frame (imgui works in immediate mode).
            Internally, the Widgets take care of rendering the actual content.
        """
        imgui.new_frame()
        if imgui.begin_main_menu_bar():
            main_menu_height = imgui.get_window_height()
            if imgui.begin_menu("File", True):

                clicked_quit, selected_quit = imgui.menu_item(
                    "Quit", 'Cmd+Q', False, True
                )

                if clicked_quit:
                    exit(1)

                imgui.end_menu()
            imgui.end_main_menu_bar()

        if self._is_reposition_imgui_menu:
            self._reposition_gui_menu(menu_width=350, main_menu_height=main_menu_height)
        imgui.begin("Scene Information", True)

        for widget in self.widgets:
            widget.paint(state)

        self._is_imgui_hovered = imgui.core.is_window_hovered(imgui.HOVERED_ANY_WINDOW |
                                                              imgui.HOVERED_ALLOW_WHEN_BLOCKED_BY_POPUP |
                                                              imgui.HOVERED_ALLOW_WHEN_BLOCKED_BY_ACTIVE_ITEM)
        self._is_imgui_hovered = self._is_imgui_hovered or \
                                 imgui.core.is_any_item_hovered() or imgui.is_any_item_active()
        imgui.end()

        self._is_imgui_focused = imgui.is_window_focused(imgui.FOCUS_ROOT_WINDOW)

        imgui.end_frame()
        imgui.render()

    def render_canvas(self, render_core, time_delta, force_render):
        """ Invoke the render-core to render all neural fields and blend into a single Renderbuffer.
        The rgb and depth channels passed on to the app.
        """
        # The render core returns a RenderBuffer
        renderbuffer = render_core.render(time_delta, force_render)
        buffer_attachment = renderbuffer.image().rgba
        buffer_attachment = buffer_attachment.flip([0])  # Flip y axis
        img = buffer_attachment.byte().contiguous()

        buffer_attachment_depth = renderbuffer.depth
        buffer_attachment_depth = buffer_attachment_depth.flip([0])
        depth_img = buffer_attachment_depth.repeat(1,1,4).contiguous()

        return img, depth_img

    def _blit_to_gl_renderbuffer(self, img, depth_img, canvas_program, cugl_rgb_handle, cugl_depth_handle, height):
        if self.blitdevice2device:
            # Device to device copy: Copy CUDA buffer to GL Texture mem
            shared_tex = canvas_program['tex']
            shared_tex_depth = canvas_program['depth_tex']

            # copy from torch into buffer
            assert shared_tex.nbytes == img.numel() * img.element_size()
            assert shared_tex_depth.nbytes == depth_img.numel() * depth_img.element_size()    # TODO: using a 4d tex
            cuda_2d_memcpy(resource_handle=cugl_rgb_handle, shared_tex=shared_tex, img=img, height=height)
            cuda_2d_memcpy(resource_handle=cugl_depth_handle, shared_tex=shared_tex_depth, img=depth_img, height=height)
            torch.cuda.synchronize()
        else:
            # Device to host to device copy: Move torch tensors to cpu and upload as texture data
            canvas_program['tex'] = img.cpu().numpy()
            canvas_program['depth_tex'] = depth_img.cpu().numpy()

        canvas_program.draw(gl.GL_TRIANGLE_STRIP)

    def update_renderer_state(self, wisp_state, dt):
        """
        Populate the scene state object with the most recent information about the interactive renderer.
        The scene state, for example, may be used by the GUI widgets to display up to date information.
        This function is invoked in the beginning of the render() function, before the gui and the canvas are drawn.
        :param wisp_state The WispState object holding shared information between components about the wisp app.
        :param dt Amount of time elapsed since the last update.
        """
        wisp_state.renderer.fps = app.clock.get_fps()
        wisp_state.renderer.dt = dt
        wisp_state.renderer.cam_controller = type(self.user_mode)

    def change_user_mode(self, user_mode: str):
        """ Changes the camera controller mode """
        if user_mode == 'Trackball':
            self.wisp_state.renderer.cam_controller = TrackballCameraMode
        elif user_mode == 'First Person View':
            self.wisp_state.renderer.cam_controller = FirstPersonCameraMode
        elif user_mode == 'Turntable':
            self.wisp_state.renderer.cam_controller = TurntableCameraMode

    @torch.no_grad()
    @enable_amp
    def redraw(self):
        """ Asks the render core to redraw the scene:
        - The scene graph will be refreshed (new objects added will create their renderers if needed)
        - Data layers will regenerate according to up-to-date state.
        render() may internally invoke redraw() when the canvas is tagged as "dirty".
        A render() call is required to display changes caused by redraw() on the canvas.
        """
        # Let the renderer redraw the data layers if needed
        self.render_core.redraw()

        # Regenerate the GL primitives according to up-to-date data layers
        layers_to_draw = self.render_core.active_data_layers()
        self.prim_painter.redraw(layers_to_draw)

    @torch.no_grad()
    @enable_amp
    def render(self):
        """ Renders a single frame. """
        dt = self.render_clock.tick()  # Tick render clock: dt is now the exact time elapsed since last render

        # Populate the scene state with the most recent information about the interactive renderer.
        # The scene state, for example, may be used by the GUI widgets to display up to date information.
        self.update_renderer_state(self.wisp_state, dt)

        # Clear color / depth buffers before rendering the next frame
        clear_color = (*self.wisp_state.renderer.clear_color_value, 1.0)    # RGBA
        self.window.clear(color=clear_color)

        # imgui renders first
        self.render_gui(self.wisp_state)

        if self.redraw_every_frame:
            self.canvas_dirty = True

        # The app was asked to redraw the scene, inform the render core
        if self.canvas_dirty:
            self.redraw()

        # Invoke the timer tick event, and let the camera controller update the state of any interactions
        # of the user which involve the time elapsed (i.e: velocity, acceleration of movements).
        self.user_mode.handle_timer_tick(dt)

        # Toggle interactive mode on or off if needed to maintain interactive FPS rate
        if self.user_mode.is_interacting():
            self.render_core.set_low_resolution()
        else:
            # Allow a fraction of a second before turning full resolution on.
            # User interactions sometimes behave like a rapid burst of short and quick interactions.
            if self._was_interacting_prev_frame:
                self.interactions_clock.tick()
            time_since_last_interaction = self.interactions_clock.time() - self.interactions_clock.last_ts
            if time_since_last_interaction > self.COOLDOWN_BETWEEN_RESOLUTION_CHANGES:
                self.render_core.set_full_resolution()

        self._was_interacting_prev_frame = self.user_mode.is_interacting()

        # render canvas: core proceeds by invoking internal renderers tracers
        # output is rendered on a Renderbuffer object, backed by torch tensors
        img, depth_img = self.render_canvas(self.render_core, dt, self.canvas_dirty)

        # glumpy code injected within the pyimgui render loop to blit the renderer core output to the actual canvas
        # The torch buffers are copied by with cuda, connected as shared resources as 2d GL textures
        self._blit_to_gl_renderbuffer(img, depth_img, self.canvas_program, self.cugl_rgb_handle,
                                      self.cugl_depth_handle, self.height)

        # Finally, render OpenGL gizmos on the canvas.
        # This may include the world grid, or vectorial lines / points belonging to data layers
        camera = self.render_core.camera
        for gizmo in self.gizmos.values():
            gizmo.render(camera)
        self.prim_painter.render(camera)
        self.canvas_dirty = False

    def register_background_task(self, hook: Callable[[], None]) -> None:
        """ Register a new callable function to run in conjunction with the rendering loop.
            The app will alternate between on_idle calls, invoking the background task, and on_draw
            calls, invoking the rendering itself, both occurring on the same thread.
        """
        if hook is not None:
            def _run_hook(dt: float):
                if not self.wisp_state.renderer.background_tasks_paused:
                    hook()
            self.window.on_idle = _run_hook

    def on_draw(self, dt=None):
        """ glumpy's event to draw the next frame. Invokes the render() function if needed. """
        # dt arg comes from the app clock, the renderer clock is maintained separately from the background tasks
        # Interactive mode on, or interaction have just started
        is_interacting = self.wisp_state.renderer.interactive_mode or self.user_mode.is_interacting()
        if is_interacting or self.is_time_to_render():
            self.render()     # Render objects uploaded to GPU

    def is_time_to_render(self):
        time_since_last_render = self.render_clock.time() - self.render_clock.last_ts
        target_fps = self.wisp_state.renderer.target_fps
        if target_fps is None or ((target_fps > 0) and time_since_last_render >= (1 / target_fps)):
            return True
        return False

    def on_resize(self, width, height):
        """ Invoked when the window is first created, or resized.
        A resize causes internal textures and buffers to regenerate according the window size.
        """
        self.width = width
        self.height = height

        # Handle deallocation of shared resources
        if self.cugl_rgb_handle is not None:
            cuda_unregister_resource(self.cugl_rgb_handle)
            self.cugl_rgb_handle = None
        if self.cugl_depth_handle is not None:
            cuda_unregister_resource(self.cugl_depth_handle)
            self.cugl_depth_handle = None
        tex = self._create_screen_texture(height, width, self.channel_depth, dtype=np.uint8)
        depth_tex = self._create_screen_texture(height, width, 4, dtype=np.float32)  # TODO: Single channel
        self.cugl_rgb_handle = self._register_cugl_shared_texture(tex)
        self.cugl_depth_handle = self._register_cugl_shared_texture(depth_tex)
        if self.canvas_program is None:
            self.canvas_program = self._create_gl_depth_billboard_program(texture=tex, depth_texture=depth_tex)
        else:
            if self.canvas_program['tex'] is not None:
                self.canvas_program['tex'].delete()
            if self.canvas_program['depth_tex'] is not None:
                self.canvas_program['depth_tex'].delete()
            self.canvas_program['tex'] = tex
            self.canvas_program['depth_tex'] = depth_tex

        self.render_core.resize_canvas(height=height, width=width)
        self.window.activate()
        gl.glViewport(0, 0, width, height)
        self._is_reposition_imgui_menu = True   # Signal menu it needs to shift after resize

    def is_canvas_event(self):
        """ Is canvas in focus or any of imgui's windows """
        return not self._is_imgui_focused

    def on_mouse_press(self, x, y, button):
        if self.is_canvas_event():
            self.user_mode.handle_mouse_press(x, y, button)

    def on_mouse_drag(self, x, y, dx, dy, button):
        if self.is_canvas_event():
            self.user_mode.handle_mouse_drag(x, y, dx, dy, button)

    def on_mouse_release(self, x, y, button):
        if self.is_canvas_event():
            self.user_mode.handle_mouse_release(x, y, button)

    def on_mouse_scroll(self, x, y, dx, dy):
        """ The mouse wheel was scrolled by (dx,dy). """
        if self.is_canvas_event():
            self.user_mode.handle_mouse_scroll(x, y, dx, dy)

    def on_mouse_motion(self, x, y, dx, dy):
        """ The mouse was moved with no buttons held down. """
        if self.is_canvas_event():
            self.user_mode.handle_mouse_motion(x, y, dx, dy)

    @property
    def width(self):
        """ Returns the canvas width """
        return self.wisp_state.renderer.canvas_width

    @width.setter
    def width(self, value: int):
        """ Sets the canvas width """
        self.wisp_state.renderer.canvas_width = value

    @property
    def height(self):
        """ Returns the canvas height """
        return self.wisp_state.renderer.canvas_height

    @height.setter
    def height(self, value: int):
        """ Sets the canvas height """
        self.wisp_state.renderer.canvas_height = value

    @property
    def channel_depth(self):
        """ Returns the number of channels the screenbuffer uses for the color attachment """
        return 4  # Assume the framebuffer keeps RGBA

    @property
    def canvas_dirty(self):
        """ Returns if the canvas is dirty,
        that is, the app requires a redraw() to stay in sync with external changes
        """
        return self.wisp_state.renderer.canvas_dirty

    @canvas_dirty.setter
    def canvas_dirty(self, value: bool):
        """ Marks the canvas as dirty,
        implying the app requires a redraw() to stay in sync with external changes
        """
        self.wisp_state.renderer.canvas_dirty = value

    def _update_imgui_keys(self, symbol):
        # Normally glfw shouldn't be explicitly imported as glumpy uses it as backend.
        # However, here we are forced to do that to take care of missing key mappings
        import glfw
        keys = [glfw.KEY_BACKSPACE, glfw.KEY_DELETE, glfw.KEY_ENTER, glfw.KEY_HOME, glfw.KEY_END,
                glfw.KEY_LEFT_SHIFT, glfw.KEY_RIGHT_SHIFT,
                glfw.KEY_RIGHT, glfw.KEY_LEFT, glfw.KEY_UP, glfw.KEY_DOWN,
                glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4, glfw.KEY_5,
                glfw.KEY_6, glfw.KEY_7, glfw.KEY_8, glfw.KEY_9,
                glfw.KEY_KP_0, glfw.KEY_KP_1, glfw.KEY_KP_2, glfw.KEY_KP_3, glfw.KEY_KP_4, glfw.KEY_KP_5,
                glfw.KEY_KP_6, glfw.KEY_KP_7, glfw.KEY_KP_8, glfw.KEY_KP_9,
                glfw.KEY_KP_ENTER, glfw.KEY_KP_ADD, glfw.KEY_KP_SUBTRACT, glfw.KEY_KP_DECIMAL
                ]
        mappings = {    # Missing keys from glumpy glfw-imgui backend
            glfw.KEY_KP_0: glfw.KEY_0,
            glfw.KEY_KP_1: glfw.KEY_1,
            glfw.KEY_KP_2: glfw.KEY_2,
            glfw.KEY_KP_3: glfw.KEY_3,
            glfw.KEY_KP_4: glfw.KEY_4,
            glfw.KEY_KP_5: glfw.KEY_5,
            glfw.KEY_KP_6: glfw.KEY_6,
            glfw.KEY_KP_7: glfw.KEY_7,
            glfw.KEY_KP_8: glfw.KEY_8,
            glfw.KEY_KP_9: glfw.KEY_9,
            glfw.KEY_KP_ENTER: glfw.KEY_ENTER,
            glfw.KEY_KP_SUBTRACT: glfw.KEY_MINUS,
            glfw.KEY_KP_DECIMAL: glfw.KEY_PERIOD
        }
        updated_symbol = symbol
        for key in keys:
            is_key_on = glfw.get_key(self.window.native_window, key)
            imgui.get_io().keys_down[key] = is_key_on

            if symbol == -1 and is_key_on and key in mappings:
                updated_symbol = mappings[key]

        # TODO: Verify imgui keys have been properly mapped during initialization..
        # imgui.get_io().key_map[imgui.KEY_BACKSPACE] = app.window.key.BACKSPACE
        # imgui.get_io().key_map[imgui.KEY_DELETE] = app.window.key.DELETE
        # imgui.get_io().key_map[imgui.KEY_ENTER] = glfw.KEY_ENTER
        # imgui.get_io().key_map[imgui.KEY_RIGHT_ARROW] = glfw.KEY_RIGHT
        # imgui.get_io().key_map[imgui.KEY_LEFT_ARROW] = glfw.KEY_LEFT
        # imgui.get_io().key_map[imgui.KEY_UP_ARROW] = glfw.KEY_UP
        # imgui.get_io().key_map[imgui.KEY_DOWN_ARROW] = glfw.KEY_DOWN
        return updated_symbol

    def on_key_press(self, symbol, modifiers):
        symbol = self._update_imgui_keys(symbol)
        if symbol > 0:
            imgui.get_io().add_input_character(symbol)

        if self.is_canvas_event():
            self.user_mode.handle_key_press(symbol, modifiers)

            # TODO: Shouldn't be here
            cam_mode = None
            if symbol in (app.window.key.T, ord('T'), ord('t')):
                cam_mode = "Trackball"
            elif symbol in (app.window.key.F, ord('F'), ord('f')):
                cam_mode = "First Person View"
            elif symbol in (app.window.key.U, ord('U'), ord('u')):
                cam_mode = "Turntable"
            if cam_mode is not None:
                self.change_user_mode(cam_mode)

    def on_key_release(self, symbol, modifiers):
        symbol = self._update_imgui_keys(symbol)

        if self.is_canvas_event():
            self.user_mode.handle_key_release(symbol, modifiers)

    def dump_framebuffer(self, path='./framebuffer'):
        # Dumps debug images of the GL screen framebuffer.
        # This framebuffer should reflect the exact content of the window.
        framebuffer = np.zeros((self.width, self.height * 3), dtype=np.uint8)
        gl.glReadPixels(0, 0, self.width, self.height,
                        gl.GL_RGB, gl.GL_UNSIGNED_BYTE, framebuffer)
        framebuffer = np.flip(framebuffer, 0)
        ext.png.from_array(framebuffer, 'RGB').save(path + '_color.png')

        framebuffer = np.zeros((self.width, self.height), dtype=np.float32)
        gl.glReadPixels(0, 0, self.width, self.height,
                        gl.GL_DEPTH_COMPONENT, gl.GL_FLOAT, framebuffer)
        framebuffer = np.flip(framebuffer, 0)
        ext.png.from_array(framebuffer, 'L').save(path + '_depth.png')

    def register_io_mappings(self):
        WispMouseButton.register_symbol(WispMouseButton.LEFT_BUTTON, app.window.mouse.LEFT)
        WispMouseButton.register_symbol(WispMouseButton.MIDDLE_BUTTON, app.window.mouse.MIDDLE)
        WispMouseButton.register_symbol(WispMouseButton.RIGHT_BUTTON, app.window.mouse.RIGHT)

        WispKey.register_symbol(WispKey.LEFT, app.window.key.LEFT)
        WispKey.register_symbol(WispKey.RIGHT, app.window.key.RIGHT)
        WispKey.register_symbol(WispKey.UP, app.window.key.UP)
        WispKey.register_symbol(WispKey.DOWN, app.window.key.DOWN)

        # TODO: Take care of remaining mappings, and verify the event handlers of glumpy were not overriden

wisp/renderer/app/wisp_app.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .cuda import *
from .render_core import *

wisp/renderer/core/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .renderers_factory import *
from .base_renderer import *
from .raytraced_renderer import *
from .decorators import *
from .scenegraph import *

wisp/renderer/core/api/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Dict, Set, Tuple, Any
import torch
from kaolin.render.camera import Camera
from wisp.core import RenderBuffer, Rays, PrimitivesPack, WispModule


@dataclass
class FramePayload:
    """This is a dataclass which holds metadata for the current frame.
    """
    camera: Camera
    visible_objects: Set[str]
    interactive_mode: bool
    render_res_x: int
    render_res_y: int
    time_delta: float   # In seconds
    clear_color: Tuple[float, float, float]
    channels: Set[str] # Channels requested for the render.


class BottomLevelRenderer(WispModule, ABC):

    def __init__(self, *args, **kwargs):
        super().__init__()
        self._data_layers = dict()

    def pre_render(self, payload: FramePayload, *args, **kwargs) -> None:
        """ General frame setup occurs here, for example:
            1. Ray Tracers - ray generation.
            2. Update shader uniforms.
        """
        pass

    def needs_refresh(self, payload: FramePayload, *args, **kwargs) -> bool:
        """ Override to optimize cases when the neural field does not require rendering from scratch. """
        return True

    @abstractmethod
    def render(self, *args, **kwargs) -> RenderBuffer:
        raise NotImplementedError('BottomLevelRenderer subclasses must implement render() logic.')

    def redraw(self):
        if self.needs_redraw():
            self._data_layers = self.regenerate_data_layers()

    def post_render(self, *args, **kwargs) -> None:
        """ General frame teardown logic takes place, for example:
            1. Cleanup of temporary information generated during the frame
            2. Caching of information relevant for the next frames.
        """
        pass

    def needs_redraw(self) -> bool:
        return True

    def regenerate_data_layers(self) -> Dict[str, PrimitivesPack]:
        return dict()

    def data_layers(self) -> Dict[str, PrimitivesPack]:
        """ Returns layers of information made of primitives, visually describing the renderer internal structures """
        return self._data_layers

    @property
    @abstractmethod
    def device(self) -> torch.device:
        raise NotImplementedError('BottomLevelRenderer subclasses must implement device')

    @property
    @abstractmethod
    def dtype(self) -> torch.dtype:
        raise NotImplementedError('BottomLevelRenderer subclasses must implement dtype')

    def acceleration_structure(self) -> Optional[str]:
        """ Returns a descriptive name of the acceleration structure used by this object, if applicable. """
        return None

    def features_structure(self) -> Optional[str]:
        """ Returns a descriptive name of the feature structure used by this object, if applicable. """
        return None

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.

        BLASGrids are generally assumed to contain a bottom level acceleration structure.
        """
        return dict()


class RasterizedRenderer(BottomLevelRenderer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @abstractmethod
    def render(self, camera: Camera) -> RenderBuffer:
        raise NotImplementedError('RasterizedRenderer subclasses must implement render() logic.')

wisp/renderer/core/api/base_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Type
from wisp.models.nefs import BaseNeuralField
from wisp.tracers import BaseTracer
from wisp.renderer.core.api.base_renderer import BottomLevelRenderer
from wisp.renderer.core.api.renderers_factory import register_neural_field_type


def field_renderer(field_type: Type[BaseNeuralField], tracer_type: Type[BaseTracer]):
    """ A decorator that registers a neural field type with a renderer.
        By registering the renderer type, the interactive renderer knows what type of renderer to create
        when dealing with this type of field.
        Essentially, this allows displaying custom types of objects on the canvas.
    """
    def _register_renderer_fn(renderer_class: Type[BottomLevelRenderer]):
        register_neural_field_type(field_type, tracer_type, renderer_class)
        return renderer_class
    return _register_renderer_fn

wisp/renderer/core/api/decorators.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Optional, Dict
import copy
import torch
from wisp.core import RenderBuffer, Rays, PrimitivesPack
from wisp.models import Pipeline
from wisp.models.nefs import BaseNeuralField
from wisp.tracers import BaseTracer
from wisp.renderer.core.api.base_renderer import BottomLevelRenderer, FramePayload
from wisp.renderer.core.api.decorators import field_renderer
from wisp.gfx.datalayers import Datalayers, OctreeDatalayers, AABBDatalayers
from wisp.accelstructs import OctreeAS, AxisAlignedBBoxAS


@field_renderer(BaseNeuralField, BaseTracer)
class RayTracedRenderer(BottomLevelRenderer):
    """ A default neural field renderers for all neural ray-based pipelines.
        The renderer is registered with the general BaseNeuralField & BaseTracer
        to make a default fallback for future combos of neural field / tracer which
        don't implement a dedicated renderer.

        Renderers represent "visualizable" objects in the interactive renderer system.
        These include anything that could potentially be painted on the canvas, how to tune it, and what data-layers
        it supports.
        This specific renderer is concerned with neural objects which rely on ray tracing.

        This class also serves as a convenience super-class for other renderers which shouldn't implement the entire
        BottomLevelRenderer functionality from scratch.
    """
    def __init__(self, nef: BaseNeuralField, tracer: BaseTracer, batch_size=None, *args, **kwargs):
        super().__init__(nef, *args, **kwargs)
        self.nef = nef.eval()
        self.tracer = tracer
        self.layers_painter = self.create_layers_painter(self.nef)

        self.batch_size = batch_size
        self._data_layers = self.regenerate_data_layers()

        # Per frame cached info
        self.render_res_x = None
        self.render_res_y = None
        self.output_width = None
        self.output_height = None
        self.far_clipping = None
        self.channels = None

    @classmethod
    def from_pipeline(cls, pipeline: Pipeline, **kwargs):
        """ Builds a bottom level renderer from the building block of a pipeline. """
        # Create a copy of the pipeline's tracer, with the same type and values
        # This allows the renderer to modify the interactive tracer without affecting the pipeline
        tracer = copy.deepcopy(pipeline.tracer)

        # Pass the kwargs to the renderer also, in case it wants to further modify the tracer with them.
        return cls(nef=pipeline.nef, tracer=tracer, **kwargs)

    @classmethod
    def create_layers_painter(cls, nef: BaseNeuralField) -> Optional[Datalayers]:
        """ NeuralRadianceFieldPackedRenderer can draw datalayers showing the occupancy status.
        These depend on the bottom level acceleration structure.
        """
        if not hasattr(nef.grid, 'blas'):
            return None
        elif isinstance(nef.grid.blas, AxisAlignedBBoxAS):
            return AABBDatalayers()
        elif isinstance(nef.grid.blas, OctreeAS):
            return OctreeDatalayers()
        else:
            return None

    def needs_refresh(self, payload: FramePayload, *args, **kwargs) -> bool:
        """ Should the neural field be "refreshed" (meaning: retraced to a new RenderBuffer).
            Retracing here generally means invoking calls that do not upload new data to the GPU.
            The general logic cannot assume anything about the state of the neural field and therefore opts
            to always refresh it.
            WARNING: This is suboptimal for neural fields which take time to retrace.
            Users are recommended to subclass and override this function by keeping track of the last state
            and what fields changed since the last invocation.
        """
        return True

    def needs_redraw(self) -> bool:
        """ Returns if a full redrawing of the object should occur.
            This includes regeneration of buffers that may get uploaded to the GPU, such as vectorial data layers.
        """
        if self.layers_painter is not None:
            return self.layers_painter.needs_redraw(self.nef.grid.blas)
        else:
            return True

    def regenerate_data_layers(self) -> Dict[str, PrimitivesPack]:
        """ Regenarates the vectorial data layers depicting additional information about the object. """
        if self.layers_painter is not None:
            return self.layers_painter.regenerate_data_layers(self.nef.grid.blas)
        else:
            return dict()

    def pre_render(self, payload: FramePayload, *args, **kwargs) -> None:
        """ Actions which should take place before rendering:
        Store information from payload, adjust resolution, and so forth.
        """
        super().pre_render(payload)
        self.render_res_x = payload.render_res_x
        self.render_res_y = payload.render_res_y
        self.output_width = payload.camera.width
        self.output_height = payload.camera.height
        self.far_clipping = payload.camera.far
        self.channels = payload.channels

    def render(self, rays: Rays) -> RenderBuffer:
        rb = RenderBuffer(hit=None)

        # Feed as single batch
        if self.batch_size is None:
            rb = self.tracer(self.nef, rays=rays, channels=self.channels)
        else:   # Apply in batches
            for ray_batch in rays.split(self.batch_size):
                rb += self.tracer(self.nef, rays=ray_batch, channels=self.channels)

        # Rescale renderbuffer to original size
        rb = rb.reshape(self.render_res_y, self.render_res_x, -1)
        if self.render_res_x != self.output_width or self.render_res_y != self.output_height:
            rb = rb.scale(size=(self.output_height, self.output_width))
        return rb

    def post_render(self) -> None:
        """ Actions which should take place after rendering:
        Store information for next frame.
        """
        pass

    def acceleration_structure(self) -> str:
        """ Returns a human readable name of the bottom level acceleration structure used by this renderer """
        if getattr(self.nef, 'grid') is None or getattr(self.nef.grid, 'blas') is None:
            return "None"
        elif hasattr(self.nef.grid.blas, 'name'):
            return self.nef.grid.blas.name()
        else:
            return "Unknown"

    def features_structure(self) -> str:
        """ Returns a human readable name of the feature structure used by this renderer """
        if getattr(self.nef, 'grid') is None:
            return "None"
        elif hasattr(self.nef.grid, 'name'):
            return self.nef.grid.name()
        else:
            return "Unknown"

    @property
    def dtype(self) -> torch.dtype:
        return torch.float32

    @property
    def device(self) -> torch.device:
        return self.nef.device

wisp/renderer/core/api/raytraced_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from collections import defaultdict, deque
from typing import Type, TYPE_CHECKING
from wisp.models import Pipeline
from wisp.models.nefs import BaseNeuralField
from wisp.tracers import BaseTracer
if TYPE_CHECKING:  # Prevent circular imports mess due to typed annotations
    from wisp.renderer.core.api.base_renderer import BottomLevelRenderer
    from wisp.renderer.core.api.raytraced_renderer import RayTracedRenderer

# All renderers supported by the interactive renderer should be registered here
# Registration can be quickly achieved via the @field_renderer decorator.
# Make sure custom modules are imported before neural renderers are constructed!
_REGISTERED_RENDERABLE_NEURAL_FIELDS = defaultdict(dict)


def register_neural_field_type(neural_field_type: Type[BaseNeuralField],
                               tracer_type: Type[BaseTracer],
                               renderer_type: Type[BottomLevelRenderer]):
    """ Register new types of neural fields with their associated bottom level renderers using this function.
        This allowes the interactive renderer to display this neural field type on the canvas.
    """
    field_name = neural_field_type.__name__
    tracer_name = tracer_type.__name__
    _REGISTERED_RENDERABLE_NEURAL_FIELDS[field_name][tracer_name] = renderer_type


def _neural_field_to_renderer_cls(pipeline: Pipeline) -> Type[RayTracedRenderer]:
    tracer_type = type(pipeline.tracer)

    # Start by iterating the current tracer type - look for renderers compatible with the current nef type
    # or any of its parents (the hierarchy of nefs take precedence over the hierarchy of tracers).
    renderer_cls = None
    while tracer_type:
        tracer_name = tracer_type.__name__

        # Look for a renderer compatible with the current tracer type and nef classes
        type_queue = deque([type(pipeline.nef)])

        renderer_cls = None
        while type_queue:
            # Query nef + tracer combo
            field_type = type_queue.popleft()
            field_name = field_type.__name__
            supported_tracers = _REGISTERED_RENDERABLE_NEURAL_FIELDS.get(field_name)
            renderer_cls = supported_tracers.get(tracer_name) if supported_tracers is not None else None

            # Current nef + tracer pair doesn't match any registered renderer
            if renderer_cls is not None:
                break
            else:   # Try querying all parent(nef) + tracer combos for compatibility
                bases = field_type.__bases__
                if len(bases) > 0:
                    type_queue.append(*bases)

        if renderer_cls is not None:
            break   # Found a renderer class
        else:
            # Didn't find a renderer class - repeat the process with the parent class of the tracer
            tracer_base_types = tracer_type.__bases__
            if len(tracer_base_types) > 0:
                # Does tracer have a single parent class which inherits from BaseTracer?
                # If so, keep looking
                tracer_base_types = [base for base in tracer_type.__bases__ if issubclass(base, BaseTracer)]
                tracer_type = tracer_base_types[0] if len(tracer_base_types) == 1 else None
            else:
                # Reached end of tracers hierarchy or it is too ambiguous, quit and fail gracefully
                tracer_type = None

    if tracer_type is None:
        raise ValueError(f'Renderer factory encountered an unknown neural pipeline: '
                         f'Neural Field {type(pipeline.nef).__name__} with tracer {type(pipeline.tracer).__name__}. '
                         'Please register the factory to reflect what kind of renderer should be created for this '
                         'type of neural field/tracer.')
    return renderer_cls


def create_neural_field_renderer(neural_object: Pipeline, **kwargs) -> RayTracedRenderer:
    # TODO (operel): support creation of non ray traced renderers

    # Fetch the neural field object. If given a pipeline, query the nef field.
    # Otherwise assume we're given the neural field explicitly.
    if isinstance(neural_object, Pipeline):
        # Obtain the appropriate renderer class, compatible with this neural field type and tracer
        renderer_cls = _neural_field_to_renderer_cls(neural_object)

        # In the case of a pipeline, we use a specialized constructor which builds the renderer from its components
        # kwargs may override the pipeline default settings
        renderer_instance = renderer_cls.from_pipeline(neural_object, **kwargs)
    # else:
        # TODO (operel): not sure we even need this path..
        #   Would a NeF ever exist without a tracer? If needed we can just choose any tracer combo
        # # Obtain the appropriate renderer class, compatible with this neural field type
        # nef = neural_object
        # renderer_cls = _neural_field_to_renderer_cls(nef)
        # renderer_args = kwargs
        #
        # # Create an instance of a renderer, wrapping the neural field and using the given arguments
        # renderer_instance = renderer_cls(nef, **renderer_args)

    return renderer_instance


# TODO (operel): Double check the interactive app can loads a model from file and call create_renderer
# Model loading should now be at the app's responsibility
# # Load nerf if model_path is available and nerf is not
# if nerf is None:
#     model_path = kwargs['model_path']
#     nerf = torch.load(model_path)

wisp/renderer/core/api/renderers_factory.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import torch
from wisp.core import ObjectTransform
from wisp.models import Pipeline
from wisp.framework import WispState, BottomLevelRendererState


def request_redraw(state):
    """ Marks the canvas as dirty,
    forcing the renderer core to refresh the object renderers on the next rendering iteration.
    """
    state.renderer.canvas_dirty = True


def add_pipeline_to_scene_graph(state: WispState,
                                name: str,
                                pipeline: Pipeline,
                                transform: ObjectTransform = None,
                                **setup_args):
    """ Adds a new object pipeline to the scene graph.
    The pipeline contains all components required to trace this object.
    The object is represented in the SceneGraph by a newly created BottomLevelRenderer.

    Args:
        state (WispState): A wisp state object, containing the scene graph information.
        name (str): Unique name of object added to the scene graph
        pipeline (Pipeline): A pipeline consisting of a field (and possibly a tracer) of the added object.
        transform (ObjectTransform): The object transform, containing a 4x4 transformation matrix which transforms
        the object from local object space to world space.
        setup_args (Dict): Optional setup args which control how the BottomLevelRenderer will be created.
    """
    if transform is None:
        transform = ObjectTransform(device=pipeline.nef.device)
    state.graph.neural_pipelines[name] = pipeline
    state.graph.bl_renderers[name] = BottomLevelRendererState(status='pending', transform=transform,
                                                              setup_args=setup_args)
    request_redraw(state)   # Let renderer core know it should refresh next frame


def add_to_scene_graph(state: WispState,
                       name: str,
                       obj,
                       transform: ObjectTransform = None,
                       **setup_args):
    """ Adds a new object to the scene graph.
    obj can be any supported object type, neural or non-neural.
    This is the most general function used to manage adding new objects to the scene graph.

    Args:
        state (WispState): A wisp state object, containing the scene graph information.
        name (str): Unique name of object added to the scene graph
        obj (object): Any object supported by the scene-graph.
        i.e: for neural fields, obj is a Pipeline.
        transform (ObjectTransform): The object transform, containing a 4x4 transformation matrix which transforms
        the object from local object space to world space.
        setup_args (Dict): Optional setup args which control how the BottomLevelRenderer will be created.
    """
    if isinstance(obj, Pipeline):
        add_pipeline_to_scene_graph(state, name, obj, transform, **setup_args)
    else:   # TODO (operel): Currently only neural pipelines are supported
        raise NotImplementedError(f'Unsupported object type added to scene graph: {obj}')


def remove_from_scene_graph(state: WispState, name: str):
    """ Removes an existing pipeline from the scene graph.

        Args:
            state (WispState): A wisp state object, containing the scene graph information.
            name (str): Unique name of object added to the scene graph
    """
    assert name in state.graph.neural_pipelines, f'Scene graph requested to remove non-existing object: {name}'
    del state.graph.neural_pipelines[name]
    del state.graph.bl_renderers[name]
    request_redraw(state)  # Let renderer core know it should refresh next frame

wisp/renderer/core/api/scenegraph.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from .io import *
from .camera_controller_mode import *
from .first_person import *
from .trackball import *
from .turntable import *

wisp/renderer/core/control/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from abc import ABC, abstractmethod
import numpy as np
import torch
import wisp.framework.state as state
from wisp.renderer.core.control.io import WispKey


class CameraControlMode(ABC):
    def __init__(self, render_core, wisp_state: state.WispState):
        self.render_core = render_core
        self.state = wisp_state

        self.PI = torch.tensor(np.pi, device=render_core.device)
        self.planes_forbidden_zooming_through = []

        # ----------------------------------
        # Sensitivity & velocity parameters
        # ----------------------------------
        # The controller uses kinematics for smooth FPS considerate transitions
        # Pan operations with the keyboard:
        self._key_pan_initial_velocity = 0.8           # start with this v0 velocity
        self._key_pan_deacceleration = 3.2             # deaccelerate with this acceleration factor a ~ 4*v0
        self._key_pan_distance_weight = 0.5            # pan amount is distance dependant, weighted by this factor

        # Pan operations with the mouse:
        self._mouse_pan_distance_weight = 0.002        # pan amount is distance dependant, weighted by this factor

        # Zoom operations with the mouse:
        self._mouse_zoom_initial_velocity = 10.0       # start with this v0 velocity
        self._mouse_zoom_deacceleration = 40.0         # deaccelerate with this acceleration factor a ~ 4*v0
        self._zoom_persp_distance_weight = 0.25        # zoom is cam-pos dependant, weighted by this factor
        self._zoom_ortho_distance_weight = 0.2         # zoom is fov-distance dependant, weighted by this factor
        self._zoom_ortho_fov_dist_range = (1e-4, 1e2)  # zoom in ortho mode is limited to this fov-distance range

        # Current state
        self._current_pan_velocity = 0
        self._current_pan_direction = None
        self._remaining_pan_time = 0

        self.interactions_stack = list()

    @classmethod
    @abstractmethod
    def name(cls):
        raise NotImplementedError("Camera modes should state a descriptive name")

    def handle_timer_tick(self, dt):
        self.progress_pan(dt)

    def handle_mouse_scroll(self, x, y, dx, dy):
        """ The mouse wheel was scrolled by (dx,dy). """
        self.stop_all_current_interactions()
        if dy < 0:
            self.start_pan(pan_direction='forward',
                           velocity=self._mouse_zoom_initial_velocity, deaccelaration=self._mouse_zoom_deacceleration)
        else:
            self.start_pan(pan_direction='backward',
                           velocity=self._mouse_zoom_initial_velocity, deaccelaration=self._mouse_zoom_deacceleration)

    def handle_key_press(self, symbol, modifiers):
        self.stop_all_current_interactions()
        if symbol == WispKey.LEFT:
            self.start_pan(pan_direction='left',
                           velocity=self._key_pan_initial_velocity, deaccelaration=self._key_pan_deacceleration)
        elif symbol == WispKey.RIGHT:
            self.start_pan(pan_direction='right',
                           velocity = self._key_pan_initial_velocity, deaccelaration=self._key_pan_deacceleration)
        elif symbol == WispKey.UP:
            self.start_pan(pan_direction='up',
                           velocity=self._key_pan_initial_velocity, deaccelaration=self._key_pan_deacceleration)
        elif symbol == WispKey.DOWN:
            self.start_pan(pan_direction='down',
                           velocity=self._key_pan_initial_velocity, deaccelaration=self._key_pan_deacceleration)

    def start_pan(self, pan_direction, velocity=None, deaccelaration=None):
        self.start_interaction(f'pan_{pan_direction}')
        self._current_pan_velocity = velocity
        self._current_pan_deacceleration = deaccelaration
        self._current_pan_direction = pan_direction
        self._remaining_pan_time = abs(self._current_pan_velocity / self._current_pan_deacceleration)
        if pan_direction in ('left', 'down', 'backward'):  # Directions that move opposite to camera axes
            self._current_pan_velocity *= -1

    def zoom(self, amount):
        if self.camera.lens_type == 'ortho':
            # Under orthographic projection, objects are not affected by distance to the camera
            amount = self._zoom_ortho_distance_weight * self.camera.fov_distance * abs(amount) * np.sign(amount)
            self.camera.zoom(amount)
            # Keep distance at reasonable range
            self.camera.fov_distance = torch.clamp(self.camera.fov_distance, *self._zoom_ortho_fov_dist_range)
        else:
            dist = self.camera.cam_pos().norm()
            amount *= self._zoom_persp_distance_weight * dist
            self.camera.move_forward(amount)

    def progress_pan(self, dt):
        if self._current_pan_direction is None or self._current_pan_velocity == 0:
            return
        dt = min(self._remaining_pan_time, dt)
        pos_delta = dt * self._current_pan_velocity
        if self._current_pan_direction in ('forward', 'backward'):
            cam_pos, cam_forward = self.camera.cam_pos().squeeze(), self.camera.cam_forward().squeeze()
            new_pos = (cam_pos + cam_forward * pos_delta)
            if ('xz' in self.planes_forbidden_zooming_through and new_pos[1].sign() * cam_pos[1].sign() == -1) or \
                ('xy' in self.planes_forbidden_zooming_through and new_pos[2].sign() * cam_pos[2].sign() == -1) or \
                ('yz' in self.planes_forbidden_zooming_through and new_pos[0].sign() * cam_pos[0].sign() == -1):
                self._remaining_pan_time = 0
            else:
                self.zoom(pos_delta)
        elif self._current_pan_direction in ('right', 'left'):
            dist = self.camera.cam_pos().norm()
            pos_delta *= self._key_pan_distance_weight * dist
            self.camera.move_right(pos_delta)
        elif self._current_pan_direction in ('up', 'down'):
            dist = self.camera.cam_pos().norm()
            pos_delta *= self._key_pan_distance_weight * dist
            self.camera.move_up(pos_delta)
        velocity_sign = np.sign(self._current_pan_velocity)
        deaccel_amount = velocity_sign * self._current_pan_deacceleration * dt
        self._current_pan_velocity -= deaccel_amount
        self._remaining_pan_time = max(0, self._remaining_pan_time - dt)
        if np.sign(self._current_pan_velocity) != velocity_sign or \
                self._current_pan_velocity == 0 or self._remaining_pan_time == 0:
            self.end_pan()

    def end_pan(self):
        self.end_interaction()
        if not self.is_interacting():   # End only if other interactions have not taken place meanwhile
            self._current_pan_velocity = 0
            self._current_pan_direction = None
            self._remaining_pan_time = 0

    def handle_key_release(self, symbol, modifiers):
        pass

    def handle_mouse_press(self, x, y, button):
        self.start_interaction('pan_withmouse')

    def handle_mouse_drag(self, x, y, dx, dy, button):
        dist_normalization = self._mouse_pan_distance_weight * self.camera.cam_pos().norm()
        self.camera.move_right(dist_normalization * -dx)
        self.camera.move_up(dist_normalization * dy)

    def handle_mouse_release(self, x, y, button):
        self.end_pan()

    def handle_mouse_motion(self, x, y, dx, dy):
        """ The mouse was moved with no buttons held down. """
        pass

    def start_interaction(self, interaction_id):
        self.interactions_stack.append(interaction_id)

    def end_interaction(self):
        # On some occassions the app may be out of focus and some interactions are not registered properly.
        # Silently ignore that.
        if len(self.interactions_stack) > 0:
            self.interactions_stack.pop()

    def stop_all_current_interactions(self):
        while self.is_interacting():
            last_interaction_started = self.get_last_interaction_started()
            if last_interaction_started.startswith('pan'):
                self.end_pan()
            else:
                self.end_interaction()

    def is_interacting(self):
        return len(self.interactions_stack) > 0

    def get_last_interaction_started(self):
        return self.interactions_stack[-1] if self.is_interacting() else None

    def has_interaction(self, interaction_id):
        return interaction_id in self.interactions_stack

    @property
    def camera(self):
        return self.render_core.camera

wisp/renderer/core/control/camera_controller_mode.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import numpy as np
import wisp.framework.state as state
from wisp.renderer.core.control.camera_controller_mode import CameraControlMode
from wisp.renderer.core.control.io import WispMouseButton


class FirstPersonCameraMode(CameraControlMode):

    def __init__(self, render_core, wisp_state: state.WispState):
        super().__init__(render_core, wisp_state)
        self._rotation_sensitivity = 0.2

    @classmethod
    def name(cls):
        return "First Person View"

    def handle_mouse_drag(self, x, y, dx, dy, button):
        if self.has_interaction('fps_rotate'):
            camera = self.render_core.camera
            w, h = camera.width, camera.height
            max_rot = 2 * self.PI * self._rotation_sensitivity
            full_dist = np.minimum(h, w)
            yaw = -max_rot * dx / full_dist  # Rotates camera in right-forward plane
            pitch = -max_rot * dy / full_dist  # Rotates camera in up-forward plane
            self.render_core.camera.rotate(yaw, pitch)
        else:
            super().handle_mouse_drag(x, y, dx, dy, button)

    def handle_mouse_press(self, x, y, button):
        self.stop_all_current_interactions()
        if button == WispMouseButton.LEFT_BUTTON:
            self.start_interaction('fps_rotate')
        elif button == WispMouseButton.MIDDLE_BUTTON:
            super().handle_mouse_press(x, y, button)

    def handle_mouse_release(self, x, y, button):
        if self.has_interaction('fps_rotate'):
            self.end_interaction()
        else:
            super().handle_mouse_release(x, y, button)

wisp/renderer/core/control/first_person.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from enum import Enum, auto

_registered_mouse_buttons = dict()
_registered_keys = dict()


class WispMouseButton(Enum):
    LEFT_BUTTON = auto()
    MIDDLE_BUTTON = auto()
    RIGHT_BUTTON = auto()

    @classmethod
    def register_symbol(cls, wisp_mb: WispMouseButton, symbol):
        _registered_mouse_buttons[wisp_mb] = symbol

    def __eq__(self, other):
        """ Perform equality test with the mapped symbol """
        mapped_symbol = _registered_mouse_buttons.get(self, None)
        assert mapped_symbol is not None, f"Error: WispMouseButton {self.name} have not been registered by the app."
        return mapped_symbol == other

    def __hash__(self):
        return hash(self.value)


class WispKey(Enum):
    LEFT = auto()
    RIGHT = auto()
    UP = auto()
    DOWN = auto()
    a = auto()
    b = auto()
    c = auto()
    d = auto()
    e = auto()
    f = auto()
    g = auto()
    h = auto()
    i = auto()
    j = auto()
    k = auto()
    l = auto()
    m = auto()
    n = auto()
    o = auto()
    p = auto()
    q = auto()
    r = auto()
    s = auto()
    t = auto()
    u = auto()
    v = auto()
    w = auto()
    x = auto()
    y = auto()
    z = auto()
    A = auto()
    B = auto()
    C = auto()
    D = auto()
    E = auto()
    F = auto()
    G = auto()
    H = auto()
    I = auto()
    J = auto()
    K = auto()
    L = auto()
    M = auto()
    N = auto()
    O = auto()
    P = auto()
    Q = auto()
    R = auto()
    S = auto()
    T = auto()
    U = auto()
    V = auto()
    W = auto()
    X = auto()
    Y = auto()
    Z = auto()
    ONE = auto()
    TWO = auto()
    THREE = auto()
    FOUR = auto()
    FIVE = auto()
    SIX = auto()
    SEVEN = auto()
    EIGHT = auto()
    NINE = auto()
    ZERO = auto()
    PLUS = auto()
    MINUS = auto()
    LEFT_BRACKET = auto()
    RIGHT_BRACKET = auto()
    LEFT_SQUARE_BRACKET = auto()
    RIGHT_SQUARE_BRACKET = auto()
    LEFT_CURLY_BRACKET = auto()
    RIGHT_CURLY_BRACKET = auto()
    DOT = auto()
    COMMA = auto()

    @classmethod
    def register_symbol(cls, wisp_key: WispKey, symbol):
        _registered_keys[wisp_key] = symbol

    def __eq__(self, other):
        mapped_symbol = _registered_keys.get(self, None)
        assert mapped_symbol is not None, f"Error: WispKey {self.name} have not been registered by the app."
        return mapped_symbol == other

    def __hash__(self):
        return hash(self.value)

wisp/renderer/core/control/io.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import copy
import torch
import wisp.framework.state as state
from wisp.renderer.core.control.camera_controller_mode import CameraControlMode
from wisp.renderer.core.control.io import WispMouseButton

# A trackball camera, allowing free orbit movement in all axes.
# Also known as Arcball, or "Orbit" camera
# Original algorithm from:
# * ARCBALL: A User Interface for Specifying Three-Dimensional Orientation Using a Mouse,
#   Shoemake, K, Proceedings of the conference on Graphics interface 1992
#
# * Animating rotation with quaternion curves, Shoemake, K, SIGGRAPH 1985

def quat_mul(Q1, Q2):
    return torch.tensor([Q1[0] * Q2[3] + Q1[3] * Q2[0] - Q1[2] * Q2[1] + Q1[1] * Q2[2],
                         Q1[1] * Q2[3] + Q1[2] * Q2[0] + Q1[3] * Q2[1] - Q1[0] * Q2[2],
                         Q1[2] * Q2[3] - Q1[1] * Q2[0] + Q1[0] * Q2[1] + Q1[3] * Q2[2],
                         Q1[3] * Q2[3] - Q1[0] * Q2[0] - Q1[1] * Q2[1] - Q1[2] * Q2[2]])


def quat_matrix(q): # True only for unit quaternions
    xx = q[0] * q[0]
    xy = q[0] * q[1]
    xz = q[0] * q[2]
    xw = q[0] * q[3]
    yy = q[1] * q[1]
    yz = q[1] * q[2]
    yw = q[1] * q[3]
    zz = q[2] * q[2]
    zw = q[2] * q[3]
    ww = q[3] * q[3]
    return torch.tensor([[ww + xx - yy - zz, 2.0 * (xy - zw), 2.0 * (xz + yw), 0.0],
                         [2.0 * (xy + zw), ww - xx + yy - zz, 2.0 * (yz - xw), 0.0],
                         [2.0 * (xz - yw), 2.0 * (yz + xw), ww - xx - yy + zz, 0.0],
                         [0.0, 0.0, 0.0, 1.0]], dtype=torch.float64)


class TrackballCameraMode(CameraControlMode):

    def __init__(self, render_core, wisp_state: state.WispState):
        super().__init__(render_core, wisp_state)
        self.radius = 1.0
        self.tb_scale = 1.1
        self.sensitivity = 1.0
        self.reset_center_of_focus(reset_radius=True)

        self.q = torch.tensor([0.0, 0.0, 0.0, 1.0])
        self.q0 = torch.tensor([0.0, 0.0, 0.0, 1.0])
        self.q1 = None
        self.v0 = None
        self.v1 = None

        self.initial_camera = None
        self.planes_forbidden_zooming_through = self.state.renderer.reference_grids

    @classmethod
    def name(cls):
        return "Trackball"
 
    def reset_center_of_focus(self, reset_radius=False):
        pos = self.camera.cam_pos().squeeze()
        forward_axis = self.camera.cam_forward().squeeze()
        forward_axis /= forward_axis.norm()
        if reset_radius:
            self.radius = torch.dot(pos, forward_axis)
        self.focus_at = pos - self.radius * forward_axis   

    def mouse2vector(self, mx, my):
        """ Converts mouse position in screen coordinates:
            [0,0] at top left of screen to [width, height] at bottom-right
            to coordinates on arcball
        """
        # First we convert the mouse coordinates to normalized device coordinates, multiplied by the arcball radius.
        # v is a vector of (x,y,z), and will contain coordinates projected to the arcball surface
        # TODO: Is x and y rotated inversely?
        half_width = 0.5 * self.camera.width
        half_height = 0.5 * self.camera.height
        v = torch.tensor([half_width - float(mx), float(my) - half_height, 0.0])
        normalization_factor = half_height if half_width >= half_height else half_width
        v *= self.tb_scale / float(normalization_factor)

        # v is now in Normalized Device Coordinates
        # Next we need to calculate the z coordinate, which is currently set to 0

        # xy_power = x^2 + y^2
        xy_power = torch.pow(v, 2).sum()

        if xy_power < 1.0:
            v[2] = -torch.sqrt(1.0 - xy_power)   # z = -sqrt(R - x^2 - y^2)
        else:
            v /= torch.sqrt(xy_power)   # x = x/sqrt(x^2 + y^2) ; y = y/sqrt(x^2 + y^2) ; z = 0.0

        return v

    def handle_mouse_press(self, x, y, button):
        self.stop_all_current_interactions()
        if button == WispMouseButton.LEFT_BUTTON:
            self.start_interaction('trackball_rotate')
            self.initial_camera = copy.deepcopy(self.camera)
            self.v0 = self.mouse2vector(x, y)
        elif button == WispMouseButton.MIDDLE_BUTTON:
            super().handle_mouse_press(x, y, button)

    def handle_mouse_release(self, x, y, button):
        if self.has_interaction('trackball_rotate'):
            self.end_interaction()
            self.q0 = torch.tensor([0.0, 0.0, 0.0, 1.0])
            self.initial_camera = None
        else:
            super().handle_mouse_release(x, y, button)

    def handle_mouse_drag(self, x, y, dx, dy, button):
        if not self.has_interaction('trackball_rotate'):
            super().handle_mouse_drag(x, y, dx, dy, button)
        else:
            camera = self.camera

            if torch.allclose(self.v0, torch.zeros_like(self.v0)):
                self.v0 = self.mouse2vector(x, y)   # Start position

            self.v1 = self.mouse2vector(x, y)       # Current position

            # find quaterion for previous to current frame rotation
            axis = torch.cross(self.v1, self.v0)
            angle = -torch.dot(self.v1, self.v0).unsqueeze(0)
            angle *= self.sensitivity
            self.q1 = torch.cat((axis, angle))

            # Apply rotation to starting rotation quaternion
            self.q = quat_mul(self.q1, self.q0)

            # To rotation matrix
            Q = quat_matrix(self.q)
            Q = Q.to(camera.device)

            # find length of vector from eye to focus at point
            pos = self.initial_camera.cam_pos().reshape(-1)
            vec = self.focus_at - pos
            length = torch.sqrt(torch.dot(vec, vec))

            # create translation in z direction to/from focua at point (in camera space)
            T = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
            T[2, 3] = length
            Tinv = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
            Tinv[2, 3] = -length

            view_matrix = self.initial_camera.view_matrix()[0]   # Unbatch

            # apply transforms
            view_matrix = Tinv @ Q @ T @ view_matrix

            # update extrinsic matrix
            camera.update(view_matrix)


wisp/renderer/core/control/trackball.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import numpy as np
import torch
import wisp.framework.state as state
from wisp.renderer.core.control.camera_controller_mode import CameraControlMode
from wisp.renderer.core.control.io import WispMouseButton


# A turntable camera with 2 degrees of freedom.
# This camera controller is more restrictive, but maintains the "horizon" orientation to
# allow for a more intuitive interface.

class TurntableCameraMode(CameraControlMode):

    def __init__(self, render_core, wisp_state: state.WispState):
        super().__init__(render_core, wisp_state)
        self.sensitivity = 0.95
        self.reset_center_of_focus(reset_radius=True)

        reference_grids = self.state.renderer.reference_grids
        self.planes_forbidden_zooming_through = reference_grids
        self.reference_plane = reference_grids[0] if len(reference_grids) > 0 else "xy"

    @classmethod
    def name(cls):
        return "Turntable"

    def reset_center_of_focus(self, reset_radius=False):
        pos = self.camera.cam_pos().squeeze()
        forward_axis = self.camera.cam_forward().squeeze()
        forward_axis /= forward_axis.norm()
        if reset_radius:
            self.radius = torch.dot(pos, forward_axis)
        self.focus_at = pos - self.radius * forward_axis

    def yaw_camera_in_world_space(self, camera, yaw):
        translate = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        translate[:3, 3] = self.focus_at
        retranslate = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        translate[:3, 3] = -self.focus_at

        rot_yaw = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        rot_yaw[0, 0] = torch.cos(yaw)
        rot_yaw[0, 2] = -torch.sin(yaw)
        rot_yaw[2, 0] = torch.sin(yaw)
        rot_yaw[2, 2] = torch.cos(yaw)
        view_matrix = camera.view_matrix()[0]   # Unbatch
        # Translate world center to focus point, rotate, and translate back
        view_matrix = view_matrix @ retranslate @ rot_yaw @ translate
        view_matrix = view_matrix.unsqueeze(0)
        camera.update(view_matrix)

    def pitch_camera_in_world_space(self, camera, pitch):
        translate = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        translate[:3, 3] = self.focus_at
        retranslate = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        translate[:3, 3] = -self.focus_at

        rot_pitch = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        rot_pitch[1, 1] = torch.cos(pitch)
        rot_pitch[1, 2] = torch.sin(pitch)
        rot_pitch[2, 1] = -torch.sin(pitch)
        rot_pitch[2, 2] = torch.cos(pitch)
        view_matrix = camera.view_matrix()[0]   # Unbatch
        # Translate world center to focus point, rotate, and translate back
        view_matrix = view_matrix @ retranslate @ rot_pitch @ translate
        view_matrix = view_matrix.unsqueeze(0)
        camera.update(view_matrix)

    def roll_camera_in_world_space(self, camera, roll):
        translate = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        translate[:3, 3] = self.focus_at
        retranslate = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        translate[:3, 3] = -self.focus_at

        rot_pitch = torch.eye(4, device=camera.device, dtype=camera.dtype).repeat(1, 1)
        rot_pitch[0, 0] = torch.cos(roll)
        rot_pitch[0, 1] = -torch.sin(roll)
        rot_pitch[1, 0] = torch.sin(roll)
        rot_pitch[1, 1] = torch.cos(roll)
        view_matrix = camera.view_matrix()[0]   # Unbatch
        # Translate world center to focus point, rotate, and translate back
        view_matrix = view_matrix @ retranslate @ rot_pitch @ translate
        view_matrix = view_matrix.unsqueeze(0)
        camera.update(view_matrix)

    def handle_mouse_drag(self, x, y, dx, dy, button):
        if self.has_interaction('turntable_rotate'):
            camera = self.camera
            w, h = camera.width, camera.height
            max_rot = 2 * self.PI * self.sensitivity
            full_dist = np.minimum(h, w)

            in_plane_amount = -max_rot * dx / full_dist  # Rotates camera in normal direction to plane (world space)
            pitch = -max_rot * dy / full_dist            # Rotates camera in up-forward direction (camera space)
            camera.t = torch.zeros(3, dtype=camera.dtype, device=camera.device)
            camera.rotate(pitch=pitch)

            if self.reference_plane == 'xz':
                self.yaw_camera_in_world_space(camera, in_plane_amount)
            elif self.reference_plane == 'xy':
                self.roll_camera_in_world_space(camera, in_plane_amount)
            elif self.reference_plane == 'yz':
                self.pitch_camera_in_world_space(camera, in_plane_amount)

            backward_dir = camera.cam_forward().squeeze() / camera.cam_forward().norm()
            camera.translate(self.radius * backward_dir)
        else:
            super().handle_mouse_drag(x, y, dx, dy, button)

    def handle_mouse_press(self, x, y, button):
        self.stop_all_current_interactions()
        if button == WispMouseButton.LEFT_BUTTON:
            self.start_interaction('turntable_rotate')
        elif button == WispMouseButton.MIDDLE_BUTTON:
            super().handle_mouse_press(x, y, button)

    def handle_mouse_release(self, x, y, button):
        if self.has_interaction('turntable_rotate'):
            self.end_interaction()
        else:
            super().handle_mouse_release(x, y, button)

    def end_pan(self):
        last_interaction = self.get_last_interaction_started()
        if last_interaction is not None:
            self.reset_center_of_focus(reset_radius=True)
        else:
            self.reset_center_of_focus(reset_radius=False)
        super().end_pan()

wisp/renderer/core/control/turntable.py
Sharath-girish/Shacira
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import os
from contextlib import contextmanager


if not os.environ.get('ENABLE_PYCUDA') == '1':
    from cuda import cuda
    import torch

    @contextmanager
    def cuda_map_resource(img):
        """Context manager simplifying use of cuda.cuGraphicsMapResources / cuGraphicsSubResourceGetMappedArray.
        Boilerplate code based in part on pytorch-glumpy.
        """
        # args: (count, resource, stream)
        mapping_result = cuda.cuGraphicsMapResources(1, img, torch.cuda.default_stream().cuda_stream)
        if mapping_result[0] != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError("Failed to map GL graphics resource to be accessed by CUDA.")

        # args (resource, arrayIndex, mipLevel)
        mapping_array = cuda.cuGraphicsSubResourceGetMappedArray(img, 0, 0)
        if mapping_array[0] != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError("Failed to get mapped array from GL resource for CUDA copy.")
        yield mapping_array[1]

        unmapping_result = cuda.cuGraphicsUnmapResources(1, img, torch.cuda.default_stream().cuda_stream)
        if unmapping_result[0] != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError("Failed to unmap GL graphics resource from being accessed by CUDA.")


    def cuda_2d_memcpy(resource_handle, shared_tex, img, height):
        cpy = cuda.CUDA_MEMCPY2D()
        with cuda_map_resource(resource_handle) as ary:
            cpy.srcDevice = img.data_ptr()
            cpy.srcMemoryType = cuda.CUmemorytype.CU_MEMORYTYPE_DEVICE
            cpy.dstArray = ary
            cpy.dstMemoryType = cuda.CUmemorytype.CU_MEMORYTYPE_ARRAY
            cpy.WidthInBytes = cpy.srcPitch = cpy.dstPitch = shared_tex.nbytes // height
            cpy.Height = height
            cpy_result = cuda.cuMemcpy2DUnaligned(cpy)
            if cpy_result[0] != cuda.CUresult.CUDA_SUCCESS:
                raise RuntimeError("Failed to memcopy cuda buffer to GL.")


    def cuda_register_gl_image(image, target):
        # Create shared GL / CUDA resource
        map_flags = cuda.CUgraphicsRegisterFlags.CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD
        register_result = cuda.cuGraphicsGLRegisterImage(image=image, target=target, Flags=map_flags)
        if register_result[0] != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError('Failed to register GL texture as a CUDA shared resource.')
        resource_handle = register_result[1]
        return resource_handle

    def cuda_unregister_resource(handle):
        unregister_result = cuda.cuGraphicsUnregisterResource(handle)
        if unregister_result[0] != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError('Failed to unregister CUDA resource.')

else:
    import pycuda
    import pycuda.gl as pycuda_gl

    window = None
    try:
        # !!! Should be called when interactive wisp loads, before any torch ops take place !!!
        # The following is a hacky workaround due to a cublas error on interfering streams:
        # pycuda.gl fails to initialize after torch performs batched matrix multiplication
        # (the bugs causes any following torch.dot invocations to fail)..
        # The solution is to initialize pycuda.gl early when wisp loads.
        # To load pycuda.gl, we require a GL context of some window,
        # so here we force glumpy-glfw to create an invisible window, which generates an opengl context.
        # Then immediately import pycuda.gl.autoint to let it initialize properly
        from glumpy import app

        # Tell glumpy to use glfw backend
        app.use("glfw_imgui")
        # Let glumpy use glfw to create an invisible window
        window = app.Window(width=10, height=10, title='dummy', visible=False)

        # pycuda initializes the default context with "cuGLCtxCreate", but this call will fail if a GL context
        # is not currently set. Therefore import is invoked only after glfw obtains a GL context.
        # See: https://documen.tician.de/pycuda/gl.html#module-pycuda.gl.autoinit
        import pycuda.gl.autoinit
        # Next tell torch to initialize the primary cuda context
        import torch
        torch.cuda.init()
        # pycuda should not create a new context, but retain the torch one
        import pycuda.driver as cuda
        pycuda_context = cuda.Device(0).retain_primary_context()

    except (ModuleNotFoundError, ImportError):
        pass  # Don't fail if interactive mode is disabled (e.g: glumpy or pycuda are unavailable)
    finally:
        if window is not None:
            window.close()

    @contextmanager
    def cuda_map_resource(img):
        """Context manager simplifying use of pycuda_gl.RegisteredImage.
        Boilerplate code based in part on pytorch-glumpy.
        """
        mapping = img.map()
        yield mapping.array(0, 0)
        mapping.unmap()

    def cuda_2d_memcpy(resource_handle, shared_tex, img, height):
        cpy = pycuda.driver.Memcpy2D()
        with cuda_map_resource(resource_handle) as ary:
            cpy.set_src_device(img.data_ptr())
            cpy.set_dst_array(ary)
            cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = shared_tex.nbytes // height
            cpy.height = height
            cpy(aligned=False)

    def cuda_register_gl_image(image, target):
        # Create shared GL / CUDA resource
        map_flags = pycuda_gl.graphics_map_flags.WRITE_DISCARD
        resource_handle = pycuda_gl.RegisteredImage(image, target, map_flags)
        return resource_handle

    def cuda_unregister_resource(handle):
        # Nothing to be done - when ref count reaches zero on proxy object in python, unregister should
        # be called automatically
        pass

wisp/renderer/core/cuda.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
import abc
import numpy as np
import torch
import copy
from collections import defaultdict
from typing import Dict, List, Iterable, Tuple
from kaolin.render.camera import Camera, PinholeIntrinsics, OrthographicIntrinsics
from wisp.framework import WispState, BottomLevelRendererState
from wisp.core import RenderBuffer, Rays, PrimitivesPack, create_default_channel, ObjectTransform
from wisp.ops.raygen import generate_pinhole_rays, generate_ortho_rays, generate_centered_pixel_coords
from wisp.renderer.core.api import BottomLevelRenderer, RayTracedRenderer, FramePayload, create_neural_field_renderer
from wisp.gfx.datalayers import CameraDatalayers


def enable_amp(func):
    """ An extension to @torch.cuda.amp.autocast which queries WispState to check if
    mixed precision should be enabled.
    """
    def _enable_amp(self: RendererCore, *args, **kwargs):
        with torch.cuda.amp.autocast(enabled=self.state.renderer.enable_amp):
            return func(self, *args, **kwargs)
    return _enable_amp


class RendererCore:
    def __init__(self, state: WispState):
        self.state = state
        self.device = state.renderer.device

        # Create a camera for user view
        self.camera = self._setup_camera(state)
        self._camera_layers = CameraDatalayers()

        # Set up the list of available bottom level object renderers, according to scene graph
        self._renderers = None
        self._tlas = None
        self.refresh_bl_renderers(state)

        self.res_x, self.res_y = None, None
        self.set_full_resolution()

        self._last_state = dict()
        self._last_renderbuffer = None

        # Minimal resolution supported by RendererCore
        self.MIN_RES = 128

    def _default_camera(self, lens="perspective"):
        # TODO: move defaults elsewhere
        if lens == 'perspective':
            return Camera.from_args(
                eye=torch.tensor([4.0, 4.0, 4.0]),
                at=torch.tensor([0.0, 0.0, 0.0]),
                up=torch.tensor([0.0, 1.0, 0.0]),
                fov=30 * np.pi / 180,  # In radians
                x0=0.0, y0=0.0,
                width=800, height=800,
                near=1e-2, far=20,
                dtype=torch.float64,
                device=self.device
            )
        elif lens == 'orthographic':
            return Camera.from_args(
                eye=torch.tensor([4.0, 4.0, 4.0]),
                at=torch.tensor([0.0, 0.0, 0.0]),
                up=torch.tensor([0.0, 1.0, 0.0]),
                width=800, height=800,
                near=-800, far=800,
                fov_distance=1.0,
                dtype=torch.float64,
                device=self.device
            )

    def change_camera_projection_mode(self, lens: str) -> None:
        # Need to update only the intrinsic component of the camera
        # TODO (operel): in the future, kaolin should implement a capability to switch lens and approximate params
        # from one type to another where applicable
        if lens == 'Perspective':
            intrinsics = PinholeIntrinsics.from_fov(
                fov=30 * np.pi / 180,       # In radians
                x0=0.0, y0=0.0,
                width=self.camera.width, height=self.camera.height,
                near=1e-2, far=1e2,
                dtype=self.camera.dtype,
                device=self.camera.device
            )
        elif lens == 'Orthographic':
            intrinsics = OrthographicIntrinsics.from_frustum(
                width=self.camera.width, height=self.camera.height,
                near=-self.camera.width, far=self.camera.height,
                fov_distance=1.0,
                dtype=self.camera.dtype,
                device=self.camera.device
            )
        else:
            raise ValueError(f"Unknown lens type: {lens} given to render_core.change_camera_projection_mode")
        self.camera.intrinsics = intrinsics
        del self._last_state['camera']  # The previous camera is now obsolete

    def _setup_camera(self, state: WispState):
        # Use selected camera to control canvas
        camera = state.renderer.selected_camera
        if camera is None:
            # Create a default camera
            lens_type = self.state.renderer.selected_camera_lens
            camera = self._default_camera(lens_type)

        camera = camera.to(self.device)
        return camera

    def refresh_bl_renderers(self, state: WispState) -> None:
        renderers = dict()
        scene_graph = state.graph

        # Remove obsolete bottom level renderers for pipelines that no longer exist
        for obj_name in list(scene_graph.bl_renderers.keys()):
            if obj_name not in scene_graph.neural_pipelines:
                del scene_graph.bl_renderers[obj_name]

        # Set up a renderer for all neural pipelines in the scene
        for renderer_id, neural_pipeline in scene_graph.neural_pipelines.items():
            # See if a descriptor for the renderer already exists.
            bl_state = scene_graph.bl_renderers.get(renderer_id)
            if bl_state is None:
                # If not, create a default one
                object_transform = ObjectTransform(device=self.device,dtype=self.camera.dtype)
                bl_state = BottomLevelRendererState(status='pending', setup_args=dict(), transform=object_transform)
                scene_graph.bl_renderers[renderer_id] = bl_state

            if bl_state.status == 'loaded':
                assert bl_state.renderer is not None, \
                       f'status of renderer {renderer_id} shows it was loaded, but renderer instance is None.'
                renderers[renderer_id] = bl_state.renderer
            elif bl_state.status == 'pending':
                bl_state.renderer = create_neural_field_renderer(neural_object=neural_pipeline, **bl_state.setup_args)
                bl_state.status = 'loaded'
                renderers[renderer_id] = bl_state.renderer
                scene_graph.visible_objects[renderer_id] = True  # Mark as visible when first loaded
            elif bl_state.status == 'ignored':
                pass    # Skip renderer on purpose
            else:
                raise ValueError(f'Invalid bottom level renderer state: {bl_state.status}')
        self._renderers = renderers

        # Refresh TLAS
        self._tlas = self._setup_tlas(state)

    def _setup_tlas(self, state: WispState):
        # Currently the top-level acceleration structure uses a straightforward ordered list stub
        return ListTLAS(state, self._renderers)

    def set_full_resolution(self):
        self.res_x = self.camera.width
        self.res_y = self.camera.height
        self.interactive_mode = False

    def set_low_resolution(self):
        self.res_x = self.camera.width // 4
        self.res_y = self.camera.height // 4
        self.interactive_mode = True

    def resize_canvas(self, width, height):
        self.camera.intrinsics.width = width
        self.camera.intrinsics.height = height
        self.set_low_resolution()

    def _update_scene_graph(self):
        """ Update scene graph information about objects and their data layers """
        # New data layers maybe have been generated, update the scene graph about their existence.
        # Some existing layers may have been regenerated, in which case we copy their previous "toggled on/off" status.
        data_layers = self._bl_renderer_data_layers()
        for object_id, bl_renderer_state in self.state.graph.bl_renderers.items():
            object_layers = data_layers[object_id]
            bl_renderer_state.data_layers = object_layers
            toggled_data_layers = defaultdict(bool)
            for layer_name, layer in object_layers.items():  # Copy over previous toggles if such exist
                toggled_data_layers[layer_name] = bl_renderer_state.toggled_data_layers.get(layer_name, False)
            bl_renderer_state.toggled_data_layers = toggled_data_layers

    @torch.no_grad()
    @enable_amp
    def redraw(self) -> None:
        """ Allow bottom-level renderers to refresh internal information, such as data layers. """
        # Read phase: sync with scene graph, create renderers for new objects added
        self.refresh_bl_renderers(self.state)

        # Invoke internal redraw() on all current visible renderers, to imply it's time to refresh content
        # (i.e. data layers may get generate here behind the scenes)
        scene_graph = self.state.graph
        for obj_id, renderer in self._renderers.items():
            if obj_id in scene_graph.visible_objects:
                renderer.redraw()

        # Write phase: update scene graph back with latest info from render core, i.e. new data layers generated
        self._update_scene_graph()

    @torch.no_grad()
    @enable_amp
    def render(self, time_delta=None, force_render=False) -> RenderBuffer:
        """Render a frame.

        Args:
            time_delta (float): The time delta from the previous frame, used to control renderer parameters
                                based on the amount of detected lag. 
            force_render (bool): If True, will always output a fresh new RenderBuffer. 
                                 Otherwise the RenderBuffer can be a stale copy of the the previous frame
                                 if no updates are detected.

        Returns: 
            (wisp.core.RenderBuffer): The rendered buffer.
        """
        payload = self._prepare_payload(time_delta)
        rb = self._render_payload(payload, force_render)
        output_rb = self._post_render(payload, rb)
        return output_rb

    def _prepare_payload(self, time_delta=None) -> FramePayload:
        """This function will prepare the FramePayload for the current frame.

        The FramePayload contains metadata for the current frame, from which the RenderBuffer will be 
        generated from. 

        Args:
            time_delta (float): The time delta from the previous frame, used to control renderer parameters
                                based on the amount of detected lag. 

        Returns:
            (wisp.renderer.core.api.FramePayload): The metadata for the frame.
        """
        # Adjust resolution of all renderers to maintain FPS
        camera = self.camera
        clear_color = self.state.renderer.clear_color_value
        res_x, res_y = self.res_x, self.res_y

        # If the FPS is slow, downscale the resolution for the render.
        is_fps_lagging = time_delta is not None and \
                         self.target_fps is not None and \
                         self.target_fps > 0 and \
                         time_delta > (1.0 / self.target_fps)
        if self.interactive_mode and is_fps_lagging:
            if res_x > self.MIN_RES and res_y > self.MIN_RES:
               res_x //= 2
               res_y //= 2

        # TODO(ttakikawa): Leaving a note here to think about whether this should be the case...
        # The renderer always needs depth, alpha, and rgb
        required_channels = {"rgb", "depth", "alpha"}
        selected_canvas_channel = self.state.renderer.selected_canvas_channel.lower()
        visible_objects = set([k for k,v in self.state.graph.visible_objects.items() if v])
        payload = FramePayload(camera=camera, interactive_mode=self.interactive_mode,
                               render_res_x=res_x, render_res_y=res_y, time_delta=time_delta,
                               visible_objects=visible_objects, clear_color=clear_color,
                               channels={selected_canvas_channel}.union(required_channels))
        for renderer_id, renderer in self._renderers.items():
            if renderer_id in payload.visible_objects:
                renderer.pre_render(payload)

        return payload

    def raygen(self, camera, res_x, res_y):
        ray_grid = generate_centered_pixel_coords(camera.width, camera.height, res_x, res_y, device=self.device)
        if camera.lens_type == 'pinhole':
            rays = generate_pinhole_rays(camera, ray_grid)
        elif camera.lens_type == 'ortho':
            rays = generate_ortho_rays(camera, ray_grid)
        else:
            raise ValueError(f'RendererCore failed to raygen on unknown camera lens type: {camera.lens_type}')
        return rays

    def _create_empty_rb(self, height, width, dtype=torch.float32) -> RenderBuffer:
        clear_color = self.state.renderer.clear_color_value
        clear_depth = self.state.renderer.clear_depth_value

        return RenderBuffer(
            rgb=torch.tensor(clear_color, dtype=dtype, device=self.device).repeat(height, width, 1),
            alpha=torch.zeros((height, width, 1), dtype=dtype, device=self.device),
            depth=torch.full((height, width, 1), fill_value=clear_depth, dtype=dtype, device=self.device),
            hit=None
        )

    def _render_payload(self, payload: FramePayload, force_render: bool) -> RenderBuffer:
        """Renders a RenderBuffer using a FramePayload which contains metadata.

        Args:
            payload (wisp.renderer.core.api.FramePayload): Metadata for the frame to be renderered.
            force_render (bool): If True, will always output a fresh new RenderBuffer. 
                                 Otherwise the RenderBuffer can be a stale copy of the the previous frame
                                 if no updates are detected.
        
        Returns:
            (wisp.core.RenderBuffer): The rendered buffer.
        """
        camera = payload.camera
        res_x, res_y = payload.render_res_x, payload.render_res_y

        visible_renderers = [r for r_id, r in self._renderers.items() if r_id in payload.visible_objects]
        renderers_to_refresh = list(filter(lambda renderer: renderer.needs_refresh(payload), visible_renderers))
        if not self.needs_refresh() and len(renderers_to_refresh) == 0 and not force_render:
            return self._last_renderbuffer  # No need to regenerate..

        # Generate rays
        rays = self.raygen(camera, res_x, res_y)
        renderers_in_view = self._tlas.traverse(rays, payload)

        rb_dtype = torch.float32
        clear_depth = self.state.renderer.clear_depth_value

        out_rb = self._create_empty_rb(height=camera.height, width=camera.width, dtype=rb_dtype)
        for renderer, hit_rays in renderers_in_view:
            if isinstance(renderer, RayTracedRenderer):
                in_rays = hit_rays.to(device=renderer.device, dtype=renderer.dtype)
                rb = renderer.render(in_rays)
            else:   # RasterizedRenderer
                # TODO (operel): Handle transformed rasterized objects
                in_cam = self.camera.to(device=renderer.device, dtype=renderer.dtype)
                rb = renderer.render(in_cam)

            rb = rb.to(device=self.device)
            rb.rgb = rb.rgb.to(dtype=rb_dtype)

            rb.alpha = rb.alpha.to(dtype=rb_dtype)
            rb.depth = rb.depth.to(dtype=rb_dtype)

            # TODO (operel): if rb.depth is None -> painters algorithm
            # Normalize ray-traced depth buffer to graphics api range
            img_dims = rb.depth.shape

            # Clip depth values which fall outside of the view frustum
            clip_mask = camera.clip_mask(rb.depth.squeeze(-1))
            rb.alpha[~clip_mask] = 0.0

            # Normalize depth from [0, inf] to NDC space according to camera settings
            # (depends on near / far and NDC space)
            ndc_depth = camera.normalize_depth(rb.depth.reshape(-1, 1))
            rb.depth = ndc_depth.reshape(img_dims)

            # Set depth of missed rays to far clipping plane, as PackedRFTracer initializes depth to 0 and writes values
            # only for hit rays
            alpha_mask = ~rb.hit[...,0]
            rb.depth[alpha_mask] = clear_depth

            rb.depth = rb.depth.to(rb_dtype)
            out_rb = out_rb.blend(rb, channel_kit=self.state.graph.channels)

        return out_rb

    def _post_render(self, payload: FramePayload, rb: RenderBuffer) -> RenderBuffer:
        # Update current resolution in case it was decreased to maintain fps
        self.res_x, self.res_y = payload.render_res_x, payload.render_res_y

        # Cache information to accelerate next frames
        self._last_renderbuffer = rb

        # Record last state, to, i.e, calculate if needs to redraw future frames
        self._last_state['camera'] = copy.deepcopy(payload.camera)
        self._last_state['res_x'] = payload.render_res_x
        self._last_state['res_y'] = payload.render_res_y
        for renderer_id, renderer in self._renderers.items():
            if renderer_id in payload.visible_objects:
                renderer.post_render()

        # Create an output renderbuffer to contain the currently viewed mode as rgba channel
        output_rb = self.map_output_channels_to_rgba(rb)
        return output_rb

    def needs_refresh(self) -> bool:
        if len(self._last_state) == 0:
            return True

        # Resolution check: if not full resolution - canvas is dirty
        if self._last_state['res_x'] != self.camera.width or self._last_state['res_y'] != self.camera.height:
            return True

        for att_name, prev_val in self._last_state.items():
            if not hasattr(self, att_name):
                continue
            curr_val = self.__getattribute__(att_name)
            if isinstance(curr_val, Camera):
                if not torch.allclose(curr_val, prev_val):
                    return True
            elif curr_val != prev_val:
                return True
        return False

    def _bl_renderer_data_layers(self) -> Dict[str, PrimitivesPack]:
        """ Returns the bottom level object data layers"""
        layers = dict()
        for renderer_id, renderer in self._renderers.items():
            layers[renderer_id] = renderer.data_layers()
        return layers

    def _cameras_data_layers(self) -> Iterable[PrimitivesPack]:
        """ Returns the available cameras data layer (all visible cameras layers are packed together) """
        cameras_to_redraw = {camera_id: camera for camera_id, camera in self.state.graph.cameras.items()
                             if self.state.graph.visible_objects.get(camera_id, False)}
        layers = self._camera_layers.regenerate_data_layers(cameras_to_redraw, self.state.renderer.clear_color_value)
        return layers.values()

    def active_data_layers(self) -> List[PrimitivesPack]:
        layers_to_draw = []
        for obj_state in self.state.graph.bl_renderers.values():
            for layer_id, layer in obj_state.data_layers.items():
                if obj_state.toggled_data_layers[layer_id]:
                    # Attach object transform reference to data layers, assumes the object maintains this transform ref
                    layer.transform = obj_state.transform
                    layers_to_draw.append(layer)
        camera_data_layers = self._cameras_data_layers()
        layers_to_draw.extend(camera_data_layers)
        return layers_to_draw

    def map_output_channels_to_rgba(self, rb: RenderBuffer):
        selected_output_channel = self.state.renderer.selected_canvas_channel.lower()
        rb_channel = rb.get_channel(selected_output_channel)

        if rb_channel is None:
            # Unknown channel type configured to view over the canvas.
            # That can happen if, i.e. no object have traced a RenderBuffer with this channel.
            # Instead of failing, create an empty rb
            height, width = rb.rgb.shape[:2]
            return self._create_empty_rb(height=height, width=width, dtype=rb.rgb.dtype)

        # Normalize channel to [0, 1]
        channels_kit = self.state.graph.channels
        channel_info = channels_kit.get(selected_output_channel, create_default_channel())
        normalized_channel = channel_info.normalize_fn(rb_channel.clone())  # Clone to protect from modifications

        # To RGB (in normalized space)
        # TODO (operel): incorporate color maps
        channel_dim = normalized_channel.shape[-1]
        if channel_dim == 1:
            rgb = torch.cat((normalized_channel, normalized_channel, normalized_channel), dim=-1)
        elif channel_dim == 2:
            rgb = torch.cat((normalized_channel, normalized_channel, torch.zeros_like(normalized_channel)), dim=-1)
        elif channel_dim == 3:
            rgb = normalized_channel
        else:
            raise ValueError('Cannot display channels with more than 3 dimensions over the canvas.')

        canvas_rb = RenderBuffer(rgb=rgb, depth=rb.depth, alpha=rb.alpha)
        return canvas_rb

    @property
    def renderers(self) -> Dict[str, BottomLevelRenderer]:
        """ All loaded bottom level renderers currently employed by the renderer core """
        return self._renderers

    @property
    def camera(self) -> Camera:
        return self.state.renderer.selected_camera

    @camera.setter
    def camera(self, camera: Camera) -> None:
        self.state.renderer.selected_camera = camera

    @property
    def target_fps(self) -> float:
        return self.state.renderer.target_fps

    @property
    def interactive_mode(self) -> bool:
        return self.state.renderer.interactive_mode

    @interactive_mode.setter
    def interactive_mode(self, mode: bool) -> None:
        self.state.renderer.interactive_mode = mode

    @property
    def selected_camera_lens(self) -> str:
        return self.state.renderer.selected_camera_lens

    @selected_camera_lens.setter
    def selected_camera_lens(self, lens: str):
        self.state.renderer.selected_camera_lens = lens


class TLAS(abc.ABC):
    def traverse(self, rays: Rays, payload: FramePayload) -> Dict[BottomLevelRenderer, Rays]:
        pass

    def transform_rays(self, rays: Rays, transform: ObjectTransform) -> Rays:
        rays_shape = rays.shape
        rays = rays.reshape((-1, 3))
        ray_origins = torch.cat((rays.origins, torch.ones_like(rays.origins[:,:1])), dim=-1)
        ray_dirs = torch.cat((rays.dirs, torch.zeros_like(rays.dirs[:,:1])), dim=-1)
        inv_model_matrix = transform.to(device=ray_dirs.device).inv_model_matrix().to(dtype=rays.origins.dtype)
        ray_origins = (inv_model_matrix @ ray_origins.T).T
        ray_dirs = (inv_model_matrix @ ray_dirs.T).T
        transformed_rays = Rays(ray_origins[:, :3], dirs=ray_dirs[:, :3],
                                dist_min=rays.dist_min, dist_max=rays.dist_max)
        return transformed_rays.reshape((*rays_shape, 3))

class ListTLAS(TLAS):
    def __init__(self, state: WispState, bl_renderers: Dict[str, BottomLevelRenderer]):
        self.state = state
        self.bl_renderers = list(bl_renderers.values())
        self.bl_renderer_ids = list(bl_renderers.keys())

    def traverse(self, rays: Rays, payload: FramePayload) -> List[Tuple[BottomLevelRenderer, Rays]]:
        # TODO (operel): invoke aabb test on bottom level renderers
        renderers_to_hit_rays = list()
        for renderer_id, renderer in zip(self.bl_renderer_ids, self.bl_renderers):
            if renderer_id in payload.visible_objects:
                object_transform = self.state.graph.bl_renderers[renderer_id].transform
                hit_rays = self.transform_rays(rays, object_transform)
                renderers_to_hit_rays.append((renderer, hit_rays))
        return renderers_to_hit_rays

wisp/renderer/core/render_core.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .radiance_pipeline_renderer import NeuralRadianceFieldPackedRenderer
from .sdf_pipeline_renderer import NeuralSDFPackedRenderer
from .spc_pipeline_renderer import SPCRenderer

wisp/renderer/core/renderers/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from wisp.core import RenderBuffer, Rays, PrimitivesPack
from wisp.renderer.core.api import RayTracedRenderer, FramePayload, field_renderer
from wisp.models.nefs.nerf import NeuralRadianceField, BaseNeuralField
from wisp.tracers import PackedRFTracer


@field_renderer(BaseNeuralField, PackedRFTracer)
class NeuralRadianceFieldPackedRenderer(RayTracedRenderer):
    """ A neural field renderers for pipelines of NeuralRadianceField + PackedRFTracer.
        The renderer is registered with the general BaseNeuralField to make a default fallback for future neural field
        subclasses which use the PackedRFTracer and don't implement a dedicated renderer.

        NeuralRadianceFieldPackedRenderer is capable of reducing the tracer settings to lower quality when interactive
        mode is on.
    """

    def __init__(self, nef: BaseNeuralField, tracer: PackedRFTracer, batch_size: int = 2**14, num_steps: int = None,
                 raymarch_type: str = None, *args, **kwargs):
        """
        Construct a new neural radiance field from the nef + tracer pipeline.
        By default, tracing will use the tracer args, unless specific values are specified for overriding those
        defaults.

        Args:
            nef (BaseNeuralField): Neural field component of the pipeline. The neural field is expected to be ray based.
            tracer (PackedRFTracer): Tracer component of the pipeline, assumed to trace radiance fields.
            batch_size (int): Amount of rays processed per batch, in inference time. A high default is in place,
                since often tracers are configured for training batch sizes, which are suboptimal for inference time.
            num_steps (int): Number of steps used for raymarching, the exact functionality depends on raymarch_type
                (see for example, OctreeAS.raymarch). By default, the tracer num_steps will be used. Specify an explicit
                value to override.
            raymarch_type (str): raymarch_type sampling strategy to use (see for example, OctreeAS.raymarch).
                By default, the tracer num_steps will be used. Specify an explicit value to override.
        """
        super().__init__(nef, tracer, batch_size, *args, **kwargs)
        if num_steps is None:
            num_steps = tracer.num_steps
        self.num_steps = num_steps
        self.num_steps_interactive = max(num_steps // 4, 1)
        if raymarch_type is None:
            raymarch_type = tracer.raymarch_type
        self.raymarch_type = raymarch_type
        self.bg_color = tracer.bg_color

        self._last_state = dict()

    def pre_render(self, payload: FramePayload, *args, **kwargs) -> None:
        super().pre_render(payload)
        self.bg_color = 'black' if payload.clear_color == (0.0, 0.0, 0.0) else 'white'
        if payload.interactive_mode:
            self.tracer.num_steps = self.num_steps_interactive
        else:
            self.tracer.num_steps = self.num_steps

    def needs_refresh(self, payload: FramePayload, *args, **kwargs) -> bool:
        return self._last_state.get('num_steps', 0) < self.num_steps or \
               self._last_state.get('channels') != self.channels

    def render(self, rays: Rays) -> RenderBuffer:
        rb = RenderBuffer(hit=None)
        for ray_batch in rays.split(self.batch_size):
            rb += self.tracer(self.nef,
                              rays=ray_batch,
                              channels=self.channels,
                              lod_idx=None,  # TODO(ttakikawa): Add a way to control the LOD in the GUI
                              raymarch_type=self.raymarch_type,
                              num_steps=self.num_steps,
                              bg_color=self.bg_color)

        # Rescale renderbuffer to original size
        rb = rb.reshape(self.render_res_y, self.render_res_x, -1)
        if self.render_res_x != self.output_width or self.render_res_y != self.output_height:
            rb = rb.scale(size=(self.output_height, self.output_width))
        return rb

    def post_render(self) -> None:
        self._last_state['num_steps'] = self.tracer.num_steps
        self._last_state['channels'] = self.channels

wisp/renderer/core/renderers/radiance_pipeline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from wisp.renderer.core.api import RayTracedRenderer, FramePayload, field_renderer
from wisp.models.nefs.neural_sdf import NeuralSDF, BaseNeuralField
from wisp.tracers import PackedSDFTracer


@field_renderer(BaseNeuralField, PackedSDFTracer)
class NeuralSDFPackedRenderer(RayTracedRenderer):
    """ A neural field renderers for pipelines of NeuralSDF + PackedSDFTracer.
        The renderer is registered with the general BaseNeuralField to make a default fallback for future neural field
        subclasses which use the PackedSDFTracer and don't implement a dedicated renderer.
    """

    def __init__(self, nef: NeuralSDF, tracer: PackedSDFTracer, num_steps=None, step_size=None, min_dis=None,
                 *args, **kwargs):
        """
        Construct a new neural signed distance field from the nef + tracer pipeline.
        By default, tracing will use the tracer args, unless specific values are specified for overriding those
        defaults.

        Args:
            nef (BaseNeuralField): Neural field component of the pipeline. The neural field is expected to be coordinate
             based.
            tracer (PackedSDFTracer): Tracer component of the pipeline, assumed to trace signed distance fields.
            num_steps (int): Number of steps used for sphere tracing.
                By default, the tracer num_steps will be used. Specify an explicit value to override.
            step_size (float): The multiplier for the sphere tracing steps. Use a value <1.0 for conservative tracing.
                By default, the tracer step_size will be used. Specify an explicit value to override.
            min_dis (float): The termination distance for sphere tracing.
                By default, the tracer min_dis will be used. Specify an explicit value to override.
        """
        super().__init__(nef, tracer, *args, **kwargs)
        self.num_steps = num_steps if num_steps is not None else tracer.num_steps
        self.step_size = step_size if step_size is not None else tracer.step_size
        self.min_dis = min_dis if min_dis is not None else tracer.min_dis
        self._last_state = dict()

    def needs_refresh(self, payload: FramePayload, *args, **kwargs) -> bool:
        return self._last_state.get('num_steps', 0) != self.num_steps or \
               self._last_state.get('step_size', 0) != self.step_size or \
               self._last_state.get('min_dis', 0) != self.min_dis or \
               self._last_state.get('channels') != self.channels

    def pre_render(self, payload: FramePayload, *args, **kwargs) -> None:
        super().pre_render(payload)
        self.tracer.num_steps = self.num_steps
        self.tracer.step_size = self.step_size
        self.tracer.min_dis = self.min_dis
        self.tracer.bg_color = 'black' if payload.clear_color == (0.0, 0.0, 0.0) else 'white'

    def post_render(self) -> None:
        self._last_state['num_steps'] = self.tracer.num_steps
        self._last_state['step_size'] = self.tracer.step_size
        self._last_state['min_dis'] = self.tracer.min_dis
        self._last_state['channels'] = self.channels

wisp/renderer/core/renderers/sdf_pipeline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from typing import Optional
from wisp.models.nefs import SPCField
from wisp.tracers import PackedSPCTracer
from wisp.gfx.datalayers import Datalayers, OctreeDatalayers
from wisp.renderer.core.api import RayTracedRenderer, field_renderer, FramePayload


@field_renderer(SPCField, PackedSPCTracer)
class SPCRenderer(RayTracedRenderer):
    """ A renderer for SPC objects, used with pipelines of SPCField + PackedSPCTracer. """

    def __init__(self, nef: SPCField, tracer: PackedSPCTracer, batch_size=4000, *args, **kwargs):
        super().__init__(nef, tracer, batch_size, *args, **kwargs)

    @classmethod
    def create_layers_painter(cls, nef: SPCField) -> Optional[Datalayers]:
        return OctreeDatalayers()   # Always assume an octree grid object exists for this field

    def needs_refresh(self, payload: FramePayload, *args, **kwargs) -> bool:
        """ SPC never requires a refresh due to internal state changes """
        return False

    def acceleration_structure(self) -> str:
        """ Returns a human readable name of the bottom level acceleration structure used by this renderer """
        return "Octree"  # Assumes to always use OctreeAS

    def features_structure(self) -> str:
        """ Returns a human readable name of the feature structure used by this renderer """
        return "Octree Grid"  # Assumes to always use OctreeGrid for storing features

wisp/renderer/core/renderers/spc_pipeline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from .gizmo import *
from .ogl.world_grid import *
from .ogl.axis_painter import *
from .ogl.primitives_painter import *

wisp/renderer/gizmos/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


from abc import ABC, abstractmethod
from kaolin.render.camera import Camera


class Gizmo(ABC):
    """
    A template for representing gizmos the interactive renderer is able to render over the canvas.
    Gizmos are entities rendered by the graphics api (i.e: OpenGL).
    Normally they're used to draw transient markups or tools over the canvas (such as a world grid or axes).
    """

    @abstractmethod
    def render(self, camera: Camera):
        """ Renders the gizmo using the graphics api. """
        raise NotImplementedError("Gizmos must implement the render function")

    @abstractmethod
    def destroy(self):
        """ Release GL resources, must be called from the rendering thread which owns the GL context """
        raise NotImplementedError("Gizmos must implement the destroy function")

wisp/renderer/gizmos/gizmo.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

wisp/renderer/gizmos/ogl/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
import copy

import numpy as np
from typing import Optional
from glumpy import gloo, gl
from kaolin.render.camera import Camera
from wisp.renderer.gizmos.gizmo import Gizmo


class AxisPainter(Gizmo):

    def __init__(self, axes_length, line_width, origin=None, axes=None, is_bidirectional=True):
        """
        :param axes_length Axes will occupy [origin, origin + axes_length] in world coordinates
        :param line_width Line width
        :param origin numpy array of 3 coordinates for the x,y,z of the origin
        :param axes Iterable of strings to specify which axes to draw. Possible values: 'x', 'y', 'z'
        :param is_bidirectional if True, will also draw [origin, origin - axes_length] in world coordinates
        """
        if origin is None:
            origin = np.array((0.0, 0.0, 0.0))
        if axes is None:
            axes = ('x', 'y', 'z')
        self.vbo, self.ibo = self.create_line_buffers(axes_length,
                                                      origin=origin, axes=axes, is_bidirectional=is_bidirectional)
        self.line_size = line_width
        self.canvas_program: Optional[gloo.Program] = self.create_gl_program()

    def destroy(self):
        """ Release GL resources, must be called from the rendering thread which owns the GL context """
        if self.vbo is not None:
            self.vbo.delete()
            self.vbo = None
        if self.ibo is not None:
            self.ibo.delete()
            self.ibo = None
        if self.canvas_program is not None:
            self.canvas_program.delete()
            self.canvas_program = None

    def create_gl_program(self):
        vertex = """
                    uniform mat4   u_viewprojection;
                    attribute vec3 position;
                    attribute vec4 color;
                    varying vec4 v_color;
                    void main()
                    {
                        v_color = color;
                        gl_Position = u_viewprojection * vec4(position, 1.0f);
                    } """

        fragment = """
                    varying vec4 v_color;
                    void main()
                    {
                        gl_FragColor = v_color;
                    } """

        # Compile GL program
        canvas = gloo.Program(vertex, fragment)
        return canvas

    def create_line_buffers(self, axes_length, origin, axes, is_bidirectional):

        if is_bidirectional:
            segment_ends = (-axes_length, axes_length)
        else:
            segment_ends = (axes_length,)

        num_lines = len(segment_ends) * len(axes)
        vertex_buffer = np.zeros(2 * num_lines, [("position", np.float32, 3), ("color", np.float32, 4)])

        blend_to_alpha = 0.6   # Makes the line appear more transparent towards the far tip
        idx = 0
        for axis in axes:
            for end_val in segment_ends:
                start = origin
                if axis == 'x':
                    end = np.array((end_val, 0.0, 0.0))
                    color_start = np.array((1.0, 0.0, 0.0, 1.0))
                    color_end = np.array((1.0, 0.0, 0.0, blend_to_alpha))
                elif axis == 'y':
                    end = np.array((0.0, end_val, 0.0))
                    color_start = np.array((0.0, 1.0, 0.0, 1.0))
                    color_end = np.array((0.0, 1.0, 0.0, blend_to_alpha))
                elif axis == 'z':
                    end = np.array((0.0, 0.0, end_val))
                    color_start = np.array((0.0, 0.0, 1.0, 1.0))
                    color_end = np.array((0.0, 0.0, 1.0, blend_to_alpha))
                vertex_buffer["position"][idx] = start
                vertex_buffer["position"][idx+1] = end
                vertex_buffer["color"][idx] = color_start
                vertex_buffer["color"][idx+1] = color_end
                idx += 2
        vertex_buffer = vertex_buffer.view(gloo.VertexBuffer)

        index_buffer = np.arange(0, 2 * num_lines).astype(np.uint32)
        index_buffer = index_buffer.view(gloo.IndexBuffer)

        return vertex_buffer, index_buffer

    def render(self, camera: Camera):
        gl.glLineWidth(self.line_size)
        self.canvas_program["u_viewprojection"] = camera.view_projection_matrix()[0].cpu().numpy().T
        self.canvas_program.bind(self.vbo)
        self.canvas_program.draw(gl.GL_LINES, self.ibo)
wisp/renderer/gizmos/ogl/axis_painter.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
import sys
from collections import defaultdict
import torch
from typing import Optional, List

import numpy as np
from glumpy import gloo, gl
from kaolin.render.camera import Camera

from wisp.core import PrimitivesPack, ObjectTransform
from wisp.renderer.gizmos.gizmo import Gizmo

# TODO (operel): release GL resources when instance is destroyed (and issue a warning if not taken care of explicitly)

class PrimitivesPainter(Gizmo):

    def __init__(self):
        # GL program used to paint GL primitives
        self.canvas_program: Optional[gloo.Program] = self.create_gl_program()
        self.lines = []
        self.points = []

    def destroy(self):
        """ Release GL resources, must be called from the rendering thread which owns the GL context """
        self.clear()
        if self.canvas_program is not None:
            self.canvas_program.delete()
            self.canvas_program = None

    def create_gl_program(self):
        vertex = """
                    uniform mat4   u_model;
                    uniform mat4   u_viewprojection;
                    attribute vec3 position;
                    attribute vec4 color;
                    varying vec4 v_color;
                    void main()
                    {
                        v_color = color;
                        gl_Position =  u_viewprojection * u_model * vec4(position, 1.0f);
                    } """

        fragment = """
                    varying vec4 v_color;
                    void main()
                    {
                        gl_FragColor = v_color;
                    } """

        # Compile GL program
        canvas = gloo.Program(vertex, fragment)
        return canvas

    def clear(self):
        for line in self.lines:
            vertex_buffer, index_buffer, _, _ = line
            vertex_buffer.delete()
            index_buffer.delete()
        for point in self.points:
            vertex_buffer, _, _ = point
            vertex_buffer.delete()

        self.lines = []
        self.points = []

    def create_points_buffers(self, points):
        # TODO (operel): gpu copy is better..
        pos, color = points
        count = len(pos)
        vertex_buffer = np.zeros(count, [("position", np.float32, 3), ("color", np.float32, 4)])
        vertex_buffer["position"] = torch.cat(pos).cpu().numpy()
        vertex_buffer["color"] = torch.cat(color).reshape(-1, 4).cpu().numpy()
        vertex_buffer = vertex_buffer.view(gloo.VertexBuffer)
        return vertex_buffer

    def create_line_buffers(self, lines):
        # TODO (operel): gpu copy is better..
        start, end, color = lines
        count = len(start)
        vertex_buffer = np.zeros(2 * count, [("position", np.float32, 3), ("color", np.float32, 4)])
        vertex_buffer["position"] = torch.cat([start, end], dim=-1).reshape(-1, 3).cpu().numpy()
        vertex_buffer["color"] = torch.cat([color, color], dim=-1).reshape(-1, 4).cpu().numpy()
        vertex_buffer = vertex_buffer.view(gloo.VertexBuffer)

        index_buffer = np.arange(0, 2 * count).astype(np.uint32)
        index_buffer = index_buffer.view(gloo.IndexBuffer)

        return vertex_buffer, index_buffer

    def redraw(self, primitives: List[PrimitivesPack]):
        self.clear()

        # Sort primitives by gl state changes - this reduces the amount of draw calls
        draw_calls = defaultdict(lambda: dict(
                lines=dict(
                    line_width=defaultdict(PrimitivesPack)
                ),
                points=dict(
                    point_size=defaultdict(PrimitivesPack)
                ),
            )
        )

        for pack in primitives:
            transform = pack.transform if pack.transform is not None else ObjectTransform()
            if (len(pack._lines_start)):
                draw_calls[transform]['lines']['line_width'][pack.line_width].append(pack)
            elif (len(pack._points_pos)):
                draw_calls[transform]['points']['point_size'][pack.point_size].append(pack)

        for transform, grouped_calls in draw_calls.items():
            for _line_width, pack in grouped_calls['lines']['line_width'].items():
                lines = pack.lines
                vertex_buffer, index_buffer = self.create_line_buffers(lines)
                self.lines.append((vertex_buffer, index_buffer, pack.line_width, transform))
            for point_size, pack in grouped_calls['points']['point_size'].items():
                points = pack.points
                vertex_buffer = self.create_points_buffers(points)
                self.points.append((vertex_buffer, point_size, transform))

    def render(self, camera: Camera):
        for line_entry in self.lines:
            vertex_buffer, index_buffer, line_width, transform = line_entry
            gl.glLineWidth(line_width)
            self.canvas_program["u_model"] = transform.model_matrix().cpu().numpy().T
            self.canvas_program["u_viewprojection"] = camera.view_projection_matrix()[0].cpu().numpy().T
            self.canvas_program.bind(vertex_buffer)
            self.canvas_program.draw(gl.GL_LINES, index_buffer)
    
        for point_entry in self.points:
            vertex_buffer, point_size, transform = point_entry
            gl.glPointSize(point_size)
            self.canvas_program["u_model"] = transform.model_matrix().cpu().numpy().T
            self.canvas_program["u_viewprojection"] = camera.view_projection_matrix()[0].cpu().numpy().T
            self.canvas_program.bind(vertex_buffer)
            self.canvas_program.draw(gl.GL_POINTS)

wisp/renderer/gizmos/ogl/primitives_painter.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import logging
import numpy as np
from typing import Optional, Tuple
from glumpy import gloo, gl
from kaolin.render.camera import Camera
from wisp.renderer.gizmos.gizmo import Gizmo


class WorldGrid(Gizmo):

    def __init__(self, squares_per_axis: int = 20, grid_size: float = 1.0,
            line_color: Tuple = None, line_size: int = 2, plane: str = 'xy'):
        """
        Args:
            squares_per_axis (int): Number of grid squares per row, column
            grid_size (float): Grid will occupy [-size, size] in world coordinates
            line_color (Tuple): Tuple of RGB values for line color, in range [0,255]. i.e: (255, 255, 255) for white.
            line_size (int): integer between [1,32], where each grid square can be thought of a 32 x 32 pixel grid,
               and line size specifies how many rows / columns each line occupies for each square.
               The actual amount of pixels each line occupies also depends on the grid_size and camera settings.
            plane (str): Plane to align the reference grid with. Valid choices: 'xy', 'xz', 'yz'.
        """
        # Size of each square texture, to be repeated, setting this one too low will decrease quality.
        # We may go above 32 x 32 and scale the line size accordingly
        self.tex_size = 64
        self.line_size = line_size * (self.tex_size / 64)
        if line_color is None:
            line_color = (255, 255, 255)
        self.line_color = np.array(line_color)
        self.squares_per_axis = squares_per_axis
        self.grid_size = grid_size
        dtype = np.uint8
        tex_data = np.zeros((self.tex_size, self.tex_size, 4), dtype=dtype)

        for j in range(0, self.tex_size):   # Height
            for i in range(0, self.tex_size):   # Width
                tex_data[j, i, :3] = self.line_color if (i < self.line_size or j < self.line_size) else 0
                tex_data[j, i, 3] = 255 if (i < self.line_size or j < self.line_size) else 0    # Alpha

        tex = tex_data.view(gloo.Texture2D)

        tex.activate()  # Binds texture
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_REPEAT)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_REPEAT)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR_MIPMAP_LINEAR)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
        try:
            gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAX_ANISOTROPY, 16)
        except AttributeError as e:
            logging.warning('GL_TEXTURE_MAX_ANISOTROPY not available; appearance may suffer')
        gl.glGenerateMipmap(gl.GL_TEXTURE_2D)
        tex.deactivate()

        # GL program used to paint a single billboard
        self.canvas_program: Optional[gloo.Program] = self.create_gl_program(tex, plane)

    def destroy(self):
        """ Release GL resources, must be called from the rendering thread which owns the GL context """
        if self.canvas_program is not None:
            if self.canvas_program['tex'] is not None:
                self.canvas_program['tex'].delete()
                self.canvas_program['tex'] = None
            self.canvas_program.delete()
            self.canvas_program = None

    def create_gl_program(self, texture: np.ndarray, plane: str):
        vertex = """
                    uniform mat4   u_viewprojection;
                    attribute vec4 position;
                    attribute vec4 texcoord;
                    varying vec4 v_texcoord;
                    void main()
                    {
                        v_texcoord = texcoord;
                        gl_Position = u_viewprojection * position;
                    } """

        fragment = """
                    uniform sampler2D tex;
                    varying vec4 v_texcoord;
                        
                    void main()
                    {
                        vec4 sample = texture2DProj(tex, v_texcoord);
                        if (sample.a <= 0)
                            discard;
                        gl_FragColor = sample;
                    } """

        # Compile GL program
        canvas = gloo.Program(vertex, fragment, count=4)

        # Upload fixed values to GPU
        if plane == 'xz':
            canvas["position"] = [(-self.grid_size, 0.0, +self.grid_size, 1.0),
                                  (+self.grid_size, 0.0, +self.grid_size, 1.0),
                                  (-self.grid_size, 0.0, -self.grid_size, 1.0),
                                  (+self.grid_size, 0.0, -self.grid_size, 1.0)]
        elif plane == 'xy':
            canvas["position"] = [(-self.grid_size, +self.grid_size, 0.0, 1.0),
                                  (+self.grid_size, +self.grid_size, 0.0, 1.0),
                                  (-self.grid_size, -self.grid_size, 0.0, 1.0),
                                  (+self.grid_size, -self.grid_size, 0.0, 1.0)]
        elif plane == 'yz':
            canvas["position"] = [(0.0, -self.grid_size, +self.grid_size, 1.0),
                                  (0.0, +self.grid_size, +self.grid_size, 1.0),
                                  (0.0, -self.grid_size, -self.grid_size, 1.0),
                                  (0.0, +self.grid_size, -self.grid_size, 1.0)]

        # Repeat square pattern with texture coordinates > 1.0
        reps = self.squares_per_axis + float(self.line_size) / self.tex_size
        canvas['texcoord'] = [(0, 0, 0.0, 1.0), (reps, 0, 0.0, 1.0), (0, reps, 0.0, 1.0), (reps, reps, 0.0, 1.0)]
        canvas['tex'] = texture
        return canvas

    def render(self, camera: Camera):
        # TODO(operel): state control should be elsewhere
        gl.glEnable(gl.GL_DEPTH_TEST)
        gl.glDepthFunc(gl.GL_LEQUAL)
        self.canvas_program["u_viewprojection"] = camera.view_projection_matrix()[0].cpu().numpy().T
        self.canvas_program.draw(gl.GL_TRIANGLE_STRIP)

wisp/renderer/gizmos/ogl/world_grid.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .imgui.widget_accelstruct import *
from .imgui.widget_cameras import *
from .imgui.widget_gpu_stats import *
from .imgui.widget_imgui import *
from .imgui.widget_object_transform import *
from .imgui.widget_optimization import *
from .imgui.widget_property_editor import *
from .imgui.widget_radiance_pipeline_renderer import *
from .imgui.widget_raytraced_pipeline_renderer import *
from .imgui.widget_renderer_properties import *
from .imgui.widget_scene_graph import *
from .imgui.widget_sdf_pipeline_renderer import *
from .imgui.widget_wisp_module import *

wisp/renderer/gui/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

wisp/renderer/gui/imgui/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import numpy as np
import imgui
from wisp.core import WispModule
from wisp.framework import WispState
from wisp.accelstructs import BaseAS
from .widget_imgui import WidgetImgui, widget, get_widget
from .widget_property_editor import WidgetPropertyEditor


@widget(BaseAS)
class WidgetAccelStruct(WidgetImgui):

    def __init__(self):
        super().__init__()
        self.properties_widget = WidgetPropertyEditor()

    def paint(self, state: WispState, module: BaseAS = None, field_name = None, *args, **kwargs):
        if module is None:
            return

        section_title = field_name if field_name is not None else "Acceleration Structure"
        if imgui.tree_node(section_title):
            imgui.text(f"Type: {module.name()}")

            imgui.text(f"Sparsity (per LOD):")
            occupancy_hist = [o / c for o, c in zip(module.occupancy(), module.capacity())]
            width, height = imgui.get_content_region_available()
            imgui.plot_histogram(label="##accelstruct", values=np.array(occupancy_hist, dtype=np.float32),
                                 graph_size=(width, 20), scale_min=0.0, scale_max=1.0)

            properties = module.public_properties()
            table_properties = {k: p for k, p in properties.items() if not isinstance(p, WispModule)}
            recursive_properties = {k: p for k, p in properties.items() if isinstance(p, WispModule)}
            self.properties_widget.paint(state=state, properties=table_properties)

            for recursive_field_name, recursive_property in recursive_properties.items():
                child_widget = get_widget(recursive_property)
                child_widget.paint(state=state, module=recursive_property, field_name=recursive_field_name)
            imgui.tree_pop()

wisp/renderer/gui/imgui/widget_accelstruct.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
import copy

import imgui
from typing import Optional
from kaolin.render.camera import Camera
from wisp.framework import WispState
from .widget_imgui import WidgetImgui, widget, get_widget
from .widget_property_editor import WidgetPropertyEditor


@widget(Camera)
class WidgetCameraProperties(WidgetImgui):
    def __init__(self, camera_id: Optional[str] = None):
        super().__init__()

        self.camera_id = camera_id
        self.available_cameras = ['Perspective', 'Orthographic']
        self.intrinsics_to_selection_idx = dict(
            PinholeIntrinsics=self.available_cameras.index('Perspective'),
            OrthographicIntrinsics=self.available_cameras.index('Orthographic')
        )
        self.properties_editor = WidgetPropertyEditor()

    def paint(self, state: WispState, camera: Camera = None, *args, **kwargs):

        if camera is None:
            camera = state.renderer.selected_camera

        is_editable = camera == state.renderer.selected_camera

        if camera != state.renderer.selected_camera:
            button_width = 100
            if imgui.button("Set View", width=button_width):
                prev_selected_cam = state.renderer.selected_camera
                device = prev_selected_cam.device
                dtype = prev_selected_cam.dtype
                state.renderer.selected_camera = copy.deepcopy(camera).to(device, dtype)

        if imgui.tree_node("Intrinsics", imgui.TREE_NODE_DEFAULT_OPEN):
            cam_type = type(camera.intrinsics).__name__
            selected_cam_idx = self.intrinsics_to_selection_idx[cam_type]
            clicked, new_cam_idx = imgui.combo("Lens", selected_cam_idx, self.available_cameras)
            if is_editable and selected_cam_idx != new_cam_idx:
                camera = copy.deepcopy(camera)  # Camera may change, but let gui finish drawing the previous cam info
                state.renderer.selected_camera_lens = self.available_cameras[new_cam_idx]

            properties = None
            if cam_type == 'PinholeIntrinsics':
                properties = self._pinhole_properties(camera, is_editable)
            elif cam_type == 'OrthographicIntrinsics':
                properties = self._ortho_properties(camera, is_editable)
            if properties is not None:
                self.properties_editor.paint(state=state, properties=properties)
            imgui.tree_pop()

        if imgui.tree_node("Extrinsics", imgui.TREE_NODE_DEFAULT_OPEN):
            flags = imgui.INPUT_TEXT_ALWAYS_INSERT_MODE | imgui.INPUT_TEXT_CHARS_DECIMAL
            rchanged, rvalues = imgui.input_float3("Right", *camera.cam_right()[0].squeeze().cpu().numpy(),
                                                   flags=flags)
            uchanged, uvalues = imgui.input_float3("Up", *camera.cam_up()[0].squeeze().cpu().numpy(),
                                                   flags=flags)
            fchanged, fvalues = imgui.input_float3("Forward", *camera.cam_forward()[0].squeeze().cpu().numpy(),
                                                   flags=flags)
            tchanged, tvalues = imgui.input_float3("Position", *camera.cam_pos()[0].squeeze().cpu().numpy(),
                                                   flags=flags)

            if rchanged or uchanged or fchanged or tchanged:
                # TODO (operel) - make sure mem layout is ok
                # TODO (operel) - don't allow to edit training cameras because that would ruin
                updated_R = camera[0].R.new_tensor([rvalues, uvalues, fvalues]).transpose(-1,-2)
                camera[0].R = updated_R
            if tchanged:
                updated_t = camera[0].t.new_tensor([tvalues]).transpose(-1,-2)
                camera[0].t = updated_t
            imgui.tree_pop()

    def _pinhole_properties(self, camera, is_editable):
        def _x0_property():
            x0_value = camera.x0.item()
            changed, new_x0_value = imgui.core.slider_int("##camera_x0",
                                                          value=x0_value, min_value=0, max_value=camera.width)
            if changed and x0_value != new_x0_value:
                camera.x0 = new_x0_value

        def _y0_property():
            y0_value = camera.y0.item()
            changed, new_y0_value = imgui.core.slider_int("##camera_y0",
                                                          value=y0_value, min_value=0, max_value=camera.height)
            if changed and y0_value != new_y0_value:
                camera.y0 = new_y0_value

        fov_x_value = camera.fov_x.item()
        fov_y_value = camera.fov_y.item()
        fov_ratio = fov_x_value / fov_y_value

        near_value = camera.near
        def _near_property():
            changed, new_near_value = imgui.core.slider_float("##camera_near",
                                                              value=near_value, min_value=1e-2, max_value=20.0)
            if changed and near_value != new_near_value:
                camera.near = new_near_value

        far_value = camera.far
        def _far_property():
            changed, new_far_value = imgui.core.slider_float("##camera_far",
                                                              value=far_value, min_value=1.0, max_value=100.0)
            if changed and far_value != new_far_value:
                camera.far = new_far_value


        def _fov_x_property():
            changed, new_fov_x_value = imgui.core.slider_float("##camera_fov_x",
                                                               value=fov_x_value, min_value=1.0, max_value=120.0)
            if changed and fov_x_value != new_fov_x_value:
                camera.fov_x = new_fov_x_value
                camera.fov_y = new_fov_x_value / fov_ratio

        def _fov_y_property():
            changed, new_fov_y_value = imgui.core.slider_float("##camera_fov_y",
                                                               value=fov_y_value, min_value=1.0, max_value=120.0)
            if changed and fov_y_value != new_fov_y_value:
                camera.fov_y = new_fov_y_value
                camera.fov_x = new_fov_y_value * fov_ratio

        properties = {
            'Near Clip Plane': _near_property if is_editable else camera.near,
            'Far Clip Plane': _far_property if is_editable else camera.far,
            'Field of view (H)': _fov_x_property if is_editable else camera.fov_x.item(),
            'Field of view (V)': _fov_y_property if is_editable else camera.fov_y.item(),
            'Principal Point x0': _x0_property if is_editable else camera.x0.item(),
            'Principal Point y0': _y0_property if is_editable else camera.y0.item(),
            'Focal Length X': camera.focal_x.item(),
            'Focal Length Y': camera.focal_y.item(),
        }
        return properties

    def _ortho_properties(self, camera, is_editable):
        def _fov_distance():
            fov_distance_value = camera.fov_distance.item()
            MAX_ORTHO_ZOOMOUT = 75  # In distance units
            max_clipping_plane = min(max(camera.far, camera.near), MAX_ORTHO_ZOOMOUT)
            changed, new_fov_dist_value = imgui.core.slider_float("##camera_fov_distance",
                                                                  value=fov_distance_value,
                                                                  min_value=1e-4, max_value=max_clipping_plane)
            if changed and fov_distance_value != new_fov_dist_value:
                camera.fov_distance = new_fov_dist_value

        properties = {
            'FOV distance': _fov_distance if is_editable else camera.fov_distance.item(),
        }
        return properties

wisp/renderer/gui/imgui/widget_cameras.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import imgui
from .widget_imgui import WidgetImgui
from wisp.framework import WispState
import pynvml
pynvml.nvmlInit()


class WidgetGPUStats(WidgetImgui):
    def __init__(self):
        super().__init__()
        self.text_color = (0.2, 1., 0.)  # RGB

    def paint(self, state: WispState, *args, **kwargs):

        imgui.text_colored(f"Frame time: {int(state.renderer.dt*1000)}ms / FPS: {int(state.renderer.fps)}", *self.text_color)
        imgui.text(f"Interactive Mode: {'on' if state.renderer.interactive_mode else 'off'}")
        imgui.separator()

        device = state.renderer.device
        if device == 'cpu':
            device_name = 'Running on CPU'
            imgui.text(f"{device_name}")
        else:
            device_name = torch.cuda.get_device_name(device=device)
            imgui.text(f"{device_name}, CUDA v{torch.version.cuda}")
            width, height = imgui.get_content_region_available()

            device_index = getattr(device, 'index', None)
            device_index = device_index if device_index is not None else torch.cuda.current_device()
            nvml_device_handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
            nvml_info = pynvml.nvmlDeviceGetMemoryInfo(nvml_device_handle)
            used_mem = nvml_info.used
            total_mem = nvml_info.total
            dec_str = '{:.2f}'
            imgui.text("Global GPU Mem: ")
            imgui.same_line()
            imgui.progress_bar(fraction=used_mem / total_mem, size=(width / 2, 0),
                               overlay=f"{dec_str.format(used_mem / 1024**3)} / "
                                       f"{dec_str.format(total_mem / 1024**3)} GB")

wisp/renderer/gui/imgui/widget_gpu_stats.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Dict, Type, Any
from wisp.framework import WispState
from collections import deque
from wisp.core.colors import colors_generator, white, black, dark_gray, gray

# All wisp block widgets supported by the interactive renderer are registered here.
# Registration can be quickly achieved via the @widget decorator.
_WIDGETS_REGISTRY: Dict[Type[Any], Type[WidgetImgui]] = dict()

# A generator to yield the next unused color is a cyclical order, useful i.e. for colorful titles of blocks
_colors_gen = colors_generator(skip_colors=[white, black, dark_gray, gray])
next_unused_color = lambda: next(_colors_gen)


class WidgetImgui(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def paint(self, state: WispState, *args, **kwargs):
        raise NotImplementedError('imgui widgets must implement the paint method.')


def widget(wisp_block: Type[Any]):
    """ A decorator that registers a gui widget to paint the contents of a given wisp block.
        By registering a widget, the gui system knows how to load this widget when it traverses the
        scene graph & properties and encounters the wisp_block type.

        Users adding new wisp blocks can directly register corresponding widgets using this decorator.
    """
    def _register_widget_fn(widget_class: Type[WidgetImgui]):
        _WIDGETS_REGISTRY[wisp_block] = widget_class
        return widget_class
    return _register_widget_fn


def _lookup_widget(wisp_block: Any) -> Type[WidgetImgui]:
    """ Searches the wisp_block class hierarchy for any matches with some registered widget.
        wisp_block can be a
    """
    # Create a queue of base classes and search if any of them has a registered widget.
    # Basically this is a BFS on the base classes, start with the concrete class type..
    if not isinstance(wisp_block, type):    # Is this a type or an instance? Fetch the type
        wisp_block = type(wisp_block)
    base_queue = deque([wisp_block])

    widget_type = None
    while widget_type is None and base_queue:
        block_type = base_queue.popleft()
        widget_type = _WIDGETS_REGISTRY.get(block_type)
        if widget_type is None:
            bases = block_type.__bases__
            if len(bases) > 0:
                base_queue.extend(bases)
    return widget_type


def get_widget(wisp_block: Any) -> WidgetImgui:
    """ Return a widget which matches the given wisp block.
    A wisp block can be of any type / subtype which was registered with @widget.
    The lookup logic will first look for a widget registered under the type of wisp_block, and if it cannot find it,
    it will start looking up the hierarchy.
    Note that multiple-inheritance may result in undefined behavior in case of more than one match - in such cases
    users are encouraged to register some widget to the concrete type.
    Most wisp interfaces should already be registered to some general widget, so concrete types are not required
    to register a dedicated widget.
    However, they may register a specialized widget returned to display other forms of information.
    For example:

    Args:
        wisp_block: any object whose type / base types were registered with @widget.

    Returns:
        (WidgetImgui) imgui widget which matches the given wisp_block type
    """
    block_type = type(wisp_block)
    widget_type = _WIDGETS_REGISTRY.get(block_type)
    if widget_type is None:
        widget_type = _lookup_widget(block_type)
        _WIDGETS_REGISTRY[block_type] = widget_type

    if widget_type is None:
        raise ValueError(f'Gui cannot find a widget for {widget_type}. Make sure to register some widget with @widget, '
                         f'or consider if {widget_type} should subclass WispModule.')
    widget_instance = widget_type()  # Component widgets are assumed to take no args during construction
    return widget_instance

wisp/renderer/gui/imgui/widget_imgui.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import imgui
from wisp.framework import WispState
from wisp.core.transforms import ObjectTransform
from .widget_imgui import WidgetImgui, widget


@widget(ObjectTransform)
class WidgetObjectTransform(WidgetImgui):

    def __init__(self):
        super().__init__()

    def paint(self, state: WispState, object_transform: ObjectTransform = None, *args, **kwargs):
        if object_transform is None:
            return

        button_width = 100
        if imgui.button("Reset", width=button_width):
            object_transform.reset()

        changed, values = imgui.slider_float3(
            "Translate",
            object_transform.tx, object_transform.ty, object_transform.tz,
            min_value=-15.0, max_value=15.0,
            format="%.2f",
            power=1.0
        )
        if changed:
            object_transform.tx = values[0]
            object_transform.ty = values[1]
            object_transform.tz = values[2]

        changed, values = imgui.slider_float3(
            "Rotate",
            object_transform.rx, object_transform.ry, object_transform.rz,
            min_value=-180.0, max_value=180.0,
            format="%.2f",
            power=1.0
        )
        if changed:
            object_transform.rx = values[0]
            object_transform.ry = values[1]
            object_transform.rz = values[2]

        changed, values = imgui.slider_float3(
            "Scale",
            object_transform.sx, object_transform.sy, object_transform.sz,
            min_value=-15.0, max_value=15.0,
            format="%.2f",
            power=1.0
        )
        if changed:
            object_transform.sx = values[0]
            object_transform.sy = values[1]
            object_transform.sz = values[2]

wisp/renderer/gui/imgui/widget_object_transform.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import numpy as np
import datetime
import imgui
from wisp.framework import WispState
from .widget_imgui import WidgetImgui


class WidgetOptimization(WidgetImgui):
    def __init__(self):
        super().__init__()
        self.active_text_color = (0.2, 1., 0.)         # RGB (green)
        self.paused_text_color = (1.0, 0.917, 0.164)   # RGB (yellow)
        self.stopped_text_color = (1.0, 0.0, 0.0)      # RGB (yellow)

    def paint(self, state: WispState, *args, **kwargs):
        expanded, _ = imgui.collapsing_header("Optimization", visible=True, flags=imgui.TREE_NODE_DEFAULT_OPEN)
        if not expanded:
            return

        is_training_stopped = not state.optimization.running or state.renderer.background_tasks_paused
        is_training_paused = state.renderer.background_tasks_paused

        if is_training_stopped:
            imgui.text_colored(f"Training stopped.", *self.stopped_text_color)
        elif is_training_paused:
            imgui.text_colored(f"Training paused.", *self.paused_text_color)
        else:
            imgui.text_colored(f"Training running.", *self.active_text_color)

        curr_iteration = state.optimization.iteration
        curr_epoch = state.optimization.epoch
        button_width = 150
        if not is_training_stopped:
            if imgui.button("Stop Training", width=button_width):
                state.optimization.running = False
                state.renderer.background_tasks_paused = True
        else:
            if curr_epoch == 1 and curr_iteration == 0:
                if imgui.button("Start Training", width=button_width):
                    state.optimization.running = True
                    state.renderer.background_tasks_paused = False
            else:
                if imgui.button("Resume Training", width=button_width):
                    state.optimization.running = True
                    state.renderer.background_tasks_paused = False

        imgui.same_line()
        elapsed_time_seconds = state.optimization.elapsed_time
        formatted_time = str(datetime.timedelta(seconds=elapsed_time_seconds))
        formatted_time = formatted_time.split(".")[0]   # Remove microseconds
        imgui.text(f"Elapsed Time: {formatted_time}")

        if is_training_stopped:
            state.renderer.target_fps = None
        else:
            _, checkbox_enabled = imgui.checkbox(label=f"Limit rendering FPS",
                                                 state=state.renderer.target_fps is not None)

            # Would be better to have a disabled widget instead of an invisible one,
            # but as of June 2022 the official pyimgui release doesn't expose push_item_flag + ITEM_DISABLED
            if checkbox_enabled:
                imgui.same_line()
                imgui.push_item_width(-30)
                target_fps = state.renderer.target_fps if state.renderer.target_fps is not None else 0
                changed, value = imgui.slider_int("##target_fps", value=target_fps, min_value=0, max_value=60)
                state.renderer.target_fps = value
                imgui.pop_item_width()
            else:
                state.renderer.target_fps = None

        if not is_training_stopped or curr_epoch > 0 or curr_iteration > 1:
            width, height = imgui.get_content_region_available()
            total_epochs = state.optimization.max_epochs
            if total_epochs > 0:
                epoch_progress = curr_epoch / total_epochs
                imgui.text("Epoch:")
                imgui.progress_bar(fraction=epoch_progress, size=(width, 0), overlay=f"{curr_epoch} / {total_epochs}")

            total_iteration = state.optimization.iterations_per_epoch
            if total_iteration > 0:
                iterations_progress = curr_iteration / total_iteration
            else:
                iterations_progress = 0
            imgui.text("Iteration:")
            imgui.progress_bar(fraction=iterations_progress, size=(width, 0),
                               overlay=f"{curr_iteration} / {total_iteration}")

            if len(state.optimization.losses) > 0:
                if imgui.tree_node("Losses", imgui.TREE_NODE_DEFAULT_OPEN):
                    for loss_name, loss_vals in state.optimization.losses.items():
                        formatted_loss_name = loss_name.replace('_', ' ').capitalize()
                        imgui.text(f"{formatted_loss_name}:")
                        imgui.plot_lines("##losses", np.array(loss_vals, dtype=np.float32),
                                         graph_size=(width - 30, 35))
                    imgui.tree_pop()

wisp/renderer/gui/imgui/widget_optimization.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from typing import Dict
import imgui
import numpy as np
from wisp.framework import WispState
from .widget_imgui import WidgetImgui


class WidgetPropertyEditor(WidgetImgui):

    running_idx = 0  # Used to distinguish multiple instances of property editor widget

    def __init__(self):
        super().__init__()

    def paint_row(self, title, value):
        imgui.text(title)
        imgui.next_column()
        imgui.separator()
        imgui.push_item_width(-1)
        if value is None:
            imgui.text('None')
        elif isinstance(value, list) or isinstance(value, np.ndarray):
            # Value is a sequence, concatenate entries as string
            if len(value) == 0:
                txt = 'None'
            else:
                txt = ', '.join([str(e) for e in value])
            imgui.text(txt)
        elif isinstance(value, float):          # Round and format float
            imgui.text("{:.3f}".format(value))
        elif callable(value):                   # Values may be editable widgets, in which case they're callable
            try:
                value()
            except Exception as e1:
                try:
                    imgui.text(type(value).__name__)    # Callable is probably not an editable widget, use typename
                except Exception as e2:
                    raise ValueError(f"Gui failed to paint property - '{title}: {value}'. "
                                     f"This error message can occur, i.e., if a wisp-module returns a public property of "
                                     f"some unsupported type.") from e1
        else:                                   # Other types are cast to str
            imgui.text(str(value))
        imgui.pop_item_width()

    def paint(self, state: WispState, properties: Dict[str, object] = None, *args, **kwargs):
        if properties is None:
            return
        try:
            imgui.begin_group()
            WidgetPropertyEditor.running_idx += 1
            running_idx = WidgetPropertyEditor.running_idx
            imgui.columns(2, f'properties_{running_idx}', border=True)

            # Draw properties table
            for row_idx, (title, value) in enumerate(properties.items()):
                self.paint_row(title, value)
                if row_idx < len(properties) - 1:
                    imgui.next_column()
        finally:
            imgui.columns(1)
            imgui.separator()
            WidgetPropertyEditor.running_idx -= 1
            imgui.end_group()

wisp/renderer/gui/imgui/widget_property_editor.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import imgui
from wisp.framework import WispState
from wisp.renderer.core.renderers import NeuralRadianceFieldPackedRenderer
from .widget_imgui import WidgetImgui, widget, get_widget
from .widget_property_editor import WidgetPropertyEditor


@widget(NeuralRadianceFieldPackedRenderer)
class WidgetNeuralRadianceFieldRenderer(WidgetImgui):

    marcher_types = ['ray', 'voxel']

    def __init__(self):
        super().__init__()
        self.properties_widget = WidgetPropertyEditor()

    def paint_tracer(self, state: WispState, renderer: NeuralRadianceFieldPackedRenderer):
        if imgui.tree_node("Tracer", imgui.TREE_NODE_DEFAULT_OPEN):
            MAX_SAMPLES = 128               # For general tracers
            MAX_SAMPLES_RAY_MODE = 2048      # For 'ray' sampling mode

            # TODO (operel): Update the ## ids below with a unique object name to avoid imgui bug
            def _num_samples_property():
                max_value = MAX_SAMPLES
                if renderer.raymarch_type == 'ray':
                    max_value = MAX_SAMPLES_RAY_MODE
                value = min(renderer.num_steps, max_value)
                changed, value = imgui.core.slider_int(f"##samples_per_ray", value=value,
                                                       min_value=2, max_value=max_value)
                if changed:
                    renderer.num_steps = value

            def _num_samples_interactive_property():
                max_value = MAX_SAMPLES
                if renderer.raymarch_type == 'ray':
                    max_value = MAX_SAMPLES_RAY_MODE
                value = min(renderer.num_steps_interactive, max_value)
                changed, value = imgui.core.slider_int(f"##samples_per_ray_interactive", value=value,
                                                       min_value=2, max_value=max_value)
                if changed:
                    renderer.num_steps_interactive = value

            def _batch_size_property():
                changed, value = imgui.core.slider_int("##batch_size", value=renderer.batch_size,
                                                       min_value=2**9, max_value=2**18)
                if changed:
                    renderer.batch_size = value

            def _marcher_type_property():
                selected_marcher_idx = self.marcher_types.index(renderer.raymarch_type)
                changed, selected_marcher_idx = imgui.combo("##marcher_type",
                                                            selected_marcher_idx, self.marcher_types)
                if changed:
                    new_marcher_mode = self.marcher_types[selected_marcher_idx]
                    if new_marcher_mode != 'ray':
                        renderer.num_steps = min(MAX_SAMPLES, renderer.num_steps)
                    renderer.raymarch_type = new_marcher_mode

            properties = {
                "Ray Samples (full-res)": _num_samples_property,                 # Samples per ray
                "Ray Samples (interact.)": _num_samples_interactive_property,    # Samples per ray
                "Batch Size (Rays)": _batch_size_property,
                "Marcher Type": _marcher_type_property,
                "Render Resolution (W x H)": f"{renderer.render_res_x} x {renderer.render_res_y}"
            }
            self.properties_widget.paint(state=state, properties=properties)
            imgui.tree_pop()

    def paint(self, state: WispState, renderer: NeuralRadianceFieldPackedRenderer = None, *args, **kwargs):
        if renderer is None:
            return
        self.paint_tracer(state, renderer)
        nef_widget = get_widget(renderer.nef)
        nef_widget.paint(state=state, module=renderer.nef)

wisp/renderer/gui/imgui/widget_radiance_pipeline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from wisp.framework import WispState
from wisp.renderer.core.api import RayTracedRenderer
from .widget_imgui import WidgetImgui, widget, get_widget
from .widget_property_editor import WidgetPropertyEditor


@widget(RayTracedRenderer)
class WidgetRayTracedRenderer(WidgetImgui):
    """ A default widget for ray tracing based bottom-level-renderers.
    If you added a new type of neural field and specialized a new type of RayTracedRenderer / BottomLevelRenderer,
    but not a widget, this default widget will be loaded as a fallback.
    The neural field information will be displayed, but no assumptions on the inference-time tracer can be made,
    therefore tracer fields will not be displayed.
    """

    def __init__(self):
        super().__init__()
        self.properties_widget = WidgetPropertyEditor()

    def paint(self, state: WispState, renderer: RayTracedRenderer = None, *args, **kwargs):
        if renderer is None:
            return
        nef_widget = get_widget(renderer.nef)
        nef_widget.paint(state=state, module=renderer.nef)

wisp/renderer/gui/imgui/widget_raytraced_pipeline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import imgui
from .widget_imgui import WidgetImgui, widget
from .widget_cameras import WidgetCameraProperties
from wisp.core.colors import black, white, dark_gray
from wisp.framework import WispState, InteractiveRendererState
from wisp.renderer.core.control import FirstPersonCameraMode, TrackballCameraMode, TurntableCameraMode


@widget(InteractiveRendererState)
class WidgetInteractiveVisualizerProperties(WidgetImgui):
    def __init__(self):
        super().__init__()
        self.camera_properties = WidgetCameraProperties(camera_id="User Camera")

        self.available_camera_modes = [FirstPersonCameraMode, TrackballCameraMode, TurntableCameraMode]
        self.available_camera_mode_names = [mode.name() for mode in self.available_camera_modes]

    def paint(self, state: WispState, *args, **kwargs):
        expanded, _ = imgui.collapsing_header("Renderer", visible=True, flags=imgui.TREE_NODE_DEFAULT_OPEN)
        if expanded:
            imgui.text(f"Canvas Size (W,H): "
                       f"{state.renderer.canvas_width} x {state.renderer.canvas_height}")

            bg_btn_width = bg_btn_height = 15
            imgui.text("Background Color: ")
            imgui.same_line()
            bg_buttons = {
                'Black': black,
                'Dark Gray': dark_gray,
                'White': white,
            }
            for btn_caption, btn_color in bg_buttons.items():
                imgui.same_line()
                if imgui.color_button(btn_caption, *btn_color, 1.0, 0, bg_btn_width, bg_btn_height):
                    state.renderer.clear_color_value = btn_color

            available_canvas_channels = state.renderer.available_canvas_channels
            selected_canvas_channel = state.renderer.selected_canvas_channel
            if selected_canvas_channel is not None and available_canvas_channels is not None and \
                    selected_canvas_channel in available_canvas_channels:
                curr_channel_idx = available_canvas_channels.index(selected_canvas_channel)
                clicked, new_channel_idx = imgui.combo("Channel", curr_channel_idx, available_canvas_channels)
                if new_channel_idx != curr_channel_idx:
                    state.renderer.selected_canvas_channel = available_canvas_channels[new_channel_idx]

            cam_mode = state.renderer.cam_controller
            if cam_mode is not None:
                curr_mode_idx = self.available_camera_mode_names.index(cam_mode.name())
                clicked, new_mode_idx = imgui.combo("Controller", curr_mode_idx, self.available_camera_mode_names)
                if new_mode_idx != curr_mode_idx:  # Update state with new mode if changed
                    new_mode_cls = self.available_camera_modes[new_mode_idx]
                    state.renderer.cam_controller = new_mode_cls

            if imgui.tree_node("User Camera", imgui.TREE_NODE_DEFAULT_OPEN):
                self.camera_properties.paint(state, camera=state.renderer.selected_camera)
                imgui.tree_pop()

wisp/renderer/gui/imgui/widget_renderer_properties.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from collections import defaultdict
import imgui
from wisp.core import WispModule
from wisp.framework import WispState
from .widget_imgui import WidgetImgui, get_widget, next_unused_color
from .widget_cameras import WidgetCameraProperties
from .widget_object_transform import WidgetObjectTransform
from wisp.renderer.core.api import request_redraw


class WidgetSceneGraph(WidgetImgui):

    def __init__(self):
        super().__init__()
        self.transform_widget = WidgetObjectTransform()
        self._object_colors = defaultdict(next_unused_color)  # Generates a new color on every query "miss"

    def get_bl_renderer_widget(self, object_type):
        return get_widget(object_type)

    def get_object_color(self, key):
        """ Generates a color per type and caches it for future invocations """
        return self._object_colors[key]

    def get_object_title(self, object_type):
        if isinstance(object_type, WispModule):
            return object_type.name()           # All wisp modules should carry meaningful names
        else:
            return type(object_type).__name__   # Not a wisp module, default to class name

    @staticmethod
    def paint_object_checkbox(state, obj_id):
        visible_objects = state.graph.visible_objects
        is_checked = visible_objects.get(obj_id, False)
        visibility_toggled, is_checked = imgui.checkbox(f"##{obj_id}", is_checked)
        visible_objects[obj_id] = is_checked
        if visibility_toggled:
            request_redraw(state)

    @staticmethod
    def paint_all_objects_checkbox(state):
        visible_objects = state.graph.visible_objects
        is_checked = any([visible_objects.get(obj_id, False) for obj_id in state.graph.bl_renderers.keys()])
        visibility_toggled, is_checked = imgui.checkbox(f"##sg_all_bl_renderers", is_checked)
        if visibility_toggled:
            for obj_id in state.graph.bl_renderers.keys():
                visible_objects[obj_id] = is_checked
            request_redraw(state)

    @staticmethod
    def paint_all_cameras_checkbox(state):
        visible_objects = state.graph.visible_objects
        is_checked = any([visible_objects.get(obj_id, False) for obj_id in state.graph.cameras.keys()])
        visibility_toggled, is_checked = imgui.checkbox(f"##sg_all_cameras", is_checked)
        if visibility_toggled:
            for cam_id in state.graph.cameras.keys():
                visible_objects[cam_id] = is_checked
            request_redraw(state)

    def paint(self, state: WispState, *args, **kwargs):
        expanded, _ = imgui.collapsing_header("Scene Objects", visible=True, flags=imgui.TREE_NODE_DEFAULT_OPEN)
        if expanded:
            bl_renderers = state.graph.bl_renderers

            if len(bl_renderers) > 0:
                self.paint_all_objects_checkbox(state)
                imgui.same_line()
                if imgui.tree_node("Objects", imgui.TREE_NODE_DEFAULT_OPEN):
                    for obj_id, obj in bl_renderers.items():
                        if obj.status != 'loaded':
                            continue
                        obj_type = type(obj.renderer).__name__
                        obj_color = self.get_object_color(obj_type)

                        self.paint_object_checkbox(state, obj_id)
                        imgui.same_line()
                        if imgui.tree_node(obj_id, imgui.TREE_NODE_DEFAULT_OPEN):
                            imgui.text(f"Type:")
                            imgui.same_line()
                            obj_title = self.get_object_title(obj.renderer)
                            imgui.text_colored(f"{obj_title}", *obj_color)

                            object_transform = obj.transform
                            if object_transform is not None:
                                if imgui.tree_node("Transform", imgui.TREE_NODE_DEFAULT_OPEN):
                                    self.transform_widget.paint(state, object_transform)
                                imgui.tree_pop()

                            if imgui.tree_node("Properties", imgui.TREE_NODE_DEFAULT_OPEN):
                                features_structure = obj.renderer.features_structure()
                                if features_structure is not None:
                                    imgui.text(f"Feature Structure:")
                                    imgui.same_line()
                                    obj_color = self.get_object_color(features_structure)
                                    imgui.text_colored(features_structure, *obj_color)
                                acceleration_structure = obj.renderer.acceleration_structure()
                                if acceleration_structure is not None:
                                    imgui.text(f"Acceleration Structure:")
                                    imgui.same_line()
                                    imgui.text(acceleration_structure)

                                bl_renderer_widget = self.get_bl_renderer_widget(obj.renderer)
                                if bl_renderer_widget is not None:
                                    bl_renderer_widget.paint(state, obj.renderer)
                                imgui.tree_pop()

                            if imgui.tree_node("Layers"):
                                obj_data_layers = obj.data_layers
                                toggled_obj_data_layers = obj.toggled_data_layers
                                for layer in obj_data_layers.keys():
                                    is_prev_selected = toggled_obj_data_layers[layer]
                                    _, is_selected = imgui.checkbox(f"{layer}", is_prev_selected)
                                    toggled_obj_data_layers[layer] = is_selected
                                    if is_prev_selected != is_selected:
                                        request_redraw(state)
                                imgui.tree_pop()
                            imgui.tree_pop()
                    imgui.tree_pop()

            if len(state.graph.cameras) > 0:
                self.paint_all_cameras_checkbox(state)
                imgui.same_line()
                if imgui.tree_node("Cameras", imgui.TREE_NODE_DEFAULT_OPEN):
                    for cam_id, cam in state.graph.cameras.items():
                        self.paint_object_checkbox(state, cam_id)
                        imgui.same_line()
                        if imgui.tree_node(f"Camera {cam_id}"):
                            obj_type = type(cam).__name__
                            obj_color = self.get_object_color(obj_type)
                            imgui.text(f"Type:")
                            imgui.same_line()
                            imgui.text_colored(f"{obj_type}", *obj_color)
                            cam_widget = WidgetCameraProperties(camera_id=cam_id)
                            cam_widget.paint(state, cam)
                            imgui.tree_pop()
                    imgui.tree_pop()
                # TODO (operel): Add proper camera widget caching

wisp/renderer/gui/imgui/widget_scene_graph.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import imgui
from wisp.framework import WispState
from wisp.renderer.core.renderers import NeuralSDFPackedRenderer
from .widget_imgui import WidgetImgui, get_widget, widget
from .widget_property_editor import WidgetPropertyEditor


@widget(NeuralSDFPackedRenderer)
class WidgetNeuralSDFRenderer(WidgetImgui):

    def __init__(self):
        super().__init__()
        self.properties_widget = WidgetPropertyEditor()

    def paint_tracer(self, state: WispState, renderer: NeuralSDFPackedRenderer):
        if imgui.tree_node("Tracer", imgui.TREE_NODE_DEFAULT_OPEN):
            def _num_samples_property():
                changed, value = imgui.core.slider_int(f"##steps_per_ray", value=renderer.num_steps,
                                                       min_value=1, max_value=256)
                renderer.num_steps = value

            def _batch_size_property():
                changed, value = imgui.core.slider_int("##batch_size", value=renderer.batch_size,
                                                       min_value=2 ** 9, max_value=2 ** 15)
                renderer.batch_size = value

            properties = {
                "Steps per ray": _num_samples_property,
                # "Batch Size (Rays)": _batch_size_property,
                "Render Resolution (W x H)": f"{renderer.render_res_x} x {renderer.render_res_y}",
            }
            self.properties_widget.paint(state=state, properties=properties)
            imgui.tree_pop()

    def paint(self, state: WispState, renderer: NeuralSDFPackedRenderer = None, *args, **kwargs):
        if renderer is None:
            return
        self.paint_tracer(state, renderer)
        nef_widget = get_widget(renderer.nef)
        nef_widget.paint(state=state, module=renderer.nef)

wisp/renderer/gui/imgui/widget_sdf_pipeline_renderer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import imgui
from wisp.core import WispModule
from wisp.accelstructs import BaseAS
from wisp.models.grids import BLASGrid
from wisp.models.nefs import BaseNeuralField
from wisp.models.decoders import BasicDecoder
from wisp.tracers import BaseTracer
from wisp.models.embedders import PositionalEmbedder
from wisp.framework import WispState
from .widget_imgui import WidgetImgui, widget, get_widget
from .widget_property_editor import WidgetPropertyEditor


@widget(WispModule)
class WidgetWispModule(WidgetImgui):

    def __init__(self):
        super().__init__()
        self.properties_widget = WidgetPropertyEditor()

    def get_type_title(self, module: WispModule):
        if isinstance(module, BaseAS):
            return "Acceleration Structure"
        elif isinstance(module, BLASGrid):
            return "Grid"
        elif isinstance(module, BaseNeuralField):
            return "Neural Field"
        elif isinstance(module, BasicDecoder):          # TODO (operel): interface decoders
            return "Decoder"
        elif isinstance(module, PositionalEmbedder):    # TODO (operel): interface embedders
            return "Embedder"
        elif isinstance(module, BaseTracer):
            return "Tracer"
        else:
            return "Module"

    def paint(self, state: WispState, module: WispModule = None, field_name=None, *args, **kwargs):
        if module is None:
            return

        type_title = self.get_type_title(module)
        section_title = field_name if field_name is not None else type_title
        if imgui.tree_node(section_title):
            imgui.text(f"Type: {module.name()}")

            properties = module.public_properties()
            recursive_properties = {k: p for k, p in properties.items() if isinstance(p, WispModule)}
            for recursive_field_name, recursive_property in recursive_properties.items():
                child_widget = get_widget(recursive_property)
                child_widget.paint(state=state, module=recursive_property, field_name=recursive_field_name)

            table_properties = {k: p for k, p in properties.items() if not isinstance(p, WispModule)}
            self.properties_widget.paint(state=state, properties=table_properties)
            imgui.tree_pop()

wisp/renderer/gui/imgui/widget_wisp_module.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .base_tracer import BaseTracer
from .packed_rf_tracer import PackedRFTracer
from .packed_sdf_tracer import PackedSDFTracer
from .packed_spc_tracer import PackedSPCTracer

wisp/tracers/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from abc import abstractmethod, ABC
from typing import Dict, Any
import inspect
import torch
import torch.nn as nn
from wisp.core import Rays
from wisp.core import WispModule

class BaseTracer(WispModule, ABC):
    """Base class for all tracers within Wisp.
    Tracers drive the mapping process which takes an input "Neural Field", and outputs a RenderBuffer of pixels.
    Different tracers may employ different algorithms for querying points, tracing and marching rays through the
    neural field.
    A common paradigm for tracers to employ is as follows:
    1. Take input in the form of rays
    2. Generate samples by tracing / marching rays, or querying coordinates over the neural field.
       Possibly make use of the neural field spatial structure for high performance.
    3. Invoke neural field's methods to decode sample features into actual channel values, such as color, density,
       signed distance, and so forth.
    4. Aggregate the sample values to decide on the final pixel value.
       The exact output may depend on the requested channel type, blending mode or other parameters.
    Wisp tracers are therefore flexible, and designed to be compatible with specific neural fields,
    depending on the forward functions they support and internal grid structures they use.
    Tracers are generally expected to be differentiable (e.g. they're part of the training loop),
    though non-differentiable tracers are also allowed.
    """

    def __init__(self):
        """Initializes the tracer class and sets the default arguments for trace.
        This should be overrided and called if you want to pass custom defaults into the renderer.
        If overridden, it should keep the arguments to `self.trace` in `self.` class variables.
        Then, if these variables exist and no function arguments are passed into forward,
        it will override them as the default.
        """
        super().__init__()

    @abstractmethod
    def get_supported_channels(self):
        """Returns the set of channel names this tracer may output.

        Implement the function to return the supported channels, e.g.       
        return set(["depth", "rgb"])

        Returns:
            (set): Set of channel strings.
        """
        pass

    @abstractmethod
    def get_required_nef_channels(self):
        """Returns the channels required by neural fields to be compatible with this tracer.
        
        Implement the function to return the required channels, e.g.
        return set(["rgb", "density"])

        Returns:
            (set): Set of channel strings.
        """
        pass

    @abstractmethod
    def trace(self, nef, rays, channels, extra_channels, *args, **kwargs):
        """Apply the forward map on the nef. 

        Tracers are required to implement this function, which commonly follows these paradigm:
        1. Take input in the form of rays
        2. Generate samples by tracing / marching rays, or querying coordinates over the neural field.
           Possibly make use of the neural field spatial structure for high performance.
        3. Invoke neural field's methods to decode sample features into actual channel values, such as color, density,
           signed distance, and so forth.
        4. Aggregate the sample values to decide on the final pixel value.
           The exact output may depend on the requested channel type, blending mode or other parameters.
        
        Args:
            nef (nn.Module): A neural field that uses a grid class.
            rays (Rays): Pack of rays to trace through the neural field.
            channels (set): The set of requested channels. The trace method can return channels that 
                            were not requested since those channels often had to be computed anyways.
            extra_channels (set): Requested extra channels, which are not first class channels supported by
                the tracer but will still be able to handle with some fallback options.

        Returns:
            (wisp.RenderBuffer): A dataclass which holds the output buffers from the tracer.
        """
        pass

    def forward(self, nef, rays: Rays, channels=None, **kwargs):
        """Queries the tracer with channels.

        Args:
            nef (BaseNeuralField): Neural field to be traced. The nef will be queried for decoded sample values.
            rays (Rays): Pack of rays to trace through the neural field.
            channels (str or list of str or set of str): Requested channel names.
            This list should include at least all channels in tracer.get_supported_channels(),
            and may include extra channels in addition.
            kwargs: Any keyword argument passed in will be passed into the respective forward functions.

        Returns:
            (wisp.RenderBuffer): A dataclass which holds the output buffers from the tracer.
        """
        nef_channels = nef.get_supported_channels()
        unsupported_inputs = self.get_required_nef_channels() - nef_channels
        if unsupported_inputs:
            raise Exception(f"The neural field class {type(nef)} does not output the required channels {unsupported_inputs}.")

        if channels is None:
            requested_channels = self.get_supported_channels()
        elif isinstance(channels, str):
            requested_channels = set([channels])
        else:
            requested_channels = set(channels)
        extra_channels = requested_channels - self.get_supported_channels()
        unsupported_outputs = extra_channels - nef_channels
        if unsupported_outputs:
            raise Exception(f"Channels {unsupported_outputs} are not supported in the tracer {type(self)} or neural field {type(nef)}.")
    
        if extra_channels is None:
            requested_extra_channels = set()
        elif isinstance(extra_channels, str):
            requested_extra_channels = set([extra_channels])
        else:
            requested_extra_channels = set(extra_channels)

        argspec = inspect.getfullargspec(self.trace)

        # Skip self, nef, rays, channel, extra_channels
        required_args = argspec.args[:-len(argspec.defaults)][5:]   # TODO (operel): this is brittle
        optional_args = argspec.args[-len(argspec.defaults):]
        
        input_args = {}
        for _arg in required_args:
            # TODO(ttakiakwa): This doesn't actually format the string, fix :) 
            if _arg not in kwargs:
                raise Exception(f"Argument {_arg} not found as input to in {type(self)}.trace()")
            input_args[_arg] = kwargs[_arg]
        for _arg in optional_args:
            if _arg in kwargs:
                # By default, the function args will take priority
                input_args[_arg] = kwargs[_arg]
            else:
                # Check if default_args are set, and use them if they are.
                default_arg = getattr(self, _arg, None)
                if default_arg is not None:
                    input_args[_arg] = default_arg
        with torch.cuda.nvtx.range("Tracer.trace"):
            rb = self.trace(nef, rays, requested_channels, requested_extra_channels, **input_args)
        return rb

    def public_properties(self) -> Dict[str, Any]:
        """ Wisp modules expose their public properties in a dictionary.
        The purpose of this method is to give an easy table of outwards facing attributes,
        for the purpose of logging, gui apps, etc.
        """
        return dict()

wisp/tracers/base_tracer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import torch.nn as nn
import kaolin.render.spc as spc_render
from wisp.core import RenderBuffer
from wisp.tracers import BaseTracer


class PackedRFTracer(BaseTracer):
    """Tracer class for sparse (packed) radiance fields.
    - Packed: each ray yields a custom number of samples, which are therefore packed in a flat form within a tensor,
     see: https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#packed
    - RF: Radiance Field
    PackedRFTracer is differentiable, and can be employed within training loops.

    This tracer class expects the neural field to expose a BLASGrid: a Bottom-Level-Acceleration-Structure Grid,
    i.e. a grid that inherits the BLASGrid class for both a feature structure and an occupancy acceleration structure).
    """
    def __init__(self, raymarch_type='voxel', num_steps=128, step_size=1.0, bg_color='white'):
        """Set the default trace() arguments.

        Args:
            raymarch_type (str): Sample generation strategy to use for raymarch.
                'voxel' - intersects the rays with the acceleration structure cells.
                    Then among the intersected cells, each cell is sampled `num_steps` times.
                'ray' - samples `num_steps` along each ray, and then filters out samples which falls outside of occupied
                    cells of the acceleration structure.
            num_steps (int): The number of steps to use for the sampling. The meaning of this parameter changes
                depending on `raymarch_type`:
                'voxel' - each acceleration structure cell which intersects a ray is sampled `num_steps` times.
                'ray' - number of samples generated per ray, before culling away samples which don't fall
                    within occupied cells.
                The exact number of samples generated, therefore, depends on this parameter but also the occupancy
                status of the acceleration structure.
            step_size (float): The step size between samples. Currently unused, but will be used for a new
                               sampling method in the future.
            bg_color (str): The background color to use.
        """
        super().__init__()
        self.raymarch_type = raymarch_type
        self.num_steps = num_steps
        self.step_size = step_size
        self.bg_color = bg_color

    def get_supported_channels(self):
        """Returns the set of channel names this tracer may output.
        
        Returns:
            (set): Set of channel strings.
        """
        return {"depth", "hit", "rgb", "alpha"}

    def get_required_nef_channels(self):
        """Returns the channels required by neural fields to be compatible with this tracer.
        
        Returns:
            (set): Set of channel strings.
        """
        return {"rgb", "density"}

    def trace(self, nef, rays, channels, extra_channels,
              lod_idx=None, raymarch_type='voxel', num_steps=64, step_size=1.0, bg_color='white'):
        """Trace the rays against the neural field.

        Args:
            nef (nn.Module): A neural field that uses a grid class.
            channels (set): The set of requested channels. The trace method can return channels that 
                            were not requested since those channels often had to be computed anyways.
            extra_channels (set): If there are any extra channels requested, this tracer will by default
                                  perform volumetric integration on those channels.
            rays (wisp.core.Rays): Ray origins and directions of shape [N, 3]
            lod_idx (int): LOD index to render at. 
            raymarch_type (str): The type of raymarching algorithm to use. Currently we support:
                                 voxel: Finds num_steps # of samples per intersected voxel
                                 ray: Finds num_steps # of samples per ray, and filters them by intersected samples
            num_steps (int): The number of steps to use for the sampling.
            step_size (float): The step size between samples. Currently unused, but will be used for a new
                               sampling method in the future.
            bg_color (str): The background color to use. TODO(ttakikawa): Might be able to simplify / remove

        Returns:
            (wisp.RenderBuffer): A dataclass which holds the output buffers from the render.
        """
        #TODO(ttakikawa): Use a more robust method
        assert nef.grid is not None and "this tracer requires a grid"

        N = rays.origins.shape[0]
        
        if "depth" in channels:
            depth = torch.zeros(N, 1, device=rays.origins.device)
        else: 
            depth = None
        
        if bg_color == 'white':
            rgb = torch.ones(N, 3, device=rays.origins.device)
        else:
            rgb = torch.zeros(N, 3, device=rays.origins.device)
        hit = torch.zeros(N, device=rays.origins.device, dtype=torch.bool)
        out_alpha = torch.zeros(N, 1, device=rays.origins.device)

        if lod_idx is None:
            lod_idx = nef.grid.num_lods - 1

        # By default, PackedRFTracer will attempt to use the highest level of detail for the ray sampling.
        # This however may not actually do anything; the ray sampling behaviours are often single-LOD
        # and is governed by however the underlying feature grid class uses the BLAS to implement the sampling.
        raymarch_results = nef.grid.raymarch(rays,
                                             level=nef.grid.active_lods[lod_idx],
                                             num_samples=num_steps,
                                             raymarch_type=raymarch_type)
        ridx = raymarch_results.ridx
        samples = raymarch_results.samples
        deltas = raymarch_results.deltas
        boundary = raymarch_results.boundary
        depths = raymarch_results.depth_samples

        # Get the indices of the ray tensor which correspond to hits
        ridx_hit = ridx[boundary]
        # Compute the color and density for each ray and their samples
        hit_ray_d = rays.dirs.index_select(0, ridx)

        # Compute the color and density for each ray and their samples
        num_samples = samples.shape[0]
        color, density = nef(coords=samples, ray_d=hit_ray_d, lod_idx=lod_idx, channels=["rgb", "density"])
        density = density.reshape(num_samples, 1)    # Protect against squeezed return shape
        del ridx

        # Compute optical thickness
        tau = density * deltas
        del density, deltas
        ray_colors, transmittance = spc_render.exponential_integration(color, tau, boundary, exclusive=True)

        if "depth" in channels:
            ray_depth = spc_render.sum_reduce(depths.reshape(num_samples, 1) * transmittance, boundary)
            depth[ridx_hit, :] = ray_depth

        alpha = spc_render.sum_reduce(transmittance, boundary)
        out_alpha[ridx_hit] = alpha
        hit[ridx_hit] = alpha[...,0] > 0.0

        # Populate the background
        if bg_color == 'white':
            color = (1.0-alpha) + ray_colors
        else:
            color = alpha * ray_colors
        rgb[ridx_hit] = color

        extra_outputs = {}
        for channel in extra_channels:
            feats = nef(coords=samples,
                        ray_d=hit_ray_d,
                        lod_idx=lod_idx,
                        channels=channel)
            num_channels = feats.shape[-1]
            ray_feats, transmittance = spc_render.exponential_integration(
                feats.view(num_samples, num_channels), tau, boundary, exclusive=True
            )
            composited_feats = alpha * ray_feats
            out_feats = torch.zeros(N, num_channels, device=feats.device)
            out_feats[ridx_hit] = composited_feats
            extra_outputs[channel] = out_feats

        return RenderBuffer(depth=depth, hit=hit, rgb=rgb, alpha=out_alpha, **extra_outputs)


wisp/tracers/packed_rf_tracer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import torch
import torch.nn.functional as F
import torch.nn as nn
import kaolin.render.spc as spc_render
from wisp.core import RenderBuffer
from wisp.ops.differential import finitediff_gradient
from wisp.ops.geometric import find_depth_bound
from wisp.tracers import BaseTracer

class PackedSDFTracer(BaseTracer):
    """Tracer class for sparse SDFs.

    - Packed: each ray yields a custom number of samples, which are therefore packed in a flat form within a tensor,
     see: https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#packed
    - SDF: Signed Distance Function
    PackedSDFTracer is non-differentiable, and follows the sphere-tracer implementation of
    Neural Geometric Level of Detail (Takikawa et al. 2021).

    This tracer class expects the neural field to expose a BLASGrid: a Bottom-Level-Acceleration-Structure Grid,
    i.e. a grid that inherits the BLASGrid class for both a feature structure and an occupancy acceleration structure).
    """
    def __init__(self, num_steps=128, step_size=1.0, min_dis=0.0003):
        """Set the default trace() arguments. """
        super().__init__()
        self.num_steps = num_steps
        self.step_size = step_size
        self.min_dis = min_dis

    def get_supported_channels(self):
        """Returns the set of channel names this tracer may output.
        
        Returns:
            (set): Set of channel strings.
        """
        return {"depth", "normal", "xyz", "hit", "rgb", "alpha"}

    def get_required_nef_channels(self):
        """Returns the channels required by neural fields to be compatible with this tracer.
        
        Returns:
            (set): Set of channel strings.
        """
        return {"sdf"}

    def trace(self, nef, rays, channels, extra_channels, lod_idx=None, num_steps=64,
              step_size=1.0, min_dis=1e-4):
        """Trace the rays against the neural field.

        Args:
            nef (nn.Module): A neural field that uses a grid class.
            rays (wisp.core.Rays): Ray origins and directions of shape [N, 3]
            channels (set): The set of requested channels. The trace method can return channels that
                            were not requested since those channels often had to be computed anyways.
            extra_channels (set): If there are any extra channels requested, this tracer will by default
                                  query those extra channels at surface intersection points.
            lod_idx (int): LOD index to render at.
            num_steps (int): The number of steps to use for sphere tracing.
            step_size (float): The multiplier for the sphere tracing steps. 
                               Use a value <1.0 for conservative tracing.
            min_dis (float): The termination distance for sphere tracing.

        Returns:
            (wisp.RenderBuffer): A dataclass which holds the output buffers from the render.
        """
        #TODO(ttakikawa): Use a more robust method
        assert nef.grid is not None and "this tracer requires a grid"
        
        if lod_idx is None:
            lod_idx = nef.grid.num_lods - 1

        invres = 1.0

        # Trace SPC
        raytrace_results = nef.grid.raytrace(rays, nef.grid.active_lods[lod_idx], with_exit=True)
        ridx = raytrace_results.ridx
        pidx = raytrace_results.pidx
        depth = raytrace_results.depth

        depth[...,0:1] += 1e-5

        first_hit = spc_render.mark_pack_boundaries(ridx)
        curr_idxes = torch.nonzero(first_hit)[...,0].int()

        first_ridx = ridx[first_hit].long()
        nug_o = rays.origins[first_ridx]
        nug_d = rays.dirs[first_ridx]

        mask = torch.ones([first_ridx.shape[0]], device=nug_o.device).bool()
        hit = torch.zeros_like(mask).bool()

        t = depth[first_hit][...,0:1]
        x = torch.addcmul(nug_o, nug_d, t)
        dist = torch.zeros_like(t)

        curr_pidx = pidx[first_hit].long()
        
        # Doing things with where is not super efficient, but we have to make do with what we have...
        with torch.no_grad():

            # Calculate SDF for current set of query points
            sdf = nef(coords=x[mask], lod_idx=lod_idx, pidx=curr_pidx[mask], channels="sdf") * invres * step_size
            dist[mask] = sdf.to(dist.dtype)
            dist[~mask] = 20
            dist_prev = dist.clone()

            for i in range(num_steps):
                # Two-stage Ray Marching
                
                # Step 1: Use SDF to march
                t += dist
                x = torch.where(mask.view(mask.shape[0], 1), torch.addcmul(nug_o, nug_d, t), x)
                hit = torch.where(mask, torch.abs(dist)[...,0] < min_dis * invres, hit)
                hit |= torch.where(mask, 
                                   torch.abs(dist+dist_prev)[...,0] * 0.5 < (min_dis*5) * invres, hit)
                mask = torch.where(mask, (t < rays.dist_max)[...,0], mask)
                mask &= ~hit
                if not mask.any():
                    break
                dist_prev = torch.where(mask.view(mask.shape[0], 1), dist, dist_prev)
                
                # Step 2: Use AABBs to march
                next_idxes = find_depth_bound(t, depth, first_hit, curr_idxes=curr_idxes)
                mask &= (next_idxes != -1)
                aabb_mask = (next_idxes != curr_idxes)
                curr_idxes = torch.where(mask, next_idxes, curr_idxes)

                t = torch.where((mask & aabb_mask).view(mask.shape[0], 1), depth[curr_idxes.long(), 0:1], t)
                x = torch.where(mask.view(mask.shape[0], 1), torch.addcmul(nug_o, nug_d, t), x)
                
                curr_pidx = torch.where(mask, pidx[curr_idxes.long()].long(), curr_pidx)
                if not mask.any():
                    break
                sdf = nef(coords=x[mask], lod_idx=lod_idx, pidx=curr_pidx[mask], channels="sdf") * invres * step_size
                dist[mask] = sdf.to(dist.dtype)
    
        x_buffer = torch.zeros_like(rays.origins)
        depth_buffer = torch.zeros_like(rays.origins[...,0:1])
        hit_buffer = torch.zeros_like(rays.origins[...,0]).bool()
        normal_buffer = torch.zeros_like(rays.origins)
        rgb_buffer = torch.zeros(*rays.origins.shape[:-1], 3, device=rays.origins.device)
        alpha_buffer = torch.zeros(*rays.origins.shape[:-1], 1, device=rays.origins.device)
        hit_buffer[first_ridx] = hit
        
        extra_outputs = {}
        for channel in extra_channels:
            feats = nef(coords=x[hit], lod_idx=lod_idx, channels=channel)
            extra_buffer = torch.zeros(*rays.origins.shape[:-1], feats.shape[-1], device=feats.device)
            extra_buffer[hit_buffer] = feats.to(extra_buffer.dtype)

        x_buffer[hit_buffer] = x[hit]
        depth_buffer[hit_buffer] = t[hit]
        
        # TODO(ttakikawa): RGB channel should _not_ be the normals. This should be only the case
        # if some shader object defines this to be the case.
        if "rgb" in channels or "normal" in channels:
            grad = finitediff_gradient(x[hit], nef.get_forward_function("sdf"))
            normal_buffer[hit_buffer] = F.normalize(grad, p=2, dim=-1, eps=1e-5)
            rgb_buffer[..., :3] = (normal_buffer + 1.0) / 2.0
        
        alpha_buffer[hit_buffer] = 1.0
        return RenderBuffer(xyz=x_buffer, depth=depth_buffer, hit=hit_buffer, normal=normal_buffer,
                            rgb=rgb_buffer, alpha=alpha_buffer, **extra_outputs)

wisp/tracers/packed_sdf_tracer.py
Sharath-girish/Shacira
import torch
import kaolin.render.spc as spc_render
from wisp.tracers import BaseTracer
from wisp.core import RenderBuffer


class PackedSPCTracer(BaseTracer):
    """Tracer class for sparse point clouds (packed rays).
    The logic of this tracer is straightforward and does not involve any neural operations:
    rays are intersected against the SPC points (cell centers).
    Each ray returns the color of the intersected cell, if such exists.

    See: https://github.com/NVIDIAGameWorks/kaolin-wisp/tree/main/examples/spc_browser
    See also: https://kaolin.readthedocs.io/en/latest/notes/spc_summary.html#spc
    """

    def __init__(self):
        """Set the default trace() arguments. """
        super().__init__()

    def get_supported_channels(self):
        """Returns the set of channel names this tracer may output.

        Returns:
            (set): Set of channel strings.
        """
        return {"depth", "hit", "rgb", "alpha"}

    def get_required_nef_channels(self):
        """Returns the channels required by neural fields to be compatible with this tracer.

        Returns:
            (set): Set of channel strings.
        """
        return {"rgb"}

    def trace(self, nef, rays, channels, extra_channels, lod_idx=None):
        """Trace the rays against the neural field.

        Args:
            nef (nn.Module): A neural field that uses a grid class.
            rays (wisp.core.Rays): Ray origins and directions of shape [N, 3]
            channels (set): The set of requested channels. The trace method can return channels that
                            were not requested since those channels often had to be computed anyways.
            lod_idx (int): LOD index to render at.

        Returns:
            (wisp.RenderBuffer): A dataclass which holds the output buffers from the render.
        """
        N = rays.origins.shape[0]

        # By default, SPCRFTracer will use the highest level of detail for the ray sampling.
        if lod_idx is None:
            lod_idx = nef.grid.blas.max_level

        raytrace_results = nef.grid.blas.raytrace(rays, lod_idx, with_exit=False)
        ridx = raytrace_results.ridx
        pidx = raytrace_results.pidx
        depths = raytrace_results.depth


        # Get the indices of the ray tensor which correspond to hits
        first_hits_mask = spc_render.mark_pack_boundaries(ridx)
        first_hits_point = pidx[first_hits_mask]
        first_hits_ray = ridx[first_hits_mask]
        first_hits_depth = depths[first_hits_mask]

        # Get the color for each ray
        color = nef(ridx_hit=first_hits_point.long(), channels="rgb")

        del ridx, pidx, rays

        # Fetch colors and depth for closest hits
        ray_colors = color.squeeze(1)
        ray_depth = first_hits_depth
        depth = torch.zeros(N, 1, device=ray_depth.device)
        depth[first_hits_ray.long(), :] = ray_depth
        alpha = torch.ones([color.shape[0], 1], device=color.device)
        hit = torch.zeros(N, device=color.device).bool()

        # Populate the background
        rgb = torch.zeros(N, 3, device=color.device)
        out_alpha = torch.zeros(N, 1, device=color.device)
        color = alpha * ray_colors

        hit[first_hits_ray.long()] = alpha[..., 0] > 0.0
        rgb[first_hits_ray.long(), :3] = color
        out_alpha[first_hits_ray.long()] = alpha

        return RenderBuffer(depth=depth, hit=hit, rgb=rgb, alpha=out_alpha)

wisp/tracers/packed_spc_tracer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .base_trainer import BaseTrainer, log_metric_to_wandb, log_images_to_wandb
from .multiview_trainer import MultiviewTrainer
from .image_trainer import ImageTrainer
from .sdf_trainer import SDFTrainer

wisp/trainers/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.  #
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import os
import time
import logging as log
from datetime import datetime
from abc import ABC, abstractmethod
import torch
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
from wisp.offline_renderer import OfflineRenderer
from wisp.framework import WispState, BottomLevelRendererState
from wisp.datasets import WispDataset, default_collate
from wisp.renderer.core.api import add_to_scene_graph
from wisp.models.grids import LatentGrid, HashGrid, CodebookOctreeGrid
from wisp.models.latent_decoders import MultiLatentDecoder, LatentDecoder
from wisp.utils.schedulers import DecayScheduler

import wandb
import numpy as np

log_metrics_only = False

def log_metric_to_wandb(key, _object, step):
    wandb.log({key: _object}, step=step, commit=False)


def log_images_to_wandb(key, image, step):
    if not log_metrics_only:
        wandb.log({key: wandb.Image(np.moveaxis(image, 0, -1))}, step=step, commit=False)


class BaseTrainer(ABC):
    """
    Base class for the trainer.

    The default overall flow of things:

    init()
    |- set_renderer()
    |- set_logger()

    train():
        pre_training()
        (i) for every epoch:
            |- pre_epoch()

            (ii) for every iteration:
                |- pre_step()
                |- step()
                |- post_step()

            post_epoch()
            |- log_tb()
            |- save_model()
            |- render_tb()
            |- resample_dataset()

            |- validate()
        post_training()

    iterate() runs a single iteration step of train() through all internal lifecycle methods,
    meaning a single run over loop (ii), and loop (i) if loop (ii) is finished.
    This is useful for cases like gui apps which run without a training loop.

    Each of these events can be overridden, or extended with super().

    """

    #######################
    # Initialization
    #######################

    def __init__(self, pipeline, train_dataset: WispDataset, num_epochs, batch_size,
                 optim_cls, lr, weight_decay, weight_decay_decoder, grid_lr, ldec_lr, optim_params, log_dir, device,
                 exp_name=None, info=None, scene_state=None, extra_args=None, validation_dataset: WispDataset = None,
                 render_tb_every=-1, save_every=-1, trainer_mode='validate', using_wandb=False,
                 enable_amp=True, metrics_only=False, use_scaler=True, writer=None):
        """Constructor.
        
        Args:
            pipeline (wisp.core.Pipeline): The pipeline with tracer and neural field to train.
            train_dataset (wisp.datasets.WispDataset): Dataset to used for generating training batches.
            num_epochs (int): The number of epochs to run the training for.
            batch_size (int): The batch size used in training.
            optim_cls (torch.optim): The Optimizer object to use
            lr (float): The learning rate to use
            weight_decay (float): The weight decay to use
            weight_decay_decoder (float): The weight decay to use for decoder parameters
            optim_params (dict): Optional params for the optimizer.
            device (device): The device to run the training on. 
            log_dir (str): The directory to save the training logs in.
            exp_name (str): The experiment name to use for logging purposes.
            info (str): The args to save to the logger.
            scene_state (wisp.core.State): Use this to inject a scene state from the outside to be synced
                                           elsewhere.
            extra_args (dict): Optional dict of extra_args for easy prototyping.
            validation_dataset (wisp.datasets.WispDataset): Validation dataset used for evaluating metrics.
            render_tb_every (int): The number of epochs between renders for tensorboard logging. -1 = no rendering.
            save_every (int): The number of epochs between model saves. -1 = no saving.
            trainer_mode (str): 'train' or 'validate' for choosing running training or validation only modes.
                Currently used only for titles within logs.
            using_wandb (bool): When True, weights & biases will be used for logging.
            enable_amp (bool): If enabled, the step() training function will use mixed precision.
            writer (SummaryWriter): Tensorboard writer for logging, initialized if None

        """
        log.info(f'Info: \n{info}')
        log.info(f'Training on {extra_args["dataset_path"]}')

        # initialize scene_state
        if scene_state is None:
            scene_state = WispState()
        self.scene_state = scene_state

        self.extra_args = extra_args
        self.info = info
        self.trainer_mode = trainer_mode

        self.pipeline = pipeline
        log.info("Total number of parameters: {}".format(
            sum(p.numel() for p in self.pipeline.nef.parameters()))
        )
        # Set device to use
        self.device = device
        device_name = torch.cuda.get_device_name(device=self.device)
        log.info(f'Using {device_name} with CUDA v{torch.version.cuda}')

        self.init_renderer()

        self.train_dataset = train_dataset
        self.validation_dataset = validation_dataset

        # Optimizer params
        self.optim_cls = optim_cls
        self.lr = lr
        self.weight_decay = weight_decay
        self.weight_decay_decoder = weight_decay_decoder
        self.grid_lr = grid_lr
        self.ldec_lr = ldec_lr
        self.optim_params = optim_params
        self.init_optimizer()

        # Latent decoder params
        if extra_args["ldecode_enabled"]:
            self.entropy_reg_sched = DecayScheduler(num_epochs, extra_args["entropy_reg_sched"], extra_args["entropy_reg"],
                                                     extra_args["entropy_reg_end"], 
                                                     params={'decay_period':extra_args['decay_period'], 'temperature':extra_args['temperature']})
            self.ldec_lr_sched = DecayScheduler(extra_args["ldec_lr_warmup"], "linear", 0.1*extra_args['ldec_lr'], extra_args['ldec_lr'])
            # self.ldec_lr_sched = DecayScheduler(10, "fix", extra_args['ldec_lr'])
        if extra_args["ldecode_enabled"] or extra_args["use_sga"]:
            self.temperature_sched = DecayScheduler(num_epochs, "exp", 1.0, extra_args["temperature"],
                                                     {'temperature': extra_args["temperature"], 'decay_period': extra_args["decay_period"]})
        # Training params
        self.epoch = 1
        self.iteration = 0
        self.max_epochs = num_epochs
        self.batch_size = batch_size
        self.exp_name = exp_name if exp_name else "unnamed_experiment"

        self.populate_scenegraph()
        if use_scaler:
            self.scaler = torch.cuda.amp.GradScaler()
        else:
            self.scaler = None

        # In-training variables
        self.train_data_loader_iter = None
        self.val_data_loader = None
        self.train_dataset_size = None
        self.log_dict = {}
        self.log_dir = log_dir
        self.init_dataloader()

        self.writer = writer
        self.render_tb_every = render_tb_every
        self.save_every = save_every
        self.using_wandb = using_wandb
        self.enable_amp = enable_amp
        global log_metrics_only
        log_metrics_only = metrics_only
        self.metrics_only = metrics_only

    def populate_scenegraph(self):
        """ Updates the scenegraph with information about available objects.
        Doing so exposes these objects to other components, like visualizers and loggers.
        """
        # Add object to scene graph: if interactive mode is on, this will make sure the visualizer can display it.
        # batch_size is an optional setup arg here which hints the visualizer how many rays can be processed at once
        # (e.g. this is the pipeline's batch_size used for inference time)
        add_to_scene_graph(state=self.scene_state, name=self.exp_name, obj=self.pipeline, batch_size=2 ** 14)

    def init_dataloader(self):
        self.train_data_loader = DataLoader(self.train_dataset,
                                            batch_size=self.batch_size,
                                            collate_fn=default_collate,
                                            shuffle=True, pin_memory=True,
                                            num_workers=self.extra_args['dataloader_num_workers'])
        self.iterations_per_epoch = len(self.train_data_loader)

    def init_optimizer(self):
        """Default initialization for the optimizer.
        """

        params_dict = { name : param for name, param in self.pipeline.nef.named_parameters()}
        
        params = []

        decoder_params = []
        grid_params = []
        ldec_params = []
        rest_params = []
        prob_model_params = []
        for name in params_dict:
            
            if 'decoder' in name:
                # If "decoder" is in the name, there's a good chance it is in fact a decoder,
                # so use weight_decay
                decoder_params.append(params_dict[name])

            elif 'grid' in name:
                # If "grid" is in the name, there's a good chance it is in fact a grid,
                # so use grid_lr_weight
                # If latent_dec is in the name it's a latent decoder parameter
                
                if 'latent_dec' in name:
                    ldec_params.append(params_dict[name])
                elif 'prob_model' in name:
                    prob_model_params.append(params_dict[name])
                else:
                    grid_params.append(params_dict[name])

            else:
                rest_params.append(params_dict[name])

        params.append({"params" : decoder_params,
                       "lr": self.lr, 
                       "weight_decay": 0.0,
                       "name":"decoder"})

        params.append({"params" : grid_params,
                       "lr": self.grid_lr,
                       "weight_decay": self.weight_decay,
                       "name":"grid"})
        
        params.append({"params" : ldec_params,
                       "lr": self.ldec_lr,
                       "weight_decay": self.weight_decay_decoder,
                       "name":"latent_dec"})
        
        params.append({"params" : prob_model_params,
                       "lr": 1.0e-4,
                       "weight_decay": self.weight_decay_decoder,
                       "name":"prob_models"})
        
        params.append({"params" : rest_params,
                       "lr": self.lr,
                       "weight_decay": 0.0,
                       "name":"rest"})

        self.optimizer = self.optim_cls(params, **self.optim_params)

    def init_renderer(self):
        """Default initalization for the renderer.
        """
        self.renderer = OfflineRenderer(**self.extra_args)

    #######################
    # Data load
    #######################

    def reset_data_iterator(self):
        """Rewind the iterator for the new epoch.
        """
        self.scene_state.optimization.iterations_per_epoch = len(self.train_data_loader)
        self.train_data_loader_iter = iter(self.train_data_loader)

    def next_batch(self):
        """Actually iterate the data loader.
        """
        return next(self.train_data_loader_iter)

    def resample_dataset(self):
        """
        Override this function if some custom logic is needed.

        Args:
            (torch.utils.data.Dataset): Training dataset.
        """
        if hasattr(self.train_dataset, 'resample'):
            log.info("Reset DataLoader")
            self.train_dataset.resample()
            self.init_dataloader()
        else:
            raise ValueError("resample=True but the dataset doesn't have a resample method")

    #######################
    # Training Life-cycle
    #######################

    def is_first_iteration(self):
        return self.total_iterations == 1

    def is_any_iterations_remaining(self):
        return self.total_iterations < self.max_iterations

    def begin_epoch(self):
        """Begin epoch.
        """
        self.reset_data_iterator()
        self.pre_epoch()
        self.init_log_dict()
        self.epoch_start_time = time.time()

    def end_epoch(self):
        """End epoch.
        """
        current_time = time.time()
        elapsed_time = current_time - self.epoch_start_time
        self.epoch_start_time = current_time
        # TODO(ttakikawa): Don't always write to TB
        self.writer.add_scalar(f'time/elapsed_ms_per_epoch', elapsed_time * 1000, self.epoch)
        if self.using_wandb:
            log_metric_to_wandb(f'time/elapsed_ms_per_epoch', elapsed_time * 1000, self.epoch)

        self.post_epoch()

        # Save every 10 epochs to not overload the disk
        if self.extra_args["resume"] and self.epoch % 10 == 0:
            self.save_state()
            
        if self.extra_args["valid_every"] > -1 and \
                self.epoch % self.extra_args["valid_every"] == 0 and \
                self.epoch != 0:
            self.validate()

        if self.epoch < self.max_epochs:
            self.iteration = 1
            self.epoch += 1
        else:
            self.is_optimization_running = False

    def grow(self):
        stage = min(self.extra_args["num_lods"],
                    (self.epoch // self.extra_args["grow_every"]) + 1)  # 1 indexed
        if self.extra_args["growth_strategy"] == 'onebyone':
            self.loss_lods = [stage - 1]
        elif self.extra_args["growth_strategy"] == 'increase':
            self.loss_lods = list(range(0, stage))
        elif self.extra_args["growth_strategy"] == 'shrink':
            self.loss_lods = list(range(0, self.extra_args["num_lods"]))[stage - 1:]
        elif self.extra_args["growth_strategy"] == 'finetocoarse':
            self.loss_lods = list(range(
                0, self.extra_args["num_lods"]
            ))[self.extra_args["num_lods"] - stage:]
        elif self.extra_args["growth_strategy"] == 'onlylast':
            self.loss_lods = list(range(0, self.extra_args["num_lods"]))[-1:]
        else:
            raise NotImplementedError

    def iterate(self):
        """Advances the training by one training step (batch).
        """
        if self.is_optimization_running:
            if self.is_first_iteration():
                self.pre_training()
            iter_start_time = time.time()
            try:
                if self.train_data_loader_iter is None:
                    self.begin_epoch()
                self.iteration += 1
                data = self.next_batch()
            except StopIteration:
                self.end_epoch()
                if self.is_any_iterations_remaining():
                    self.begin_epoch()
                    data = self.next_batch()
            
            if self.is_any_iterations_remaining():
                self.pre_step()
                with torch.cuda.amp.autocast(self.enable_amp):
                    self.step(data)
                self.post_step()
                iter_end_time = time.time()
            else:
                iter_end_time = time.time()
                if self.using_wandb and (not self.is_optimization_running):
                    wandb.run.summary["train_time"] = (self.scene_state.optimization.elapsed_time+
                                                       (iter_end_time-iter_start_time))
                if not self.is_optimization_running:
                    self.post_training()
            self.scene_state.optimization.elapsed_time += iter_end_time - iter_start_time

    def save_state(self):
        """Save the current state of the training.
        """
        state = {
            "epoch": self.epoch,
            "model": self.pipeline.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            'scene_state': self.scene_state
        }
        torch.save(state, os.path.join(self.log_dir, "resume_state.pth"))

    def resume_state(self):
        """Resume the training from the saved state.
        """
        if os.path.exists(os.path.join(self.log_dir, "resume_state.pth")):
            state = torch.load(os.path.join(self.log_dir, "resume_state.pth"))
            self.epoch = state["epoch"]+1
            self.pipeline.load_state_dict(state["model"])
            self.optimizer.load_state_dict(state["optimizer"])
            self.scene_state = state['scene_state']
            log.info("Found saved state. Resuming training from epoch {}".format(self.epoch))
        else:
            log.info("No saved state found!")

    def save_model(self):
        """
        Override this function to change model saving.
        """

        if self.extra_args["save_as_new"]:
            model_fname = os.path.join(self.log_dir, f'model-ep{self.epoch}-it{self.iteration}.pth')
        else:
            model_fname = os.path.join(self.log_dir, f'model.pth')

        log.info(f'Saving model checkpoint to: {model_fname}')
        if self.extra_args["model_format"] == "full":
            torch.save(self.pipeline, model_fname)
        else:
            torch.save(self.pipeline.state_dict(), model_fname)

        if self.using_wandb and not self.metrics_only and False: # Don't log model artifacts in wandb
            name = wandb.util.make_artifact_name_safe(f"{wandb.run.name}-model")
            model_artifact = wandb.Artifact(name, type="model")
            model_artifact.add_file(model_fname)
            wandb.run.log_artifact(model_artifact, aliases=["latest", f"ep{self.epoch}_it{self.iteration}"])

    def train(self):
        """
        Override this if some very specific training procedure is needed.
        """
        with torch.autograd.profiler.emit_nvtx(enabled=self.extra_args["profile"]):
            self.is_optimization_running = True
            while self.is_optimization_running:
                self.iterate()

    #######################
    # Training Events
    #######################

    def pre_training(self):
        """
        Override this function to change the logic which runs before the first training iteration.
        This function runs once before training starts.
        """
        # Default TensorBoard Logging
        self.writer = SummaryWriter(self.log_dir, purge_step=0)
        self.writer.add_text('Info', self.info)
        
        if self.extra_args["resume"]:
            self.resume_state()

        if self.using_wandb:
            wandb_project = self.extra_args["wandb_project"]
            wandb_run_name = self.extra_args.get("wandb_run_name")
            wandb_entity = self.extra_args.get("wandb_entity")
            wandb_mode = self.extra_args.get("wandb_mode")
            wandb.init(
                project=wandb_project,
                name=self.exp_name if wandb_run_name is None else wandb_run_name,
                entity=wandb_entity,
                job_type=self.trainer_mode,
                config=self.extra_args,
                sync_tensorboard=not self.metrics_only,
                dir=self.log_dir,
                mode=wandb_mode,
                resume=self.extra_args["resume"],
            )

    def post_training(self):
        """
        Override this function to change the logic which runs after the last training iteration.
        This function runs once after training ends.
        """
        self.writer.close()
        if self.extra_args["resume"]:
            if os.path.exists(os.path.join(self.log_dir, "resume_state.pth")):
                os.remove(os.path.join(self.log_dir, "resume_state.pth"))
                print(f"Finished training. Removed resume state file at {os.path.join(self.log_dir, 'resume_state.pth')}")
        if self.using_wandb:
            wandb.finish()

    def pre_epoch(self):
        """
        Override this function to change the pre-epoch preprocessing.
        This function runs once before the epoch.
        """
        # The DataLoader is refreshed before every epoch, because by default, the dataset refreshes
        # (resamples) after every epoch.

        self.loss_lods = list(range(0, self.extra_args["num_lods"]))
        if self.extra_args["grow_every"] > 0:
            self.grow()

        if self.extra_args["only_last"]:
            self.loss_lods = self.loss_lods[-1:]

        if self.extra_args["resample"] and self.epoch % self.extra_args["resample_every"] == 0 and self.epoch > 1:
            self.resample_dataset()

        if self.extra_args["ldecode_enabled"]:
            grid = self.pipeline.nef.grid
            assert isinstance(grid, LatentGrid)
            self.entropy_reg_lambda = self.entropy_reg_sched(self.epoch)
            if self.extra_args["use_sga"] and isinstance(grid.latent_dec, LatentDecoder):
                grid.latent_dec.diff_sampling = self.extra_args["diff_sampling"]
            if isinstance(grid.latent_dec, MultiLatentDecoder) or self.extra_args["use_sga"]:
                grid.latent_dec.temperature = self.temperature_sched(self.epoch)
            if self.extra_args["use_sga"] and (self.epoch)/self.max_epochs>self.extra_args['decay_period']:
                grid.latent_dec.use_sga = False
            if isinstance(grid.latent_dec, MultiLatentDecoder):
                grid.latent_dec.straight_through = (self.epoch/self.max_epochs)>self.extra_args['decay_period']
            
        self.pipeline.train()

    def post_epoch(self):
        """
        Override this function to change the post-epoch post processing.

        By default, this function logs to Tensorboard, renders images to Tensorboard, saves the model,
        and resamples the dataset.

        To keep default behaviour but also augment with other features, do

          super().post_epoch()

        in the derived method.
        """
        self.pipeline.eval()

        total_loss = self.log_dict['total_loss'] / len(self.train_data_loader)
        self.scene_state.optimization.losses['total_loss'].append(total_loss)

        if isinstance(self.pipeline.nef.grid, LatentGrid) or isinstance(self.pipeline.nef.grid, HashGrid) or isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            grid = self.pipeline.nef.grid
            size_ldec, size_latents = grid.size(use_torchac=False, use_prob_model=False)
            size_remainder = sum([p.numel()*torch.finfo(p.dtype).bits for n,p in self.pipeline.named_parameters() if 'grid' not in n])
            if size_ldec>0:
                self.log_dict['ldec_size'] = size_ldec/8e3
            self.log_dict['latent_size'] = size_latents/8e3
            self.log_dict['remainder_size'] = size_remainder/8e3
            self.log_dict['total_size'] = (size_latents+size_remainder+size_ldec)/8e3
            if isinstance(self.pipeline.nef.grid, LatentGrid):
                self.log_dict['rounding_loss'] = torch.mean(torch.abs(grid.codebook - torch.round(grid.codebook))).item()

        if self.extra_args["log_tb_every"] > -1 and self.epoch % self.extra_args["log_tb_every"] == 0:
            self.log_cli()
            self.log_tb()

        # Render visualizations to tensorboard
        if self.render_tb_every > -1 and self.epoch % self.render_tb_every == 0 and not self.metrics_only:
            self.render_tb()

       # Save model
        if self.save_every > -1 and self.epoch % self.save_every == 0 and self.epoch != 0 and not self.metrics_only:
            self.save_model()

    def pre_step(self):
        """
        Override this function to change the pre-step preprocessing (runs per iteration).
        """
        pass

    def post_step(self):
        """
        Override this function to change the pre-step preprocessing (runs per iteration).
        """
        pass

    @abstractmethod
    def step(self, data):
        """Advance the training by one step using the batched data supplied.

        data (dict): Dictionary of the input batch from the DataLoader.
        """
        pass

    @abstractmethod
    def validate(self):
        pass

    #######################
    # Logging
    #######################

    def init_log_dict(self):
        """
        Override this function to use custom logs.
        """
        self.log_dict['total_loss'] = 0.0
        self.log_dict['total_iter_count'] = 0

    def log_model_details(self):
        # TODO (operel): Brittle
        log.info(f"Position Embed Dim: {self.pipeline.nef.pos_embed_dim}")
        log.info(f"View Embed Dim: {self.pipeline.nef.view_embed_dim}")

    def log_cli(self):
        """
        Override this function to change CLI logging.

        By default, this function only runs every epoch.
        """
        # Average over iterations
        log_text = 'EPOCH {}/{}'.format(self.epoch, self.max_epochs)
        log_text += ' | total loss: {:>.3E}'.format(self.log_dict['total_loss'] / len(self.train_data_loader))

    def log_tb(self):
        """
        Override this function to change loss / other numeric logging to TensorBoard / Wandb.
        """
        for key in self.log_dict:
            if 'loss' in key:
                self.writer.add_scalar(f'loss/{key}', self.log_dict[key] / len(self.train_data_loader), self.epoch)
                if self.using_wandb:
                    log_metric_to_wandb(f'loss/{key}', self.log_dict[key] / len(self.train_data_loader), self.epoch)
            if 'size' in key:
                self.writer.add_scalar(f'size/{key}', self.log_dict[key], self.epoch)
                if self.using_wandb:
                    log_metric_to_wandb(f'size/{key}', self.log_dict[key], self.epoch)

    def render_tb(self):
        """
        Override this function to change render logging to TensorBoard / Wandb.
        """
        self.pipeline.eval()
        for d in [self.extra_args["num_lods"] - 1]:
            out = self.renderer.shade_images(self.pipeline,
                                             f=self.extra_args["camera_origin"],
                                             t=self.extra_args["camera_lookat"],
                                             fov=self.extra_args["camera_fov"],
                                             lod_idx=d,
                                             camera_clamp=self.extra_args["camera_clamp"])

            # Premultiply the alphas since we're writing to PNG (technically they're already premultiplied)
            if self.extra_args["bg_color"] == 'black' and out.rgb.shape[-1] > 3:
                bg = torch.ones_like(out.rgb[..., :3])
                out.rgb[..., :3] += bg * (1.0 - out.rgb[..., 3:4])

            out = out.image().byte().numpy_dict()

            log_buffers = ['depth', 'hit', 'normal', 'rgb', 'alpha']

            for key in log_buffers:
                if out.get(key) is not None:
                    self.writer.add_image(f'{key}/{d}', out[key].T, self.epoch)
                    if self.using_wandb:
                        log_images_to_wandb(f'{key}/{d}', out[key].T, self.epoch)

    #######################
    # Properties
    #######################

    @property
    def is_optimization_running(self) -> bool:
        return self.scene_state.optimization.running

    @is_optimization_running.setter
    def is_optimization_running(self, is_running: bool):
        self.scene_state.optimization.running = is_running

    @property
    def epoch(self) -> int:
        """ Epoch counter, starts at 1 and ends at max epochs"""
        return self.scene_state.optimization.epoch

    @epoch.setter
    def epoch(self, epoch: int):
        self.scene_state.optimization.epoch = epoch

    @property
    def iteration(self) -> int:
        """ Iteration counter, for current epoch. Starts at 1 and ends at iterations_per_epoch """
        return self.scene_state.optimization.iteration

    @iteration.setter
    def iteration(self, iteration: int):
        """ Iteration counter, for current epoch """
        self.scene_state.optimization.iteration = iteration

    @property
    def iterations_per_epoch(self) -> int:
        """ How many iterations should run per epoch """
        return self.scene_state.optimization.iterations_per_epoch

    @iterations_per_epoch.setter
    def iterations_per_epoch(self, iterations: int):
        """ How many iterations should run per epoch """
        self.scene_state.optimization.iterations_per_epoch = iterations

    @property
    def total_iterations(self) -> int:
        """ Total iteration steps the trainer took so far, for all epochs.
            Starts at 1 and ends at max_iterations
        """
        return (self.epoch - 1) * self.iterations_per_epoch + self.iteration

    @property
    def max_epochs(self) -> int:
        """ Total number of epochs set for this optimization task.
        The first epoch starts at 1 and the last epoch ends at the returned `max_epochs` value.
        """
        return self.scene_state.optimization.max_epochs

    @max_epochs.setter
    def max_epochs(self, num_epochs):
        """ Total number of epochs set for this optimization task.
        The first epoch starts at 1 and the last epoch ends at `num_epochs`.
        """
        self.scene_state.optimization.max_epochs = num_epochs

    @property
    def max_iterations(self) -> int:
        """ Total number of iterations set for this optimization task. """
        return self.max_epochs * self.iterations_per_epoch

wisp/trainers/base_trainer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import logging as log
import numpy as np
import os
import cv2
import copy
import torch
import random
import wandb
import time
import json
import math
from tqdm import tqdm
from dataclasses import dataclass, field
from torch.utils.tensorboard import SummaryWriter
from wisp.trainers import BaseTrainer, log_metric_to_wandb, log_images_to_wandb
from wisp.datasets import BaseImageDataset,ImageBatch, default_collate
from wisp.ops.image import hwc_to_chw, psnr, clamped_psnr, clamped_mse
from wisp.utils.schedulers import DecayScheduler
from wisp.models.grids import LatentGrid, HashGrid, CodebookOctreeGrid
from wisp.models.latent_decoders import LatentDecoder, MultiLatentDecoder

os.environ["OPENCV_IO_MAX_IMAGE_PIXELS"] = pow(2,40).__str__()


class ImageTrainer(BaseTrainer):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.ldec_lr_sched = DecayScheduler(self.extra_args["ldec_lr_warmup"], "fix", self.extra_args['ldec_lr'])
        
        # Update optimization state about the current train set and val set used
        self.scene_state.data.train_data.append(self.train_dataset)

        # self.log_fname = f'{datetime.now().strftime("%Y%m%d-%H%M%S")}'
        self.log_dir = os.path.join(self.log_dir, self.exp_name)
        self.log_dir_cur = os.path.join(
            self.log_dir,
            os.path.splitext(os.path.basename(self.train_dataset.image_path))[0]
            # self.log_fname
        )
        os.makedirs(self.log_dir_cur,exist_ok=True)
        if os.path.exists(os.path.join(self.log_dir,'complete')):
            print(f'Experiment already exists at {self.log_dir}, exiting...')
            exit(1)

        log.info(f'Logging to {self.log_dir_cur}')
 

    def populate_scenegraph(self):
        """ Disable scenegraph update for images
        """
        pass

    def init_renderer(self):
        """Disable renderer for images
        """
        pass

    def init_dataloader(self):
        assert self.batch_size==1, "Image loader should be with batch size 1 only"
        self.train_data_loader = torch.utils.data.DataLoader(self.train_dataset,
                                            batch_size=self.batch_size,
                                            shuffle=False, pin_memory=False,
                                            num_workers=self.extra_args['dataloader_num_workers'])
        self.iterations_per_epoch = len(self.train_data_loader)

    def init_log_dict(self):
        """Custom logging dictionary.
        """
        super().init_log_dict()
        self.log_dict['rgb_loss'] = 0.0
        self.log_dict['total_loss'] = 0.0
        self.log_dict['PSNR'] = 0.0
        if self.extra_args["ldecode_enabled"] and self.entropy_reg_lambda > 0:
            self.log_dict['ent_loss'] = 0.0
            self.log_dict['net_kbytes_codebook'] = 0.0

    #######################
    # Data load
    #######################

    def reset_data_iterator(self):
        """Rewind the iterator for the new epoch if coordinates not static.
        """
        if not (self.train_dataset.static_coords and not self.is_first_iteration()):
            self.scene_state.optimization.iterations_per_epoch = len(self.train_data_loader)
            self.train_data_loader_iter = iter(self.train_data_loader)

    #######################
    # Training Events
    #######################
        
    def pre_training(self):
        """
        Override this function to change the logic which runs before the first training iteration.
        This function runs once before training starts.
        """
        self.best_state = {'rgb_loss':np.inf, 'PSNR': 0.0}
        if self.train_dataset.static_coords:
            #  Pointer to hold static data when initialized
            self.data = None

    def pre_epoch(self):
        """
        This function runs once before the epoch.
        """

        # Data resampling if not static every epoch
        if self.extra_args["resample"] and self.epoch % self.extra_args["resample_every"] == 0 \
            and self.epoch > 1 and not self.train_dataset.static_coords:
            self.resample_dataset()

        # If using latent decoder
        if self.extra_args["ldecode_enabled"]:
            grid = self.pipeline.nef.grid
            assert isinstance(grid, LatentGrid)

            # Entropy regularization loss schedule over training
            self.entropy_reg_lambda = self.entropy_reg_sched(self.epoch)

            # Temperature schedule for SGA while training
            if isinstance(grid.latent_dec, MultiLatentDecoder) or self.extra_args["use_sga"]:
                grid.latent_dec.temperature = self.temperature_sched(self.epoch)

            # Disable SGA after the decay period. Uses simple Straight Through Estimator after
            if self.extra_args["use_sga"] and (self.epoch)/self.max_epochs>self.extra_args['decay_period']:
                grid.latent_dec.use_sga = False
            
            # When straight through enabled, use only single decoder from MultiLatentDecoder
            if isinstance(grid.latent_dec, MultiLatentDecoder):
                grid.latent_dec.straight_through = (self.epoch/self.max_epochs)>self.extra_args['decay_period']
        
        # For static coordinates, pipeline is set to train only in the first epoch/iteration
        if not (self.train_dataset.static_coords and not self.is_first_iteration()):
            self.pipeline.train()

    def post_epoch(self):
        """
        Runs once after the epoch
        """

        if not self.train_dataset.static_coords:
            self.p_bar.close()

        self.log_dict['PSNR'] = self.log_dict['PSNR'] / self.iterations_per_epoch
        if isinstance(self.pipeline.nef.grid, LatentGrid) or \
            isinstance(self.pipeline.nef.grid, HashGrid) or \
            isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            grid = self.pipeline.nef.grid
            use_torchac = self.total_iterations % self.extra_args["log_every"] == 0
            size_ldec, size_latents = grid.size(use_torchac=use_torchac, use_prob_model=False)
            size_remainder = sum([p.numel()*torch.finfo(p.dtype).bits for n,p in self.pipeline.named_parameters() if 'grid' not in n])
            if size_ldec>0:
                self.log_dict['ldec_size'] = size_ldec/8e3
            self.log_dict['latent_size'] = size_latents/8e3
            self.log_dict['remainder_size'] = size_remainder/8e3
            self.log_dict['total_size'] = (size_latents+size_remainder+size_ldec)/8e3
            self.log_dict['BPP'] = (size_latents+size_remainder+size_ldec)/math.prod(self.train_dataset.image_size)
            if isinstance(self.pipeline.nef.grid, LatentGrid):
                self.log_dict['rounding_loss'] = torch.mean(torch.abs(grid.codebook - torch.round(grid.codebook))).item()


        if self.train_dataset.static_coords and self.log_dict['rgb_loss'] < self.best_state['rgb_loss']:
            self.best_state['PSNR'] = self.log_dict['PSNR']
            self.best_state['BPP'] = self.log_dict['BPP']
            self.best_state['total_size'] = self.log_dict['total_size']
            self.best_state['rgb_loss'] = self.log_dict['rgb_loss']
            self.best_state['state_dict'] = copy.deepcopy(self.pipeline.state_dict())

        for k,v in self.log_dict.items():
            if 'loss' not in k:
                self.scene_state.optimization.metrics[k].append(v)
            else:
                self.scene_state.optimization.losses[k].append(v)
            
        if self.extra_args["log_every"] > -1 and self.epoch % self.extra_args["log_every"] == 0:
            self.log_cli()
            self.log_tb()

        # Render visualizations to tensorboard
        if self.render_tb_every > -1 and self.epoch % self.render_tb_every == 0 and not self.metrics_only:
            self.render_tb()


    def begin_epoch(self):
        """Begin epoch.
        """
        self.reset_data_iterator()
        self.pre_epoch()
        self.init_log_dict()
        self.epoch_start_time = time.time()
        if not self.train_dataset.static_coords:
            self.p_bar = tqdm(total=len(self.train_data_loader))

    def end_epoch(self):
        """End epoch.
        """
        current_time = time.time()
        elapsed_time = current_time - self.epoch_start_time
        self.epoch_start_time = current_time
        self.writer.add_scalar(f'time/elapsed_per_epoch', elapsed_time, self.epoch)
        if self.using_wandb:
            log_metric_to_wandb(f'time/elapsed_per_epoch', elapsed_time, self.epoch)

        self.post_epoch()

        # Save every N epochs to not overload the disk
        if self.extra_args["resume"] and self.epoch % self.extra_args["save_every"] == 0:
            self.save_state()
            
        # Validate every N epochs if enabled and coords not static (as train same as val then)
        if  self.extra_args["valid_every"] > -1 and \
            self.epoch % self.extra_args["valid_every"] == 0 and \
            self.epoch != 0 and not self.train_dataset.static_coords:
            with torch.no_grad():
                self.validate()

        if self.epoch < self.max_epochs:
            self.iteration = 0
            self.epoch += 1
        else:
            self.is_optimization_running = False

    def iterate_static(self):
        """Advances the training by one training step (batch) for static coordinates (one iteration)
        """
        assert self.iterations_per_epoch == 1
        if self.is_optimization_running:
            self.iteration += 1
            if self.is_first_iteration():
                self.pre_training()
            
            iter_start_time = time.time()
            self.begin_epoch()

            # Load data only first time since data is static
            if self.is_first_iteration():
                assert self.data is None
                self.data = self.next_batch()
            data = self.data
        
            # if self.is_any_iterations_remaining():
            self.pre_step()
            with torch.cuda.amp.autocast(self.enable_amp):
                self.step(data)
            self.post_step()
            iter_end_time = time.time()
            self.end_epoch()

            if not self.is_any_iterations_remaining():
                if self.using_wandb and (not self.is_optimization_running):
                    wandb.run.summary[f"image{self.train_dataset.image_idx}/train_time"] = (self.scene_state.optimization.elapsed_time+
                                                       (iter_end_time-iter_start_time))
                if not self.is_optimization_running:
                    self.post_training()
            self.scene_state.optimization.elapsed_time += iter_end_time - iter_start_time

    @torch.cuda.nvtx.range("ImageTrainer.step")
    def step(self, data):
        """Implement the optimization over image-space loss.
        """
        # Move to GPU only the first time if coords static
        if not (self.train_dataset.static_coords and not self.is_first_iteration()):
            data['coords'] = data['coords'].to(self.device)
            data['rgb'] = data['rgb'].to(self.device)
        coords = data['coords'].squeeze(0)
        coords = self.train_dataset.transform_coords(coords, self.device)
        img_gts = data['rgb'].squeeze(0)
        
        self.optimizer.zero_grad(set_to_none=True)
            
        loss = rgb_loss = 0

        with torch.no_grad():
            if (isinstance(self.pipeline.nef.grid, LatentGrid) and 
                self.extra_args["ldecode_enabled"] and 
                isinstance(self.pipeline.nef.grid.latent_dec, LatentDecoder) and
                self.extra_args["norm"]!='none' and 
                self.extra_args["norm_every"]%self.total_iterations==0):
                decoder = self.pipeline.nef.grid.latent_dec
                latents = self.pipeline.nef.grid.codebook
                if self.extra_args["norm"] == "max":
                    decoder.div = torch.max(torch.abs(latents.min(dim=0,keepdim=True)[0]),\
                                            torch.abs(latents.max(dim=0,keepdim=True)[0]))
                elif self.extra_args["norm"] == "std":
                    decoder.div = latents.std(dim=0,keepdim=True)

        pred = self.pipeline.nef(coords=coords, channels=["rgb"])[0]
        res = 1.0
        rgb_loss += ((pred - res * img_gts)**2).mean()

        self.log_dict['PSNR'] += clamped_psnr(pred, img_gts)
        # If coordinates static (single step per epoch), train is same as validation data
        # Best state is stored at end of validation otherwise
        if self.train_dataset.static_coords and rgb_loss < self.best_state['rgb_loss'] and not self.metrics_only:
            H,W = self.train_dataset.image_size
            pred_image = torch.clamp(pred.detach().cpu().reshape(H,W,3)*255,0.,255.)\
                                     .numpy().astype(np.uint8)
            self.best_state['pred'] = pred_image

        loss += self.extra_args["rgb_loss"] * rgb_loss
        if self.extra_args["ldecode_enabled"] and self.entropy_reg_lambda > 0:
            avg_bits, net_bits = self.pipeline.nef.grid.ent_loss(
                                    self.total_iterations-1, 
                                    is_val=not self.pipeline.training
                                    )
            ent_loss = self.entropy_reg_lambda * avg_bits
            loss += ent_loss
            self.log_dict['ent_loss'] = ent_loss.item()
            self.log_dict['net_kbytes_codebook'] = net_bits/8/1024

        self.log_dict['rgb_loss'] += rgb_loss.item()
        self.log_dict['total_loss'] += loss.item()
        
        if (isinstance(self.pipeline.nef.grid, LatentGrid) and 
            self.extra_args["ldecode_enabled"] and 
            isinstance(self.pipeline.nef.grid.latent_dec, LatentDecoder) and
            self.extra_args["scale_grid_lr"]!='none'):
            with torch.no_grad():
                for group in self.optimizer.param_groups:
                    if group['name'] == 'grid':
                        norm = self.pipeline.nef.grid.latent_dec.scale_norm().item()
                        if self.extra_args["scale_grid_lr"]=='mul':
                            group['lr'] = self.extra_args["grid_lr"]* norm
                        elif self.extra_args["scale_grid_lr"]=='div':
                            group['lr'] = self.extra_args["grid_lr"]/ norm
                        elif self.extra_args["scale_grid_lr"]=='none':
                            group['lr'] = self.extra_args["grid_lr"]
                        else:
                            raise Exception(f"Unknown grid_lr scaling mode {self.extra_args['scale_grid_lr']}")
                        if self.total_iterations %100 == 0:
                            print("Latent learning rate: ", group['lr'], norm)
                    elif group['name'] == 'latent_dec':
                        group['lr'] = self.ldec_lr_sched(self.epoch)
                        # if self.total_iterations %100 == 0:
                        #     print("Latent decoder learning rate: ", group['lr'])

        with torch.cuda.nvtx.range("ImageTrainer.backward"):
            if self.scaler:
                self.scaler.scale(loss).backward()
            else:
                loss.backward()

            if self.scaler:
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                self.optimizer.step()

    def post_step(self):
        if not self.train_dataset.static_coords:
            self.p_bar.update(1)

    def train(self):
        """
        Override this if some very specific training procedure is needed.
        """
        with torch.autograd.profiler.emit_nvtx(enabled=self.extra_args["profile"]):
            self.is_optimization_running = True
            while self.is_optimization_running:
                if self.train_dataset.static_coords:
                    self.iterate_static()
                else:
                    self.iterate()

    def validate(self):
        """Implement validation. 
        """
        val_dict = {'MSE':[], 'num_samples': [], 'clamped_MSE': [], 'PSNR': 0.0}
        val_dataset = self.train_dataset
        sample_mode_state = self.train_dataset.sample_mode
        # val_dataset.sample_mode = 'eval'
        # val_dataset.load()
        
        image_idx = val_dataset.image_idx
        val_data_loader = torch.utils.data.DataLoader(val_dataset,
                                            batch_size=self.batch_size,
                                            shuffle=False, pin_memory=False,
                                            num_workers=self.extra_args['dataloader_num_workers'])
        iterations_per_epoch = len(val_data_loader)

        # For validation, disable SGA annealing if enabled
        if self.extra_args["ldecode_enabled"]:
            grid = self.pipeline.nef.grid
            assert isinstance(grid, LatentGrid)
            if self.extra_args["use_sga"]:
                use_sga_state = grid.latent_dec.use_sga
                grid.latent_dec.use_sga = False
            if isinstance(grid.latent_dec, MultiLatentDecoder):
                straight_through_state = grid.latent_dec.straight_through
                grid.latent_dec.straight_through = True

        self.pipeline.eval()
        H,W = self.train_dataset.image_size
        pred_image = torch.zeros(H*W,3)
        coord_idx = 0
        log.info('Starting validation for image {}/{}'.format(self.train_dataset.image_idx, self.train_dataset.num_images))
        p_bar = tqdm(total=len(val_data_loader))
        
        for iteration, data in enumerate(val_data_loader):
            # Map to device
            coords = data['coords'].to(self.device).squeeze(0)
            coords = val_dataset.transform_coords(coords, self.device)
            img_gts = data['rgb'].to(self.device).squeeze(0)
            pred = self.pipeline.nef(coords=coords, channels=["rgb"])[0]
            res = 1.0
            rgb_loss = ((pred - res * img_gts)**2).sum()
            
            val_dict['PSNR'] += [clamped_psnr(pred, img_gts)]
            val_dict['MSE'] += [rgb_loss.item()]
            val_dict['num_samples'] += [img_gts.size(0)]
            pred_image[val_dataset.shuffle_idx[coord_idx:coord_idx+coords.size(0)]] = pred.detach().cpu()
            coord_idx += coords.size(0)
            p_bar.update(1)
        p_bar.close()
        pred_image = torch.clamp(pred_image.reshape(H,W,3)*255,0.,255.).numpy().astype(np.uint8)
        gt_image = torch.clamp(val_dataset.data['orig_rgb'].reshape(H,W,3)*255,0.,255.).numpy().astype(np.uint8)
        val_dict['PSNR'] = sum(val_dict['PSNR'])/sum(val_dict['num_samples'])
        val_dict['MSE'] = sum(val_dict['MSE'])/sum(val_dict['num_samples'])
        if not self.metrics_only:
            val_dict['pred'] = pred_image
        
        if not self.train_dataset.static_coords and val_dict['MSE'] < self.best_state['rgb_loss']:
            self.best_state['rgb_loss'] = val_dict['MSE']
            self.best_state['PSNR'] = val_dict['PSNR']
            self.best_state['state_dict'] = copy.deepcopy(self.pipeline.state_dict())
            if not self.metrics_only:
                self.best_state['pred'] = val_dict['pred']
            if isinstance(self.pipeline.nef.grid, LatentGrid) or isinstance(self.pipeline.nef.grid, HashGrid) or isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
                self.best_state['total_size'] = self.log_dict['total_size']
                self.best_state['BPP'] = self.log_dict['total_size']*8000/math.prod(val_dataset.image_size)

        log_text = 'Validation Image {}/ {} EPOCH {}/{}'.format(self.train_dataset.image_idx, self.train_dataset.num_images, self.epoch, self.max_epochs)
        log_text += ' | PSNR: {:>.2E}'.format(val_dict['PSNR'])
        if isinstance(self.pipeline.nef.grid, LatentGrid) or isinstance(self.pipeline.nef.grid, HashGrid) or isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            val_dict['BPP'] = self.log_dict['total_size']*8000/math.prod(val_dataset.image_size)
            log_text += ' | BPP: {:>.2E}'.format(self.log_dict['BPP'])
            log_text += ' | total size (kB): {:>.2E}'.format(self.log_dict['total_size'])
        log_text += ' | MSE: {:>.2E}'.format(val_dict['MSE'])
        if 'ent_loss' in self.log_dict:
            log_text += ' | ent loss: {:>.2E}'.format(self.log_dict['ent_loss'] / len(self.train_data_loader))
            # log_text += ' | codebook size (kB): {:>.3E}'.format(self.log_dict['net_kbytes_codebook'] / len(self.train_data_loader))
        
        log.info(log_text)

        if self.using_wandb:
            log_metric_to_wandb(f'image{image_idx}/validation/PSNR',val_dict['PSNR'], self.epoch)
            log_metric_to_wandb(f'image{image_idx}/validation/BPP',val_dict['BPP'], self.epoch)
            log_metric_to_wandb(f'image{image_idx}/validation/MSE',val_dict['MSE'], self.epoch)

        self.train_dataset.sample_mode = sample_mode_state
        # Resume state for SGA/Straight-through
        if self.extra_args["ldecode_enabled"]:
            latent_dec = self.pipeline.nef.grid.latent_dec
            if self.extra_args["use_sga"]:
                latent_dec.use_sga = use_sga_state
            if isinstance(latent_dec, MultiLatentDecoder):
                latent_dec.straight_through = straight_through_state
    
    def post_training(self):
        """
        Override this function to change the logic which runs after the last training iteration.
        This function runs once after training ends.
        """

        self.pipeline.load_state_dict(self.best_state['state_dict'])
        model_fname = os.path.join(self.log_dir_cur, f'model_best.pth')
        log.info(f'Saving best model checkpoint to: {model_fname}')
        if self.extra_args["model_format"] == "full":
            torch.save(self.pipeline, model_fname)
        else:
            torch.save(self.pipeline.state_dict(), model_fname)

        if not self.metrics_only:
            best_image = self.best_state['pred']
            H,W = self.train_dataset.image_size
            cv2.imwrite(os.path.join(self.log_dir_cur,'predicted.png'), cv2.cvtColor(best_image, cv2.COLOR_BGR2RGB))

        # Obtain model sizes and BPP for the best PSNR prediction 
        if isinstance(self.pipeline.nef.grid, LatentGrid) or \
            isinstance(self.pipeline.nef.grid, HashGrid) or \
            isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            grid = self.pipeline.nef.grid
            size_ldec, size_latents = grid.size(use_torchac=True, use_prob_model=False)
            size_remainder = sum([p.numel()*torch.finfo(p.dtype).bits for n,p in self.pipeline.named_parameters() if 'grid' not in n])
            self.best_state['BPP'] = (size_latents+size_remainder+size_ldec)/math.prod(self.train_dataset.image_size)
            if size_ldec>0:
                self.best_state['ldec_size'] = size_ldec/8e3
            self.best_state['latent_size'] = size_latents/8e3
            self.best_state['remainder_size'] = size_remainder/8e3
            self.best_state['total_size'] = (size_latents+size_remainder+size_ldec)/8e3
            if isinstance(self.pipeline.nef.grid, LatentGrid):
                self.best_state['rounding_loss'] = torch.mean(torch.abs(grid.codebook - torch.round(grid.codebook))).item()

        state = self.scene_state.optimization
        torch.save(state, os.path.join(self.log_dir_cur, "scene_state.pth"))

        # Store only metrics
        if not self.metrics_only:
            del self.best_state['pred']
        del self.best_state['state_dict']
        with open(os.path.join(self.log_dir_cur,'metrics.json'), 'w') as f:
            json.dump(self.best_state, f)

    #######################
    # Logging
    #######################
        
    def log_cli(self):
        log_text = 'Image {}/ {} EPOCH {}/{}'.format(self.train_dataset.image_idx, self.train_dataset.num_images, self.epoch, self.max_epochs)
        log_text += ' | PSNR: {:>.2E}'.format(self.log_dict['PSNR'])
        if  isinstance(self.pipeline.nef.grid, LatentGrid) or \
            isinstance(self.pipeline.nef.grid, HashGrid) or \
            isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            log_text += ' | BPP: {:>.2E}'.format(self.log_dict['BPP'])
            log_text += ' | total size (kB): {:>.2E}'.format(self.log_dict['total_size'])
        log_text += ' | total loss: {:>.2E}'.format(self.log_dict['total_loss'] / len(self.train_data_loader))
        log_text += ' | rgb loss: {:>.2E}'.format(self.log_dict['rgb_loss'] / len(self.train_data_loader))
        if 'ent_loss' in self.log_dict:
            log_text += ' | ent loss: {:>.2E}'.format(self.log_dict['ent_loss'] / len(self.train_data_loader))
            # log_text += ' | codebook size (kB): {:>.3E}'.format(self.log_dict['net_kbytes_codebook'] / len(self.train_data_loader))
        if isinstance(self.pipeline.nef.grid, LatentGrid) or isinstance(self.pipeline.nef.grid, HashGrid) or isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            # if 'dec_size' in self.log_dict and isinstance(self.pipeline.nef.grid.latent_dec, MultiLatentDecoder):
            #     log_text += ' | dec size (kB): {:>.2E}'.format(self.log_dict['dec_size'])
            # log_text += ' | latent size (kB): {:>.2E}'.format(self.log_dict['latent_size'])
            log_text += ' | total size (kB): {:>.2E}'.format(self.log_dict['total_size'])
            if isinstance(self.pipeline.nef.grid, LatentGrid) and isinstance(self.pipeline.nef.grid.latent_dec, MultiLatentDecoder) or self.extra_args["use_sga"]:
                log_text += ' | temp: {:>.2E}'.format(self.pipeline.nef.grid.latent_dec.temperature)
            if self.extra_args["use_sga"]:
                log_text += ' | sga: {}'.format(self.pipeline.nef.grid.latent_dec.use_sga)
        
        log.info(log_text)

    def log_tb(self):
        """
        Override this function to change loss / other numeric logging to TensorBoard / Wandb.
        """
        image_idx = self.train_dataset.image_idx
        for key in self.log_dict:
            if 'loss' in key:
                self.writer.add_scalar(f'image{image_idx}/loss/{key}', self.log_dict[key] / len(self.train_data_loader), self.epoch)
                if self.using_wandb:
                    log_metric_to_wandb(f'image{image_idx}/loss/{key}', self.log_dict[key] / len(self.train_data_loader), self.epoch)
            if 'size' in key:
                self.writer.add_scalar(f'image{image_idx}/size/{key}', self.log_dict[key], self.epoch)
                if self.using_wandb:
                    log_metric_to_wandb(f'image{image_idx}/size/{key}', self.log_dict[key], self.epoch)

    def render_tb(self):
        image_idx = self.train_dataset.image_idx
        self.writer.add_image(f'image{image_idx}/pred', hwc_to_chw(self.best_state['pred']), self.epoch)
        if self.using_wandb:
            log_images_to_wandb(f'image{image_idx}/pred', hwc_to_chw(self.best_state['pred']), self.epoch)

    def save_state(self):
        """Save the current state of the training.
        """
        state = {
            "epoch": self.epoch,
            "image_idx": self.train_dataset.image_idx,
            "model": self.pipeline.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            'opt_state': self.scene_state.optimization
        }
        torch.save(state, os.path.join(self.log_dir, "resume_state.pth"))

    def resume_state(self):
        """Resume the training from the saved state.
        """
        if os.path.exists(os.path.join(self.log_dir, "resume_state.pth")):
            state = torch.load(os.path.join(self.log_dir, "resume_state.pth"))
            self.epoch = state["epoch"]+1
            self.pipeline.load_state_dict(state["model"])
            self.optimizer.load_state_dict(state["optimizer"])
            self.scene_state.optimization = state['opt_state']
            log.info("Found saved state. Resuming training from epoch {}".format(self.epoch))
        else:
            log.info("No saved state found!")

wisp/trainers/image_trainer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.


import os
import logging as log
from tqdm import tqdm
import random
import pandas as pd
import torch
from wisp.trainers import BaseTrainer, log_metric_to_wandb, log_images_to_wandb
from wisp.ops.image import write_png, write_exr
from wisp.ops.image.metrics import psnr, lpips, ssim
from wisp.datasets import MultiviewDataset
from wisp.core import Rays, RenderBuffer
from wisp.models.latent_decoders import MultiLatentDecoder, LatentDecoder
from wisp.models.grids import LatentGrid, HashGrid, CodebookOctreeGrid

import wandb
import numpy as np
import wisp.trainers.base_trainer as base_trainer
from PIL import Image


class MultiviewTrainer(BaseTrainer):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Update optimization state about the current train set and val set used
        self.scene_state.optimization.train_data.append(self.train_dataset)
        self.scene_state.optimization.validation_data.append(self.validation_dataset)

        # self.log_fname = f'{datetime.now().strftime("%Y%m%d-%H%M%S")}'
        self.log_dir = os.path.join(
            self.log_dir,
            self.exp_name,
            # self.log_fname
        )

        if os.path.exists(os.path.join(self.log_dir, 'val')):
            print(f'Experiment already exists at {self.log_dir}, exiting...')
            exit(1)

        log.info(f'Logging to {self.log_dir}')
        
    def populate_scenegraph(self):
        """ Updates the scenegraph with information about available objects.
        Doing so exposes these objects to other components, like visualizers and loggers.
        """
        super().populate_scenegraph()
        self.scene_state.graph.cameras = self.train_dataset.cameras

    def pre_step(self):
        """Override pre_step to support pruning.
        """
        super().pre_step()
        
        if self.extra_args["prune_every"] > -1 and \
           self.total_iterations > 1 and \
           self.total_iterations % self.extra_args["prune_every"] == 0:
            self.pipeline.nef.prune()

    def init_log_dict(self):
        """Custom log dict.
        """
        super().init_log_dict()
        self.log_dict['rgb_loss'] = 0.0
        if self.extra_args["ldecode_enabled"] and self.entropy_reg_lambda > 0:
            self.log_dict['ent_loss'] = 0.0
            self.log_dict['net_kbytes_codebook'] = 0.0

    @torch.cuda.nvtx.range("MultiviewTrainer.step")
    def step(self, data):
        """Implement the optimization over image-space loss.
        """
        # Map to device
        rays = data['rays'].to(self.device).squeeze(0)
        img_gts = data['rgb'].to(self.device).squeeze(0)

        self.optimizer.zero_grad(set_to_none=True)
            
        loss = 0
        
        if self.extra_args["random_lod"]:
            # Sample from a geometric distribution
            population = [i for i in range(self.pipeline.nef.grid.num_lods)]
            weights = [2**i for i in range(self.pipeline.nef.grid.num_lods)]
            weights = [i/sum(weights) for i in weights]
            lod_idx = random.choices(population, weights)[0]
        else:
            # Sample only the max lod (None is max lod by default)
            lod_idx = None

        rb = self.pipeline(rays=rays, lod_idx=lod_idx, channels=["rgb"])

        # RGB Loss
        #rgb_loss = F.mse_loss(rb.rgb, img_gts, reduction='none')
        rgb_loss = torch.abs(rb.rgb[..., :3] - img_gts[..., :3])

        rgb_loss = rgb_loss.mean()
        loss += self.extra_args["rgb_loss"] * rgb_loss
        if self.extra_args["ldecode_enabled"] and self.entropy_reg_lambda > 0:
            avg_bits, net_bits = self.pipeline.nef.grid.ent_loss(self.total_iterations-1, is_val=self.pipeline.training)
            ent_loss = self.entropy_reg_lambda * avg_bits
            loss += ent_loss
            self.log_dict['ent_loss'] += ent_loss.item()
            self.log_dict['net_kbytes_codebook'] += net_bits/8/1024

        self.log_dict['rgb_loss'] += rgb_loss.item()

        self.log_dict['total_loss'] += loss.item()
        
        if (isinstance(self.pipeline.nef.grid, LatentGrid) and 
            self.extra_args["ldecode_enabled"] and 
            isinstance(self.pipeline.nef.grid.latent_dec, LatentDecoder) and
            self.extra_args["scale_grid_lr"]!='none'):
            with torch.no_grad():
                for group in self.optimizer.param_groups:
                    if group['name'] == 'grid':
                        norm = self.pipeline.nef.grid.latent_dec.scale_norm().item()
                        if self.extra_args["scale_grid_lr"]=='mul':
                            group['lr'] = self.extra_args["grid_lr"]* norm
                        elif self.extra_args["scale_grid_lr"]=='div':
                            group['lr'] = self.extra_args["grid_lr"]/ norm
                        else:
                            raise Exception(f"Unknown grid_lr scaling mode {self.extra_args['scale_grid_lr']}")
                        if self.total_iterations %100 == 0:
                            print("Latent learning rate: ", group['lr'], norm)
                    elif group['name'] == 'latent_dec':
                        group['lr'] = self.ldec_lr_sched(self.epoch)
                        # if self.total_iterations %100 == 0:
                        #     print("Latent decoder learning rate: ", group['lr'])

        with torch.cuda.nvtx.range("MultiviewTrainer.backward"):
            if self.scaler:
                self.scaler.scale(loss).backward()
                with torch.no_grad():
                    if (isinstance(self.pipeline.nef.grid, LatentGrid) and 
                        self.extra_args["ldecode_enabled"] and 
                        isinstance(self.pipeline.nef.grid.latent_dec, LatentDecoder)) and self.total_iterations %100 == 0:
                        latent_dec = self.pipeline.nef.grid.latent_dec
                        grad_norm = list(latent_dec.layers.children())[0].scale.grad.norm().item()
                        if grad_norm == 0.0:
                            raise Exception("Gradient norm is 0! Exiting ...")
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                loss.backward()
                self.optimizer.step()
        
    def log_cli(self):
        log_text = 'EPOCH {}/{}'.format(self.epoch, self.max_epochs)
        log_text += ' | total loss: {:>.2E}'.format(self.log_dict['total_loss'] / len(self.train_data_loader))
        log_text += ' | rgb loss: {:>.2E}'.format(self.log_dict['rgb_loss'] / len(self.train_data_loader))
        if 'ent_loss' in self.log_dict:
            log_text += ' | ent loss: {:>.2E}'.format(self.log_dict['ent_loss'] / len(self.train_data_loader))
            # log_text += ' | codebook size (kB): {:>.3E}'.format(self.log_dict['net_kbytes_codebook'] / len(self.train_data_loader))
        if isinstance(self.pipeline.nef.grid, LatentGrid) or isinstance(self.pipeline.nef.grid, HashGrid) or isinstance(self.pipeline.nef.grid, CodebookOctreeGrid):
            if 'dec_size' in self.log_dict and isinstance(self.pipeline.nef.grid.latent_dec, MultiLatentDecoder):
                log_text += ' | dec size (kB): {:>.2E}'.format(self.log_dict['dec_size'])
            log_text += ' | latent size (kB): {:>.2E}'.format(self.log_dict['latent_size'])
            log_text += ' | total size (kB): {:>.2E}'.format(self.log_dict['total_size'])
            if isinstance(self.pipeline.nef.grid, LatentGrid) and isinstance(self.pipeline.nef.grid.latent_dec, MultiLatentDecoder) or self.extra_args["use_sga"]:
                log_text += ' | temp: {:>.2E}'.format(self.pipeline.nef.grid.latent_dec.temperature)
            if self.extra_args["use_sga"]:
                log_text += ' | sga: {}'.format(self.pipeline.nef.grid.latent_dec.use_sga)
        
        log.info(log_text)

    def evaluate_metrics(self, dataset: MultiviewDataset, lod_idx, name=None, lpips_model=None):

        img_count = len(dataset)
        img_shape = dataset.img_shape

        psnr_total = 0.0
        lpips_total = 0.0
        ssim_total = 0.0
        with torch.no_grad():
            for idx, full_batch in tqdm(enumerate(dataset)):
                gts = full_batch['rgb'].to('cuda')
                rays = full_batch['rays'].to('cuda')
                rb = self.renderer.render(self.pipeline, rays, lod_idx=lod_idx)

                gts = gts.reshape(*img_shape, -1)
                rb = rb.reshape(*img_shape, -1)

                psnr_total += psnr(rb.rgb[...,:3], gts[...,:3])
                if lpips_model:
                    lpips_total += lpips(rb.rgb[...,:3], gts[...,:3], lpips_model)
                
                ssim_total += ssim(rb.rgb[...,:3], gts[...,:3])
                
                out_rb = RenderBuffer(rgb=rb.rgb, depth=rb.depth, alpha=rb.alpha,
                                      gts=gts, err=(gts[..., :3] - rb.rgb[..., :3])**2)
                exrdict = out_rb.reshape(*img_shape, -1).cpu().exr_dict()
                
                out_name = f"{idx}"
                if name is not None:
                    out_name += "-" + name

                if not base_trainer.log_metrics_only:
                    try:
                            write_exr(os.path.join(self.valid_log_dir, out_name + ".exr"), exrdict)
                    except:
                        if hasattr(self, "exr_exception"):
                            pass
                        else:
                            self.exr_exception = True
                            log.info("Skipping EXR logging since pyexr is not found.")
                    write_png(os.path.join(self.valid_log_dir, out_name + ".png"), rb.cpu().image().byte().rgb)
                    # write_png(os.path.join('val_lego/gt', f"gt_{idx}.png"), (gts*255.0).cpu().byte())

        psnr_total /= img_count
        lpips_total /= img_count
        ssim_total /= img_count

        metrics_dict = {"psnr": psnr_total, "ssim": ssim_total}

        log_text = 'EPOCH {}/{}'.format(self.epoch, self.max_epochs)
        log_text += ' | {}: {:.2f}'.format(f"{name} PSNR", psnr_total)
        log_text += ' | {}: {:.6f}'.format(f"{name} SSIM", ssim_total)

        if lpips_model:
            log_text += ' | {}: {:.6f}'.format(f"{name} LPIPS", lpips_total)
            metrics_dict["lpips"] = lpips_total
        log.info(log_text)
 
        return metrics_dict

    def render_final_view(self, num_angles, camera_distance):
        angles = np.pi * 0.1 * np.array(list(range(num_angles + 1)))
        x = -camera_distance * np.sin(angles)
        y = self.extra_args["camera_origin"][1]
        z = -camera_distance * np.cos(angles)
        os.environ["RENDERING_FINAL"] = "1"
        for d in range(self.extra_args["num_lods"]):
            out_rgb = []
            for idx in tqdm(range(num_angles + 1), desc=f"Generating 360 Degree of View for LOD {d}"):
                log_metric_to_wandb(f"LOD-{d}-360-Degree-Scene/step", idx, step=idx)
                out = self.renderer.shade_images(
                    self.pipeline,
                    f=[x[idx], y, z[idx]],
                    t=self.extra_args["camera_lookat"],
                    fov=self.extra_args["camera_fov"],
                    lod_idx=d,
                    camera_clamp=self.extra_args["camera_clamp"]
                )
                out = out.image().byte().numpy_dict()
                if out.get('rgb') is not None:
                    log_images_to_wandb(f"LOD-{d}-360-Degree-Scene/RGB", out['rgb'].T, idx)
                    out_rgb.append(Image.fromarray(np.moveaxis(out['rgb'].T, 0, -1)))
                if out.get('rgba') is not None:
                    log_images_to_wandb(f"LOD-{d}-360-Degree-Scene/RGBA", out['rgba'].T, idx)
                if out.get('depth') is not None:
                    log_images_to_wandb(f"LOD-{d}-360-Degree-Scene/Depth", out['depth'].T, idx)
                if out.get('normal') is not None:
                    log_images_to_wandb(f"LOD-{d}-360-Degree-Scene/Normal", out['normal'].T, idx)
                if out.get('alpha') is not None:
                    log_images_to_wandb(f"LOD-{d}-360-Degree-Scene/Alpha", out['alpha'].T, idx)
                wandb.log({})
        
            if not base_trainer.log_metrics_only:
                rgb_gif = out_rgb[0]
                gif_path = os.path.join(self.log_dir, "rgb.gif")
                rgb_gif.save(gif_path, save_all=True, append_images=out_rgb[1:], optimize=False, loop=0)
                wandb.log({f"360-Degree-Scene/RGB-Rendering/LOD-{d}": wandb.Video(gif_path)})
    
    def validate(self):
        self.pipeline.eval()

        # record_dict contains trainer args, but omits torch.Tensor fields which were not explicitly converted to
        # numpy or some other format. This is required as parquet doesn't support torch.Tensors
        # (and also for output size considerations)
        record_dict = {k: v for k, v in self.extra_args.items() if not isinstance(v, torch.Tensor)}
        dataset_name = os.path.splitext(os.path.basename(self.validation_dataset.dataset_path))[0]
        model_fname = os.path.abspath(os.path.join(self.log_dir, f'model.pth'))
        record_dict.update({"dataset_name" : dataset_name, "epoch": self.epoch, 
                            "exp_name" : self.exp_name, "model_fname": model_fname})
        parent_log_dir = os.path.dirname(self.log_dir)

        log.info("Beginning validation...")
        img_shape = self.validation_dataset.img_shape
        log.info(f"Running validation on dataset with {len(self.validation_dataset)} images "
                 f"at resolution {img_shape[0]}x{img_shape[1]}")

        self.valid_log_dir = os.path.join(self.log_dir, "val")
        log.info(f"Saving validation result to {self.valid_log_dir}")
        if not os.path.exists(self.valid_log_dir):
            os.makedirs(self.valid_log_dir)

        lods = list(range(self.pipeline.nef.grid.num_lods))
        try:
            from lpips import LPIPS
            lpips_model = LPIPS(net='vgg').cuda()
        except:
            lpips_model = None
            if hasattr(self, "lpips_exception"):
                pass
            else:
                self.lpips_exception = True
                log.info("Skipping LPIPS since lpips is not found.")
        evaluation_results = self.evaluate_metrics(self.validation_dataset, lods[-1],
                                                   f"lod{lods[-1]}", lpips_model=lpips_model)
        record_dict.update(evaluation_results)
        if self.using_wandb:
            for key in evaluation_results:
                log_metric_to_wandb(f"Validation/{key}", evaluation_results[key], self.epoch)
        
        if not base_trainer.log_metrics_only:
            df = pd.DataFrame.from_records([record_dict])
            df['lod'] = lods[-1]
            fname = os.path.join(parent_log_dir, f"logs.parquet")
            if os.path.exists(fname):
                df_ = pd.read_parquet(fname)
                df = pd.concat([df_, df])
            df.to_parquet(fname, index=False)

    def pre_training(self):
        """
        Override this function to change the logic which runs before the first training iteration.
        This function runs once before training starts.
        """
        super().pre_training()
        if self.using_wandb and not base_trainer.log_metrics_only:
            for d in range(self.extra_args["num_lods"]):
                wandb.define_metric(f"LOD-{d}-360-Degree-Scene")
                wandb.define_metric(
                    f"LOD-{d}-360-Degree-Scene",
                    step_metric=f"LOD-{d}-360-Degree-Scene/step"
                )

    def post_training(self):
        """
        Override this function to change the logic which runs after the last training iteration.
        This function runs once after training ends.
        """
        wandb_viz_nerf_angles = self.extra_args.get("wandb_viz_nerf_angles", 0)
        wandb_viz_nerf_distance = self.extra_args.get("wandb_viz_nerf_distance")
        if self.using_wandb and wandb_viz_nerf_angles != 0 and not base_trainer.log_metrics_only:
            self.render_final_view(
                num_angles=wandb_viz_nerf_angles,
                camera_distance=wandb_viz_nerf_distance
            )
        super().post_training()

wisp/trainers/multiview_trainer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import logging as log
import torch
from wisp.trainers import BaseTrainer, log_metric_to_wandb, log_images_to_wandb
from wisp.datasets import MeshSampledSDFDataset, OctreeSampledSDFDataset, SDFBatch
from wisp.ops.sdf import compute_sdf_iou
from wisp.ops.image import hwc_to_chw


class SDFTrainer(BaseTrainer):

    def init_log_dict(self):
        """Custom logging dictionary.
        """
        super().init_log_dict()
        self.log_dict['rgb_loss'] = 0
        self.log_dict['l2_loss'] = 0

    @torch.cuda.nvtx.range("SDFTrainer.step")
    def step(self, data: SDFBatch):
        """Implement training from ground truth TSDF.
        """
        # Map to device
        pts = data['coords'].to(self.device)
        gts = data['sdf'].to(self.device)
        if self.extra_args["sample_tex"]:
            rgb = data['rgb'].to(self.device)

        # Prepare for inference
        batch_size = pts.shape[0]
        self.pipeline.zero_grad()

        # Calculate loss
        loss = 0

        l2_loss = 0.0
        _l2_loss = 0.0
        rgb_loss = 0.0

        preds = []

        if self.extra_args["sample_tex"]:
            for lod_idx in self.loss_lods:
                preds.append(*self.pipeline.nef(coords=pts, lod_idx=lod_idx, channels=["rgb", "sdf"]))
            for i, pred in zip(self.loss_lods, preds):
                _rgb_loss = ((pred[0] - rgb[...,:3])**2).sum()

                rgb_loss += _rgb_loss     

                res = 1.0
                _l2_loss = ((pred[1] - res * gts)**2).sum()
                l2_loss += _l2_loss
            
                loss += rgb_loss
                self.log_dict['rgb_loss'] += rgb_loss.item()
        else:
            for lod_idx in self.loss_lods:
                preds.append(*self.pipeline.nef(coords=pts, lod_idx=lod_idx, channels=["sdf"]))
            for i, pred in zip(self.loss_lods, preds):
                res = 1.0
                _l2_loss = ((pred - res * gts)**2).sum()
                l2_loss += _l2_loss
        
        loss += l2_loss

        loss /= batch_size

        # Update logs
        self.log_dict['l2_loss'] += _l2_loss.item()
        self.log_dict['total_loss'] += loss.item() 

        # Backpropagate
        with torch.cuda.nvtx.range("SDFTrainer.backward"):
            loss.backward()
            self.optimizer.step()

    def log_cli(self):
        """Override logging.
        """
        log_text = 'EPOCH {}/{}'.format(self.epoch, self.max_epochs)
        log_text += ' | total loss: {:>.3E}'.format(self.log_dict['total_loss'] / len(self.train_data_loader))
        log_text += ' | l2 loss: {:>.3E}'.format(self.log_dict['l2_loss'] / len(self.train_data_loader))
        log_text += ' | rgb loss: {:>.3E}'.format(self.log_dict['rgb_loss'] / len(self.train_data_loader))
        log.info(log_text)

    def render_tb(self):
        super().render_tb()

        self.pipeline.eval()
        for d in [self.pipeline.nef.grid.num_lods - 1]:
            if self.extra_args["log_2d"]:
                out_x = self.renderer.sdf_slice(self.pipeline.nef.get_forward_function("sdf"), dim=0)
                out_y = self.renderer.sdf_slice(self.pipeline.nef.get_forward_function("sdf"), dim=1)
                out_z = self.renderer.sdf_slice(self.pipeline.nef.get_forward_function("sdf"), dim=2)
                out_x = torch.FloatTensor(out_x)
                out_y = torch.FloatTensor(out_y)
                out_z = torch.FloatTensor(out_z)
                self.writer.add_image(f'Cross-section/X/{d}', hwc_to_chw(out_x), self.epoch)
                self.writer.add_image(f'Cross-section/Y/{d}', hwc_to_chw(out_y), self.epoch)
                self.writer.add_image(f'Cross-section/Z/{d}', hwc_to_chw(out_z), self.epoch)
                if self.using_wandb:
                    log_images_to_wandb(f'Cross-section/X/{d}', hwc_to_chw(out_x), self.epoch)
                    log_images_to_wandb(f'Cross-section/Y/{d}', hwc_to_chw(out_y), self.epoch)
                    log_images_to_wandb(f'Cross-section/Z/{d}', hwc_to_chw(out_z), self.epoch)

    def validate(self):
        """Implement validation. Just computes IOU.
        """
        # Same as training since we're overfitting
        if isinstance(self.train_dataset, MeshSampledSDFDataset):
            metric_name = "volumetric_iou"
        elif isinstance(self.train_dataset, OctreeSampledSDFDataset):
            metric_name = "narrowband_iou"
        else:
            raise NotImplementedError
        
        val_dict = {}
        val_dict[metric_name] = []
    
        # Uniform points metrics
        for n_iter, data in enumerate(self.train_data_loader):

            pts = data['coords'].to(self.device)
            gts = data['sdf'].to(self.device)
            nrm = data['normals'].to(self.device) if data.get('normals') is not None else None

            for lod_idx in self.loss_lods:
                pred = self.pipeline.nef(coords=pts, lod_idx=lod_idx, channels="sdf")
                val_dict[metric_name] += [float(compute_sdf_iou(pred, gts))]

        log_text = 'EPOCH {}/{}'.format(self.epoch, self.max_epochs)

        for k, v in val_dict.items():
            score_total = 0.0
            for lod, score in zip(self.loss_lods, v):
                self.writer.add_scalar(f'Validation/{k}/{lod}', score, self.epoch)
                score_total += score
            log_text += ' | {}: {:.4f}'.format(k, score_total / len(v))
        log.info(log_text)

wisp/trainers/sdf_trainer.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

from .helper_classes import DotDict
from .debug import PsDebugger
from .perf import PerfTimer, print_gpu_memory

wisp/utils/__init__.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import pdb
import polyscope as ps
import wisp.ops.mesh as mesh_ops

""" This module contains utility function and classes for debugging purposes. """


class PsDebugger:
    """This is a small convenience class to use Polyscope.
    
    You can basically inject Polyscope visualization for any function using something like:

    psd = PsDebugger()
    psd.register_point_cloud("test", point_tensor)
    psd.add_vector_quantity("test", "test_vectors", vector_tensor)
    psd.show()

    These functions mostly just follow polyscope, so for more details check out: https://polyscope.run
    """
    def __init__(self):
        ps.init()
        self.pcls = {}

    def register_curve_network(self, name, tensor, **kwargs):
        if 'cpu' not in str(tensor.device):
            tensor = tensor.cpu().detach()
        self.pcls[name] = ps.register_curve_network(name, tensor.reshape(-1, 3).numpy(), 'line', **kwargs)

    def register_point_cloud(self, name, tensor, **kwargs):
        if 'cpu' not in str(tensor.device):
            tensor = tensor.cpu().detach()
        self.pcls[name] = ps.register_point_cloud(name, tensor.reshape(-1, 3).numpy(), **kwargs)

    def add_vector_quantity(self, pcl_name, vec_name, tensor, **kwargs):
        if 'cpu' not in str(tensor.device):
            tensor = tensor.cpu().detach()
        self.pcls[pcl_name].add_vector_quantity(vec_name, tensor.reshape(-1, 3).numpy(), **kwargs)
    
    def add_scalar_quantity(self, pcl_name, s_name, tensor, **kwargs):
        if 'cpu' not in str(tensor.device):
            tensor = tensor.cpu().detach()
        self.pcls[pcl_name].add_scalar_quantity(s_name, tensor.reshape(-1).numpy(), **kwargs)
    
    def add_color_quantity(self, pcl_name, c_name, tensor, **kwargs):
        if 'cpu' not in str(tensor.device):
            tensor = tensor.cpu().detach()
        self.pcls[pcl_name].add_color_quantity(c_name, tensor.reshape(-1, 3).numpy(), **kwargs)

    def add_surface_mesh(self, name, obj_path, **kwargs):
        verts, faces = mesh_ops.load_obj(obj_path)
        ps.register_surface_mesh(name, verts.numpy(), faces.numpy(), **kwargs)
    
    def show(self):
        ps.show()
        pdb.set_trace()


wisp/utils/debug.py
Sharath-girish/Shacira
class DotDict(dict):
    """
    from https://stackoverflow.com/questions/13520421/recursive-dotdict

    a dictionary that supports dot notation
    as well as dictionary access notation
    usage: d = DotDict() or d = DotDict({'val1':'first'})
    set attributes: d.val2 = 'second' or d['val2'] = 'second'
    get attributes: d.val2 or d['val2']
    """

    __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

    def __init__(self, dct=None):
        if dct is not None:
            for key, value in dct.items():
                if hasattr(value, "keys"):
                    value = DotDict(value)
                self[key] = value

wisp/utils/helper_classes.py
Sharath-girish/Shacira
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.

import time
import torch


""" This module contains utility function and classes for measuring latency / memory. """


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def colorize_time(elapsed):
    """Returns colors based on the significance of the time elapsed.
    """
    if elapsed > 1e-3:
        return bcolors.FAIL + "{:.3e}".format(elapsed) + bcolors.ENDC
    elif elapsed > 1e-4:
        return bcolors.WARNING + "{:.3e}".format(elapsed) + bcolors.ENDC
    elif elapsed > 1e-5:
        return bcolors.OKBLUE + "{:.3e}".format(elapsed) + bcolors.ENDC
    else:
        return "{:.3e}".format(elapsed)

def print_gpu_memory():
    """Prints GPU memory used.
    """
    torch.cuda.empty_cache()
    print(f"{torch.cuda.memory_allocated()//(1024*1024)} mb")


class PerfTimer():
    """Super simple performance timer.
    """
    def __init__(self, activate=False, show_memory=False, print_mode=True):
        self.prev_time = time.process_time()
        self.start = torch.cuda.Event(enable_timing=True)
        self.end = torch.cuda.Event(enable_timing=True)
        self.prev_time_gpu = self.start.record()
        self.counter = 0
        self.activate = activate
        self.show_memory = show_memory
        self.print_mode = print_mode

    def reset(self):
        self.counter = 0
        self.prev_time = time.process_time()
        self.start = torch.cuda.Event(enable_timing=True)
        self.end = torch.cuda.Event(enable_timing=True)
        self.prev_time_gpu = self.start.record()

    def check(self, name=None):
        if self.activate:
            cpu_time = time.process_time() - self.prev_time
          
            self.end.record()
            torch.cuda.synchronize()

            gpu_time = self.start.elapsed_time(self.end) / 1e3
            
            if self.print_mode:
                cpu_time_disp = colorize_time(cpu_time)
                gpu_time_disp = colorize_time(gpu_time)
                if name:
                    print("CPU Checkpoint {}: {} s".format(name, cpu_time_disp))
                    print("GPU Checkpoint {}: {} s".format(name, gpu_time_disp))
                else:
                    print("CPU Checkpoint {}: {} s".format(self.counter, cpu_time_disp))
                    print("GPU Checkpoint {}: {} s".format(self.counter, gpu_time_disp))
                if self.show_memory:
                    print(f"{torch.cuda.memory_allocated()//1048576}MB")
                

            self.prev_time = time.process_time()
            self.prev_time_gpu = self.start.record()
            self.counter += 1
            return cpu_time, gpu_time



wisp/utils/perf.py
Sharath-girish/Shacira
import math
import numpy as np

class DecayScheduler(object):
    '''A simple class for decaying schedule of various hyperparameters.'''

    def __init__(self, total_steps, decay_name='fix', start=0, end=0, params=None):
        self.decay_name = decay_name
        self.start = start
        self.end = end
        self.total_steps = total_steps
        self.params = params

    def __call__(self, step):
        if self.decay_name == 'fix':
            return self.start
        elif self.decay_name == 'linear':
            if step>self.total_steps:
                return self.end
            return self.start + (self.end - self.start) * step / self.total_steps
        elif self.decay_name == 'exp':
            return max(self.end, self.start*(np.exp(-np.log(1/self.params['temperature'])*step/self.total_steps/self.params['decay_period'])))
            # return self.start * (self.end / self.start) ** (step / self.total_steps)
        elif self.decay_name == 'inv_sqrt':
            return self.start * (self.total_steps / (self.total_steps + step)) ** 0.5
        elif self.decay_name == 'cosine':
            return self.end + 0.5 * (self.start - self.end) * (1 + math.cos(step / self.total_steps * math.pi))
        else:
            raise ValueError('Unknown decay name: {}'.format(self.decay_name))
wisp/utils/schedulers.py
Sharath-girish/Shacira