# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from argparse import ArgumentParser


parser = ArgumentParser()
parser.add_argument("--config", type=str)
parser.add_argument("--path", type=str, default='./data/img/pluto.png')
parser.add_argument("--alias", type=str)
parser.add_argument("--ds_device", type=str, default='auto')
parser.add_argument("--log2_hashmap_size_ref", type=int, default=24)
config = parser.parse_args()
config.config_fp = __file__

if config.alias is None:
    config.alias = config.path.split('/')[-1].split('.')[0]

config.task = 'image'
config.workspace = f'log/img'
config.test = False
config.seed = None
config.cfactor = 128
config.train_shuffle_mode = 2
config.train_num_samples = 2**18
config.record_training = False

config.pe_lc0_freq = [1e1, 1e2]
config.pe_lc0_rbf_freq = [2**-3, 2**12]
config.num_levels = 2
config.level_dim = 2

# Model
config.arch = 'rbf'
config.rbf_type = 'ivq_a'
config.rbf_lc0_normalize = True
config.n_kernel = 'auto'
config.point_nn_kernel = 4
config.ks_alpha = 1
config.n_hidden_fl = 32

config.num_levels_ref = 16
config.level_dim_ref = 2
config.base_resolution_ref = 16

config.log2_hashmap_size = config.log2_hashmap_size_ref
config.base_resolution = 16
config.levels_omit = []

config.pe_freqs = []
config.pe_hg0_freq = []
config.pe_lc0_rbf_keep = 0

config.num_layers = 3
config.hidden_dim = 64
config.lc_act = 'none'
config.act = 'relu'
config.lc_init = [1e-4]
config.lcb_init = [1e-4]
config.w_init = [None, None, None]
config.b_init = [None, None, None]
config.a_init = [9, 30, 30, 30]

config.rbf_suffixes = ['0']

config.fix_params = ['kc0', 'ks0']
config.fix_params += ['a0', 'a1', 'a2', 'a3']

# Special
config.ema_decay = None
config.fp16 = False

# Training
config.max_steps = 35000
config.val_freq = 1.0
config.val_first = False
config.log_train = True
config.log_img = False
config.log_kernel_img = False
config.log_sdf_slice = False
config.save_pred = True
config.save_pred_pt = False
config.save_ckpt = False
config.train_metric_list = ['psnr', 'mae']
config.val_metric_list = ['psnr', 'mae', 'mse', 'ssim_ski', 'ms-ssim', 'lpips_alex', 'lpips_vgg']

# Optimizers
eps = 1e-15
lr = 1e-2/2
config.optims = {}
config.optims["dec"] = {'type': 'Adam', 'lr': 1e-2, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}
config.optims["hg0"] = {'type': 'Adam', 'lr': lr, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}
config.optims["lc0"] = {'type': 'Adam', 'lr': lr, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}
config.optims["lcb0"] = {'type': 'Adam', 'lr': lr, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}

# lr schedulers
T_max = config.max_steps
lr_gamma = 100
config.lr_schs = {}
config.lr_schs["dec"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}
config.lr_schs["hg0"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}
config.lr_schs["lc0"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}
config.lr_schs["lcb0"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}

# RBF params init
config.kc_init_config = {}
config.kc_init_config['0'] = {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0, 'n_iter': 10, 'weight_src': 'gt_grad'}

config.kw_init_config = {}
config.kw_init_config['0'] = {'type': 'v3', 'points_sampling': 1, 'alpha': 0.5, 'weight_exp': 1, 'weight_thres': 0, 'weight_src': 'gt_grad'}

config.save_fn = f'{config.alias}-{config.arch}'

config.exp_name = f'{config.alias}-{config.arch}'

configs/img.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

config = {
    "expname": "rbf_kplanes",
    "logdir": "./log/llff",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 4,
    "data_dirs": ["data/nerf_llff_data/fern"],
    # Data settings for LLFF
    "hold_every": 8,
    "contract": False,
    "ndc": True,
    "near_scaling": 0.89,
    "ndc_far": 2.6,

    # Optimization settings
    "num_steps": 40_001,
    "batch_size": 4096,
    "eval_batch_size": 4096,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr_config": {
        "base": 2e-2,
        "lc0": 2e-2,
        "kc0": 3e-5,
        "ks0": 1e-2,
 },

    # Regularization
    "plane_tv_weight": 0e-4,
    "plane_tv_weight_proposal_net": 0e-4,
    "l1_proposal_net_weight": 0,
    "histogram_loss_weight": 1.0, 
    "distortion_loss_weight": 0.001,

    # Training settings
    "train_fp16": True,
    "save_every": 40000,
    "valid_every": 40000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 48,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    "multiscale_res": [1, 2, 4, 8],
    "density_activation": "trunc_exp",
    "concat_features_across_scales": True,
    "linear_decoder": False,
    "grid_config": [{
        "input_coordinate_dim": 3,
        "output_coordinate_dim": [16, 16, 16, 15],
        "grid_dimensions": 2,
        "resolution": [64, 64, 64],
    }],
    # RBF related
 'field_type': 'kplane_rbf',
 'rbf_config': {
    'get_init_data': False,
    'init_steps': 2000,
    'init_data_fp': 'init_data/data_name',
    'init_rbf': True,
    's_dims': [512, 512, 512],
    'kp_ref_config': {
      'multiscale_res': [1, 2, 4, 8],
      'output_coordinate_dim': [16, 16, 16, 16],
      'resolution': [64, 64, 64],
    },
    'rbf_type': 'ivq_a',
    'rbf_lc0_normalize': True,
    'n_kernel': 'auto',
    'point_nn_kernel': 8,
    'ks_alpha': 1,
    'lc0_dim': 16,
    'pe_lc0_freq': [],
    'pe_lc0_rbf_freq': [],
    'pe_lc0_rbf_keep': 'all',
    'pe_kp_freq': [],
    'lc_init': [0.1, 0.5],
    'lcb_init': [0.1, 0.5],
    'kp_init': [0.1, 0.5],
    'kpb_init': [0, 0],
    'rbf_suffixes': ['0'],
    'fix_params': ['kc0', 'ks0'],
    'kc_init_config': {
      '0': {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0.5, 'n_iter': 10}
    },
    'kw_init_config': {
      '0': {'type': 'v3', 'points_sampling': 1, 'alpha': 0.5, 'weight_exp': 1, 'weight_thres': 0.5}
    },
 },
}


# Process config
rbf_config = config['rbf_config']
rbf_config['init_data_fp'] += f"_{rbf_config['init_steps']}.pt"
config['num_steps'] -= rbf_config['init_steps']
config['save_every'] -= rbf_config['init_steps']
config['valid_every'] -= rbf_config['init_steps']

config['expname'] = f"data_name_{config['expname']}"

configs/nerf_kplanes/llff.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

config = {
    "expname": "rbf_kplanes",
    "logdir": "./log/llff_init",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 4,
    "data_dirs": ["data/nerf_llff_data/fern"],
    # Data settings for LLFF
    "hold_every": 8,
    "contract": False,
    "ndc": True,
    "near_scaling": 0.89,
    "ndc_far": 2.6,

    # Optimization settings
    "num_steps": 40_001,
    "batch_size": 4096,
    "eval_batch_size": 4096,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr_config": {
        "base": 2e-2,
        "lc0": 2e-2,
        "kc0": 3e-5,
        "ks0": 1e-2,
 },

    # Regularization
    "plane_tv_weight": 1e-4,
    "plane_tv_weight_proposal_net": 1e-4,
    "l1_proposal_net_weight": 0,
    "histogram_loss_weight": 1.0, 
    "distortion_loss_weight": 0.001,

    # Training settings
    "train_fp16": True,
    "save_every": 40000,
    "valid_every": 40000,
    "save_outputs": False,

    # Raymarching settings
    "num_samples": 48,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    "multiscale_res": [1, 2, 4, 8],
    "density_activation": "trunc_exp",
    "concat_features_across_scales": True,
    "linear_decoder": False,
    "grid_config": [{
        "input_coordinate_dim": 3,
        "output_coordinate_dim": [16, 16, 16, 16],
        "grid_dimensions": 2,
        "resolution": [64, 64, 64],
    }],
    # RBF related
 'field_type': 'kplane_rbf',
 'rbf_config': {
    'get_init_data': True,
    'init_steps': 2000,
    'init_data_fp': 'init_data/data_name',
    'init_rbf': False,
    's_dims': [512, 512, 512],
    'kp_ref_config': {
      'multiscale_res': [1, 2, 4, 8],
      'output_coordinate_dim': [16, 16, 16, 16],
      'resolution': [64, 64, 64],
    },
    'rbf_type': 'ivq_a',
    'rbf_lc0_normalize': True,
    'n_kernel': 'auto',
    'point_nn_kernel': 8,
    'ks_alpha': 1,
    'lc0_dim': 0,
    'pe_lc0_freq': [],
    'pe_lc0_rbf_freq': [],
    'pe_lc0_rbf_keep': 'all',
    'pe_kp_freq': [],
    'lc_init': [-1e-4, 1e-4],
    'lcb_init': [0.1, 0.5],
    'kp_init': [0.1, 0.5],
    'kpb_init': None,
    'rbf_suffixes': ['0'],
    'fix_params': ['kc0', 'ks0'],
    'kc_init_config': {
      '0': {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0.5, 'n_iter': 10}
    },
    'kw_init_config': {
      '0': {'type': 'v3', 'points_sampling': 1, 'alpha': 0.5, 'weight_exp': 1, 'weight_thres': 0.5}
    },
 },
}


# Process config
rbf_config = config['rbf_config']
rbf_config['init_data_fp'] += f"_{rbf_config['init_steps']}.pt"
config['num_steps'] = rbf_config['init_steps'] + 1

config['expname'] = f"data_name_{config['expname']}"

configs/nerf_kplanes/llff_init.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

config = dict(
    dataset_name = 'blender',
    datadir = './data/nerf_synthetic',
    basedir = './log/nerf_synthetic',
    expname = 'rbf_tensorf',
    add_version_id = True,

    data_name = 'chair',
    # data_name = 'drums',
    # data_name = 'ficus',
    # data_name = 'hotdog',
    # data_name = 'lego',
    # data_name = 'materials',
    # data_name = 'mic',
    # data_name = 'ship',

    n_iters = 30000,
    batch_size = 4096,
    batch_size_init = 4096,

    N_voxel_init = int(128**3),
    N_voxel_final = int(200**3),
    N_voxel_real_init = int(128**3),
    N_voxel_real_final = int(350**3),
    upsamp_list = [2000,3000,4000,5500,7000],
    update_AlphaMask_list = [2000,4000],
    shrink_0 = True,
    no_upsample = False,
    scale_reso = True,
    fp16 = True,

    N_vis = 5,
    vis_every = 10000,

    render_test = 1,
    save_img = True,
    save_video = True,
    save_ckpt = True,

    model_name = 'nerf_tensorf.network.TensorVMSplitRBF',
    resol_min = 16,
    n_level = 2,
    level_types = ['vm'],
    n_lamb_sigma = [14, 16],
    n_lamb_sh = [28, 32],
    data_dim_color = None,  # None means no basis_mat

    shadingMode = 'ASG_Fea',
    fea2denseAct = 'softplus',

    view_pe = 3,
    fea_pe = -1,
    btn_freq = [3e-1, 1e1],
    featureC = 256,

    lr_basis = 1e-3,
    lr_basis_mat = 1e-3,
    lr_init = 0.02,
    lr_g_factor = 0.02,

    Ro_weight = 0.3,
    L1_weight_inital = 8e-5 / 6,
    L1_weight_rest = 0.,
    Ortho_weight = 0.,
    TV_weight_density = 0.,
    TV_weight_app = 0.,
    rm_weight_mask_thre = 1e-4,
)

# RBF related
config['rbf_config'] = {
        'get_init_data': False,
        'init_steps': 1000,
        'init_data_fp': f"init_data/data_name",
        'init_rbf': True,
        's_dims': 'aabb',
        'ref_config': {
          'N_voxel_final': int(300**3),
          'n_lamb_sigma': [16],
          'n_lamb_sh': [48],
        },
        'rbf_type': 'ivq_a',
        'rbf_lc0_normalize': False,
        'n_kernel': 'auto',
        'point_nn_kernel': 5,
        'ks_alpha': 1,
        'lc0_dim': 32,
        'lcd0_dim': 0,
        'pe_lc0_freq': [],
        'pe_lc0_rbf_freq': [],
        'pe_lc0_rbf_keep': 0,
        'lc_init': [-1e-6, 1e-6],
        'lcd_init': [-0e-6, 0e-6],
        'lcb_init': None,
        'rbf_suffixes': ['0'],
        'fix_params': ['kc0', 'ks0'],
        'kc_init_config': {
        '0': {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0, 'n_iter': 10,
              'weight_src': 'alpha_feat_grad'},
        },
        'kw_init_config': {
        '0': {'type': 'v3', 'points_sampling': 1, 'alpha': 0.3, 'weight_exp': 1, 'weight_thres': 0,
              'weight_src': 'alpha_feat_grad'},
        },
        'lr_config': {
          'lc0': 2e-2,
          'lcd0': 2e-2,
          'lcb0': 2e-2,
          'kc0': 3e-5,
          'ks0': 1e-2,
        },
    }


import os

# Process config
config['datadir'] = os.path.join(config['datadir'], 'data_name')
config['expname'] = f"data_name_{config['expname']}"

rbf_config = config['rbf_config']
rbf_config['init_data_fp'] += f"_{rbf_config['init_steps']}.pt"

if config['no_upsample']:
  config['N_voxel_real_init'] = config['N_voxel_real_final']

configs/nerf_tensorf/nerf_synthetic.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

config = dict(
    dataset_name = 'blender',
    datadir = './data/nerf_synthetic',
    basedir = './log/nerf_synthetic_3m',
    expname = 'rbf_tensorf',
    add_version_id = True,

    data_name = 'chair',
    # data_name = 'drums',
    # data_name = 'ficus',
    # data_name = 'hotdog',
    # data_name = 'lego',
    # data_name = 'materials',
    # data_name = 'mic',
    # data_name = 'ship',

    n_iters = 30000,
    batch_size = 4096,
    batch_size_init = 4096,

    N_voxel_init = int(128**3),
    N_voxel_final = int(200**3),
    N_voxel_real_init = int(128**3),
    N_voxel_real_final = int(235**3),
    upsamp_list = [2000,3000,4000,5500,7000],
    update_AlphaMask_list = [2000,4000],
    shrink_0 = True,
    no_upsample = False,
    scale_reso = True,
    fp16 = True,

    N_vis = 5,
    vis_every = 10000,

    render_test = 1,
    save_img = True,
    save_video = True,
    save_ckpt = True,

    model_name = 'nerf_tensorf.network.TensorVMSplitRBF',
    resol_min = 16,
    n_level = 2,
    level_types = ['vm'],
    n_lamb_sigma = [6, 16],
    n_lamb_sh = [12, 32],
    data_dim_color = None,  # None means no basis_mat

    shadingMode = 'ASG_Fea',
    fea2denseAct = 'softplus',

    view_pe = 3,
    fea_pe = -1,
    btn_freq = [3e-1, 1e1],
    featureC = 256,

    lr_basis = 1e-3,
    lr_basis_mat = 1e-3,
    lr_init = 0.02,
    lr_g_factor = 0.02,

    Ro_weight = 0.3,
    L1_weight_inital = 8e-5 / 6,
    L1_weight_rest = 0.,
    Ortho_weight = 0.,
    TV_weight_density = 0.,
    TV_weight_app = 0.,
    rm_weight_mask_thre = 1e-4,
)

# RBF related
config['rbf_config'] = {
        'get_init_data': False,
        'init_steps': 1000,
        'init_data_fp': f"init_data/data_name",
        'init_rbf': True,
        's_dims': 'aabb',
        'ref_config': {
          'N_voxel_final': int(128**3),
          'n_lamb_sigma': [16],
          'n_lamb_sh': [48],
        },
        'rbf_type': 'ivq_a',
        'rbf_lc0_normalize': False,
        'n_kernel': 'auto',
        'point_nn_kernel': 5,
        'ks_alpha': 1,
        'lc0_dim': 32,
        'lcd0_dim': 0,
        'pe_lc0_freq': [],
        'pe_lc0_rbf_freq': [],
        'pe_lc0_rbf_keep': 0,
        'lc_init': [-1e-6, 1e-6],
        'lcd_init': [-0e-6, 0e-6],
        'lcb_init': None,
        'rbf_suffixes': ['0'],
        'fix_params': ['kc0', 'ks0'],
        'kc_init_config': {
        '0': {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0, 'n_iter': 10,
              'weight_src': 'alpha_feat_grad'},
        },
        'kw_init_config': {
        '0': {'type': 'v3', 'points_sampling': 1, 'alpha': 0.3, 'weight_exp': 1, 'weight_thres': 0,
              'weight_src': 'alpha_feat_grad'},
        },
        'lr_config': {
          'lc0': 2e-2,
          'lcd0': 2e-2,
          'lcb0': 2e-2,
          'kc0': 3e-5,
          'ks0': 1e-2,
        },
    }


import os

# Process config
config['datadir'] = os.path.join(config['datadir'], 'data_name')
config['expname'] = f"data_name_{config['expname']}"

rbf_config = config['rbf_config']
rbf_config['init_data_fp'] += f"_{rbf_config['init_steps']}.pt"

if config['no_upsample']:
  config['N_voxel_real_init'] = config['N_voxel_real_final']

configs/nerf_tensorf/nerf_synthetic_3m.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

config = dict(
    dataset_name = 'blender',
    datadir = './data/nerf_synthetic',
    basedir = './log/nerf_synthetic_init',
    expname = 'rbf_tensorf',
    add_version_id = True,

    data_name = 'chair',
    # data_name = 'drums',
    # data_name = 'ficus',
    # data_name = 'hotdog',
    # data_name = 'lego',
    # data_name = 'materials',
    # data_name = 'mic',
    # data_name = 'ship',

    n_iters = 30000,
    batch_size = 4096,
    batch_size_init = 4096,

    N_voxel_init = int(128**3),
    N_voxel_final = int(200**3),
    N_voxel_real_init = int(128**3),
    N_voxel_real_final = int(350**3),
    upsamp_list = [2000,3000,4000,5500,7000],
    update_AlphaMask_list = [2000,4000],
    shrink_0 = True,
    no_upsample = False,
    scale_reso = True,
    fp16 = True,

    N_vis = 5,
    vis_every = 10000,

    render_test = 1,
    save_img = False,
    save_video = False,
    save_ckpt = False,

    model_name = 'nerf_tensorf.network.TensorVMSplitRBF',
    resol_min = 16,
    n_level = 1,
    level_types = ['vm'],
    n_lamb_sigma = [16],
    n_lamb_sh = [48],
    data_dim_color = 27,  # None means no basis_mat

    shadingMode = 'ASG_Fea',
    fea2denseAct = 'softplus',

    view_pe = -1,
    fea_pe = -1,
    btn_freq = [],
    featureC = 256,

    lr_basis = 1e-3,
    lr_basis_mat = 1e-3,
    lr_init = 0.02,
    lr_g_factor = 0.02,

    Ro_weight = 0.3,
    L1_weight_inital = 8e-5 / 6,
    L1_weight_rest = 0.,
    Ortho_weight = 0.,
    TV_weight_density = 0.,
    TV_weight_app = 0.,
    rm_weight_mask_thre = 1e-4,
)

# RBF related
config['rbf_config'] = {
        'get_init_data': True,
        'init_steps': 1000,
        'init_data_fp': f"init_data/data_name",
        'init_rbf': True,
        's_dims': 'aabb',
        'ref_config': {
          'N_voxel_final': int(300**3),
          'n_lamb_sigma': [16],
          'n_lamb_sh': [48],
        },
        'rbf_type': 'ivq_a',
        'rbf_lc0_normalize': False,
        'n_kernel': 'auto',
        'point_nn_kernel': 5,
        'ks_alpha': 1,
        'lc0_dim': 32,
        'lcd0_dim': 0,
        'pe_lc0_freq': [],
        'pe_lc0_rbf_freq': [],
        'pe_lc0_rbf_keep': 0,
        'lc_init': [-1e-6, 1e-6],
        'lcd_init': [-0e-6, 0e-6],
        'lcb_init': None,
        'rbf_suffixes': ['0'],
        'fix_params': ['kc0', 'ks0'],
        'kc_init_config': {
        '0': {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0, 'n_iter': 10,
              'weight_src': 'alpha_feat_grad'},
        },
        'kw_init_config': {
        '0': {'type': 'v3', 'points_sampling': 1, 'alpha': 0.3, 'weight_exp': 1, 'weight_thres': 0,
              'weight_src': 'alpha_feat_grad'},
        },
        'lr_config': {
          'lc0': 2e-2,
          'lcd0': 2e-2,
          'lcb0': 2e-2,
          'kc0': 3e-5,
          'ks0': 1e-2,
        },
    }


import os

# Process config
config['datadir'] = os.path.join(config['datadir'], 'data_name')
config['expname'] = f"data_name_{config['expname']}"

rbf_config = config['rbf_config']
rbf_config['init_data_fp'] += f"_{rbf_config['init_steps']}.pt"

if config['no_upsample']:
  config['N_voxel_real_init'] = config['N_voxel_real_final']

configs/nerf_tensorf/nerf_synthetic_init.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from argparse import ArgumentParser


parser = ArgumentParser()
parser.add_argument("--config", type=str)
parser.add_argument("--path", type=str, default='./data/sdf/armadillo_nrml.obj')
parser.add_argument("--alias", type=str)
parser.add_argument("--ds_device", type=str, default='auto')
parser.add_argument("--log2_hashmap_size_ref", type=int, default=15)
config = parser.parse_args()
config.config_fp = __file__

if config.alias is None:
    config.alias = config.path.split('/')[-1].split('.')[0].split('_nrml')[0]

config.task = 'sdf'
config.workspace = f'log/sdf'
config.test = False
config.seed = None
config.cfactor = 128
config.train_shuffle_mode = 1
config.train_num_samples = 2**18
config.train_presample = True
config.train_epoch_size = 1000
config.val_resolution = 1024
config.clip_sdf = None
config.record_training = False

config.pe_lc0_freq = [30, 300]
config.pe_lc0_rbf_freq = [2**0, 2**3]
config.num_levels = 16
config.level_dim = 1

# Model
config.arch = 'rbf'
config.rbf_type = 'nlin_s'
config.rbf_lc0_normalize = True
config.n_kernel = 'auto'
config.point_nn_kernel = 8
config.ks_alpha = 1
config.n_hidden_fl = 16

config.num_levels_ref = 16
config.level_dim_ref = 2
config.base_resolution_ref = 16

config.log2_hashmap_size = config.log2_hashmap_size_ref
config.base_resolution = 16
config.desired_resolution = 2048
config.levels_omit = []

config.pe_freqs = []
config.pe_hg0_freq = []
config.pe_lc0_rbf_keep = 0

config.num_layers = 3
config.hidden_dim = 64
config.lc_act = 'none'
config.act = 'relu'
config.lc_init = [1e-4]
config.lcb_init = [1e-4]
config.w_init = [None, None, None]
config.b_init = [None, None, None]
config.a_init = [9, 30, 30, 30]

config.rbf_suffixes = ['0']

config.fix_params = ['kc0', 'ks0']
config.fix_params += ['a0', 'a1', 'a2', 'a3']

# Special
config.ema_decay = None
config.fp16 = False

# Training
config.max_steps = 20000
config.val_freq = 1.0
config.val_first = False
config.log_train = True
config.log_img = False
config.log_kernel_img = False
config.log_sdf_slice = False
config.save_pred = True
config.save_pred_pt = False
config.save_ckpt = False
config.train_metric_list = ['psnr', 'mae']
config.val_metric_list = ['psnr', 'iou', 'mae', 'mesh']
config.metric_config = {}
config.metric_config['n_surface_samples'] = 1000000
config.metric_config['fscore_tau'] = [1e-2, 2e-3, 1e-3, 2e-4, 1e-4]

# Optimizers
eps = 1e-15
lr = 1e-4
config.optims = {}
config.optims["dec"] = {'type': 'Adam', 'lr': 1e-4, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}
config.optims["hg0"] = {'type': 'Adam', 'lr': lr, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}
config.optims["lc0"] = {'type': 'Adam', 'lr': lr, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}
config.optims["lcb0"] = {'type': 'Adam', 'lr': lr, 'betas': (0.9, 0.99), 'eps': eps, 'wd': 0}

# lr schedulers
T_max = config.max_steps
lr_gamma = 10
config.lr_schs = {}
config.lr_schs["dec"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}
config.lr_schs["hg0"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}
config.lr_schs["lc0"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}
config.lr_schs["lcb0"] = {'type': 'cosine', 'T_max': T_max, 'gamma': lr_gamma}

# RBF params init
config.kc_init_config = {}
config.kc_init_config['0'] = {'type': 'v3', 'points_sampling': 1, 'reg_sampling': 0, 'weight_exp': 1, 'weight_thres': 0, 'n_iter': 10, 'weight_src': 'gt_inv', 'weight_src_thres': 1e-9, 'weight_src_npts': 26214400}

config.kw_init_config = {}
config.kw_init_config['0'] = {'type': 'v3', 'points_sampling': 1, 'alpha': 0.5, 'weight_exp': 1, 'weight_thres': 0, 'weight_src': 'gt_inv', 'weight_src_thres': 1e-9, 'weight_src_npts': 26214400}

config.save_fn = f'{config.alias}-{config.arch}'

config.exp_name = f'{config.alias}-{config.arch}'

configs/sdf.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import time
import math
import numpy as np
import trimesh
import torch
import torch.nn as nn
import torch.nn.functional as F
from pykdtree.kdtree import KDTree

from thirdparty.torch_ngp.encoding import get_encoder
import util_misc
import util_clustering
import util_network


act_dict = {'none': 0, 'relu': 1, 'sine': 2, 'gauss': 3}


class IMGNetwork(nn.Module):
    def __init__(self,
                 cmin,
                 cmax,
                 encoding="hashgrid",
                 num_layers=3,
                 skips=[],
                 hidden_dim=64,
                 in_dim=2,
                 out_dim=3,
                 num_levels=16,
                 level_dim=2,
                 base_resolution=16,
                 log2_hashmap_size=24,
                 desired_resolution=2048,
                 act='relu',
                 lc_act='relu',
                 lc_init=1e-4,
                 lca_init=None,
                 w_init=None,
                 b_init=None,
                 a_init=None,
                 pe_freqs=[],
                 levels_omit=[],
                 ):
        super().__init__()
        if not isinstance(cmin, list):
            cmin = [cmin] * in_dim
        if not isinstance(cmax, list):
            cmax = [cmax] * in_dim
        self.register_buffer('cmin', torch.tensor(cmin))  # ... y x
        self.register_buffer('cmax', torch.tensor(cmax))  # ... y x

        self.num_layers = num_layers
        self.skips = skips
        self.hidden_dim = hidden_dim
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.act = act
        self.lc_act = lc_act
        self.lc_init = lc_init
        self.lca_init = lca_init
        self.w_init = w_init
        self.b_init = b_init
        self.a_init = a_init

        self.encoder, self.backbone_in_dim = get_encoder(encoding, input_dim=in_dim, num_levels=num_levels, 
            level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, 
            desired_resolution=desired_resolution, levels_omit=levels_omit)
        self.encoder.embeddings.data.uniform_(-self.lc_init, self.lc_init)
        self.out_lc0_dim = self.backbone_in_dim

        # PE
        self.pe = None
        if len(pe_freqs) > 0:
            P = torch.cat([torch.eye(2)*2**i for i in pe_freqs], 1)
            self.pe = util_network.PE(P)
            self.backbone_in_dim += self.pe.out_dim
        
        backbone = []
        self.w_dims = []
        for l in range(num_layers):
            if l == 0:
                in_dim = self.backbone_in_dim
            elif l in self.skips:
                in_dim = self.hidden_dim + self.backbone_in_dim
            else:
                in_dim = self.hidden_dim
            if l == num_layers - 1:
                out_dim = self.out_dim
            else:
                out_dim = self.hidden_dim
            backbone.append(nn.Linear(in_dim, out_dim, bias=True))
            self.w_dims.append([out_dim, in_dim])
        self.backbone = nn.ModuleList(backbone)

        self.init_w_b()

    def count_params(self):
        params = {}
        params['hg'] = util_misc.count_parameters(self.encoder)
        params['mlp'] = util_misc.count_parameters(self.backbone)
        params['total'] = sum([v for _, v in params.items()])
        for k, v in params.items():
            print(f'{k} params: {v}')
        return params

    def init_w_b(self):
        with torch.no_grad():
            for i in range(self.num_layers):
                wi = self.backbone[i].weight
                bi = self.backbone[i].bias
                out_dim, in_dim = self.w_dims[i]

                if len(self.w_init) > i and self.w_init[i] is not None:
                    val = self.w_init[i]
                    nn.init.uniform_(wi, -val, val)
                elif self.act == 'sine':
                    if i == 0:
                        nn.init.uniform_(wi, -1 / in_dim, 1 / in_dim)
                    else:
                        nn.init.uniform_(wi, -np.sqrt(6 / in_dim) / 30, np.sqrt(6 / in_dim) / 30)
                
                if len(self.b_init) > i and self.b_init[i] is not None:
                    val = self.b_init[i]
                    nn.init.uniform_(bi, -val, val)

    def clip_rbf_params(self):
        pass

    def forward(self, x, **kwargs):
        # x: [B, 3]
        out_other = {}

        h = self.encoder(x/self.cmax.flip(-1)[None])
        out_other['out_lc0'] = h

        if self.lc_act == 'none':
            pass
        elif self.lc_act == 'relu':
            h = F.relu(h, inplace=True)
        elif self.lc_act == 'sine':
            h = util_network.scaledsin_activation(h, torch.tensor(self.lca_init).to(h.device))
        else:
            raise NotImplementedError

        if self.pe is not None:
            h = torch.cat((h, self.pe(x)), dim=-1)

        for l in range(self.num_layers):
            if l in self.skips:
                h = torch.cat([h, x], dim=-1)
            h = self.backbone[l](h)
            if l != self.num_layers - 1:
                if self.act is None:
                    pass
                elif self.act == 'relu':
                    h = F.relu(h, inplace=True)
                elif self.act == 'sine':
                    h = util_network.scaledsin_activation(h, torch.tensor(30).to(h.device))
                elif self.act == 'gauss':
                    h = util_network.gaussian_activation(h, self.a_init[l])
                else:
                    raise NotImplementedError

        return h, out_other


class rbf(nn.Module):
    def __init__(self, cmin, cmax, s_dims, in_dim=2, out_dim=3, 
                num_layers=3, skips=[], hidden_dim=64, n_hidden_fl=20,
                num_levels_ref=16, level_dim_ref=2, base_resolution_ref=16, log2_hashmap_size_ref=24,
                num_levels=2, level_dim=2, base_resolution=16, log2_hashmap_size=24, desired_resolution=2048,
                rbf_type='nlin_s', n_kernel=64, point_nn_kernel=4, ks_alpha=1, 
                lc_init=[3e-1], lcb_init=None,
                w_init=None, b_init=None, a_init=None,
                sparse_embd_grad=True, act='relu', lc_act=None, 
                rbf_suffixes=None, kc_init_config=None, 
                rbf_lc0_normalize=True, pe_freqs=[], pe_lc0_freq=None, pe_hg0_freq=None,
                pe_lc0_rbf_freq=None, pe_lc0_rbf_keep=None, 
                **kwargs):
        super().__init__()
        if not isinstance(cmin, list):
            cmin = [cmin] * in_dim
        if not isinstance(cmax, list):
            cmax = [cmax] * in_dim
        self.cmin = torch.tensor(cmin)  # ... y x
        self.cmax = torch.tensor(cmax)  # ... y x
        self.register_buffer('cmin_gpu', self.cmin.clone())  # ... y x
        self.register_buffer('cmax_gpu', self.cmax.clone())  # ... y x
        self.s_dims = torch.tensor(s_dims)  # ... h w

        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_levels = num_levels
        self.rbf_type = rbf_type
        self.point_nn_kernel = point_nn_kernel
        self.lc_init = lc_init
        self.lcb_init = lcb_init
        self.w_init = w_init
        self.b_init = b_init
        self.a_init = a_init
        self.act = act_dict[act]
        self.lc_act = act_dict[lc_act]
        self.rbf_suffixes = rbf_suffixes
        self.rbf_lc0_normalize = rbf_lc0_normalize
        self.pe_lc0_rbf_freq = pe_lc0_rbf_freq
        self.pe_lc0_rbf_keep = pe_lc0_rbf_keep
        self.pe_lc0_freq = pe_lc0_freq
        self.pe_hg0_freq = pe_hg0_freq
        self.kc_mult = 1  # kc_mult=2 means extra copy of initial kc

        lc0_dim = n_hidden_fl
        hg0_dim = num_levels*level_dim
        lcb0_dim = lc0_dim + hg0_dim
        self.backbone_in_dim = lcb0_dim
        lcb0_dim = self.backbone_in_dim

        if len(self.pe_lc0_freq) >= 2:
            self.pe_lc0_freqs = torch.linspace(np.log2(self.pe_lc0_freq[0]), np.log2(self.pe_lc0_freq[1]), 
                                               hidden_dim, device=0)
            self.pe_lc0_freqs = torch.exp2(self.pe_lc0_freqs)
        else:
            self.pe_lc0_freqs = None
        if len(self.pe_lc0_rbf_freq) >= 2 and self.pe_lc0_rbf_keep < lc0_dim:
            self.pe_lc0_rbf_freqs = torch.linspace(np.log2(self.pe_lc0_rbf_freq[0]), np.log2(self.pe_lc0_rbf_freq[1]), 
                                                   lc0_dim - self.pe_lc0_rbf_keep, device=0)
            self.pe_lc0_rbf_freqs = torch.exp2(self.pe_lc0_rbf_freqs)
        else:
            self.pe_lc0_rbf_freqs = None

        self.kc_init_regular = {}
        for k, v in kc_init_config.items():
            if v['type'] == 'none':
                self.kc_init_regular[k] = True
            else:
                self.kc_init_regular[k] = False

        # Get reference number of params
        temp, _ = get_encoder('hashgrid', input_dim=in_dim, num_levels=num_levels_ref, 
            level_dim=level_dim_ref, base_resolution=base_resolution_ref, log2_hashmap_size=log2_hashmap_size_ref, 
            desired_resolution=desired_resolution, allocate_params=False)
        n_params_ref = temp.n_params+(32+hidden_dim*(num_layers-2)+out_dim)*hidden_dim+(hidden_dim*(num_layers-1)+out_dim)

        # PE
        self.pe_x = None
        if len(pe_freqs) > 0:
            P = torch.cat([torch.eye(self.in_dim)*2**i for i in pe_freqs], 1)
            self.pe_x = util_network.PE(P)
            self.backbone_in_dim += self.pe_x.out_dim

        # Hash grid
        if num_levels > 0:
            self.hg0, _ = get_encoder('hashgrid', input_dim=in_dim, num_levels=num_levels, 
                level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, 
                desired_resolution=desired_resolution)
            self.hg0.embeddings.data.uniform_(-self.lc_init[0], self.lc_init[0])

        # Backbone
        backbone = []
        self.w_dims = []
        for l in range(num_layers):
            if l == 0:
                in_dim_i = self.backbone_in_dim
            else:
                in_dim_i = self.hidden_dim
            if l == num_layers - 1:
                out_dim_i = self.out_dim
            else:
                out_dim_i = self.hidden_dim
            self.w_dims.append([out_dim_i, in_dim_i])
            backbone.append(nn.Linear(in_dim_i, out_dim_i, bias=True))
        self.backbone = nn.ModuleList(backbone)
        self.lcb0 = torch.nn.Parameter(torch.zeros(lcb0_dim))

        # Activation parameters
        for l in range(num_layers+1):
            setattr(self, f'a{l}', torch.nn.Parameter(torch.ones(1) * self.a_init[l]))

        # RBF parameters
        if n_kernel == 'auto':
            n_params_rbf = n_params_ref - self.count_params(exclude_rbf_lc0=True, verbose=False)['total'][-1]
            n_kernel = int(n_params_rbf // (lc0_dim + util_misc.get_rbf_params_per_kernel(
                self.rbf_type, in_dim, self.kc_mult)))
        if self.kc_init_regular['0']:
            n_kernel = util_misc.get_lower_int_power(n_kernel, in_dim)
        print('n_kernel:', n_kernel)
        self.n_kernel = n_kernel
        self.lc_dims = [[n_kernel, lc0_dim]]
        self.kc0, self.ks0, self.lc0, self.ks_dims, k_dims, kc_interval = self.create_rbf_params(
            self.rbf_type, n_kernel, in_dim, lc0_dim, sparse_embd_grad, 
            self.cmin, self.cmax, ks_alpha, is_bag=False)
        self.register_buffer(f'k_dims_0', k_dims)
        self.register_buffer(f'kci0', kc_interval)
        
        self.init_lc()
        self.init_w_b()

        self.rbf_fn = eval(f'util_network.rbf_{self.rbf_type}_fb')

    def create_rbf_params(self, rbf_type, n_kernel, in_dim, lc_dim, sparse_embd_grad, 
                          cmin, cmax, ks_alpha, scale_grad_by_freq=False, is_bag=True):
        if rbf_type.endswith('_a'):
            ks_dims = [in_dim, in_dim]
        elif rbf_type.endswith('_d'):
            ks_dims = [in_dim]
        elif rbf_type.endswith('_s'):
            ks_dims = [1]
        else:
            raise NotImplementedError
        kc = torch.nn.Embedding(n_kernel, in_dim, scale_grad_by_freq=False, sparse=sparse_embd_grad)
        ks = torch.nn.Embedding(n_kernel, np.prod(ks_dims), scale_grad_by_freq=False, 
            sparse=sparse_embd_grad)
        if rbf_type.endswith('_a'):
            ks.weight.data = torch.eye(in_dim)[None, ...].repeat(n_kernel, 1, 1).reshape(n_kernel, -1)
        elif rbf_type.endswith('_d'):
            ks.weight.data[:] = 1
        elif rbf_type.endswith('_s'):
            ks.weight.data[:] = 1
        else:
            raise NotImplementedError

        if is_bag:
            lc = torch.nn.EmbeddingBag(n_kernel, lc_dim, scale_grad_by_freq=scale_grad_by_freq, 
                                       sparse=sparse_embd_grad, mode='sum')
        else:
            lc = torch.nn.Embedding(n_kernel, lc_dim, scale_grad_by_freq=scale_grad_by_freq, 
                                    sparse=sparse_embd_grad)

        # Initialize
        n_kernel_per_dim = int(math.ceil(n_kernel**(1/in_dim)))
        k_dims = [n_kernel_per_dim] * in_dim
        points, _, side_interval = util_misc.get_grid_points(k_dims, align_corners=True, 
            vmin=cmin.tolist(), vmax=cmax.tolist())
        points = points.flip(-1)  # x y ...
        side_interval = side_interval.flip(-1)  # x y ...
        kc.weight.data = points[:n_kernel]
        if rbf_type.startswith('nlin_'):
            if rbf_type.endswith('_s'):
                ks.weight.data *= 1 / side_interval.mean()
            elif rbf_type.endswith('_d'):
                ks.weight.data *= 1 / side_interval[None]
            else:
                raise NotImplementedError
        else:
            if rbf_type.endswith('_s'):
                ks.weight.data *= 1 / (side_interval.mean() / 2) ** 2
            elif rbf_type.endswith('_d'):
                ks.weight.data *= 1 / (side_interval[None] / 2) ** 2
            elif rbf_type.endswith('_a'):
                ks.weight.data *= (1 / (side_interval / 2) ** 2)[:, None].repeat(1, side_interval.shape[0]).view(1, -1)
            else:
                raise NotImplementedError
        ks.weight.data *= ks_alpha

        return kc, ks, lc, ks_dims, torch.tensor(k_dims).flip(-1), side_interval

    def count_params_rbf_i(self, suffix, rbf_type, kc_mult=1):
        params = util_misc.count_parameters(getattr(self, f'lc{suffix}') if suffix=='0' else getattr(self, suffix))
        kc = getattr(self, f'kc{suffix}')
        ks = getattr(self, f'ks{suffix}')
        if suffix[0] == 'w' or suffix[0] == 'b' or suffix[0] == 'a':
            pass
        else:
            params += util_misc.count_parameters(kc) * kc_mult
            if rbf_type.endswith('_a'):
                params += (util_misc.count_parameters(ks)*(self.in_dim+1)/2/self.in_dim).astype(np.int64)
            elif rbf_type.endswith('_s') or rbf_type.endswith('_d'):
                params += (util_misc.count_parameters(ks)).astype(np.int64)
            else:
                raise NotImplementedError
        return params

    def count_params(self, exclude_rbf_lc0=False, verbose=True):
        params = {}
        params['hg0'] = util_misc.count_parameters(self.hg0) if self.num_levels > 0 else np.array([0, 0, 0], dtype=np.int64)
        params['mlp'] = util_misc.count_parameters(self.backbone)
        params['rbf_lcb0'] = util_misc.count_parameters(self.lcb0)

        if not exclude_rbf_lc0:
            params['rbf_lc0'] = self.count_params_rbf_i('0', self.rbf_type, self.kc_mult)

        params['total'] = sum([v for _, v in params.items()])
        if verbose:
            for k, v in params.items():
                print(f'{k} params: {v}')

        return params

    def init_lc(self):
        with torch.no_grad():
            for i in range(1):
                lci = getattr(self, f'lc{i}') if hasattr(self, f'lc{i}') else None
                lcbi = getattr(self, f'lcb{i}') if hasattr(self, f'lcb{i}') else None
                if isinstance(lci, torch.nn.Embedding) or isinstance(lci, torch.nn.EmbeddingBag):
                    lci = lci.weight
                n_in_i, n_out_i = self.lc_dims[i][-2], self.lc_dims[i][-1]
                if lci is not None:
                    val = self.lc_init[i]
                    nn.init.uniform_(lci, -val, val)
                if lcbi is not None:
                    val = self.lcb_init[i]
                    nn.init.uniform_(lcbi, -val, val)

    def init_w_b(self):
        with torch.no_grad():
            for l in range(self.num_layers):
                out_dim, in_dim = self.w_dims[l]

                wi = self.backbone[l].weight
                if self.w_init[l] is None:
                    pass
                elif self.w_init[l] == 'act':
                    if self.act == 2:
                        if l == 0:
                            nn.init.uniform_(wi, -1 / in_dim, 1 / in_dim)
                        else:
                            nn.init.uniform_(wi, -np.sqrt(6 / in_dim) / 30, np.sqrt(6 / in_dim) / 30)
                else:
                    val = self.w_init[l]
                    nn.init.uniform_(wi, -val, val)
                
                bi = self.backbone[l].bias
                if self.b_init[l] is not None:
                    val = self.b_init[l]
                    nn.init.uniform_(bi, -val, val)

    def get_kc(self, suffix):
        return getattr(self, 'kc' + suffix).weight

    def get_ks(self, suffix):
        return getattr(self, 'ks' + suffix).weight.view(-1, *self.ks_dims)

    def clip_rbf_params(self):
        for k in self.rbf_suffixes:
            if getattr(self, 'ks'+k).weight.requires_grad:
                if k == '0':
                    vmin_factor = 100
                    vmax_factor = 10
                else:
                    vmin_factor = 1
                    vmax_factor = 2
                util_misc.clip_kw_sq(getattr(self, 'ks'+k).weight, self.rbf_type, 
                                     self.cmin.flip(-1), self.cmax.flip(-1), 
                                     self.s_dims.flip(-1), is_ks=True, is_flat=True, 
                                     vmin_factor=vmin_factor, vmax_factor=vmax_factor,
                                     vmin_min=None)

    def update_point_kernel_idx_i(self, x_g, kc, device=0):
        '''
        x_g: [p d], tensor
        '''
        kd_tree = KDTree(kc.detach().cpu().numpy())
        _, kernel_idx = util_clustering.query_chunked(kd_tree, x_g.cpu().numpy(), 
            k=self.point_nn_kernel, sqr_dists=True, chunk_size=int(2e8), return_dist=False)
        point_kernel_idx = torch.tensor(kernel_idx.astype(np.int32), device=device)  # [p topk]
        if point_kernel_idx.ndim == 1:
            point_kernel_idx = point_kernel_idx[:, None]
        return point_kernel_idx
    
    def update_point_kernel_idx(self, x_g, device=0):
        t = time.time()
        for k in self.rbf_suffixes:
            if not self.kc_init_regular[k] and self.point_nn_kernel > 0:
                point_kernel_idx = self.update_point_kernel_idx_i(x_g, getattr(self, f'kc{k}').weight, device)
                setattr(self, f'point_kernel_idx_{k}', point_kernel_idx)
        print('Update point kernel idx', time.time() - t)

    def forward_kernel_idx(self, x_g, point_idx, suffix):
        if self.use_train_knn and hasattr(self, 'point_kernel_idx_' + suffix):
            # Use pre-computed knn
            point_kernel_idx = getattr(self, 'point_kernel_idx_' + suffix)
            kernel_idx = point_kernel_idx[point_idx.to(point_kernel_idx.device)].to(x_g.device)  # [p topk]
        elif self.kc_init_regular[suffix]:
            # Find first ring neighbors in regular grid
            kernel_idx = util_misc.get_multi_index_first_ring(x_g, 
                getattr(self, f'k_dims_{suffix}'), getattr(self, f'kci{suffix}'), 
                self.cmin_gpu.flip(-1), self.cmax_gpu.flip(-1), True)  # (p, n_ngb, n_in)
            kernel_idx = util_misc.ravel_multi_index(kernel_idx, getattr(self, f'k_dims_{suffix}'))  # (p, n_ngb)
        else:
            # Knn on the fly
            kd_tree = KDTree(getattr(self, 'kc' + suffix).weight.detach().cpu().numpy())
            _, kernel_idx = util_clustering.query_chunked(kd_tree, x_g.cpu().numpy(), 
                k=self.point_nn_kernel, sqr_dists=True, chunk_size=int(2e8), return_dist=False)
            kernel_idx = torch.tensor(kernel_idx.astype(np.int32), device=x_g.device)  # [p topk]
            if kernel_idx.ndim == 1:
                kernel_idx = kernel_idx[:, None]
            
        return kernel_idx

    def forward_rbf(self, x_g, kernel_idx, suffix):
        if kernel_idx is None:  # Use all kernels for each point
            kc = getattr(self, 'kc' + suffix).weight[None]  # [1 k d]
            ks = getattr(self, 'ks' + suffix).weight[None]
            ks = ks.view(*ks.shape[:2], *self.ks_dims)  # [1 k d d] or [1 k 1]
            rbf_out, _ = self.rbf_fn(x_g, kc, ks)  # [p k_topk]
        else:
            kc = getattr(self, 'kc' + suffix)(kernel_idx)  # [p k d]
            ks = getattr(self, 'ks' + suffix)(kernel_idx).view(*kernel_idx.shape, *self.ks_dims)  # [p k d d] or [p k d] or [p k 1]
            rbf_out, _ = self.rbf_fn(x_g, kc, ks)  # [p k_topk]
        return rbf_out

    def forward(self, x_g, point_idx, **kwargs):
        """
        x_g: (p, d), coords within entire domain
        point_idx: (p), idx within block
        return: (p, n_out)
        """
        out_other = {}

        suffix = '0'
        if self.point_nn_kernel <= 0:  # Use all kernels for each point
            rbf_out = self.forward_rbf(x_g, None, suffix)  # [p nk]
            out = rbf_out @ self.lc0.weight  # [p hfl]
        else:
            kernel_idx = self.forward_kernel_idx(x_g, point_idx, suffix)
            rbf_out = self.forward_rbf(x_g, kernel_idx, suffix)  # [p k_topk]
            if self.rbf_lc0_normalize:
                rbf_out = rbf_out / (rbf_out.detach().sum(-1, keepdim=True) + 1e-8)

            out = self.lc0(kernel_idx)  # [p k_topk d_lc0]
            rbf_out = rbf_out[..., None]  # [p k_topk 1]
            if len(self.pe_lc0_rbf_freq) >= 2 and self.pe_lc0_rbf_keep < out.shape[-1]:
                if self.pe_lc0_rbf_keep > 0:
                    rbf_out = torch.cat([rbf_out.expand(-1, -1, self.pe_lc0_rbf_keep), 
                        torch.sin(rbf_out * self.pe_lc0_rbf_freqs[None, None])], -1)  # [p k_topk d_lc0]
                else:
                    rbf_out = torch.sin(rbf_out * self.pe_lc0_rbf_freqs[None, None])  # [p k_topk d_lc0]
            out = (out * rbf_out).sum(1)  # [p d_lc0]

        if self.num_levels > 0:
            out_hg = self.hg0(x_g / self.cmax_gpu.flip(-1)[None])  # [p d_hg0]
        else:
            out_hg = None
            
        if out_hg is not None:
            out = torch.cat([out_hg, out], -1)
        out = out + self.lcb0[None]

        h = out
        if self.lc_act == 0:
            pass
        elif self.lc_act == 1:
            h = F.relu(h, inplace=True)
        elif self.lc_act == 2:
            h = util_network.scaledsin_activation(h, self.a0[None])
        else:
            raise NotImplementedError

        if self.pe_x is not None:
            h = torch.cat((h, self.pe_x(x_g)), dim=-1)

        # MLP
        for l in range(self.num_layers):
            h = self.backbone[l](h)
            
            if l == 0 and self.pe_lc0_freqs is not None:
                h = h + torch.sin(h * self.pe_lc0_freqs[None])
            elif l != self.num_layers - 1:
                if self.act == 0:
                    pass
                elif self.act == 1:
                    h = F.relu(h, inplace=True)
                elif self.act == 2:
                    h = util_network.scaledsin_activation(h, getattr(self, f'a{l+1}')[None])
                elif self.act == 3:
                    h = util_network.gaussian_activation(h, 30)
                else:
                    raise NotImplementedError

        return h, out_other

img_sdf/network.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import numpy as np
import kornia
import math
import time
from tqdm import tqdm

import torch
from torch.utils.data import Dataset

import trimesh
import pysdf

import util_misc


class IMGDataset(Dataset):
    def __init__(self, gt, cmin, cmax, s_dims, num_samples=2**18, ns_per_block=1, shuffle_mode=1, device=0):
        super().__init__()
        self.gt = gt
        if self.gt.dtype is torch.uint8:
            self.gt_cpu = self.gt.to(dtype=torch.float32) / 255.  # [h w c]
        else:
            raise NotImplementedError
        self.device = device
        self.device_cuda = 0 if device == 'cpu' else device
        self.shuffle_mode = shuffle_mode

        if not isinstance(cmin, list):
            cmin = [cmin] * self.s_dims.shape[0]
        if not isinstance(cmax, list):
            cmax = [cmax] * self.s_dims.shape[0]
        self.cmin = torch.tensor(cmin)  # ... y x
        self.cmax = torch.tensor(cmax)  # ... y x
        self.s_dims = torch.tensor(s_dims)  # ... h w
        self.n_point = self.s_dims.prod().item()
        self.d = self.s_dims.shape[0]

        if self.shuffle_mode in [0, 1]:
            self.num_samples = num_samples
            self.size = math.ceil(self.n_point / self.num_samples)
            self.b_dims = torch.ones_like(self.s_dims)
            self.p_dims = self.s_dims.clone()
        elif self.shuffle_mode in [2]:
            self.ns_per_block = ns_per_block
            n_blocks = math.ceil(num_samples / self.ns_per_block)
            self.p_dims = ((self.s_dims.prod() / n_blocks)**(1 / self.d)).floor().repeat(self.d).to(torch.long)
            self.b_dims = (self.s_dims / self.p_dims).ceil().to(torch.long)
            self.num_samples = self.ns_per_block * self.b_dims.prod().item()
            self.size = math.ceil(self.p_dims.prod().item() / self.ns_per_block)
        else:
            raise NotImplementedError

        # get spatial gradients
        self.gt_grad = kornia.filters.SpatialGradient(mode='sobel', order=1, normalized=True)(self.gt_cpu.movedim(-1, 0)[None])  # [1 c 2 h w]
        self.gt_grad = self.gt_grad[0].movedim((0, 1), (-2, -1)).pow(2).sum(dim=[-2, -1]).sqrt()[..., None]  # [h w 1]
        self.gt_grad_cpu = self.gt_grad  # [h w 1]

        # reshape
        self.gt = self.gt.reshape(-1, self.gt.shape[-1])
        self.gt_grad = self.gt_grad.reshape(-1, self.gt_grad.shape[-1])

        # get points
        self.points = util_misc.get_grid_points(self.s_dims.tolist(), align_corners=False, 
            vmin=self.cmin.tolist(), vmax=self.cmax.tolist())[0].flip(-1)  # (... * h * w, d), x y ...
    
    def __len__(self):
        return self.size

    def prepare(self):
        self.points = self.points.to(self.device_cuda)
        self.gt = self.gt.to(self.device_cuda)

        if self.shuffle_mode > 0:
            self.generate_point_order(True)
        else:
            self.generate_point_order(False)

    def generate_point_order(self, shuffle=False):
        if shuffle:
            if self.shuffle_mode == 0:
                pass
            elif self.shuffle_mode == 1:
                self.point_order = util_misc.random_shuffle_chunked(self.n_point, chunk_size=10000*10000, device=self.device)
            elif self.shuffle_mode == 2:
                self.point_order, self.point_mask = util_misc.stratified_shuffle_chunked(self.s_dims, self.b_dims, 
                    chunk_size=10000*5000, device=self.device)
            else:
                raise NotImplementedError

            if self.device == 'cpu': torch.cuda.empty_cache()
        else:
            self.point_order = torch.arange(0, self.n_point, dtype=torch.long, device=self.device_cuda)

    def shuffle_order(self, batch_idx):
        if batch_idx == 0 and self.shuffle_mode > 0:
            self.generate_point_order(True)

    def __getitem__(self, idx):
        if self.shuffle_mode < 2:
            point_idx = self.point_order[idx*self.num_samples:(idx+1)*self.num_samples]
        elif self.shuffle_mode == 2:
            mask = self.point_mask[:, idx*self.ns_per_block:(idx+1)*self.ns_per_block].reshape(-1)
            point_idx = self.point_order[:, idx*self.ns_per_block:(idx+1)*self.ns_per_block].reshape(-1)[mask]
        else:
            raise NotImplementedError
        point_idx = point_idx.to(torch.long)

        return {'point_idx': point_idx, 'points': self.points[point_idx], 
            'gt': self.gt[point_idx].to(torch.float32)/255.}


class SDFDataset(Dataset):
    def __init__(self, mesh, cmin, cmax, s_dims, num_samples=2**18, size=100, presample=False, is_grid=False, shuffle_mode=0, clip_sdf=None, mesh_fp='', normalize_mesh=False, device=0):
        super().__init__()
        self.mesh = mesh.copy()
        self.s_dims = torch.tensor(s_dims)  # ... h w
        self.n_point = self.s_dims.prod().item()
        self.size = size
        self.presample = presample
        self.is_grid = is_grid
        self.device = device
        self.device_cuda = device
        self.shuffle_mode = shuffle_mode
        self.ratio = [4, 3, 1]
        self.ratio_total = sum(self.ratio)

        self.d = self.s_dims.shape[0]
        if not isinstance(cmin, list):
            cmin = [cmin] * self.d
        if not isinstance(cmax, list):
            cmax = [cmax] * self.d
        self.cmin = torch.tensor(cmin)  # ... y x
        self.cmax = torch.tensor(cmax)  # ... y x

        if normalize_mesh:
            # normalize to [-1, 1] (different from instant-sdf where is [0, 1])
            vs = self.mesh.vertices
            vmin = vs.min(0)
            vmax = vs.max(0)
            v_center = (vmin + vmax) / 2
            vs = (vs - v_center[None, :])
            v_scale = 1. / (np.sqrt(np.sum(vs**2, -1).max()) / 0.99)
            vs *= v_scale
            self.mesh.vertices = vs

        if not self.mesh.is_watertight:
            print(f"[WARN] mesh is not watertight! SDF maybe incorrect.")

        self.sdf_fn = pysdf.SDF(self.mesh.vertices, self.mesh.faces)
        
        self.num_samples = num_samples
        assert self.num_samples % 8 == 0, "num_samples must be divisible by 8."
        self.clip_sdf = clip_sdf

        if self.is_grid:
            print('Presample grid points and SDF')
            self.presample = True
            self.points = util_misc.get_grid_points(self.s_dims.tolist(), align_corners=True, 
                vmin=self.cmin.tolist(), vmax=self.cmax.tolist())[0].flip(-1)  # (... * h * w, d), x y ...
            self.size = math.ceil(self.points.shape[0] / self.num_samples)
            self.gt = torch.zeros([self.points.shape[0], 1])
            chunk_size = int(2**24)
            for i in tqdm(range(0, self.points.shape[0], chunk_size)):
                self.gt[i:i+chunk_size] = torch.from_numpy(-self.sdf_fn(self.points[i:i+chunk_size].numpy())[:, None].astype(np.float32))
        elif self.presample:
            save_fp = mesh_fp.rsplit('.', 1)[0] + '_train_points.pt'
            if os.path.exists(save_fp):
                print('Load random points and SDF')
                temp = torch.load(save_fp)
                self.points = temp['points']
                self.gt = temp['gt']
                self.size = self.points.shape[0]
                
                if self.points.shape[1] != self.num_samples:
                    num_samples = self.points.shape[1]
                    size_new = []
                    start = 0
                    for i in self.ratio:
                        end = start + int(num_samples * i / self.ratio_total)
                        size_new.append((self.size * (end - start)) // int(self.num_samples * i / self.ratio_total))
                        start = end
                    self.size = min(size_new)
                    
                    self.points = []
                    self.gt = []
                    start = 0
                    for i in self.ratio:
                        end = start + int(num_samples * i / self.ratio_total)
                        num_samples_i = int(self.num_samples * i / self.ratio_total)
                        self.points.append(temp['points'][:, start:end].reshape(-1, 3)[:self.size*num_samples_i, :].reshape(self.size, -1, 3))
                        self.gt.append(temp['gt'][:, start:end].reshape(-1, 1)[:self.size*num_samples_i].reshape(self.size, -1, 1))
                        start = end
                    self.points = torch.cat(self.points, 1)
                    self.gt = torch.cat(self.gt, 1)
            else:
                print('Presample random points and SDF')
                self.points = torch.zeros([self.size, self.num_samples, 3])
                self.gt = torch.zeros([self.size, self.num_samples, 1])
                for i in tqdm(range(self.size)):
                    points, _ = self.sample_points(compute_sdf=False)
                    self.points[i] = torch.from_numpy(points)
                chunk_size = int(2**24) // (self.num_samples * self.ratio[0] // self.ratio_total)
                for i in tqdm(range(0, self.size, chunk_size)):
                    points = self.points[i:i+chunk_size, self.num_samples * self.ratio[0] // self.ratio_total:]
                    npt = points.shape[1]
                    self.gt[i:i+chunk_size, self.num_samples * self.ratio[0] // self.ratio_total:] = \
                        torch.from_numpy(-self.sdf_fn(points.reshape(-1, 3).numpy()).astype(np.float32)).reshape(-1, npt, 1)
                print('Save random points and SDF')
                torch.save({'points': self.points, 'gt': self.gt}, save_fp)

        if self.presample and self.clip_sdf is not None:
            self.gt.clip_(-self.clip_sdf, self.clip_sdf)

        if self.is_grid:
            # get spatial gradients
            self.gt_grad = kornia.filters.SpatialGradient3d(mode='diff', order=1)(self.gt.reshape(*s_dims, -1).movedim(-1, 0)[None])  # [1 c 3 d h w]
            self.gt_grad = self.gt_grad[0].movedim((0, 1), (-2, -1)).pow(2).sum(dim=[-2, -1]).sqrt()[..., None]  # [... h w 1]
            self.gt_grad = self.gt_grad.reshape(-1, self.gt_grad.shape[-1])
            self.gt_grad_cpu = self.gt_grad  # [d h w 1]
        
    def prepare(self):
        if hasattr(self, 'points'):
            self.points = self.points.to(self.device_cuda)
        if hasattr(self, 'gt'):
            self.gt = self.gt.to(self.device_cuda)
            
        if self.shuffle_mode > 0:
            self.generate_point_order(True)
        else:
            self.generate_point_order(False)

    def generate_point_order(self, shuffle=False):
        if self.presample:
            if shuffle:
                if self.shuffle_mode == 0:
                    pass
                elif self.shuffle_mode == 1:
                    if self.is_grid:
                        self.point_order = torch.randperm(self.n_point, dtype=torch.int32, device=self.device)
                    else:
                        self.point_order = torch.randperm(self.size, dtype=torch.int32, device=self.device)
                else:
                    raise NotImplementedError
            else:
                if self.is_grid:
                    pass
                else:
                    self.point_order = torch.arange(0, self.size, dtype=torch.int32, device=self.device_cuda)

    def shuffle_order(self, batch_idx):
        if batch_idx == 0 and self.shuffle_mode > 0:
            self.generate_point_order(True)

    def sample_points(self, compute_sdf=True):
        ratio = self.ratio
        total = self.ratio_total

        # surface
        points_surface = self.mesh.sample(self.num_samples * (ratio[0] + ratio[1]) // total)
        # perturb surface
        points_surface[self.num_samples * ratio[0] // total:] += 0.01 * np.random.randn(self.num_samples * ratio[1] // total, 3)
        # random
        points_uniform = np.random.rand(self.num_samples * ratio[2] // total, 3) * 2 - 1
        points = np.concatenate([points_surface, points_uniform], axis=0).astype(np.float32)

        if compute_sdf:
            sdfs = np.zeros((self.num_samples, 1))
            sdfs[self.num_samples * ratio[0] // total:] = -self.sdf_fn(points[self.num_samples * ratio[0] // total:])[:,None].astype(np.float32)
        else:
            sdfs = None

        return points, sdfs

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        if not self.is_grid:
            if not self.presample:
                point_idx = np.zeros([1], dtype=np.int32)
                points, sdfs = self.sample_points()
                if self.clip_sdf is not None:
                    sdfs = sdfs.clip(-self.clip_sdf, self.clip_sdf)
            else:
                idx_use = self.point_order[idx].to(torch.long)
                point_idx = torch.arange(self.num_samples*idx_use, self.num_samples*(idx_use+1), device=self.points.device)
                points = self.points[idx_use]
                sdfs = self.gt[idx_use]
        else:
            if self.shuffle_mode > 0:
                point_idx = self.point_order[self.num_samples*idx:self.num_samples*(idx+1)].to(torch.long)
            else:
                point_idx = torch.arange(self.num_samples*idx, self.num_samples*(idx+1), device=self.points.device)
            points = self.points[point_idx]
            sdfs = self.gt[point_idx]

        results = {
            'gt': sdfs,
            'points': points,
            'point_idx': point_idx,
        }

        return results

img_sdf/provider.py
oppo-us-research/NeuRBF
# Modified from https://github.com/ashawkey/torch-ngp/blob/main/sdf/utils.py

import os
import glob
import tqdm
import random
import warnings
import tensorboardX
import imageio.v3 as iio
import math
import shutil

import numpy as np
import pandas as pd

import time
from datetime import datetime

import cv2
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader

from rich.console import Console
from torch_ema import ExponentialMovingAverage

import packaging
import util_metric
import util_network
import util_misc
import util_summary


def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    #torch.backends.cudnn.deterministic = True
    #torch.backends.cudnn.benchmark = True


class Trainer(object):
    def __init__(self, 
                 name, # name of this experiment
                 model, # network 
                 hparams,
                 criterion=None, # loss function, if None, assume inline implementation in train_step
                 ema_decay=None, # if use EMA, set the decay
                 local_rank=0, # which GPU am I
                 world_size=1, # total num of GPUs
                 device=None, # device to use, usually setting to None is OK. (auto choose device)
                 mute=False, # whether to mute all print
                 fp16=False, # amp optimize level
                 eval_interval=1, # eval once every $ epoch
                 workspace='workspace', # workspace to save logs & ckpts
                 use_checkpoint="scratch", # which ckpt to use at init time
                 use_tensorboardX=True, # whether to use tensorboard for logging
                 scheduler_update_every_step=True, # whether to call scheduler.step() after every train step
                 ):
        self.name = name
        self.hparams = hparams
        self.mute = mute
        self.local_rank = local_rank
        self.world_size = world_size
        self.workspace = workspace
        self.eval_interval = eval_interval
        self.use_checkpoint = use_checkpoint
        self.use_tensorboardX = use_tensorboardX
        self.time_stamp = time.strftime("%Y-%m-%d_%H-%M-%S")
        self.scheduler_update_every_step = scheduler_update_every_step
        self.device = device if device is not None else torch.device(f'cuda:{local_rank}' if torch.cuda.is_available() else 'cpu')
        self.console = Console()
        self.train_metric_list = self.hparams.train_metric_list
        self.val_metric_list = self.hparams.val_metric_list
        self.log_train = self.hparams.log_train
        if not self.log_train: self.use_tensorboardX = False
        self.t_train = 0

        model.to(self.device)
        if self.world_size > 1:
            raise NotImplementedError
        self.model = model

        if isinstance(criterion, nn.Module):
            criterion.to(self.device)
        self.criterion = criterion

        param_groups = util_network.get_param_groups(self.model)
        self.optims = util_network.configure_optimizers(param_groups, self.hparams.optims)
        self.lr_schs = util_network.configure_lr_schedulers(self.optims, self.hparams.lr_schs)

        # variable init
        self.epoch = 0
        self.global_step = 0
        self.local_step = 0

        # workspace prepare
        self.log_ptr = None
        if self.workspace is not None:
            os.makedirs(self.workspace, exist_ok=True)        
            self.log_path = os.path.join(workspace, f"log_{self.name}.txt")
            self.log_ptr = open(self.log_path, "a+")

            self.ckpt_path = os.path.join(self.workspace, 'checkpoints')
            self.best_path = f"{self.ckpt_path}/{self.name}.pth.tar"
            os.makedirs(self.ckpt_path, exist_ok=True)
            
        self.log(f'[INFO] Trainer: {self.name} | {self.time_stamp} | {self.device} | {"fp32"} | {self.workspace}')

        if self.workspace is not None:
            if self.use_checkpoint == "scratch":
                self.log("[INFO] Training from scratch ...")
            else:
                raise NotImplementedError

    def __del__(self):
        if hasattr(self, 'log_ptr') and self.log_ptr:
            self.log_ptr.close()

    def log(self, *args, **kwargs):
        if self.local_rank == 0:
            if not self.mute: 
                self.console.print(*args, **kwargs)
            if self.log_ptr: 
                print(*args, file=self.log_ptr)
                self.log_ptr.flush() # write immediately to file

    def save_img(self, pred, save_path=None):
        if save_path is None:
            save_path = os.path.join(self.workspace, 'validation', f'{self.name}_{self.epoch}.png')
        self.log(f"==> Saving image to {save_path}")
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        img = (pred.clip(self.hparams.vmin, self.hparams.vmax).cpu().numpy() * 255).astype(np.uint8)
        iio.imwrite(save_path, img)
        self.log(f"==> Finished saving image.")

    def save_mesh(self, pred, save_path=None):
        if save_path is None:
            save_path = os.path.join(self.workspace, 'validation', f'{self.name}_{self.epoch}.obj')
        self.log(f"==> Saving mesh to {save_path}")
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        mesh = util_misc.extract_mesh(pred, 0, self.extract_mesh_cmin, self.extract_mesh_cmax)
        mesh.export(save_path)
        self.log(f"==> Finished saving mesh.")

    ### ------------------------------	

    def train_step(self, data):
        X = data["points"][0]  # [B, 2]
        y = data["gt"][0]  # [B, 3]
        point_idx = data["point_idx"][0]  # [B]
        
        pred, out_other = self.model(X, point_idx=point_idx)
        loss = self.criterion(pred, y)

        return pred, y, loss, out_other

    def eval_step(self, data):
        return self.train_step(data)

    def test_step(self, data):  
        X = data["points"][0]
        point_idx = data["point_idx"][0]
        pred = self.model(X, point_idx=point_idx)
        return pred        

    ### ------------------------------

    def train(self, train_loader, valid_loader, max_steps):
        if self.local_rank == 0:
            pbar = tqdm.tqdm(total=max_steps, bar_format='{desc} {percentage:3.0f}% {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]')

            shutil.copy(self.hparams.config_fp, 
                os.path.join(self.workspace, "run", self.name, os.path.basename(self.hparams.config_fp)))
            
            if self.use_tensorboardX:
                self.writer = tensorboardX.SummaryWriter(os.path.join(self.workspace, "run", self.name))

            if self.hparams.task == 'sdf':
                self.gt_mesh = train_loader.dataset.mesh
                self.extract_mesh_cmin = valid_loader.dataset.cmin.numpy()
                self.extract_mesh_cmax = valid_loader.dataset.cmax.numpy()
        
        self.max_steps = max_steps
        self.max_epochs = math.ceil(self.max_steps / len(train_loader))

        if self.hparams.val_first:
            self.pred_latest = self.evaluate_one_epoch(valid_loader)

        # Save initial pred
        if self.hparams.record_training:
            pred_latest = self.evaluate_one_epoch(valid_loader, val_metric_list=['psnr', 'mae', 'mse', 'ssim_ski', 'lpips_alex'], full=False)
            self.val_metrics_proc = {k: [v] for k, v in self.val_metrics.items()}
            save_dir = os.path.join(self.workspace, 'results', self.name)
            os.makedirs(save_dir, exist_ok=True)
            save_fn = os.path.join(save_dir, f'step_{str(self.global_step).zfill(6)}')
            iio.imwrite(f'{save_fn}.png', (pred_latest.clip(self.hparams.vmin, self.hparams.vmax).detach().cpu().numpy() * 255).astype(np.uint8))
            
        for epoch in range(self.epoch + 1, self.max_epochs + 1):
            self.epoch = epoch
            self.train_one_epoch(train_loader, pbar, valid_loader)

            if self.epoch % self.eval_interval == 0 or self.epoch == self.max_epochs:
                self.pred_latest = self.evaluate_one_epoch(valid_loader)

        if self.local_rank == 0:
            pbar.close()
            if self.use_tensorboardX:
                self.writer.close()

    def prepare_data(self, data):
        for k, v in data.items():
            if torch.is_tensor(v):
                data[k] = v.to(self.device, non_blocking=True)
            else:
                raise NotImplementedError
        return data

    def train_one_epoch(self, loader, pbar, valid_loader):
        t = time.time()
        self.model.train()
        self.model.use_train_knn = True
        loader.dataset.shuffle_order(0)
        self.local_step = 0
        total_loss = 0
        self.t_train += time.time() - t
        metrics_str = ''

        for data in loader:
            t = time.time()
            if self.global_step == self.max_steps:
                break
            self.local_step += 1
            self.global_step += 1
            
            data = self.prepare_data(data)

            for k, v in self.optims.items():
                v.zero_grad()

            preds, truths, loss, _ = self.train_step(data)
            loss.backward()

            for k, v in self.optims.items():
                v.step()

            if self.scheduler_update_every_step:
                for k, v in self.lr_schs.items():
                    v.step()

            loss_val = loss.item()
            total_loss += loss_val
            self.t_train += time.time() - t

            if self.local_rank == 0:
                loss_avg = total_loss/self.local_step

                # Save preds during training
                if self.hparams.record_training:
                    pred_latest = self.evaluate_one_epoch(valid_loader, val_metric_list=['psnr', 'mae', 'mse', 'ssim_ski', 'lpips_alex'], full=False)
                    for k, v in self.val_metrics.items():
                        self.val_metrics_proc[k].append(v)
                    if (self.global_step % 10 == 0 or self.global_step == self.max_steps):
                        save_dir = os.path.join(self.workspace, 'results', self.name)
                        save_fn = os.path.join(save_dir, f'step_{str(self.global_step).zfill(6)}')
                        iio.imwrite(f'{save_fn}.png', (pred_latest.clip(self.hparams.vmin, self.hparams.vmax).detach().cpu().numpy() * 255).astype(np.uint8))

                if self.log_train and ((self.global_step - 1) % 10 == 0 or self.local_step == 1 or self.global_step == self.max_steps):
                    self.train_metrics = self.compute_metrics(
                        preds.detach().clip(self.hparams.vmin, self.hparams.vmax), truths.detach(), 
                        self.train_metric_list)
                    if self.use_tensorboardX:
                        for k, v in self.optims.items():
                            self.writer.add_scalar(f"train/lr_{k}", v.param_groups[0]['lr'], self.global_step)
                        self.writer.add_scalar("train/loss", loss_val, self.global_step)
                        self.writer.add_scalar("train/loss_avg", loss_avg, self.global_step)
                        for k, v in self.train_metrics.items():
                            self.writer.add_scalar(f"train/{k}", v, self.global_step)
                    metrics_str = f"psnr={self.train_metrics['psnr']:.4f}"
                
                pbar.set_description(f"Epoch {self.epoch}/{self.max_epochs} training: loss={loss_val:.6f} ({loss_avg:.6f}), {metrics_str}, lr_dec={self.optims['dec'].param_groups[0]['lr']:.6f}")
                pbar.update(1)

        if not self.scheduler_update_every_step:
            raise NotImplementedError

    def compute_metrics(self, pred, gt, metrics):
        out = {}
        has_mesh_metrics = False
        for metric in metrics:
            if metric in ['mesh']:
                has_mesh_metrics = True
                break
        if has_mesh_metrics:
            pred_mesh = util_misc.extract_mesh(pred, 0, self.extract_mesh_cmin, self.extract_mesh_cmax)

        for metric in metrics:
            if metric == 'iou':
                out['iou'] = util_metric.iou(pred, gt, 0).item()
            elif metric == 'psnr':
                out['psnr'] = util_metric.psnr(pred, gt, vmin=self.hparams.vmin, vmax=self.hparams.vmax).item()
            elif metric == 'mae':
                out['mae'] = util_metric.mae(pred, gt).item()
            elif metric == 'mse':
                out['mse'] = ((pred - gt)**2).mean().item()
            elif metric == 'ssim':
                out['ssim'] = util_metric.ssim_func(pred, gt)
            elif metric == 'ssim_ski':
                out['ssim_ski'] = util_metric.ssim_ski_func(pred, gt, vmin=self.hparams.vmin, vmax=self.hparams.vmax)
            elif metric == 'ms-ssim':
                out['ms-ssim'] = util_metric.msssim(pred, gt)
            elif metric == 'lpips_alex':
                try:
                    out['lpips_alex'] = util_metric.rgb_lpips(pred, gt, net_name='alex', device=pred.device)
                except:
                    self.log("Failed to compute lpips_alex")
                    out['lpips_alex'] = None
            elif metric == 'lpips_vgg':
                try:
                    out['lpips_vgg'] = util_metric.rgb_lpips(pred, gt, net_name='vgg', device=pred.device)
                except:
                    self.log("Failed to compute lpips_vgg")
                    out['lpips_vgg'] = None
            elif metric == 'mesh':
                out.update(util_metric.mesh_metrics(pred_mesh, self.gt_mesh, 
                    self.hparams.metric_config['n_surface_samples'], self.hparams.metric_config['fscore_tau']))
        return out

    def evaluate_one_epoch(self, loader, val_metric_list=None, full=True):
        self.model.eval()
        if not self.hparams.train_val_same_points:
            self.model.use_train_knn = False
        else:
            self.model.use_train_knn = True
        total_loss = 0

        if self.local_rank == 0:
            pbar = tqdm.tqdm(total=len(loader) * loader.batch_size, bar_format='{desc} {percentage:3.0f}% {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]                                               ')

        pred = torch.zeros([loader.dataset.n_point, self.model.out_dim])
        gt = torch.zeros([loader.dataset.n_point, self.model.out_dim])

        with torch.no_grad():
            self.local_step = 0
            for data in loader:    
                self.local_step += 1
                
                data = self.prepare_data(data)

                preds, truths, loss, out_other = self.eval_step(data)
                pred[data['point_idx']] = preds.cpu()
                gt[data['point_idx']] = truths.cpu()

                loss_val = loss.item()
                total_loss += loss_val

                # only rank = 0 will perform evaluation.
                if self.local_rank == 0:
                    pbar.set_description(f"Epoch {self.epoch}/{self.max_epochs} validation: loss={loss_val:.6f} ({total_loss/self.local_step:.6f})")
                    pbar.update(loader.batch_size)

        loss_avg = total_loss / self.local_step
        pred.clip_(self.hparams.vmin, self.hparams.vmax)
        pred = pred.reshape(*loader.dataset.s_dims, -1)
        gt = gt.reshape(*loader.dataset.s_dims, -1)

        if self.local_rank == 0:
            if val_metric_list is None:
                val_metric_list = self.val_metric_list
            self.val_metrics = self.compute_metrics(pred, gt, val_metric_list)

            if full:
                metrics_str = f"psnr={self.val_metrics['psnr']:.4f}"
                if 'iou' in self.val_metrics:
                    metrics_str += f", iou={self.val_metrics['iou']:.6f}"
                pbar.set_description(f"Epoch {self.epoch}/{self.max_epochs} validation: loss={loss_avg:.6f}, {metrics_str}")
                pbar.close()

                if self.use_tensorboardX:
                    for k, v in self.val_metrics.items():
                        self.writer.add_scalar(f"val/{k}", v, self.global_step)

        if full:
            self.write_img_summary(loader.dataset, pred, gt)

        return pred

    def inference_step(self, points):
        return self.model(points, point_idx=None)

    def evaluate_points(self, points_dict):
        self.model.eval()
        self.model.use_train_knn = False
        metrics = {}

        for k in points_dict['points']:
            with torch.no_grad():
                pts = torch.tensor(points_dict['points'][k])
                gt = torch.tensor(points_dict['sdfs'][k])
                chunks = torch.split(pts, 2**18)
                pred = []
                for chunk_pts in tqdm.tqdm(chunks):
                    pred_i, _ = self.inference_step(chunk_pts.to(self.device))
                    pred.append(pred_i.cpu())
            pred = torch.cat(pred, dim=0)
            pred.clip_(self.hparams.vmin, self.hparams.vmax)
            gt.clip_(self.hparams.vmin, self.hparams.vmax)

            metrics[k] = {}
            metrics[k]['iou'] = util_metric.iou(pred, gt, 0).item()
            metrics[k]['psnr'] = util_metric.psnr(pred, gt, vmin=self.hparams.vmin, vmax=self.hparams.vmax).item()
            metrics[k]['mae'] = util_metric.mae(pred, gt).item()

        return metrics

    def write_img_summary(self, dataset, pred, gt):
        cmin, cmax = dataset.cmin.flip(-1).numpy(), dataset.cmax.flip(-1).numpy()  # x y ...
        if self.hparams.log_img:
            util_summary.write_image_summary_new(self.writer, self.global_step, 
                dataset.gt_cpu.movedim(-1, 0)[None], 
                dataset.gt_grad_cpu.movedim(-1, 0)[None], 
                pred.cpu().movedim(-1, 0)[None], 
                cmin, cmax, suffix='')
        if self.hparams.log_kernel_img and hasattr(self.model, 'get_kc'):
            for k in self.model.rbf_suffixes:
                kc = self.model.get_kc(k).detach().cpu().numpy()
                kw_sq = util_misc.ks_to_kw_sq(self.model.get_ks(k).detach().cpu().numpy(), self.model.rbf_type)
                util_summary.write_kernel_summary_2d(
                    self.writer, self.global_step, 
                    dataset.gt_cpu.permute(-1, 0, 1)[None], 
                    dataset.gt_grad_cpu.permute(-1, 0, 1)[None], 
                    pred.cpu().permute(-1, 0, 1)[None], 
                    kc, kw_sq, cmin, cmax, suffix='_' + k)
        if self.hparams.log_sdf_slice:
            util_summary.write_sdf_slice_summary(self.writer, self.global_step, 
                gt.cpu().movedim(-1, 0)[None], 
                pred.cpu().movedim(-1, 0)[None], 
                suffix='')

    def save_checkpoint(self, metrics):
        state = metrics
        state['epoch'] = self.epoch
        state['global_step'] = self.global_step
        state['model'] = self.model.state_dict()
        file_path = f"{self.ckpt_path}/{self.name}_ep{self.epoch:04d}.pth"
        torch.save(state, file_path)
            
    def load_checkpoint(self, checkpoint=None):
        raise NotImplementedError
            
img_sdf/utils.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
os.environ["OMP_NUM_THREADS"] = "256"

import sys
import importlib
import time
import glob
import math
from PIL import Image
Image.MAX_IMAGE_PIXELS = None # enable reading large image
import GPUtil
import trimesh

import numpy as np
import torch

import img_sdf.utils as utils
import util_misc
import util_network
import util_metric
import pprint
import pickle


def main(opt):
    # Select devices
    opt.device = util_misc.select_devices('1#', force_reselect=True, excludeID=[])[0]
    print(f'GPU to be used: {opt.device}')
    os.environ["CUDA_VISIBLE_DEVICES"] = f"{opt.device}"
    temp = torch.ones([100, 100], device=0)
    del temp

    if opt.seed is not None:
        utils.seed_everything(opt.seed)

    import img_sdf.network as network

    # Process experiment name
    exp_list = glob.glob(f'{opt.workspace}/run/{opt.exp_name}-v*'.replace('[', '[[]'))
    if len(exp_list) == 0:
        vid = 0
    else:
        vid = max([int(v.split(f'{opt.exp_name}-v')[-1]) for v in exp_list]) + 1
    opt.exp_name += f'-v{vid}'
    os.makedirs(os.path.join(opt.workspace, "run", opt.exp_name), exist_ok=True)        
    print(f"[INFO] exp_name: {opt.exp_name}")

    # Load data
    if opt.task == 'image':
        gt = Image.open(opt.path).convert('RGB')
        if opt.path.endswith('.jpg') or opt.path.endswith('.png') or opt.path.endswith('.tif'):
            gt = np.uint8(gt)  # [h w c]
        else:
            raise NotImplementedError
        gt = torch.tensor(gt)  # [h w c]
        print(f"[INFO] image: {gt.shape} {opt.path}")

        s_dims = gt.shape[:-1]
        n_train_point = torch.tensor(s_dims).prod()
        if n_train_point > torch.iinfo(torch.int).max: raise NotImplementedError
        if n_train_point >= 50000*20000:
            opt.kc_init_config['0']['points_sampling'] = 2
            if 'lpips_alex' in opt.val_metric_list: opt.val_metric_list.remove('lpips_alex')
        if n_train_point >= 3000*3000:
            if 'lpips_vgg' in opt.val_metric_list: opt.val_metric_list.remove('lpips_vgg')

        if opt.cfactor is None:
            opt.cmax = [1, 1]  # ... y x
        else:
            opt.cmax = (torch.tensor(gt.shape[:2]) / opt.cfactor).tolist()  # ... y x
        opt.cmin = [-i for i in opt.cmax]  # ... y x
        desired_resolution = max(gt.shape[0:2]) // 2
        in_dim = 2
        out_dim = 3
        opt.vmin = 0
        opt.vmax = 1
    elif opt.task == 'sdf':
        gt_mesh = trimesh.load(opt.path, force='mesh')
        print(f"[INFO] mesh: {gt_mesh.vertices.shape} {gt_mesh.faces.shape} {opt.path}")
        opt.cmax = [1, 1, 1]  # ... y x
        opt.cmin = [-i for i in opt.cmax]  # ... y x
        s_dims = [opt.val_resolution] * 3
        n_train_point = opt.train_num_samples * opt.train_epoch_size
        desired_resolution=opt.desired_resolution
        in_dim = 3
        out_dim = 1
        opt.vmin = -(2**2*3)**0.5
        opt.vmax = (2**2*3)**0.5
    else:
        raise NotImplementedError
    
    print('Prepare dataset and dataloader')
    if opt.ds_device == 'auto':
        device = 'cpu' if n_train_point > int(5e8) else 0
    elif opt.ds_device == 'cpu':
        device = 'cpu'
    else:
        device = int(opt.ds_device)
    if opt.task == 'image':
        from img_sdf.provider import IMGDataset
        train_dataset = IMGDataset(gt, opt.cmin, opt.cmax, s_dims, num_samples=opt.train_num_samples, ns_per_block=1, 
            shuffle_mode=opt.train_shuffle_mode, device=device)
        valid_dataset = train_dataset
        opt.train_val_same_points = True
        train_pin_memory = False
        valid_pin_memory = False
    elif opt.task == 'sdf':
        from img_sdf.provider import SDFDataset
        train_dataset = SDFDataset(gt_mesh, opt.cmin, opt.cmax, s_dims, num_samples=opt.train_num_samples, 
            size=opt.train_epoch_size, presample=opt.train_presample, shuffle_mode=opt.train_shuffle_mode, 
            clip_sdf=opt.clip_sdf, mesh_fp=opt.path, device=device)
        valid_dataset = SDFDataset(gt_mesh, opt.cmin, opt.cmax, s_dims, num_samples=2**18, is_grid=True,
            clip_sdf=opt.clip_sdf, mesh_fp=opt.path, device=device)
        opt.train_val_same_points = False
        train_pin_memory = True if not train_dataset.presample else False
        valid_pin_memory = True if not valid_dataset.presample else False
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=False, 
        num_workers=0, pin_memory=train_pin_memory)
    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=1, shuffle=False, 
        num_workers=0, pin_memory=valid_pin_memory)

    print('Init model')
    if opt.arch == 'ngp':
        t = time.time()
        model = network.IMGNetwork(opt.cmin, opt.cmax, encoding=opt.encoding, num_layers=opt.num_layers, 
            hidden_dim=opt.hidden_dim, in_dim=in_dim, out_dim=out_dim,
            num_levels=opt.num_levels, level_dim=opt.level_dim, base_resolution=opt.base_resolution,
            log2_hashmap_size=opt.log2_hashmap_size, desired_resolution=desired_resolution, 
            act=opt.act, lc_act=opt.lc_act,
            lc_init=opt.lc_init, lca_init=opt.lca_init, w_init=opt.w_init, b_init=opt.b_init, a_init=opt.a_init,
            pe_freqs=opt.pe_freqs, levels_omit=opt.levels_omit)
        t_init_model = time.time() - t
        t_init_rbf = 0
        t_knn = 0
    else:
        t = time.time()
        net = eval(f'network.{opt.arch}')
        model = net(opt.cmin, opt.cmax, s_dims, in_dim=in_dim, out_dim=out_dim,
            num_layers=opt.num_layers, hidden_dim=opt.hidden_dim, n_hidden_fl=opt.n_hidden_fl, 
            num_levels_ref=opt.num_levels_ref, level_dim_ref=opt.level_dim_ref, 
            base_resolution_ref=opt.base_resolution_ref, log2_hashmap_size_ref=opt.log2_hashmap_size_ref, 
            num_levels=opt.num_levels, level_dim=opt.level_dim, base_resolution=opt.base_resolution,
            log2_hashmap_size=opt.log2_hashmap_size, desired_resolution=desired_resolution, 
            rbf_type=opt.rbf_type, n_kernel=opt.n_kernel, point_nn_kernel=opt.point_nn_kernel, ks_alpha=opt.ks_alpha, 
            lc_init=opt.lc_init, lcb_init=opt.lcb_init, 
            w_init=opt.w_init, b_init=opt.b_init, a_init=opt.a_init,
            sparse_embd_grad=False, act=opt.act, lc_act=opt.lc_act, rbf_suffixes=opt.rbf_suffixes, 
            kc_init_config=opt.kc_init_config, rbf_lc0_normalize=opt.rbf_lc0_normalize, 
            pe_freqs=opt.pe_freqs, pe_lc0_freq=opt.pe_lc0_freq, pe_hg0_freq=opt.pe_hg0_freq,
            pe_lc0_rbf_freq=opt.pe_lc0_rbf_freq, pe_lc0_rbf_keep=opt.pe_lc0_rbf_keep)
        t_init_model = time.time() - t
        t = time.time()
        util_network.init_rbf_params(model, train_dataset, opt.kc_init_config, opt.kw_init_config, device=0)
        t_init_rbf = time.time() - t
        t = time.time()
        if hasattr(train_dataset, 'points'):
            model.update_point_kernel_idx(train_dataset.points.view(-1, train_dataset.points.shape[-1]), device=device)
        t_knn = time.time() - t
    util_network.fix_params(model, opt.fix_params)
    model.count_params()

    if opt.task == 'image':
        train_dataset.prepare()
    elif opt.task == 'sdf':
        train_dataset.prepare()
        valid_dataset.prepare()

    if opt.task == 'image':
        criterion = torch.nn.MSELoss(reduction='mean')
    elif opt.task == 'sdf':
        from thirdparty.torch_ngp.loss import mape_loss
        criterion = mape_loss

    max_steps = opt.max_steps
    max_epochs = math.ceil(max_steps / len(train_loader))
    eval_interval = math.ceil(max_epochs * opt.val_freq)

    print('Set up trainer')
    trainer = utils.Trainer(opt.exp_name, model, hparams=opt, workspace=opt.workspace, criterion=criterion,
        ema_decay=opt.ema_decay, fp16=opt.fp16, use_checkpoint='scratch', eval_interval=eval_interval, local_rank=0)

    # Train
    trainer.train(train_loader, valid_loader, max_steps=max_steps)
    t_train = trainer.t_train
    metrics = {'grid': trainer.val_metrics}

    # Eval
    if opt.task == 'sdf':
        eval_data = torch.load(opt.path.rsplit('.', 1)[0] + '_eval_points.pt')
        metrics.update(trainer.evaluate_points(eval_data))

    # Print stats
    stats = {}
    stats['time'] = {'t_init_model': t_init_model, 't_init_rbf': t_init_rbf, 't_knn': t_knn, 't_train': t_train, 
                        't_total': t_init_model + t_init_rbf + t_knn + t_train}
    stats['misc'] = {'n_epoch': trainer.epoch, 'n_iter': trainer.global_step, 'batch_size': opt.train_num_samples}
    stats['n_param'] = {f'n_param_{k}': v.tolist() for k, v in model.count_params().items()}
    stats['metrics'] = metrics
    pprint.pprint(stats, sort_dicts=False)

    # Write stats to file
    os.makedirs(os.path.join(opt.workspace, "results"), exist_ok=True)        
    save_fn = os.path.join(opt.workspace, 'results', opt.exp_name)
    with open(f'{save_fn}.txt', 'w') as f:
        pprint.pprint(stats, f, sort_dicts=False)

    if opt.save_pred:
        if opt.task == 'image':
            trainer.save_img(trainer.pred_latest, f'{save_fn}.png')
        elif opt.task == 'sdf':
            trainer.save_mesh(trainer.pred_latest, f'{save_fn}.ply')

    if opt.save_pred_pt:
        torch.save(trainer.pred_latest.detach().cpu(), f'{save_fn}.pt')
    
    if opt.save_ckpt:
        trainer.save_checkpoint(metrics)

    if hasattr(opt, 'save_kcs0') and opt.save_kcs0:
        torch.save({'kc0': trainer.model.kc0.weight.detach().cpu(), 'ks0': trainer.model.ks0.weight.detach().cpu()}, 
                    f'{save_fn}_kcs0.pt')

    # Save process data
    if opt.record_training:
        with open(f'{save_fn}.pkl', 'wb') as fp:
            pickle.dump(trainer.val_metrics_proc, fp)


if __name__ == '__main__':
    config_fp = sys.argv[2].replace('./', '').replace('.py', '').replace('/', '.')
    opt = importlib.import_module(config_fp).config
    main(opt)

main.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import importlib
import main


if __name__ == '__main__':
    config_fp = 'configs/img.py'

    config_fp = config_fp.replace('./', '').replace('.py', '').replace('/', '.')
    config = importlib.import_module(config_fp).config

    for i in range(801, 901):
        print(f'div2k {i}')

        config.workspace = f'log/img_div2k'
        config.path = f'./data/img/div2k/DIV2K_valid_HR/{str(i).zfill(4)}.png'

        config.pe_lc0_freq = [1e0, 1e3]
        config.pe_lc0_rbf_freq = [2**-3, 2**12]
        config.num_levels = 2
        config.level_dim = 2
        config.max_steps = 35000
        config.val_freq = 1.0
        config.val_first = False
        config.log_train = False
        config.log_img = False
        config.log_kernel_img = False
        config.log_sdf_slice = False
        config.save_pred = True
        config.save_pred_pt = False
        config.save_ckpt = False
        config.save_fn = str(i).zfill(4)
        config.exp_name = str(i).zfill(4)

        main.main(config)

main_img_div2k.py
oppo-us-research/NeuRBF
# Modified from https://github.com/apchenstu/TensoRF/blob/main/train.py

import os
from tqdm.auto import tqdm
from thirdparty.tensorf.opt import config_parser

import json, random
from thirdparty.tensorf.renderer import *
from thirdparty.tensorf.utils import *
from torch.utils.tensorboard import SummaryWriter
import datetime

from thirdparty.tensorf.dataLoader import dataset_dict
import sys
import pprint
import numpy as np
import torch

from thirdparty.tensorf.models.tensorBase import count_params
import thirdparty.tensorf.models.tensorBase
import thirdparty.tensorf.models.tensoRF
import thirdparty.nrff.nrff
import nerf_tensorf.network
import util_misc

import time
import pprint


gpu = util_misc.select_devices('1#', force_reselect=True, excludeID=[])[0]
if gpu is not None:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
    print(f"CUDA_VISIBLE_DEVICES set to {gpu}")
else:
    print(f"Did not set GPU.")

temp = torch.ones([100, 100], device=0)
del temp

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

renderer = OctreeRender_trilinear_fast


class SimpleSampler:
    def __init__(self, total, batch):
        self.total = total
        self.batch = batch
        self.curr = total
        self.ids = None

    def nextids(self):
        self.curr+=self.batch
        if self.curr + self.batch > self.total:
            self.ids = torch.LongTensor(np.random.permutation(self.total))
            self.curr = 0
        return self.ids[self.curr:self.curr+self.batch]


@torch.no_grad()
def export_mesh(args):

    ckpt = torch.load(args.ckpt, map_location=device)
    kwargs = ckpt['kwargs']
    kwargs.update({'device': device})
    tensorf = eval(args.model_name)(**kwargs)
    tensorf.load(ckpt)

    alpha,_ = tensorf.getDenseAlpha()
    convert_sdf_samples_to_ply(alpha.cpu(), f'{args.ckpt[:-3]}.ply',bbox=tensorf.aabb.cpu(), level=0.005)


@torch.no_grad()
def render_test(args):
    # init dataset
    dataset = dataset_dict[args.dataset_name]
    test_dataset = dataset(args.datadir, split='test', downsample=args.downsample_train, is_stack=True)
    white_bg = test_dataset.white_bg
    ndc_ray = args.ndc_ray

    if not os.path.exists(args.ckpt):
        print('the ckpt path does not exists!!')
        return

    ckpt = torch.load(args.ckpt, map_location=device)
    kwargs = ckpt['kwargs']
    kwargs.update({'device': device})
    tensorf = eval(args.model_name)(**kwargs)
    tensorf.load(ckpt)

    logfolder = os.path.dirname(args.ckpt)
    if args.render_train:
        os.makedirs(f'{logfolder}/imgs_train_all', exist_ok=True)
        train_dataset = dataset(args.datadir, split='train', downsample=args.downsample_train, is_stack=True)
        PSNRs_test, _, _, _ = evaluation(
            train_dataset,tensorf, args, renderer, f'{logfolder}/imgs_train_all/',
            N_vis=-1, N_samples=-1, white_bg = white_bg, ndc_ray=ndc_ray,device=device)
        print(f'======> {args.expname} train all psnr: {np.mean(PSNRs_test)} <========================')

    if args.render_test:
        os.makedirs(f'{logfolder}/{args.expname}/imgs_test_all', exist_ok=True)
        evaluation(test_dataset, tensorf, args, renderer, f'{logfolder}/{args.expname}/imgs_test_all/',
                                N_vis=-1, N_samples=-1, white_bg = white_bg, ndc_ray=ndc_ray,device=device)

    if args.render_path:
        c2ws = test_dataset.render_path
        os.makedirs(f'{logfolder}/{args.expname}/imgs_path_all', exist_ok=True)
        evaluation_path(test_dataset,tensorf, c2ws, renderer, f'{logfolder}/{args.expname}/imgs_path_all/',
                                N_vis=-1, N_samples=-1, white_bg = white_bg, ndc_ray=ndc_ray,device=device)


def reconstruction(args, init_data=None):
    if args.add_timestamp:
        logfolder = f'{args.basedir}/{args.expname}{datetime.datetime.now().strftime("-%Y%m%d-%H%M%S")}'
    else:
        logfolder = f'{args.basedir}/{args.expname}'
    os.makedirs(logfolder, exist_ok=True)

    # init dataset
    dataset = dataset_dict[args.dataset_name]
    train_dataset = dataset(args.datadir, split='train', downsample=args.downsample_train, is_stack=False)
    test_dataset = dataset(args.datadir, split='test', downsample=args.downsample_train, is_stack=True)
    white_bg = train_dataset.white_bg
    near_far = train_dataset.near_far
    ndc_ray = args.ndc_ray

    # init resolution
    upsamp_list = args.upsamp_list
    update_AlphaMask_list = args.update_AlphaMask_list
    n_lamb_sigma = args.n_lamb_sigma
    n_lamb_sh = args.n_lamb_sh

    # init log file
    os.makedirs(f'{logfolder}/imgs_vis', exist_ok=True)
    os.makedirs(f'{logfolder}/imgs_rgba', exist_ok=True)
    os.makedirs(f'{logfolder}/rgba', exist_ok=True)
    summary_writer = SummaryWriter(logfolder)

    # init parameters
    aabb = train_dataset.scene_bbox.to(device)
    reso_cur = N_to_reso(args.N_voxel_init, aabb)
    reso_real_cur = N_to_reso(args.N_voxel_real_init, aabb)
    nSamples = min(args.nSamples, cal_n_samples(reso_cur,args.step_ratio))


    t = time.time()
    if args.ckpt is not None:
        ckpt = torch.load(args.ckpt, map_location=device)
        kwargs = ckpt['kwargs']
        kwargs.update({'device':device})
        tensorf = eval(args.model_name)(**kwargs)
        tensorf.load(ckpt)
    else:
        tensorf = eval(args.model_name)(aabb, reso_cur, device,
                    density_n_comp=n_lamb_sigma, appearance_n_comp=n_lamb_sh, app_dim=args.data_dim_color, near_far=near_far,
                    shadingMode=args.shadingMode, 
                    alphaMask_thres=args.alpha_mask_thre, rayMarch_weight_thres=args.rm_weight_mask_thre,
                    density_shift=args.density_shift, distance_scale=args.distance_scale,
                    pos_pe=args.pos_pe, view_pe=args.view_pe, fea_pe=args.fea_pe, featureC=args.featureC, step_ratio=args.step_ratio, fea2denseAct=args.fea2denseAct, 
                    args=args, gridSize_real=reso_real_cur, init_data=init_data)
    tensorf.to(device)
    torch.cuda.synchronize()
    t_init = time.time() - t

    n_params = count_params(tensorf)

    grad_vars = tensorf.get_optparam_groups(args.lr_init, args.lr_basis)
    if args.lr_decay_iters > 0:
        lr_factor = args.lr_decay_target_ratio**(1/args.lr_decay_iters)
    else:
        args.lr_decay_iters = args.n_iters
        lr_factor = args.lr_decay_target_ratio**(1/args.n_iters)

    print("lr decay", args.lr_decay_target_ratio, args.lr_decay_iters)

    optimizer = torch.optim.Adam(grad_vars, betas=(0.9,0.99))


    #linear in logrithmic space
    N_voxel_list = (torch.round(torch.exp(torch.linspace(np.log(args.N_voxel_init), np.log(args.N_voxel_final), len(upsamp_list)+1))).long()).tolist()[1:]
    N_voxel_real_list = (torch.round(torch.exp(torch.linspace(np.log(args.N_voxel_real_init), np.log(args.N_voxel_real_final), len(upsamp_list)+1))).long()).tolist()[1:]


    torch.cuda.empty_cache()
    PSNRs,PSNRs_test = [],[0]
    batch_size = args.batch_size_init

    allrays, allrgbs = train_dataset.all_rays, train_dataset.all_rgbs
    if not args.ndc_ray:
        allrays, allrgbs = tensorf.filtering_rays(allrays, allrgbs, bbox_only=True)
    trainingSampler = SimpleSampler(allrays.shape[0], batch_size)

    Ortho_reg_weight = args.Ortho_weight
    print("initial Ortho_reg_weight", Ortho_reg_weight)

    L1_reg_weight = args.L1_weight_inital
    print("initial L1_reg_weight", L1_reg_weight)
    TV_weight_density, TV_weight_app = args.TV_weight_density, args.TV_weight_app
    tvreg = TVLoss()
    print(f"initial TV_weight density: {TV_weight_density} appearance: {TV_weight_app}")

    scaler = torch.cuda.amp.GradScaler(enabled=args.fp16)

    t_train = 0
    pbar = tqdm(range(args.n_iters), miniters=args.progress_refresh_rate, file=sys.stdout)
    for iteration in pbar:
        t = time.time()
        ray_idx = trainingSampler.nextids()
        rays_train, rgb_train = allrays[ray_idx], allrgbs[ray_idx].to(device)

        optimizer.zero_grad()

        with torch.cuda.amp.autocast(enabled=args.fp16):
            rgb_map, alphas_map, depth_map, weights, others = renderer(rays_train, tensorf, chunk=batch_size,
                                    N_samples=nSamples, white_bg = white_bg, ndc_ray=ndc_ray, device=device, is_train=True)

            total_loss = torch.mean((rgb_map - rgb_train) ** 2)
            loss = total_loss.detach().item()

            # loss
            if others['normals'] is not None and args.Ro_weight > 0:
                Ro = torch.sum(others['normals'] * others['valid_viewdirs'], dim=-1)
                Ro = F.relu(Ro).pow(2) * others['valid_weights']
                Ro  = Ro.mean()
                loss_Ro = args.Ro_weight * Ro
                total_loss += loss_Ro
                summary_writer.add_scalar('train/reg_Ro', loss_Ro.detach().item(), global_step=iteration)

            if Ortho_reg_weight > 0:
                loss_reg = tensorf.vector_comp_diffs()
                total_loss += Ortho_reg_weight*loss_reg
                summary_writer.add_scalar('train/reg', loss_reg.detach().item(), global_step=iteration)
            if L1_reg_weight > 0:
                loss_reg_L1 = tensorf.density_L1()
                total_loss += L1_reg_weight*loss_reg_L1
                summary_writer.add_scalar('train/reg_l1', loss_reg_L1.detach().item(), global_step=iteration)

            if TV_weight_density>0:
                TV_weight_density *= lr_factor
                loss_tv = tensorf.TV_loss_density(tvreg) * TV_weight_density
                total_loss = total_loss + loss_tv
                summary_writer.add_scalar('train/reg_tv_density', loss_tv.detach().item(), global_step=iteration)
            if TV_weight_app>0:
                TV_weight_app *= lr_factor
                loss_tv = tensorf.TV_loss_app(tvreg)*TV_weight_app
                total_loss = total_loss + loss_tv
                summary_writer.add_scalar('train/reg_tv_app', loss_tv.detach().item(), global_step=iteration)

        if args.fp16:
            scaler.scale(total_loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            total_loss.backward()
            optimizer.step()

        torch.cuda.synchronize()
        t_train += time.time() - t

        PSNRs.append(-10.0 * np.log(loss) / np.log(10.0))
        summary_writer.add_scalar('train/PSNR', PSNRs[-1], global_step=iteration)
        summary_writer.add_scalar('train/mse', loss, global_step=iteration)
        summary_writer.add_scalar('train/stepSize', tensorf.stepSize, global_step=iteration)
        summary_writer.add_scalar('train/nSamples', tensorf.nSamples, global_step=iteration)

        # Update lr
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * lr_factor

        # Print the current values of the losses.
        if iteration % args.progress_refresh_rate == 0:
            pbar.set_description(
                f'Iteration {iteration:05d}:'
                + f' train_psnr = {float(np.mean(PSNRs)):.2f}'
                + f' test_psnr = {float(np.mean(PSNRs_test)):.2f}'
                + f' mse = {loss:.6f}'
            )
            PSNRs = []


        if iteration % args.vis_every == args.vis_every - 1 and args.N_vis!=0:
            PSNRs_test, _, _, _ = evaluation(
                test_dataset,tensorf, args, renderer, f'{logfolder}/imgs_vis/', N_vis=args.N_vis,
                prtx=f'{iteration:06d}_', N_samples=nSamples, white_bg = white_bg, ndc_ray=ndc_ray,
                compute_extra_metrics=False, save_img=args.save_img, save_video=args.save_video)
            summary_writer.add_scalar('test/psnr', np.mean(PSNRs_test), global_step=iteration)


        t = time.time()
        if iteration in update_AlphaMask_list:

            if reso_cur[0] * reso_cur[1] * reso_cur[2]<256**3:# update volume resolution
                reso_mask = reso_cur
            new_aabb = tensorf.updateAlphaMask(tuple(reso_mask))
            
            do_shrink = True
            if args.shrink_0:
                do_shrink = iteration == update_AlphaMask_list[0]
            if do_shrink:
                tensorf.shrink(new_aabb)
                L1_reg_weight = args.L1_weight_rest
                print("continuing L1_reg_weight", L1_reg_weight)

            if not args.ndc_ray and iteration == update_AlphaMask_list[1]:
                # filter rays outside the bbox
                allrays,allrgbs = tensorf.filtering_rays(allrays,allrgbs)
                batch_size = args.batch_size
                trainingSampler = SimpleSampler(allrgbs.shape[0], batch_size)


        if iteration in upsamp_list:
            n_voxels = N_voxel_list.pop(0)
            reso_cur = N_to_reso(n_voxels, tensorf.aabb)
            nSamples = min(args.nSamples, cal_n_samples(reso_cur,args.step_ratio))
            n_voxels_real = N_voxel_real_list.pop(0)
            reso_real_cur = N_to_reso(n_voxels_real, tensorf.aabb)
            if args.scale_reso:
                reso_real_cur = scale_reso(reso_real_cur, n_voxels_real)
            tensorf.upsample_volume_grid(reso_cur, reso_real_cur)

            if args.lr_upsample_reset:
                print("reset lr to initial")
                lr_scale = 1 #0.1 ** (iteration / args.n_iters)
            else:
                lr_scale = args.lr_decay_target_ratio ** (iteration / args.n_iters)
            grad_vars = tensorf.get_optparam_groups(args.lr_init*lr_scale, args.lr_basis*lr_scale)
            optimizer = torch.optim.Adam(grad_vars, betas=(0.9, 0.99))

            if args.no_upsample:
                if iteration == upsamp_list[0]:
                    batch_size = args.batch_size
                    trainingSampler = SimpleSampler(allrgbs.shape[0], batch_size)

        torch.cuda.synchronize()
        t_train += time.time() - t

        init_data = tensorf.step_after_iter(iteration + 1)
        if init_data is not None:
            return init_data

        n_params = count_params(tensorf, False)
        for k, v in n_params.items():
            summary_writer.add_scalar(f'n_params/{k}', v, global_step=iteration)
            
        summary_writer.add_scalar(f'time/train', t_train, global_step=iteration)

    summary_writer.add_scalar(f'time/init', t_init, global_step=iteration)
    summary_writer.add_scalar(f'time/total', t_init + t_train, global_step=iteration)

    if args.render_train:
        os.makedirs(f'{logfolder}/imgs_train_all', exist_ok=True)
        train_dataset = dataset(args.datadir, split='train', downsample=args.downsample_train, is_stack=True)
        PSNRs_test, _, _, _ = evaluation(
            train_dataset,tensorf, args, renderer, f'{logfolder}/imgs_train_all/',
            N_vis=-1, N_samples=-1, white_bg = white_bg, ndc_ray=ndc_ray,device=device,
            save_img=args.save_img, save_video=args.save_video)
        print(f'======> test all psnr: {np.mean(PSNRs_test)} <========================')

    if args.render_test:
        os.makedirs(f'{logfolder}/imgs_test_all', exist_ok=True)
        PSNRs_test, ssims_test, l_vgg_test, l_alex_test = evaluation(
            test_dataset,tensorf, args, renderer, f'{logfolder}/imgs_test_all/',
            N_vis=-1, N_samples=-1, white_bg = white_bg, ndc_ray=ndc_ray,device=device,
            save_img=args.save_img, save_video=args.save_video)
        summary_writer.add_scalar('test/psnr_all', np.mean(PSNRs_test), global_step=iteration)
        summary_writer.add_scalar('test/ssim_all', np.mean(ssims_test), global_step=iteration)
        summary_writer.add_scalar('test/lpips_vgg_all', np.mean(l_vgg_test), global_step=iteration)
        summary_writer.add_scalar('test/lpips_alex_all', np.mean(l_alex_test), global_step=iteration)
        print(f'======> test all psnr: {np.mean(PSNRs_test)} <========================')

    if args.render_path:
        c2ws = test_dataset.render_path
        print('========>',c2ws.shape)
        os.makedirs(f'{logfolder}/imgs_path_all', exist_ok=True)
        evaluation_path(test_dataset,tensorf, c2ws, renderer, f'{logfolder}/imgs_path_all/',
                                N_vis=-1, N_samples=-1, white_bg = white_bg, ndc_ray=ndc_ray,device=device)
        
    n_params = count_params(tensorf, False)
    print({args.expname})

    if args.save_ckpt:
        tensorf.save(f'{logfolder}/ckpt.th')

    stats_t = {'t_init': t_init, 't_train': t_train, 't_total': t_init + t_train}
    stats_misc = {'n_iter': args.n_iters, 'batch_size': args.batch_size, 'n_param': n_params}
    stats_metric = {'psnr': np.mean(PSNRs_test), 'ssim': np.mean(ssims_test), 'lpips_vgg': np.mean(l_vgg_test), 'lpips_alex': np.mean(l_alex_test)}
    print(stats_t)
    print(stats_misc)
    print(stats_metric)
    
    save_fn = os.path.join(logfolder, 'stats')
    with open(f'{save_fn}.txt', 'w') as f:
        pprint.pprint(stats_t, f, sort_dicts=False)
        pprint.pprint(stats_misc, f, sort_dicts=False)
        pprint.pprint(stats_metric, f, sort_dicts=False)


if __name__ == '__main__':

    torch.set_default_dtype(torch.float32)
    torch.manual_seed(20211202)
    np.random.seed(20211202)

    args = config_parser()
    pprint.pprint(args)
    print(args.expname)

    if args.export_mesh:
        export_mesh(args)

    if args.render_only and (args.render_test or args.render_path):
        render_test(args)
    else:
        init_data = None
        if args.config_init is not None:
            print('Get init data...')
            args = config_parser(mode='config_init')
            pprint.pprint(args)
            print(args.expname)
            if os.path.exists(args.rbf_config.init_data_fp):
                init_data = torch.load(args.rbf_config.init_data_fp)
            else:
                init_data = reconstruction(args)
                torch.cuda.empty_cache()

        print('Training...')
        args = config_parser()
        pprint.pprint(args)
        print(args.expname)
        reconstruction(args, init_data)

main_nerf.py
oppo-us-research/NeuRBF
# Modified from https://github.com/sarafridov/K-Planes/blob/main/plenoxels/main.py

import os

import time
import argparse
import importlib.util
import logging
import pprint
import sys
from typing import List, Dict, Any
import tempfile
import glob

import numpy as np

import util_misc


def get_freer_gpu():
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_fname = os.path.join(tmpdir, "tmp")
        os.system(f'nvidia-smi -q -d Memory |grep -A5 GPU|grep Free >"{tmp_fname}"')
        if os.path.isfile(tmp_fname):
            memory_available = [int(x.split()[2]) for x in open(tmp_fname, 'r').readlines()]
            if len(memory_available) > 0:
                return np.argmax(memory_available)
    return None  # The grep doesn't work with all GPUs. If it fails we ignore it.

gpu = util_misc.select_devices('1#', force_reselect=True, excludeID=[])[0]
if gpu is not None:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
    print(f"CUDA_VISIBLE_DEVICES set to {gpu}")
else:
    print(f"Did not set GPU.")


import torch
import torch.utils.data
from thirdparty.kplanes.runners import video_trainer
from thirdparty.kplanes.runners import phototourism_trainer
from thirdparty.kplanes.runners import static_trainer
from thirdparty.kplanes.utils.create_rendering import render_to_path, decompose_space_time
from thirdparty.kplanes.utils.parse_args import parse_optfloat


temp = torch.ones([100, 100], device=0)
del temp


def setup_logging(log_level=logging.INFO):
    handlers = [logging.StreamHandler(sys.stdout)]
    logging.basicConfig(level=log_level,
                        format='%(asctime)s|%(levelname)8s| %(message)s',
                        handlers=handlers,
                        force=True)


def load_data(model_type: str, data_downsample, data_dirs, validate_only: bool, render_only: bool, **kwargs):
    data_downsample = parse_optfloat(data_downsample, default_val=1.0)

    if model_type == "video":
        return video_trainer.load_data(
            data_downsample, data_dirs, validate_only=validate_only,
            render_only=render_only, **kwargs)
    elif model_type == "phototourism":
        return phototourism_trainer.load_data(
            data_downsample, data_dirs, validate_only=validate_only,
            render_only=render_only, **kwargs
        )
    else:
        return static_trainer.load_data(
            data_downsample, data_dirs, validate_only=validate_only,
            render_only=render_only, **kwargs)


def init_trainer(model_type: str, **kwargs):
    if model_type == "video":
        from thirdparty.kplanes.runners import video_trainer
        return video_trainer.VideoTrainer(**kwargs)
    elif model_type == "phototourism":
        from thirdparty.kplanes.runners import phototourism_trainer
        return phototourism_trainer.PhototourismTrainer(**kwargs)
    else:
        from thirdparty.kplanes.runners import static_trainer
        return static_trainer.StaticTrainer(**kwargs)


def save_config(config):
    log_dir = os.path.join(config['logdir'], config['expname'])
    os.makedirs(log_dir, exist_ok=True)

    with open(os.path.join(log_dir, 'config.py'), 'wt') as out:
        out.write('config = ' + pprint.pformat(config))

    with open(os.path.join(log_dir, 'config.csv'), 'w') as f:
        for key in config.keys():
            f.write("%s\t%s\n" % (key, config[key]))


def main(args, init_data=None):
    setup_logging()

    # Set random seed
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # Import config
    spec = importlib.util.spec_from_file_location(os.path.basename(args.config), args.config)
    cfg = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(cfg)
    config: Dict[str, Any] = cfg.config

    # Process overrides from argparse into config
    if args.data_name is not None:
        config['data_dirs'] = [config['data_dirs'][0].rsplit('/', 1)[0] + '/' + args.data_name]
    overrides: List[str] = args.override
    overrides_dict = {ovr.split("=")[0]: ovr.split("=")[1] for ovr in overrides}
    config.update(overrides_dict)

    # Replace placeholders with actual value
    config['data_fn'] = config['data_dirs'][0].split('/')[-1]
    config['expname'] = config['expname'].replace('data_name', config['data_fn'])
    if 'rbf_config' in config:
        if 'init_data_fp' in config['rbf_config']:
            config['rbf_config']['init_data_fp'] = config['rbf_config']['init_data_fp'].replace('data_name', config['data_fn'])

    # Add version id
    exp_list = glob.glob(f'{config["logdir"]}/{config["expname"]}-v*'.replace('[', '[[]'))
    if len(exp_list) == 0:
        vid = 0
    else:
        vid = max([int(v.split(f'{config["expname"]}-v')[-1]) for v in exp_list]) + 1
    config["expname"] += f'-v{vid}'

    # If will load trained model
    if args.log_dir is not None:
        config['expname'] = args.log_dir.split('/')[-1]
        if 'rbf_config' in config:
            config['rbf_config']['init_rbf'] = False

    if "keyframes" in config:
        model_type = "video"
    elif "appearance_embedding_dim" in config:
        model_type = "phototourism"
    else:
        model_type = "static"
    validate_only = args.validate_only
    render_only = args.render_only
    spacetime_only = args.spacetime_only
    if validate_only and render_only:
        raise ValueError("render_only and validate_only are mutually exclusive.")
    if render_only and spacetime_only:
        raise ValueError("render_only and spacetime_only are mutually exclusive.")
    if validate_only and spacetime_only:
        raise ValueError("validate_only and spacetime_only are mutually exclusive.")

    pprint.pprint(config)
    print(config["expname"])

    if validate_only or render_only:
        assert args.log_dir is not None and os.path.isdir(args.log_dir)
    else:
        save_config(config)

    data = load_data(model_type, validate_only=validate_only, render_only=render_only or spacetime_only, **config)
    config.update(data)
    config['init_data'] = init_data
    trainer = init_trainer(model_type, **config)
    if args.log_dir is not None:
        checkpoint_path = os.path.join(args.log_dir, "model.pth")
        training_needed = not (validate_only or render_only or spacetime_only)
        trainer.load_model(torch.load(checkpoint_path), training_needed=training_needed)

    if validate_only:
        trainer.validate()
    elif render_only:
        render_to_path(trainer, extra_name=config['expname'].split('-')[0])
    elif spacetime_only:
        decompose_space_time(trainer, extra_name="")
    else:
        t_train = time.time()
        init_data = trainer.train()
        t_train = time.time() - t_train

        stats_t = {'t_init': trainer.model.t_init, 't_train': t_train, 't_total': trainer.model.t_init + t_train}
        stats_misc = {'n_iter': config["num_steps"], 'batch_size': config["batch_size"], 'n_param': trainer.model.n_params}
        if hasattr(trainer, 'val_metrics'):
            stats_metric = trainer.val_metrics[0]
        else:
            stats_metric = {}
        print(stats_t)
        print(stats_misc)
        print(stats_metric)
        
        save_fn = os.path.join(config['logdir'], config['expname'], 'stats')
        with open(f'{save_fn}.txt', 'w') as f:
            pprint.pprint(stats_t, f, sort_dicts=False)
            pprint.pprint(stats_misc, f, sort_dicts=False)
            pprint.pprint(stats_metric, f, sort_dicts=False)

        return init_data


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="")
    p.add_argument('--data_name', type=str, default=None)
    p.add_argument('--render-only', action='store_true')
    p.add_argument('--validate-only', action='store_true')
    p.add_argument('--spacetime-only', action='store_true')
    p.add_argument('--config_init', type=str, default=None)
    p.add_argument('--config', type=str, required=True)
    p.add_argument('--log-dir', type=str, default=None)
    p.add_argument('--seed', type=int, default=0)
    p.add_argument('override', nargs=argparse.REMAINDER)
    args = p.parse_args()

    init_data = None
    if args.config_init is not None:
        args_config = args.config
        args.config = args.config_init
        print('Get init data...')
        init_data = main(args)
        torch.cuda.empty_cache()
        args.config = args_config

    main(args, init_data)

main_nerf_kplanes.py
oppo-us-research/NeuRBF
# Modified from https://github.com/sarafridov/K-Planes/blob/main/plenoxels/models/kplane_field.py

import itertools
import logging as log
from typing import Optional, Union, List, Dict, Sequence, Iterable, Collection, Callable

import os
import time
import math
import numpy as np
import torch
import torch.nn as nn
import tinycudann as tcnn

from thirdparty.kplanes.ops.interpolation import grid_sample_wrapper
from thirdparty.kplanes.raymarching.spatial_distortions import SpatialDistortion

import mcubes
import trimesh
from pykdtree.kdtree import KDTree
import util_misc
import util_network
import util_init
import util_clustering


def get_normalized_directions(directions):
    """SH encoding must be in the range [0, 1]

    Args:
        directions: batch of directions
    """
    return (directions + 1.0) / 2.0


def normalize_aabb(pts, aabb):
    return (pts - aabb[0]) * (2.0 / (aabb[1] - aabb[0])) - 1.0


def init_grid_param(
        grid_nd: int,
        in_dim: int,
        out_dim: int,
        reso: Sequence[int],
        a: float = 0.1,
        b: float = 0.5):
    assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
    has_time_planes = in_dim == 4
    assert grid_nd <= in_dim
    coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
    grid_coefs = nn.ParameterList()
    for ci, coo_comb in enumerate(coo_combs):
        new_grid_coef = nn.Parameter(torch.empty(
            [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
        ))
        if has_time_planes and 3 in coo_comb:  # Initialize time planes to 1
            nn.init.ones_(new_grid_coef)
        else:
            nn.init.uniform_(new_grid_coef, a=a, b=b)
        grid_coefs.append(new_grid_coef)

    return grid_coefs


def interpolate_ms_features(pts: torch.Tensor,
                            ms_grids: Collection[Iterable[nn.Module]],
                            grid_dimensions: int,
                            concat_features: bool,
                            num_levels: Optional[int],
                            ) -> torch.Tensor:
    coo_combs = list(itertools.combinations(
        range(pts.shape[-1]), grid_dimensions)
    )
    if num_levels is None:
        num_levels = len(ms_grids)
    multi_scale_interp = [] if concat_features else 0.
    grid: nn.ParameterList
    for scale_id, grid in enumerate(ms_grids[:num_levels]):
        interp_space = 1.
        for ci, coo_comb in enumerate(coo_combs):
            # interpolate in plane
            feature_dim = grid[ci].shape[1]  # shape of grid[ci]: 1, out_dim, *reso
            interp_out_plane = (
                grid_sample_wrapper(grid[ci], pts[..., coo_comb])
                .view(-1, feature_dim)
            )
            # compute product over planes
            interp_space = interp_space * interp_out_plane

        # combine over scales
        if concat_features:
            multi_scale_interp.append(interp_space)
        else:
            multi_scale_interp = multi_scale_interp + interp_space

    if concat_features:
        multi_scale_interp = torch.cat(multi_scale_interp, dim=-1)
    return multi_scale_interp


class KPlaneRBFField(nn.Module):
    def __init__(
        self,
        aabb,
        grid_config: Union[str, List[Dict]],
        concat_features_across_scales: bool,
        multiscale_res: Optional[Sequence[int]],
        use_appearance_embedding: bool,
        appearance_embedding_dim: int,
        spatial_distortion: Optional[SpatialDistortion],
        density_activation: Callable,
        linear_decoder: bool,
        linear_decoder_layers: Optional[int],
        num_images: Optional[int],
        rbf_config: Dict,
        init_data: Optional[Dict],
    ) -> None:
        super().__init__()

        self.aabb = nn.Parameter(aabb, requires_grad=False)
        self.spatial_distortion = spatial_distortion
        self.grid_config = grid_config

        self.multiscale_res_multipliers: List[int] = multiscale_res or [1]
        self.concat_features = concat_features_across_scales
        self.density_activation = density_activation
        self.linear_decoder = linear_decoder

        self.in_dim = 3
        self.cmin = torch.tensor([-1, -1, -1])  # ... y x, aabb
        self.cmax = torch.tensor([1, 1, 1])  # ... y x, aabb
        self.register_buffer('cmin_gpu', self.cmin.clone())  # ... y x, aabb
        self.register_buffer('cmax_gpu', self.cmax.clone())  # ... y x, aabb
        self.get_init_data = rbf_config['get_init_data']
        self.init_steps = rbf_config['init_steps']
        self.init_data_fp = rbf_config['init_data_fp']
        self.init_rbf = rbf_config['init_rbf']
        self.s_dims = torch.tensor(rbf_config['s_dims'])  # ... h w, used for initializing rbf
        self.kp_ref_config = rbf_config['kp_ref_config']
        self.rbf_type = rbf_config['rbf_type']
        self.rbf_lc0_normalize = rbf_config['rbf_lc0_normalize']
        self.n_kernel = rbf_config['n_kernel']
        self.point_nn_kernel = rbf_config['point_nn_kernel']
        ks_alpha = rbf_config['ks_alpha']
        self.lc0_dim = rbf_config['lc0_dim']
        self.lc_init = rbf_config['lc_init']
        self.lcb_init = rbf_config['lcb_init']
        self.kp_init = rbf_config['kp_init']
        self.kpb_init = rbf_config['kpb_init']
        self.rbf_suffixes = rbf_config['rbf_suffixes']
        self.pe_lc0_freq = rbf_config['pe_lc0_freq']
        self.pe_lc0_rbf_freq = rbf_config['pe_lc0_rbf_freq']
        self.pe_lc0_rbf_keep = rbf_config['pe_lc0_rbf_keep']
        if self.pe_lc0_rbf_keep == 'all': self.pe_lc0_rbf_keep = self.lc0_dim

        n_freq_level = self.lc0_dim - self.pe_lc0_rbf_keep
        if n_freq_level > 0 and len(self.pe_lc0_rbf_freq) >= 1:
            if len(self.pe_lc0_rbf_freq) == 1:
                freqs = torch.ones([n_freq_level]) * self.pe_lc0_rbf_freq[0]
            elif len(self.pe_lc0_rbf_freq) >= 2:
                amin = np.log2(self.pe_lc0_rbf_freq[0])
                amax = np.log2(self.pe_lc0_rbf_freq[1])
                freqs = torch.linspace(amin, amax, n_freq_level)
                freqs = torch.exp2(freqs)
            self.register_buffer('freqs', freqs)
        else:
            self.freqs = None

        self.pe_kp_freq = rbf_config['pe_kp_freq']
        fix_params = rbf_config['fix_params']
        self.kc_init_config = rbf_config['kc_init_config']
        self.kw_init_config = rbf_config['kw_init_config']
        self.kc_mult = 1
        self.kc_init_regular = {}
        for k, v in self.kc_init_config.items():
            if v['type'] == 'none':
                self.kc_init_regular[k] = True
            else:
                self.kc_init_regular[k] = False
        sparse_embd_grad = False
        self.rbf_train_start = False

        n_params_ref = 0
        for i, res in enumerate(self.kp_ref_config['multiscale_res']):
            n_params_ref += int(self.kp_ref_config['resolution'][0] * res)**2 * 3 * \
                self.kp_ref_config['output_coordinate_dim'][i]

        # 1. Init planes
        self.grids = nn.ModuleList()
        self.kp_dim = 0
        for i, res in enumerate(self.multiscale_res_multipliers):
            # initialize coordinate grid
            config = self.grid_config[0].copy()
            # Resolution fix: multi-res only on spatial planes
            config["resolution"] = [
                int(r * res) for r in config["resolution"][:3]
            ] + config["resolution"][3:]
            gp = init_grid_param(
                grid_nd=config["grid_dimensions"],
                in_dim=config["input_coordinate_dim"],
                out_dim=config["output_coordinate_dim"][i],
                reso=config["resolution"],
                a=self.kp_init[0],
                b=self.kp_init[1],
            )
            # shape[1] is out-dim - Concatenate over feature len for each scale
            if self.concat_features:
                self.kp_dim += gp[-1].shape[1]
            else:
                self.kp_dim = gp[-1].shape[1]
            self.grids.append(gp)
        log.info(f"Initialized model grids: {self.grids}")
        n_params_kplanes = int(util_misc.count_parameters(self.grids)[-1])
        self.feature_dim = self.kp_dim

        # 2. Init rbf
        self.kc0 = nn.ModuleList()
        self.ks0 = nn.ModuleList()
        self.lc0 = nn.ModuleList()
        if self.n_kernel == 'auto':
            n_params_rbf = n_params_ref - n_params_kplanes
            self.n_kernel = int(n_params_rbf // (self.lc0_dim + util_misc.get_rbf_params_per_kernel(
                self.rbf_type, self.in_dim, self.kc_mult)))
        if self.n_kernel > 0 and self.lc0_dim > 0 and not self.get_init_data:
            if self.kc_init_regular['0']:
                self.n_kernel = util_misc.get_lower_int_power(self.n_kernel, self.in_dim)
            print('n_kernel:', self.n_kernel)
            self.lc_dims = [[self.n_kernel, self.lc0_dim]]
            self.kc0, self.ks0, self.lc0, self.ks_dims, k_dims, kc_interval = self.create_rbf_params(
                self.rbf_type, self.n_kernel, self.in_dim, self.lc0_dim, sparse_embd_grad, self.cmin, self.cmax, ks_alpha, is_bag=False)
            nn.init.uniform_(self.lc0.weight, self.lc_init[0], self.lc_init[1])
            self.register_buffer(f'k_dims_0', k_dims)
            self.register_buffer(f'kci0', kc_interval)
            self.rbf_fn = eval(f'util_network.rbf_{self.rbf_type}_fb')
            if self.init_rbf:
                self.init_rbf_params(init_data)
            self.rbf_train_start = True
            self.feature_dim += self.lc0_dim * (len(self.pe_lc0_freq) + 1)

        self.feature_dim *= len(self.pe_kp_freq) + 1

        # kpb0
        if self.kpb_init is not None:
            self.kpb = torch.nn.Parameter(torch.zeros(self.feature_dim))
            nn.init.uniform_(self.kpb, self.kpb_init[0], self.kpb_init[1])

        # 3. Init appearance code-related parameters
        self.use_average_appearance_embedding = True  # for test-time
        self.use_appearance_embedding = use_appearance_embedding
        self.num_images = num_images
        self.appearance_embedding = None
        if use_appearance_embedding:
            assert self.num_images is not None
            self.appearance_embedding_dim = appearance_embedding_dim
            # this will initialize as normal_(0.0, 1.0)
            self.appearance_embedding = nn.Embedding(self.num_images, self.appearance_embedding_dim)
        else:
            self.appearance_embedding_dim = 0

        # 4. Init decoder params
        self.direction_encoder = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "SphericalHarmonics",
                "degree": 4,
            },
        )

        # 4. Init decoder network
        if self.linear_decoder:
            assert linear_decoder_layers is not None
            # The NN learns a basis that is used instead of spherical harmonics
            # Input is an encoded view direction, output is weights for
            # combining the color features into RGB
            # This architecture is based on instant-NGP
            self.color_basis = tcnn.Network(
                n_input_dims=3 + self.appearance_embedding_dim,
                n_output_dims=3 * self.feature_dim,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 128,
                    "n_hidden_layers": linear_decoder_layers,
                },
            )
            # sigma_net just does a linear transformation on the features to get density
            self.sigma_net = tcnn.Network(
                n_input_dims=self.feature_dim,
                n_output_dims=1,
                network_config={
                    "otype": "CutlassMLP",
                    "activation": "None",
                    "output_activation": "None",
                    "n_neurons": 128,
                    "n_hidden_layers": 0,
                },
            )
        else:
            self.geo_feat_dim = 15
            self.sigma_net = tcnn.Network(
                n_input_dims=self.feature_dim,
                n_output_dims=self.geo_feat_dim + 1,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 64,
                    "n_hidden_layers": 1,
                },
            )
            self.in_dim_color = (
                    self.direction_encoder.n_output_dims
                    + self.geo_feat_dim
                    + self.appearance_embedding_dim
            )
            self.color_net = tcnn.Network(
                n_input_dims=self.in_dim_color,
                n_output_dims=3,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "Sigmoid",
                    "n_neurons": 64,
                    "n_hidden_layers": 2,
                },
            )

        # Fix params
        util_network.fix_params(self, fix_params)

    def create_rbf_params(self, rbf_type, n_kernel, in_dim, lc_dim, sparse_embd_grad, cmin, cmax, ks_alpha, scale_grad_by_freq=False, is_bag=True):
        if rbf_type.endswith('_a') or rbf_type.endswith('_f'):
            ks_dims = [in_dim, in_dim]
        elif rbf_type.endswith('_d'):
            ks_dims = [in_dim]
        elif rbf_type.endswith('_s'):
            ks_dims = [1]
        else:
            raise NotImplementedError
        kc = torch.nn.Embedding(n_kernel, in_dim, scale_grad_by_freq=False, sparse=sparse_embd_grad)
        ks = torch.nn.Embedding(n_kernel, np.prod(ks_dims), scale_grad_by_freq=False, 
            sparse=sparse_embd_grad)
        if rbf_type.endswith('_a') or rbf_type.endswith('_f'):
            ks.weight.data = torch.eye(in_dim)[None, ...].repeat(n_kernel, 1, 1).reshape(n_kernel, -1)
        elif rbf_type.endswith('_d'):
            ks.weight.data[:] = 1
        elif rbf_type.endswith('_s'):
            ks.weight.data[:] = 1
        else:
            raise NotImplementedError

        if is_bag:
            lc = torch.nn.EmbeddingBag(n_kernel, lc_dim, scale_grad_by_freq=scale_grad_by_freq, 
                                       sparse=sparse_embd_grad, mode='sum')
        else:
            lc = torch.nn.Embedding(n_kernel, lc_dim, scale_grad_by_freq=scale_grad_by_freq, sparse=sparse_embd_grad)

        # Initialize
        n_kernel_per_dim = int(math.ceil(n_kernel**(1/in_dim)))
        k_dims = [n_kernel_per_dim] * in_dim
        points, _, side_interval = util_misc.get_grid_points(k_dims, align_corners=True, 
            vmin=cmin.tolist(), vmax=cmax.tolist())
        points = points.flip(-1)  # x y ...
        side_interval = side_interval.flip(-1)  # x y ...
        kc.weight.data = points[:n_kernel]
        if rbf_type.startswith('nlin_'):
            if rbf_type.endswith('_s'):
                ks.weight.data *= 1 / side_interval.mean()
            elif rbf_type.endswith('_d'):
                ks.weight.data *= 1 / side_interval[None]
            elif rbf_type.endswith('_f'):
                ks.weight.data *= (1 / side_interval)[:, None].repeat(1, side_interval.shape[0]).view(1, -1)
            else:
                raise NotImplementedError
        else:
            if rbf_type.endswith('_s'):
                ks.weight.data *= 1 / (side_interval.mean() / 2) ** 2
            elif rbf_type.endswith('_d'):
                ks.weight.data *= 1 / (side_interval[None] / 2) ** 2
            elif rbf_type.endswith('_a'):
                ks.weight.data *= (1 / (side_interval / 2) ** 2)[:, None].repeat(1, side_interval.shape[0]).view(1, -1)
            else:
                raise NotImplementedError
        ks.weight.data *= ks_alpha

        return kc, ks, lc, ks_dims, torch.tensor(k_dims).flip(-1), side_interval

    def get_kc(self, suffix):
        return getattr(self, 'kc' + suffix).weight

    def get_ks(self, suffix):
        return getattr(self, 'ks' + suffix).weight.view(-1, *self.ks_dims)

    def clip_rbf_params(self):
        for k in self.rbf_suffixes:
            if hasattr(self, f'ks{k}.weight') and getattr(self, 'ks'+k).weight.requires_grad:
                if k == '0':
                    vmin_factor = 100
                    vmax_factor = 10
                else:
                    vmin_factor = 1
                    vmax_factor = 2
                util_misc.clip_kw_sq(
                    getattr(self, 'ks'+k).weight, self.rbf_type, self.cmin.flip(-1), self.cmax.flip(-1), 
                    self.s_dims.flip(-1), is_ks=True, is_flat=True, vmin_factor=vmin_factor, vmax_factor=vmax_factor)

    def init_rbf_params(self, init_data=None):
        # Get grid points within aabb
        t = time.time()
        pts, _, _ = util_misc.get_grid_points(self.s_dims.tolist(), align_corners=True, 
            vmin=self.cmin.tolist(), vmax=self.cmax.tolist(), device=0)
        pts.flip(-1)  # x y z
        print('Generate grid points:', time.time() - t)

        # Get density and features at these points
        if self.get_init_data:
            t = time.time()
            with torch.no_grad():
                chunk_size = 2**21
                density_all = []
                features_all = []
                for i in range(0, pts.shape[0], chunk_size):
                    features = interpolate_ms_features(
                        pts[i:i+chunk_size], ms_grids=self.grids,
                        grid_dimensions=self.grid_config[0]["grid_dimensions"],
                        concat_features=self.concat_features, num_levels=None)

                    if self.linear_decoder:
                        density_before_activation = self.sigma_net(features)  # [batch, 1]
                    else:
                        features = self.sigma_net(features)  # float16
                        features, density_before_activation = torch.split(
                            features, [self.geo_feat_dim, 1], dim=-1)  # float16

                    density = self.density_activation(density_before_activation.to(pts))  # float32
                    
                    density_all.append(density.cpu())
                    features_all.append(features.cpu())
                density_all = torch.cat(density_all, 0)
                features_all = torch.cat(features_all, 0)
                del density, features
                torch.cuda.empty_cache()
            print('Get density and features at grid points:', time.time() - t)
            density_all, features_all = density_all.cpu(), features_all.cpu()
            init_data = {'density': density_all, 'features': features_all}

            # # Save init data
            # os.makedirs(os.path.dirname(self.init_data_fp), exist_ok=True)
            # torch.save(init_data, self.init_data_fp)
            # density_all = density_all.reshape(*self.s_dims.tolist(), -1)
            # vertices, triangles = mcubes.marching_cubes(density_all.numpy()[..., 0], 50.)
            # trimesh.Trimesh(vertices, triangles, process=False).export(self.init_data_fp.split('.')[0] + '.ply')
            
            torch.cuda.empty_cache()
            return init_data

        density_all = init_data['density']
        features_all = init_data['features'].to(torch.float32)
        del init_data

        # Init rbf parameters
        util_network.init_nerf_rbf_params(self, pts, density_all, features_all, self.kc_init_config, 
                                          self.kw_init_config, device=0)
        del density_all, features_all
        torch.cuda.empty_cache()

        # Build kd tree
        self.kdtree0 = KDTree(getattr(self, 'kc0').weight.detach().cpu().numpy())

    def step_after_iter(self, step):
        if step == self.init_steps and self.get_init_data:
            return self.init_rbf_params()

    def forward_kernel_idx(self, x_g, point_idx, suffix):
        if hasattr(self, 'point_kernel_idx_' + suffix):
            # Use pre-computed knn
            point_kernel_idx = getattr(self, 'point_kernel_idx_' + suffix)
            kernel_idx = point_kernel_idx[point_idx.to(point_kernel_idx.device)].to(x_g.device)  # [p topk]
        elif self.kc_init_regular[suffix]:
            # Find first ring neighbors in regular grid
            kernel_idx = util_misc.get_multi_index_first_ring(x_g, 
                getattr(self, f'k_dims_{suffix}'), getattr(self, f'kci{suffix}'), 
                self.cmin_gpu.flip(-1), self.cmax_gpu.flip(-1), True)  # (p, n_ngb, n_in)
            kernel_idx = util_misc.ravel_multi_index(kernel_idx, getattr(self, f'k_dims_{suffix}'))  # (p, n_ngb)
        else:
            # Knn on the fly
            if not hasattr(self, 'kdtree' + suffix):
                setattr(self, 'kdtree' + suffix, KDTree(getattr(self, 'kc' + suffix).weight.detach().cpu().numpy()))
            _, kernel_idx = util_clustering.query_chunked(getattr(self, 'kdtree' + suffix), x_g.cpu().numpy(), 
                k=self.point_nn_kernel, sqr_dists=True, chunk_size=int(2e8), return_dist=False)
            kernel_idx = torch.tensor(kernel_idx.astype(np.int32), device=x_g.device)  # [p topk]
            if kernel_idx.ndim == 1:
                kernel_idx = kernel_idx[:, None]
            
        return kernel_idx

    def forward_rbf(self, x_g, kernel_idx, suffix):
        if kernel_idx is None:  # Use all kernels for each point
            kc = getattr(self, 'kc' + suffix).weight[None]  # [1 k d]
            ks = getattr(self, 'ks' + suffix).weight[None]
            ks = ks.view(*ks.shape[:2], *self.ks_dims)  # [1 k d d] or [1 k 1]
            rbf_out, _ = self.rbf_fn(x_g, kc, ks)  # [p k_topk]
        else:
            kc = getattr(self, 'kc' + suffix)(kernel_idx)  # [p k d]
            ks = getattr(self, 'ks' + suffix)(kernel_idx).view(*kernel_idx.shape, *self.ks_dims)  # [p k d d] or [p k 1]
            rbf_out, _ = self.rbf_fn(x_g, kc, ks)  # [p k_topk]
        return rbf_out

    def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
        """Computes and returns the densities."""
        if self.spatial_distortion is not None:
            pts = self.spatial_distortion(pts)
            pts = pts / 2  # from [-2, 2] to [-1, 1]
        else:
            pts = normalize_aabb(pts, self.aabb)
        n_rays, n_samples = pts.shape[:2]
        if timestamps is not None:
            timestamps = timestamps[:, None].expand(-1, n_samples)[..., None]  # [n_rays, n_samples, 1]
            pts = torch.cat((pts, timestamps), dim=-1)  # [n_rays, n_samples, 4]

        pts = pts.reshape(-1, pts.shape[-1])  # [n_rays * n_samples, 3]
        features = interpolate_ms_features(
            pts, ms_grids=self.grids,  # noqa
            grid_dimensions=self.grid_config[0]["grid_dimensions"],
            concat_features=self.concat_features, num_levels=None)
        if len(features) < 1:
            features = torch.zeros((0, 1)).to(features.device)

        # Forward rbf lc0
        if len(features) >= 1 and self.rbf_train_start:
            suffix = '0'
            if self.point_nn_kernel <= 0:  # Use all kernels for each point
                rbf_out = self.forward_rbf(pts, None, suffix)  # [p nk]
                if self.rbf_lc0_normalize:
                    rbf_out = rbf_out / (rbf_out.detach().sum(-1, keepdim=True) + 1e-8)
                out = rbf_out @ self.lc0.weight  # [p hfl]
            else:
                kernel_idx = self.forward_kernel_idx(pts, None, suffix)
                rbf_out = self.forward_rbf(pts, kernel_idx, suffix)  # [p k_topk]
                if self.rbf_lc0_normalize:
                    rbf_out = rbf_out / (rbf_out.detach().sum(-1, keepdim=True) + 1e-8)
                out = self.lc0(kernel_idx)  # [p k_topk d_lc0]
                rbf_out = rbf_out[..., None]  # [p k_topk 1]
                if self.freqs is not None:
                    if self.pe_lc0_rbf_keep > 0:
                        rbf_out = torch.cat([rbf_out.expand(-1, -1, self.pe_lc0_rbf_keep), 
                            torch.sin(rbf_out * self.freqs[None, None])], -1)  # [p k_topk d_lc0]
                    else:
                        rbf_out = torch.sin(rbf_out * self.freqs[None, None])  # [p k_topk d_lc0]
                out = (out * rbf_out).sum(1)  # [p d_lc0]

            features = torch.cat([features, out], -1)

        if len(self.pe_kp_freq) > 0:
            features = torch.cat([features] + [features*i for i in self.pe_kp_freq], -1)
        if self.kpb_init is not None:
            features = features + self.kpb[None]
        features = torch.sin(features)

        if self.linear_decoder:
            density_before_activation = self.sigma_net(features)  # [batch, 1]
        else:
            features = self.sigma_net(features)
            features, density_before_activation = torch.split(
                features, [self.geo_feat_dim, 1], dim=-1)

        density = self.density_activation(
            density_before_activation.to(pts)
        ).view(n_rays, n_samples, 1)
        return density, features

    def forward(self,
                pts: torch.Tensor,
                directions: torch.Tensor,
                timestamps: Optional[torch.Tensor] = None):
        camera_indices = None
        if self.use_appearance_embedding:
            if timestamps is None:
                raise AttributeError("timestamps (appearance-ids) are not provided.")
            camera_indices = timestamps
            timestamps = None
        density, features = self.get_density(pts, timestamps)
        n_rays, n_samples = pts.shape[:2]

        directions = directions.view(-1, 1, 3).expand(pts.shape).reshape(-1, 3)
        if not self.linear_decoder:
            directions = get_normalized_directions(directions)
            encoded_directions = self.direction_encoder(directions)

        if self.linear_decoder:
            color_features = [features]
        else:
            color_features = [encoded_directions, features.view(-1, self.geo_feat_dim)]

        if self.use_appearance_embedding:
            if camera_indices.dtype == torch.float32:
                # Interpolate between two embeddings. Currently they are hardcoded below.
                #emb1_idx, emb2_idx = 100, 121  # trevi
                emb1_idx, emb2_idx = 11, 142  # sacre
                emb_fn = self.appearance_embedding
                emb1 = emb_fn(torch.full_like(camera_indices, emb1_idx, dtype=torch.long))
                emb1 = emb1.view(emb1.shape[0], emb1.shape[2])
                emb2 = emb_fn(torch.full_like(camera_indices, emb2_idx, dtype=torch.long))
                emb2 = emb2.view(emb2.shape[0], emb2.shape[2])
                embedded_appearance = torch.lerp(emb1, emb2, camera_indices)
            elif self.training:
                embedded_appearance = self.appearance_embedding(camera_indices)
            else:
                if hasattr(self, "test_appearance_embedding"):
                    embedded_appearance = self.test_appearance_embedding(camera_indices)
                elif self.use_average_appearance_embedding:
                    embedded_appearance = torch.ones(
                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                    ) * self.appearance_embedding.mean(dim=0)
                else:
                    embedded_appearance = torch.zeros(
                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                    )

            # expand embedded_appearance from n_rays, dim to n_rays*n_samples, dim
            ea_dim = embedded_appearance.shape[-1]
            embedded_appearance = embedded_appearance.view(-1, 1, ea_dim).expand(n_rays, n_samples, -1).reshape(-1, ea_dim)
            if not self.linear_decoder:
                color_features.append(embedded_appearance)

        color_features = torch.cat(color_features, dim=-1)

        if self.linear_decoder:
            if self.use_appearance_embedding:
                basis_values = self.color_basis(torch.cat([directions, embedded_appearance], dim=-1))
            else:
                basis_values = self.color_basis(directions)  # [batch, color_feature_len * 3]
            basis_values = basis_values.view(color_features.shape[0], 3, -1)  # [batch, 3, color_feature_len]
            rgb = torch.sum(color_features[:, None, :] * basis_values, dim=-1)  # [batch, 3]
            rgb = rgb.to(directions)
            rgb = torch.sigmoid(rgb).view(n_rays, n_samples, 3)
        else:
            rgb = self.color_net(color_features).to(directions).view(n_rays, n_samples, 3)

        return {"rgb": rgb, "density": density}

    def get_params(self):
        lc0_params = {k: v for k, v in self.lc0.named_parameters(prefix="lc0")}
        kc0_params = {k: v for k, v in self.kc0.named_parameters(prefix="kc0")}
        ks0_params = {k: v for k, v in self.ks0.named_parameters(prefix="ks0")}
        field_params = {k: v for k, v in self.grids.named_parameters(prefix="grids")}
        nn_params = [
            self.sigma_net.named_parameters(prefix="sigma_net"),
            self.direction_encoder.named_parameters(prefix="direction_encoder"),
        ]
        if self.linear_decoder:
            nn_params.append(self.color_basis.named_parameters(prefix="color_basis"))
        else:
            nn_params.append(self.color_net.named_parameters(prefix="color_net"))
        nn_params = {k: v for plist in nn_params for k, v in plist}
        other_params = {k: v for k, v in self.named_parameters() if (
            k not in nn_params.keys() and k not in field_params.keys() and 
            k not in lc0_params.keys() and k not in kc0_params.keys() and k not in ks0_params.keys()
        )}
        return {
            "nn": list(nn_params.values()),
            "field": list(field_params.values()),
            "other": list(other_params.values()),
            "lc0": list(lc0_params.values()),
            "kc0": list(kc0_params.values()),
            "ks0": list(ks0_params.values()),
        }

nerf_kplanes/network.py
oppo-us-research/NeuRBF
import os
import time
import math
import numpy as np
import torch.nn as nn
import mcubes
import trimesh
from pykdtree.kdtree import KDTree
import kornia

from thirdparty.tensorf.models.tensorBase import *
import thirdparty.tensorf.utils as utils
import util_misc
import util_network
import util_clustering


class TensorVMSplitRBF(TensorBase):
    """
    Modified from https://github.com/apchenstu/TensoRF/blob/4ec894dc1341a2201fe13ae428631b58458f105d/models/tensoRF.py#L139
    """
    def __init__(self, aabb, gridSize, device, **kargs):
        self.n_level = kargs['args'].n_level
        self.level_types = list(kargs['args'].level_types)
        if len(self.level_types) == 1:
            self.level_types *= self.n_level
        assert(len(self.level_types) == self.n_level)
        for x in self.level_types:
            assert(x in ['vm', 'g'])
        
        self.N_voxel_min = kargs['args'].resol_min ** 3

        self.density_n_comp = list(kargs['density_n_comp'])
        self.app_n_comp = list(kargs['appearance_n_comp'])
        if len(self.density_n_comp) == 1:
            self.density_n_comp *= self.n_level
        if len(self.app_n_comp) == 1:
            self.app_n_comp *= self.n_level
        assert(len(self.density_n_comp) == self.n_level)
        assert(len(self.app_n_comp) == self.n_level)
        self.density_n_comp_total = sum(self.density_n_comp) * 3
        self.app_n_comp_total = sum(self.app_n_comp) * 3

        self.matMode = [[0,1], [0,2], [1,2]]
        self.vecMode =  [2, 1, 0]

        args = kargs['args']
        rbf_config = args.rbf_config
        self.lc0_dim = rbf_config['lc0_dim']
        self.lcd0_dim = rbf_config['lcd0_dim']
        self.get_init_data = rbf_config['get_init_data']
        self.n_kernel = rbf_config['n_kernel']
        self.rbf_type = rbf_config['rbf_type']
        self.ref_config = rbf_config['ref_config']
        self.in_dim = 3
        self.kc_mult = 1  # kc_mult=2 means extra copy of initial kc

        gridSize_ref = [round(self.ref_config['N_voxel_final']**(1/3)) - 1] * 3  # Not actual value, just for init rbf params
        gridSize_ref = np.array(gridSize_ref).astype(np.int64)
        gridSize_final = [round(args.N_voxel_real_final**(1/3))] * 3
        gridSize_final = np.array(gridSize_final).astype(np.int64)
        if self.n_level > 1:
            down_factor = np.exp((np.log(np.prod(gridSize_final)) - np.log(self.N_voxel_min)) / 3 / (self.n_level - 1))
        else:
            down_factor = 1
        gridSizes = []
        for i in range(self.n_level):
            gridSize_i = np.round(gridSize_final / down_factor**i).astype(np.int64)
            gridSizes.append(gridSize_i)

        n_params_ref = self.count_tensorf_params(self.ref_config['n_lamb_sigma'], [gridSize_ref]) + \
            self.count_tensorf_params(self.ref_config['n_lamb_sh'], [gridSize_ref])
        n_params_tensorf = self.count_tensorf_params(self.density_n_comp, gridSizes) + \
            self.count_tensorf_params(self.app_n_comp, gridSizes)
        if self.n_kernel == 'auto':
            n_params_rbf = n_params_ref - n_params_tensorf
            self.n_kernel = int(n_params_rbf // (self.lc0_dim + self.lcd0_dim + util_misc.get_rbf_params_per_kernel(
                self.rbf_type, self.in_dim, self.kc_mult)))
        print('n_kernel:', self.n_kernel)

        self.rbf_train_start = False
        if self.n_kernel > 0 and not self.get_init_data:
            self.rbf_train_start = True

        super(TensorVMSplitRBF, self).__init__(aabb, gridSize, device, **kargs)

        self.args = args
        self.cmin = torch.tensor([-1, -1, -1])  # ... y x, aabb
        self.cmax = torch.tensor([1, 1, 1])  # ... y x, aabb
        self.register_buffer('cmin_gpu', self.cmin.clone())  # ... y x, aabb
        self.register_buffer('cmax_gpu', self.cmax.clone())  # ... y x, aabb
        self.init_steps = rbf_config['init_steps']
        self.init_data_fp = rbf_config['init_data_fp']
        self.init_rbf = rbf_config['init_rbf']
        self.s_dims = rbf_config['s_dims']
        if self.s_dims != 'aabb':
            self.s_dims = torch.tensor(self.s_dims)  # ... h w, used for initializing rbf
            self.rbf_normalize_pts = True
        else:
            self.rbf_normalize_pts = False
        self.rbf_lc0_normalize = rbf_config['rbf_lc0_normalize']
        self.point_nn_kernel = rbf_config['point_nn_kernel']
        ks_alpha = rbf_config['ks_alpha']
        self.lc_init = rbf_config['lc_init']
        self.lcd_init = rbf_config['lcd_init']
        self.lcb_init = rbf_config['lcb_init']
        self.rbf_suffixes = rbf_config['rbf_suffixes']
        self.pe_lc0_freq = rbf_config['pe_lc0_freq']
        self.pe_lc0_rbf_freq = rbf_config['pe_lc0_rbf_freq']
        self.pe_lc0_rbf_keep = rbf_config['pe_lc0_rbf_keep']
        if self.pe_lc0_rbf_keep == 'all': self.pe_lc0_rbf_keep = self.lc0_dim

        if len(self.pe_lc0_freq) >= 2:
            self.pe_lc0_freqs = torch.linspace(np.log2(self.pe_lc0_freq[0]), np.log2(self.pe_lc0_freq[1]), 
                                              self.app_dim, device=0)
            self.pe_lc0_freqs = torch.exp2(self.pe_lc0_freqs)
        else:
            self.pe_lc0_freqs = None

        if len(self.pe_lc0_rbf_freq) >= 2 and self.pe_lc0_rbf_keep < self.lc0_dim:
            self.pe_lc0_rbf_freqs = torch.linspace(np.log2(self.pe_lc0_rbf_freq[0]), np.log2(self.pe_lc0_rbf_freq[1]), 
                                                   self.lc0_dim - self.pe_lc0_rbf_keep, device=0)
            self.pe_lc0_rbf_freqs = torch.exp2(self.pe_lc0_rbf_freqs)
        else:
            self.pe_lc0_rbf_freqs = None

        fix_params = rbf_config['fix_params']
        self.kc_init_config = rbf_config['kc_init_config']
        self.kw_init_config = rbf_config['kw_init_config']
        self.kc_init_regular = {}
        for k, v in self.kc_init_config.items():
            if v['type'] == 'none':
                self.kc_init_regular[k] = True
            else:
                self.kc_init_regular[k] = False
        sparse_embd_grad = False

        # Init rbf
        self.kc0 = nn.ModuleList()
        self.ks0 = nn.ModuleList()
        self.lc0 = nn.ModuleList()
        if self.rbf_train_start:
            if self.kc_init_regular['0']:
                self.n_kernel = util_misc.get_lower_int_power(self.n_kernel, self.in_dim)
            self.rbf_fn = eval(f'util_network.rbf_{self.rbf_type}_fb')
            if self.lc0_dim > 0:
                self.lc_dims = [[self.n_kernel, self.lc0_dim]]
                self.kc0, self.ks0, self.lc0, self.ks_dims, k_dims, kc_interval = self.create_rbf_params(
                    self.rbf_type, self.n_kernel, self.in_dim, self.lc0_dim, sparse_embd_grad, self.cmin, self.cmax, ks_alpha, is_bag=False)
                nn.init.uniform_(self.lc0.weight, self.lc_init[0], self.lc_init[1])
                self.register_buffer(f'k_dims_0', k_dims)
                self.register_buffer(f'kci0', kc_interval)
                if self.init_rbf:
                    self.init_rbf_params(kargs['init_data'])
            if self.lcd0_dim > 0:
                self.lcd_dims = [[self.n_kernel, self.lcd0_dim]]
                _, _, self.lcd0, _, _, _ = self.create_rbf_params(
                    self.rbf_type, self.n_kernel, self.in_dim, self.lcd0_dim, sparse_embd_grad, self.cmin, self.cmax, ks_alpha, is_bag=False)
                nn.init.uniform_(self.lcd0.weight, self.lcd_init[0], self.lcd_init[1])
        
        if self.lcb_init is not None:
            self.lcb0 = torch.nn.Parameter(torch.zeros((1, self.app_n_comp_total)))
            nn.init.uniform_(self.lcb0.data, self.lcb_init[0], self.lcb_init[1])

        # Init decoder network
        self.init_render_func(self.shadingMode, self.app_dim, self.pos_pe, self.view_pe, self.fea_pe, 
                              self.featureC, self.device, args=args)

        # Fix params
        util_network.fix_params(self, fix_params)


    def count_tensorf_params(self, n_component, gridSizes):
        n_params = 0
        for j in range(len(gridSizes)):
            gridSize = gridSizes[j]
            for i in range(len(self.vecMode)):
                vec_id = self.vecMode[i]
                mat_id_0, mat_id_1 = self.matMode[i]
                n_params += n_component[j] * gridSize[mat_id_1] * gridSize[mat_id_0]
                n_params += n_component[j] * gridSize[vec_id]
        return n_params


    def init_svd_volume(self, res, device):
        gridSize = np.array(res)

        if self.n_level > 1:
            down_factor = np.exp((np.log(np.prod(gridSize)) - np.log(self.N_voxel_min)) / 3 / (self.n_level - 1))
        else:
            down_factor = 1
        self.gridSizes = []
        self.density_plane = torch.nn.ParameterList()
        self.density_line = torch.nn.ParameterList()
        self.app_plane = torch.nn.ParameterList()
        self.app_line = torch.nn.ParameterList()
        for i in range(self.n_level):
            gridSize_i = np.round(gridSize / down_factor**i).astype(np.int64)
            self.gridSizes.append(gridSize_i)
            level_type = self.level_types[i]
            density_plane, density_line = self.init_one_svd(
                self.density_n_comp[i], gridSize_i, 0.1, device, level_type)
            app_plane, app_line = self.init_one_svd(
                self.app_n_comp[i], gridSize_i, 0.1, device, level_type)
            self.density_plane += density_plane
            self.density_line += density_line
            self.app_plane += app_plane
            self.app_line += app_line

        if self.rbf_train_start and self.lcd0_dim > 0:
            self.basis_mat_ds = torch.nn.Linear(2, 1, bias=False).to(device)
        
        if self.rbf_train_start:
            self.app_n_comp_total += self.lc0_dim
        if self.app_dim is not None:
            self.basis_mat = torch.nn.Linear(self.app_n_comp_total, self.app_dim, bias=False).to(device)
        else:
            self.app_dim = self.app_n_comp_total


    def init_one_svd(self, n_component, gridSize, scale, device, level_type):
        plane_coef, line_coef = [], []
        if level_type == 'vm':
            for i in range(len(self.vecMode)):
                vec_id = self.vecMode[i]
                mat_id_0, mat_id_1 = self.matMode[i]
                plane_coef.append(torch.nn.Parameter(
                    scale * torch.randn((1, n_component, gridSize[mat_id_1], gridSize[mat_id_0]))))
                line_coef.append(
                    torch.nn.Parameter(scale * torch.randn((1, n_component, gridSize[vec_id], 1))))
        else:
            plane_coef.append(torch.nn.Parameter(scale * torch.randn((1, n_component * 3, gridSize[2], gridSize[1], gridSize[0]))))
            plane_coef += [torch.nn.Parameter(torch.zeros([0])), torch.nn.Parameter(torch.zeros([0]))]
            line_coef += [torch.nn.Parameter(torch.zeros([0])), torch.nn.Parameter(torch.zeros([0])), 
                          torch.nn.Parameter(torch.zeros([0]))]

        return torch.nn.ParameterList(plane_coef).to(device), torch.nn.ParameterList(line_coef).to(device)
    

    def get_optparam_groups(self, lr_init_spatialxyz = 0.02, lr_init_network = 0.001):
        params_vm = []
        params_g = []
        for i in range(self.n_level):
            level_type = self.level_types[i]
            if level_type == 'vm':
                for j in range(3):
                    params_vm.append(self.density_line[i * 3 + j])
                    params_vm.append(self.density_plane[i * 3 + j])
                    params_vm.append(self.app_line[i * 3 + j])
                    params_vm.append(self.app_plane[i * 3 + j])
            elif level_type == 'g':
                params_g.append(self.density_plane[i * 3])
                params_g.append(self.app_plane[i * 3])

        grad_vars = []
        if len(params_vm) > 0:
            grad_vars += [{'params': params_vm, 'lr': lr_init_spatialxyz}]
        if len(params_g) > 0:
            grad_vars += [{'params': params_g, 'lr': lr_init_spatialxyz * self.args.lr_g_factor}]
        if hasattr(self, 'lc0'):
            grad_vars.append({'params': self.lc0.parameters(), 'lr': self.args.rbf_config['lr_config']['lc0']})
        if hasattr(self, 'lcd0'):
            grad_vars.append({'params': self.lcd0.parameters(), 'lr': self.args.rbf_config['lr_config']['lcd0']})
        if hasattr(self, 'lcb0'):
            grad_vars.append({'params': [self.lcb0], 'lr': self.args.rbf_config['lr_config']['lcb0']})
        if hasattr(self, 'basis_mat'):
            grad_vars += [{'params': self.basis_mat.parameters(), 'lr': self.args.lr_basis_mat}]
        if hasattr(self, 'basis_mat_ds'):
            grad_vars.append({'params': self.basis_mat_ds.parameters(), 'lr': lr_init_network})

        if isinstance(self.renderModule, torch.nn.Module):
            grad_vars += [{'params':self.renderModule.parameters(), 'lr':lr_init_network}]
        return grad_vars


    def vectorDiffs(self, vector_comps):
        total = 0
        
        for idx in range(len(vector_comps)):
            n_comp, n_size = vector_comps[idx].shape[1:-1]
            
            dotp = torch.matmul(vector_comps[idx].view(n_comp,n_size), vector_comps[idx].view(n_comp,n_size).transpose(-1,-2))
            non_diagonal = dotp.view(-1)[1:].view(n_comp-1, n_comp+1)[...,:-1]
            total = total + torch.mean(torch.abs(non_diagonal))
        return total

    def vector_comp_diffs(self):
        return self.vectorDiffs(self.density_line) + self.vectorDiffs(self.app_line)
    
    def density_L1(self):
        total = 0
        for idx in range(len(self.density_plane)):
            if self.density_plane[idx].shape[0] != 0:
                total = total + torch.mean(torch.abs(self.density_plane[idx]))
            if self.density_line[idx].shape[0] != 0:
                total = total + torch.mean(torch.abs(self.density_line[idx]))
        return total
    
    def TV_loss_density(self, reg):
        total = 0
        for idx in range(len(self.density_plane)):
            if self.density_plane[idx].shape[0] != 0:
                total = total + reg(self.density_plane[idx]) * 1e-2
        return total
        
    def TV_loss_app(self, reg):
        total = 0
        for idx in range(len(self.app_plane)):
            if self.app_plane[idx].shape[0] != 0:
                total = total + reg(self.app_plane[idx]) * 1e-2
        return total


    def compute_densityfeature(self, xyz_sampled, others={}):
        # xyz_sampled: [n_pts, 3], normalized to range [-1, 1] based on self.aabb

        # plane + line basis
        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach().view(3, -1, 1, 2)  # [3, n_pts, 1, 2]
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))  # [3, n_pts]
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)  # [3, n_pts, 1, 2]

        sigma_feature = torch.zeros((xyz_sampled.shape[0],), device=xyz_sampled.device)  # [n_pts]
        for idx_plane in range(len(self.density_plane)):
            level_type = self.level_types[idx_plane // 3]
            if level_type == 'vm':
                idx_coord = idx_plane % 3
                # density_plane: [1, c, h, w]
                plane_coef_point = F.grid_sample(self.density_plane[idx_plane], coordinate_plane[[idx_coord]],
                                                    align_corners=True).view(-1, xyz_sampled.shape[0])  # [c, n_pts]
                # density_line: [1, c, d, 1]
                line_coef_point = F.grid_sample(self.density_line[idx_plane], coordinate_line[[idx_coord]],
                                                align_corners=True).view(-1, xyz_sampled.shape[0])  # [c, n_pts]
                sigma_feature = sigma_feature + torch.sum(plane_coef_point * line_coef_point, dim=0)
            else:
                if self.density_plane[idx_plane].shape[0] != 0:
                    out = F.grid_sample(self.density_plane[idx_plane], xyz_sampled[None, :, None, None],
                                        align_corners=True).view(-1, xyz_sampled.shape[0])  # [c, n_pts]
                    sigma_feature = sigma_feature + torch.sum(out, dim=0)
        
        if self.rbf_train_start and self.lcd0_dim > 0:
            kernel_idx, rbf_out = others['kernel_idx'], others['rbf_out']
            out = self.lcd0(kernel_idx)  # [p k_topk d_lcd0]
            out = (out * rbf_out).sum(1)  # [p d_lcd0]
            sigma_feature = torch.stack([sigma_feature, out.sum(-1)], -1)
            sigma_feature = self.basis_mat_ds(sigma_feature)[..., 0]

        return sigma_feature


    def compute_appfeature(self, xyz_sampled, others={}):
        # xyz_sampled: [n_pts, 3], normalized to range [-1, 1] based on self.aabb

        # plane + line basis
        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach().view(3, -1, 1, 2)
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)

        plane_coef_point, line_coef_point = [], []
        grid_coef_point = []
        for idx_plane in range(len(self.app_plane)):
            level_type = self.level_types[idx_plane // 3]
            if level_type == 'vm':
                idx_coord = idx_plane % 3
                plane_coef_point.append(F.grid_sample(self.app_plane[idx_plane], coordinate_plane[[idx_coord]],
                                                    align_corners=True).view(-1, *xyz_sampled.shape[:1]))  # [c, n_pts]
                line_coef_point.append(F.grid_sample(self.app_line[idx_plane], coordinate_line[[idx_coord]],
                                                align_corners=True).view(-1, *xyz_sampled.shape[:1]))  # [c, n_pts]
            else:
                if self.app_plane[idx_plane].shape[0] != 0:
                    grid_coef_point.append(F.grid_sample(self.app_plane[idx_plane], xyz_sampled[None, :, None, None],
                                                        align_corners=True).view(-1, xyz_sampled.shape[0]).T)  # [n_pts, c]
        plane_coef_point, line_coef_point = torch.cat(plane_coef_point), torch.cat(line_coef_point)
        app_feature = (plane_coef_point * line_coef_point).T  # [n_pts, c]
        if len(grid_coef_point) > 0:
            app_feature = torch.cat([app_feature] + grid_coef_point, -1)

        # Forward rbf lc0
        if self.rbf_train_start:
            kernel_idx, rbf_out = others['kernel_idx'], others['rbf_out']
            if self.pe_lc0_rbf_freqs is not None:
                if self.pe_lc0_rbf_keep > 0:
                    rbf_out = torch.cat([rbf_out.expand(-1, -1, self.pe_lc0_rbf_keep), 
                        torch.sin(rbf_out * self.pe_lc0_rbf_freqs[None, None])], -1)  # [p k_topk d_lc0]
                else:
                    rbf_out = torch.sin(rbf_out * self.pe_lc0_rbf_freqs[None, None])  # [p k_topk d_lc0]
            out = self.lc0(kernel_idx)  # [p k_topk d_lc0]
            out = (out * rbf_out).sum(1)  # [p d_lc0]
            app_feature = torch.cat([app_feature, out], -1)

        if hasattr(self, 'lcb0'):
            app_feature = app_feature + self.lcb0

        if hasattr(self, 'basis_mat'):
            app_feature = self.basis_mat(app_feature)  # [n_pts, app_dim]

        if self.pe_lc0_freqs is not None:
            app_feature = app_feature + torch.sin(app_feature * self.pe_lc0_freqs[None])

        return app_feature


    def create_rbf_params(self, rbf_type, n_kernel, in_dim, lc_dim, sparse_embd_grad, cmin, cmax, ks_alpha, scale_grad_by_freq=False, is_bag=True):
        if rbf_type.endswith('_a') or rbf_type.endswith('_f'):
            ks_dims = [in_dim, in_dim]
        elif rbf_type.endswith('_d'):
            ks_dims = [in_dim]
        elif rbf_type.endswith('_s'):
            ks_dims = [1]
        else:
            raise NotImplementedError
        kc = torch.nn.Embedding(n_kernel, in_dim, scale_grad_by_freq=False, sparse=sparse_embd_grad)
        ks = torch.nn.Embedding(n_kernel, np.prod(ks_dims), scale_grad_by_freq=False, 
            sparse=sparse_embd_grad)
        if rbf_type.endswith('_a') or rbf_type.endswith('_f'):
            ks.weight.data = torch.eye(in_dim)[None, ...].repeat(n_kernel, 1, 1).reshape(n_kernel, -1)
        elif rbf_type.endswith('_d'):
            ks.weight.data[:] = 1
        elif rbf_type.endswith('_s'):
            ks.weight.data[:] = 1
        else:
            raise NotImplementedError

        if is_bag:
            lc = torch.nn.EmbeddingBag(n_kernel, lc_dim, scale_grad_by_freq=scale_grad_by_freq, 
                                       sparse=sparse_embd_grad, mode='sum')
        else:
            lc = torch.nn.Embedding(n_kernel, lc_dim, scale_grad_by_freq=scale_grad_by_freq, sparse=sparse_embd_grad)

        # Initialize
        n_kernel_per_dim = int(math.ceil(n_kernel**(1/in_dim)))
        k_dims = [n_kernel_per_dim] * in_dim
        points, _, side_interval = util_misc.get_grid_points(k_dims, align_corners=True, 
            vmin=cmin.tolist(), vmax=cmax.tolist())
        points = points.flip(-1)  # x y ...
        side_interval = side_interval.flip(-1)  # x y ...
        kc.weight.data = points[:n_kernel]
        if rbf_type.startswith('nlin_'):
            if rbf_type.endswith('_s'):
                ks.weight.data *= 1 / side_interval.mean()
            elif rbf_type.endswith('_d'):
                ks.weight.data *= 1 / side_interval[None]
            elif rbf_type.endswith('_f'):
                ks.weight.data *= (1 / side_interval)[:, None].repeat(1, side_interval.shape[0]).view(1, -1)
            else:
                raise NotImplementedError
        else:
            if rbf_type.endswith('_s'):
                ks.weight.data *= 1 / (side_interval.mean() / 2) ** 2
            elif rbf_type.endswith('_d'):
                ks.weight.data *= 1 / (side_interval[None] / 2) ** 2
            elif rbf_type.endswith('_a'):
                ks.weight.data *= (1 / (side_interval / 2) ** 2)[:, None].repeat(1, side_interval.shape[0]).view(1, -1)
            else:
                raise NotImplementedError
        ks.weight.data *= ks_alpha

        return kc, ks, lc, ks_dims, torch.tensor(k_dims).flip(-1), side_interval


    def get_kc(self, suffix):
        return getattr(self, 'kc' + suffix).weight


    def get_ks(self, suffix):
        return getattr(self, 'ks' + suffix).weight.view(-1, *self.ks_dims)


    def clip_rbf_params(self):
        for k in self.rbf_suffixes:
            if hasattr(self, f'ks{k}.weight') and getattr(self, 'ks'+k).weight.requires_grad:
                if k == '0':
                    vmin_factor = 100
                    vmax_factor = 10
                else:
                    vmin_factor = 1
                    vmax_factor = 2
                util_misc.clip_kw_sq(
                    getattr(self, 'ks'+k).weight, self.rbf_type, self.cmin.flip(-1), self.cmax.flip(-1), 
                    self.s_dims.flip(-1), is_ks=True, is_flat=True, vmin_factor=vmin_factor, vmax_factor=vmax_factor)
                

    def init_rbf_params(self, init_data=None):
        # Get grid points within aabb
        if self.s_dims == 'aabb':
            if self.get_init_data:
                self.aabb_rbf = self.aabb
                aabb_size = self.aabb_rbf[1] - self.aabb_rbf[0]  # w h ...
                sl = aabb_size.prod() ** (1/3) / 512
                self.s_dims = (aabb_size / sl).round().long().flip(-1)  # ... h w
            else:
                points_weight = init_data['points_weight']
                self.aabb_rbf = init_data['aabb'].to(self.device)
                self.s_dims = init_data['s_dims']
                del init_data

            self.cmin = self.aabb_rbf[0].flip(-1).cpu()  # ... y x
            self.cmax = self.aabb_rbf[1].flip(-1).cpu()  # ... y x
            self.register_buffer('cmin_gpu', self.cmin.clone())  # ... y x
            self.register_buffer('cmax_gpu', self.cmax.clone())  # ... y x

            t = time.time()
            pts, _, _ = util_misc.get_grid_points(self.s_dims.tolist(), align_corners=True, 
                vmin=self.cmin.tolist(), vmax=self.cmax.tolist(), device=0)
            pts = pts.flip(-1)  # x y z
            print('Generate grid points:', time.time() - t)

            # Get density and features at these points
            if self.get_init_data:
                t = time.time()
                density_all, features_all = self.get_pts_data(pts, normalize_pts=True)
                if features_all.shape[-1] > 27:
                    shuffle = torch.randperm(features_all.shape[-1]).long()[:27]
                    features_all = features_all[:, shuffle]
                print('Get density and features at grid points:', time.time() - t)
                density, features = density_all.cpu().view(-1, 1), features_all.cpu()
                del density_all, features_all

                # Compute point weight
                density = 1 - torch.exp(-density)
                features = (features - features.min()) / (features.max() - features.min())
                features = features.reshape(*self.s_dims.tolist(), -1)
                features = kornia.filters.SpatialGradient3d(mode='diff', order=1)(features.movedim(-1, 0)[None])[0]  # [c 3 d h w]
                features = features.movedim((0, 1), (-2, -1)).pow(2).sum(dim=[-2, -1]).sqrt()[..., None]  # [d h w 1]
                points_weight = density * features.reshape(-1, 1)
                del density, features

                init_data = {'points_weight': points_weight, 
                             'aabb': self.aabb_rbf.cpu(), 's_dims': self.s_dims.cpu()}

                # # Save init data
                # os.makedirs(os.path.dirname(self.init_data_fp), exist_ok=True)
                # torch.save(init_data, self.init_data_fp)
                # density_all = density_all.reshape(*self.s_dims.tolist(), -1).permute(2, 1, 0, 3)
                # vertices, triangles = mcubes.marching_cubes(density_all.numpy()[..., 0], 1.)
                # vertices = vertices / (np.array(density_all.shape[:3]) - 1.0) * \
                #     aabb_size.cpu().numpy()[None] + self.aabb_rbf[0].cpu().numpy()[None]
                # trimesh.Trimesh(vertices, triangles, process=False).export(self.init_data_fp.split('.')[0] + '.ply')

                torch.cuda.empty_cache()
                return init_data
        else:
            raise NotImplementedError
            t = time.time()
            pts, _, _ = util_misc.get_grid_points(self.s_dims.tolist(), align_corners=True, 
                vmin=self.cmin.tolist(), vmax=self.cmax.tolist(), device=0)
            pts = pts.flip(-1)  # x y z
            print('Generate grid points:', time.time() - t)

            # Get density and features at these points
            if self.get_init_data:
                t = time.time()
                density_all, features_all = self.get_pts_data(pts)
                if features_all.shape[-1] > 27:
                    shuffle = torch.randperm(features_all.shape[-1]).long()[:27]
                    features_all = features_all[:, shuffle]
                print('Get density and features at grid points:', time.time() - t)
                density_all, features_all = density_all.cpu(), features_all.cpu()
                init_data = {'density': density_all, 'features': features_all, 
                            'aabb': self.aabb.cpu()}

                # # Save init data
                # os.makedirs(os.path.dirname(self.init_data_fp), exist_ok=True)
                # torch.save(init_data, self.init_data_fp)
                # density_all = density_all.reshape(*self.s_dims.tolist(), -1).permute(2, 1, 0, 3)
                # vertices, triangles = mcubes.marching_cubes(density_all.numpy()[..., 0], 1.)
                # vertices = vertices / (density_all.shape[0] - 1.0) * 2 - 1
                # trimesh.Trimesh(vertices, triangles, process=False).export(self.init_data_fp.split('.')[0] + '.ply')

                torch.cuda.empty_cache()
                return init_data
        
            density_all = init_data['density'].view(-1, 1)
            features_all = init_data['features']
            self.aabb_rbf = init_data['aabb'].to(self.device)
            del init_data

        # Init rbf parameters
        util_network.init_nerf_rbf_params(self, pts, points_weight, None, self.kc_init_config, 
                                          self.kw_init_config, device=0)
        del points_weight
        torch.cuda.empty_cache()

        # Build kd tree
        self.kdtree0 = KDTree(getattr(self, 'kc0').weight.detach().cpu().numpy())


    def step_after_iter(self, step):
        if step == self.init_steps and self.get_init_data:
            return self.init_rbf_params()


    def compute_rbf(self, xyz_sampled):
        out = {}

        # Forward rbf lc0
        if self.rbf_train_start:
            # Un-normalize
            pts = (xyz_sampled + 1) / 2 * (self.aabb[1] - self.aabb[0]) + self.aabb[0]
            # Re-normalize based on self.aabb_rbf
            if self.rbf_normalize_pts:
                pts = (pts - self.aabb_rbf[0]) * 2.0 / (self.aabb_rbf[1] - self.aabb_rbf[0]) - 1

            suffix = '0'
            if self.point_nn_kernel <= 0:  # Use all kernels for each point
                raise NotImplementedError
            else:
                kernel_idx = self.forward_kernel_idx(pts, None, suffix)
                rbf_out = self.forward_rbf(pts, kernel_idx, suffix)  # [p k_topk]
                if self.rbf_lc0_normalize:
                    rbf_out = rbf_out / (rbf_out.detach().sum(-1, keepdim=True) + 1e-8)
                rbf_out = rbf_out[..., None]  # [p k_topk 1]

            out['kernel_idx'] = kernel_idx
            out['rbf_out'] = rbf_out

        return out
    

    def forward_kernel_idx(self, x_g, point_idx, suffix):
        if hasattr(self, 'point_kernel_idx_' + suffix):
            # Use pre-computed knn
            point_kernel_idx = getattr(self, 'point_kernel_idx_' + suffix)
            kernel_idx = point_kernel_idx[point_idx.to(point_kernel_idx.device)].to(x_g.device)  # [p topk]
        elif self.kc_init_regular[suffix]:
            # Find first ring neighbors in regular grid
            kernel_idx = util_misc.get_multi_index_first_ring(x_g, 
                getattr(self, f'k_dims_{suffix}'), getattr(self, f'kci{suffix}'), 
                self.cmin_gpu.flip(-1), self.cmax_gpu.flip(-1), True)  # (p, n_ngb, n_in)
            kernel_idx = util_misc.ravel_multi_index(kernel_idx, getattr(self, f'k_dims_{suffix}'))  # (p, n_ngb)
        else:
            # Knn on the fly
            if not hasattr(self, 'kdtree' + suffix):
                setattr(self, 'kdtree' + suffix, KDTree(getattr(self, 'kc' + suffix).weight.detach().cpu().numpy()))
            _, kernel_idx = util_clustering.query_chunked(getattr(self, 'kdtree' + suffix), x_g.cpu().numpy(), 
                k=self.point_nn_kernel, sqr_dists=True, chunk_size=int(2e8), return_dist=False)
            kernel_idx = torch.tensor(kernel_idx.astype(np.int32), device=x_g.device)  # [p topk]
            if kernel_idx.ndim == 1:
                kernel_idx = kernel_idx[:, None]
            
        return kernel_idx


    def forward_rbf(self, x_g, kernel_idx, suffix):
        if kernel_idx is None:  # Use all kernels for each point
            raise NotImplementedError
        else:
            kc = getattr(self, 'kc' + suffix)(kernel_idx)  # [p k d]
            ks = getattr(self, 'ks' + suffix)(kernel_idx).view(*kernel_idx.shape, *self.ks_dims)  # [p k d d] or [p k 1]
            rbf_out, _ = self.rbf_fn(x_g, kc, ks)  # [p k_topk]
        return rbf_out


    @torch.no_grad()
    def up_sampling_VM(self, plane_coef, line_coef, res_target):
        res_target = np.array(res_target)
        if self.n_level > 1:
            down_factor = np.exp((np.log(np.prod(res_target)) - np.log(self.N_voxel_min)) / 3 / (self.n_level - 1))
        else:
            down_factor = 1

        for idx_level in range(self.n_level):
            level_type = self.level_types[idx_level]
            res_target_i = np.round(res_target / down_factor**idx_level).astype(np.int64)
            self.gridSizes[idx_level] = res_target_i
            res_target_i = res_target_i.tolist()
            if level_type == 'vm':
                for i in range(len(self.vecMode)):
                    vec_id = self.vecMode[i]
                    mat_id_0, mat_id_1 = self.matMode[i]
                    idx = idx_level * 3 + i
                    plane_coef[idx] = torch.nn.Parameter(
                        F.interpolate(plane_coef[idx].data, size=(res_target_i[mat_id_1], res_target_i[mat_id_0]), mode='bilinear',
                                    align_corners=True))
                    line_coef[idx] = torch.nn.Parameter(
                        F.interpolate(line_coef[idx].data, size=(res_target_i[vec_id], 1), mode='bilinear', align_corners=True))
            else:
                idx = idx_level * 3
                plane_coef[idx] = torch.nn.Parameter(
                    F.interpolate(plane_coef[idx].data, size=(res_target_i[2], res_target_i[1], res_target_i[0]), mode='trilinear', align_corners=True))

        return plane_coef, line_coef

    @torch.no_grad()
    def upsample_volume_grid(self, res_target, res_target_real):
        if not self.args.no_upsample:
            self.app_plane, self.app_line = self.up_sampling_VM(self.app_plane, self.app_line, res_target_real)
            self.density_plane, self.density_line = self.up_sampling_VM(self.density_plane, self.density_line, res_target_real)

        self.update_stepSize(res_target)

    @torch.no_grad()
    def shrink(self, new_aabb):
        print("====> shrinking ...")
        aabb_size = self.aabb[1] - self.aabb[0]
        xyz_min, xyz_max = new_aabb
        gridSize = torch.from_numpy(self.gridSizes[0]).to(aabb_size.device)
        units = aabb_size / (gridSize - 1)

        t_l, b_r = (xyz_min - self.aabb[0]) / units, (xyz_max - self.aabb[0]) / units
        t_l, b_r = torch.floor(t_l).long(), torch.ceil(b_r).long() + 1
        b_r = torch.stack([b_r, gridSize]).amin(0)

        t_l_np = t_l.cpu().numpy()
        b_r_np = b_r.cpu().numpy()
        for idx_level in range(self.n_level):
            level_type = self.level_types[idx_level]
            if idx_level == 0:
                t_l_i = t_l_np
                b_r_i = b_r_np
            else:
                factor = self.gridSizes[idx_level] / self.gridSizes[0]
                t_l_i = np.floor(t_l_np * factor).astype(np.int64)
                b_r_i = np.ceil(b_r_np * factor).astype(np.int64)

            if level_type == 'vm':
                for i in range(len(self.vecMode)):
                    idx = idx_level * 3 + i
                    mode0 = self.vecMode[i]
                    self.density_line[idx] = torch.nn.Parameter(
                        self.density_line[idx].data[...,t_l_i[mode0]:b_r_i[mode0],:]
                    )
                    self.app_line[idx] = torch.nn.Parameter(
                        self.app_line[idx].data[...,t_l_i[mode0]:b_r_i[mode0],:]
                    )
                    mode0, mode1 = self.matMode[i]
                    self.density_plane[idx] = torch.nn.Parameter(
                        self.density_plane[idx].data[...,t_l_i[mode1]:b_r_i[mode1],t_l_i[mode0]:b_r_i[mode0]]
                    )
                    self.app_plane[idx] = torch.nn.Parameter(
                        self.app_plane[idx].data[...,t_l_i[mode1]:b_r_i[mode1],t_l_i[mode0]:b_r_i[mode0]]
                    )
            else:
                idx = idx_level * 3
                self.density_plane[idx] = torch.nn.Parameter(
                    self.density_plane[idx].data[...,t_l_i[2]:b_r_i[2],t_l_i[1]:b_r_i[1],t_l_i[0]:b_r_i[0]]
                    )
                self.app_plane[idx] = torch.nn.Parameter(
                    self.app_plane[idx].data[...,t_l_i[2]:b_r_i[2],t_l_i[1]:b_r_i[1],t_l_i[0]:b_r_i[0]]
                    )

        xyz_min, xyz_max = new_aabb
        t_l, b_r = (xyz_min - self.aabb[0]) / self.units, (xyz_max - self.aabb[0]) / self.units
        t_l, b_r = torch.floor(t_l).long(), torch.ceil(b_r).long() + 1
        b_r = torch.stack([b_r, self.gridSize]).amin(0)

        if not torch.all(self.alphaMask.gridSize == self.gridSize):
            t_l_r, b_r_r = t_l / (self.gridSize-1), (b_r-1) / (self.gridSize-1)
            correct_aabb = torch.zeros_like(new_aabb)
            correct_aabb[0] = (1-t_l_r)*self.aabb[0] + t_l_r*self.aabb[1]
            correct_aabb[1] = (1-b_r_r)*self.aabb[0] + b_r_r*self.aabb[1]
            print("aabb", new_aabb, "\ncorrect aabb", correct_aabb)
            new_aabb = correct_aabb

        newSize = b_r - t_l
        self.aabb = new_aabb

        voxels_orig = np.prod(self.gridSizes[0])
        target_res = utils.N_to_reso(voxels_orig, new_aabb)
        if self.args.scale_reso:
            target_res = utils.scale_reso(target_res, voxels_orig)
        self.app_plane, self.app_line = self.up_sampling_VM(self.app_plane, self.app_line, target_res)
        self.density_plane, self.density_line = self.up_sampling_VM(self.density_plane, self.density_line, target_res)

        self.update_stepSize((newSize[0], newSize[1], newSize[2]))

nerf_tensorf/network.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import trimesh
import pysdf
import numpy as np
import torch
from argparse import ArgumentParser


parser = ArgumentParser()
parser.add_argument("--path", type=str)
config = parser.parse_args()

path = config.path

print(path)
mesh = trimesh.load(path, force='mesh')
if not mesh.is_watertight:
    print(f"[WARN] mesh is not watertight! SDF maybe incorrect.")

# Normalize to [-1, 1] (different from instant-sdf where is [0, 1])
vs = mesh.vertices
vmin = vs.min(0)
vmax = vs.max(0)
v_center = (vmin + vmax) / 2
vs = (vs - v_center[None, :])
v_scale = 1. / (np.sqrt(np.sum(vs**2, -1).max()) / 0.99)
vs *= v_scale
mesh.vertices = vs
mesh.export(path.rsplit('.', 1)[0] + '_nrml.obj')

# Generate evaluation points
sdf_fn = pysdf.SDF(mesh.vertices, mesh.faces)
ns_taus = []
n_points = {f'ns_{ns_tau:.0e}': int(1e7) for ns_tau in ns_taus}
n_points.update({'unif': int(1e7)})
points = {}
sdfs = {}
# surface
if 'surf' in n_points:
    points['surf'] = mesh.sample(n_points['surf']).view(np.ndarray)
    sdfs['surf'] = np.zeros([points['surf'].shape[0], 1])
# near surface
for ns_tau in ns_taus:
    k = f'ns_{ns_tau:.0e}'
    points[k] = mesh.sample(n_points[k]).view(np.ndarray)
    points[k] += ns_tau * np.random.rand(*points[k].shape)
# uniform random
if 'unif' in n_points:
    points['unif'] = (np.random.rand(n_points['unif'], 3) * 2 - 1)
# Compute SDF
for k in points:
    if k != 'surf':
        sdfs[k] = -sdf_fn(points[k])[:, None]
    points[k] = points[k].astype(np.float32)
    sdfs[k] = sdfs[k].astype(np.float32)
torch.save({'points': points, 'sdfs': sdfs}, path.rsplit('.', 1)[0] + '_nrml_eval_points.pt')

preproc_mesh.py
oppo-us-research/NeuRBF
from .config_helpers import parse_config

__all__ = (
    "parse_config",
)

thirdparty/kplanes/configs/__init__.py
oppo-us-research/NeuRBF
import argparse
import os
from datetime import datetime


def parse_config(defaults):
    # Build experiment configuration
    parser = argparse.ArgumentParser("Train + evaluate kernel model")
    parser.add_argument("--config", default=None)
    parser.add_argument("--config-updates", default=[], nargs='*')
    parser.add_argument("--logdir", default=None)
    args = parser.parse_args()
    # Allow up to two configs, one for reloading and one for training
    reload_cfg = None
    #train_cfg = defaults
    train_cfg = None
    # assert args.logdir is not None or args.config is not None, "Must specify at least one config"
    # Passing both a logdir and a config means train new scenes using pretrained dicts
    if args.logdir is not None:
        logged_config_file = os.path.join(args.logdir, "config.yaml")
        if not os.path.isfile(logged_config_file):
            raise RuntimeError(f"logdir {args.logdir} doesn't specify a config-file")
        print(f"Loading configuration from logs at {logged_config_file}")
        reload_cfg = defaults.clone()
        reload_cfg.merge_from_file(logged_config_file)
        reload_cfg.merge_from_list(args.config_updates)
    if args.config is not None:
        train_cfg = defaults.clone()
        # Reuse the same config as was reloaded, but make updates for datasets and logdir
        if reload_cfg is not None:
            train_cfg.merge_from_file(logged_config_file)
        train_cfg.merge_from_file(args.config)
        train_cfg.merge_from_list(args.config_updates)
    print(f"[{datetime.now()}] Starting")
    return train_cfg, reload_cfg


thirdparty/kplanes/configs/config_helpers.py
oppo-us-research/NeuRBF
from .llff_dataset import LLFFDataset
from .synthetic_nerf_dataset import SyntheticNerfDataset
from .video_datasets import Video360Dataset
from .phototourism_dataset import PhotoTourismDataset

__all__ = (
    "LLFFDataset",
    "SyntheticNerfDataset",
    "Video360Dataset",
    "PhotoTourismDataset",
)

thirdparty/kplanes/datasets/__init__.py
oppo-us-research/NeuRBF
from abc import ABC
import os
from typing import Optional, List, Union

import torch
from torch.utils.data import Dataset

from .intrinsics import Intrinsics


class BaseDataset(Dataset, ABC):
    def __init__(self,
                 datadir: str,
                 scene_bbox: torch.Tensor,
                 split: str,
                 is_ndc: bool,
                 is_contracted: bool,
                 rays_o: Optional[torch.Tensor],
                 rays_d: Optional[torch.Tensor],
                 intrinsics: Union[Intrinsics, List[Intrinsics]],
                 batch_size: Optional[int] = None,
                 imgs: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
                 sampling_weights: Optional[torch.Tensor] = None,
                 weights_subsampled: int = 1,
                 ):
        self.datadir = datadir
        self.name = os.path.basename(self.datadir)
        self.scene_bbox = scene_bbox
        self.split = split
        self.is_ndc = is_ndc
        self.is_contracted = is_contracted
        self.weights_subsampled = weights_subsampled
        self.batch_size = batch_size
        if self.split == 'train':
            assert self.batch_size is not None
        self.rays_o = rays_o
        self.rays_d = rays_d
        self.imgs = imgs
        if self.imgs is not None:
            self.num_samples = len(self.imgs)
        elif self.rays_o is not None:
            self.num_samples = len(self.rays_o)
        else:
            self.num_samples = None
            #raise RuntimeError("Can't figure out num_samples.")
        self.intrinsics = intrinsics
        self.sampling_weights = sampling_weights
        if self.sampling_weights is not None:
            assert len(self.sampling_weights) == self.num_samples, (
                f"Expected {self.num_samples} sampling weights but given {len(self.sampling_weights)}."
            )
        self.sampling_batch_size = 2_000_000  # Increase this?
        if self.num_samples is not None:
            self.use_permutation = self.num_samples < 100_000_000  # 64M is static
        else:
            self.use_permutation = True
        self.perm = None

    @property
    def img_h(self) -> Union[int, List[int]]:
        if isinstance(self.intrinsics, list):
            return [i.height for i in self.intrinsics]
        return self.intrinsics.height

    @property
    def img_w(self) -> Union[int, List[int]]:
        if isinstance(self.intrinsics, list):
            return [i.width for i in self.intrinsics]
        return self.intrinsics.width

    def reset_iter(self):
        if self.sampling_weights is None and self.use_permutation:
            self.perm = torch.randperm(self.num_samples)
        else:
            del self.perm
            self.perm = None

    def get_rand_ids(self, index):
        assert self.batch_size is not None, "Can't get rand_ids for test split"
        if self.sampling_weights is not None:
            batch_size = self.batch_size // (self.weights_subsampled ** 2)
            num_weights = len(self.sampling_weights)
            if num_weights > self.sampling_batch_size:
                # Take a uniform random sample first, then according to the weights
                subset = torch.randint(
                    0, num_weights, size=(self.sampling_batch_size,),
                    dtype=torch.int64, device=self.sampling_weights.device)
                samples = torch.multinomial(
                    input=self.sampling_weights[subset], num_samples=batch_size)
                return subset[samples]
            return torch.multinomial(
                input=self.sampling_weights, num_samples=batch_size)
        else:
            batch_size = self.batch_size
            if self.use_permutation:
                return self.perm[index * batch_size: (index + 1) * batch_size]
            else:
                return torch.randint(0, self.num_samples, size=(batch_size, ))

    def __len__(self):
        if self.split == 'train':
            return (self.num_samples + self.batch_size - 1) // self.batch_size
        else:
            return self.num_samples

    def __getitem__(self, index, return_idxs: bool = False):
        if self.split == 'train':
            index = self.get_rand_ids(index)
        out = {}
        if self.rays_o is not None:
            out["rays_o"] = self.rays_o[index]
        if self.rays_d is not None:
            out["rays_d"] = self.rays_d[index]
        if self.imgs is not None:
            out["imgs"] = self.imgs[index]
        else:
            out["imgs"] = None
        if return_idxs:
            return out, index
        return out

thirdparty/kplanes/datasets/base_dataset.py
oppo-us-research/NeuRBF
import collections
import struct

import numpy as np


Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"])
BaseImage = collections.namedtuple("Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"])
CAMERA_MODELS = {
    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
    CameraModel(model_id=7, model_name="FOV", num_params=5),
    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
}
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS])


class Image(BaseImage):
    def qvec2rotmat(self):
        return qvec2rotmat(self.qvec)


def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
    """Read and unpack the next bytes from a binary file.
    :param fid:
    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
    :param endian_character: Any of {@, =, <, >, !}
    :return: Tuple of read and unpacked values.
    """
    data = fid.read(num_bytes)
    return struct.unpack(endian_character + format_char_sequence, data)


def read_cameras_text(path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasText(const std::string& path)
        void Reconstruction::ReadCamerasText(const std::string& path)
    """
    cameras = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                camera_id = int(elems[0])
                model = elems[1]
                width = int(elems[2])
                height = int(elems[3])
                params = np.array(tuple(map(float, elems[4:])))
                cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params)
    return cameras


def read_cameras_binary(path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasBinary(const std::string& path)
        void Reconstruction::ReadCamerasBinary(const std::string& path)
    """
    cameras = {}
    with open(path_to_model_file, "rb") as fid:
        num_cameras = read_next_bytes(fid, 8, "Q")[0]
        for camera_line_index in range(num_cameras):
            camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ")
            camera_id = camera_properties[0]
            model_id = camera_properties[1]
            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name  # only pinhole (model_id = 1) supported
            width = camera_properties[2]
            height = camera_properties[3]
            num_params = CAMERA_MODEL_IDS[model_id].num_params
            params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params)
            cameras[camera_id] = Camera(
                id=camera_id, model=model_name, width=width, height=height, params=np.array(params)
            )
        assert len(cameras) == num_cameras
    return cameras


def read_images_text(path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesText(const std::string& path)
        void Reconstruction::WriteImagesText(const std::string& path)
    """
    images = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                image_id = int(elems[0])
                qvec = np.array(tuple(map(float, elems[1:5])))
                tvec = np.array(tuple(map(float, elems[5:8])))
                camera_id = int(elems[8])
                image_name = elems[9]
                elems = fid.readline().split()
                xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))])
                point3D_ids = np.array(tuple(map(int, elems[2::3])))
                images[image_id] = Image(
                    id=image_id,
                    qvec=qvec,
                    tvec=tvec,
                    camera_id=camera_id,
                    name=image_name,
                    xys=xys,
                    point3D_ids=point3D_ids,
                )
    return images


def read_images_binary(path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesBinary(const std::string& path)
        void Reconstruction::WriteImagesBinary(const std::string& path)
    """
    images = {}
    with open(path_to_model_file, "rb") as fid:
        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
        for image_index in range(num_reg_images):
            binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi")
            image_id = binary_image_properties[0]
            qvec = np.array(binary_image_properties[1:5])
            tvec = np.array(binary_image_properties[5:8])
            camera_id = binary_image_properties[8]
            image_name = ""
            current_char = read_next_bytes(fid, 1, "c")[0]
            while current_char != b"\x00":  # look for the ASCII 0 entry
                image_name += current_char.decode("utf-8")
                current_char = read_next_bytes(fid, 1, "c")[0]
            num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0]
            x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2D, format_char_sequence="ddq" * num_points2D)
            xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))])
            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
            images[image_id] = Image(
                id=image_id,
                qvec=qvec,
                tvec=tvec,
                camera_id=camera_id,
                name=image_name,
                xys=xys,
                point3D_ids=point3D_ids,
            )
    return images


def qvec2rotmat(qvec):
    return np.array(
        [
            [
                1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
                2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
            ],
            [
                2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
            ],
            [
                2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
                2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
            ],
        ]
    )


def rotmat2qvec(R):
    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
    K = (
        np.array(
            [
                [Rxx - Ryy - Rzz, 0, 0, 0],
                [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
                [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
                [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],
            ]
        )
        / 3.0
    )
    eigvals, eigvecs = np.linalg.eigh(K)
    qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
    if qvec[0] < 0:
        qvec *= -1
    return

thirdparty/kplanes/datasets/colmap_utils.py
oppo-us-research/NeuRBF
from typing import Tuple, Optional, Dict, Any, List
import logging as log
import os
import resource

import torch
from torch.multiprocessing import Pool
import torchvision.transforms
from PIL import Image
import imageio.v3 as iio

from thirdparty.kplanes.utils.my_tqdm import tqdm

pil2tensor = torchvision.transforms.ToTensor()
# increase ulimit -n (number of open files) otherwise parallel loading might fail
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (16192, rlimit[1]))


def _load_phototourism_image(idx: int,
                             paths: List[str],
                             out_h: List[int],
                             out_w: List[int]) -> torch.Tensor:
    f_path = paths[idx]
    img = Image.open(f_path).convert('RGB')
    img.resize((out_w[idx], out_h[idx]), Image.LANCZOS)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]
    return img


def _parallel_loader_phototourism_image(args):
    torch.set_num_threads(1)
    return _load_phototourism_image(**args)


def _load_llff_image(idx: int,
                     paths: List[str],
                     data_dir: str,
                     out_h: int,
                     out_w: int,
                     ) -> torch.Tensor:
    f_path = os.path.join(data_dir, paths[idx])
    img = Image.open(f_path).convert('RGB')

    img = img.resize((out_w, out_h), Image.LANCZOS)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]
    return img


def _parallel_loader_llff_image(args):
    torch.set_num_threads(1)
    return _load_llff_image(**args)


def _load_nerf_image_pose(idx: int,
                          frames: List[Dict[str, Any]],
                          data_dir: str,
                          out_h: Optional[int],
                          out_w: Optional[int],
                          downsample: float,
                          ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
    # Fix file-path
    f_path = os.path.join(data_dir, frames[idx]['file_path'])
    if '.' not in os.path.basename(f_path):
        f_path += '.png'  # so silly...
    if not os.path.exists(f_path):  # there are non-exist paths in fox...
        return None
    img = Image.open(f_path)
    if out_h is None:
        out_h = int(img.size[0] / downsample)
    if out_w is None:
        out_w = int(img.size[1] / downsample)
    # Now we should downsample to out_h, out_w and low-pass filter to resolution * 2.
    # We only do the low-pass filtering if resolution * 2 is lower-res than out_h, out_w
    if out_h != out_w:
        log.warning("360 non-square")
    img = img.resize((out_w, out_h), Image.LANCZOS)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]

    pose = torch.tensor(frames[idx]['transform_matrix'], dtype=torch.float32)

    return (img, pose)


def _parallel_loader_nerf_image_pose(args):
    torch.set_num_threads(1)
    return _load_nerf_image_pose(**args)


def _load_video_1cam(idx: int,
                     paths: List[str],
                     poses: torch.Tensor,
                     out_h: int,
                     out_w: int,
                     load_every: int = 1
                     ):  # -> Tuple[List[torch.Tensor], torch.Tensor, List[int]]:
    filters = [
        ("scale", f"w={out_w}:h={out_h}")
    ]
    all_frames = iio.imread(
        paths[idx], plugin='pyav', format='rgb24', constant_framerate=True, thread_count=2,
        filter_sequence=filters,)
    imgs, timestamps = [], []
    for frame_idx, frame in enumerate(all_frames):
        if frame_idx % load_every != 0:
            continue
        if frame_idx >= 300:  # Only look at the first 10 seconds
            break
        # Frame is np.ndarray in uint8 dtype (H, W, C)
        imgs.append(
            torch.from_numpy(frame)
        )
        timestamps.append(frame_idx)
    imgs = torch.stack(imgs, 0)
    med_img, _ = torch.median(imgs, dim=0)  # [h, w, 3]
    return (imgs,
            poses[idx].expand(len(timestamps), -1, -1),
            med_img,
            torch.tensor(timestamps, dtype=torch.int32))


def _parallel_loader_video(args):
    torch.set_num_threads(1)
    return _load_video_1cam(**args)


def parallel_load_images(tqdm_title,
                         dset_type: str,
                         num_images: int,
                         **kwargs) -> List[Any]:
    max_threads = 10
    if dset_type == 'llff':
        fn = _parallel_loader_llff_image
    elif dset_type == 'synthetic':
        fn = _parallel_loader_nerf_image_pose
    elif dset_type == 'phototourism':
        fn = _parallel_loader_phototourism_image
    elif dset_type == 'video':
        fn = _parallel_loader_video
        # giac: Can increase to e.g. 10 if loading 4x subsampled images. Otherwise OOM.
        max_threads = 8
    else:
        raise ValueError(dset_type)
    p = Pool(min(max_threads, num_images))

    iterator = p.imap(fn, [{"idx": i, **kwargs} for i in range(num_images)])
    outputs = []
    for _ in tqdm(range(num_images), desc=tqdm_title):
        out = next(iterator)
        if out is not None:
            outputs.append(out)
    return outputs

thirdparty/kplanes/datasets/data_loading.py
oppo-us-research/NeuRBF
from dataclasses import dataclass


@dataclass
class Intrinsics:
    width: int
    height: int
    focal_x: float
    focal_y: float
    center_x: float
    center_y: float

    def scale(self, factor: float):
        nw = round(self.width * factor)
        nh = round(self.height * factor)
        sw = nw / self.width
        sh = nh / self.height
        self.focal_x *= sw
        self.focal_y *= sh
        self.center_x *= sw
        self.center_y *= sh
        self.width = int(nw)
        self.height = int(nh)

    def __repr__(self):
        return (f"Intrinsics(width={self.width}, height={self.height}, "
                f"focal_x={self.focal_x}, focal_y={self.focal_y}, "
                f"center_x={self.center_x}, center_y={self.center_y})")

thirdparty/kplanes/datasets/intrinsics.py
oppo-us-research/NeuRBF
import glob
import os
import logging as log
from typing import Tuple, Optional, List

import numpy as np
import torch

from .data_loading import parallel_load_images
from .ray_utils import (
    center_poses, generate_spiral_path, create_meshgrid, stack_camera_dirs, get_rays
)
from .intrinsics import Intrinsics
from .base_dataset import BaseDataset


class LLFFDataset(BaseDataset):
    def __init__(self,
                 datadir,
                 split: str,
                 batch_size: Optional[int] = None,
                 downsample: int = 4,
                 hold_every: int = 8,
                 contraction: bool = False,
                 ndc: bool = True,
                 near_scaling: float = 0.9,
                 ndc_far: float = 1.0,
                 ):
        if (not contraction) and (not ndc):
            raise ValueError("LLFF dataset expects either contraction or NDC to be enabled.")
        self.downsample = downsample
        self.hold_every = hold_every
        self.near_scaling = near_scaling
        self.ndc_far = ndc_far

        if split == 'render':
            # For rendering path we load all poses, use them to generate spiral poses.
            # No gt images exist!
            assert ndc, "Unable to generate render poses without ndc: don't know near-far."
            image_paths, poses, near_fars, intrinsics = load_llff_poses(
                datadir, downsample=downsample, split='test', hold_every=1,
                near_scaling=self.near_scaling)
            render_poses = generate_spiral_path(
                poses.numpy(), near_fars.numpy(), n_frames=120, n_rots=2, zrate=0.5,
                dt=self.near_scaling)
            self.poses = torch.from_numpy(render_poses).float()
            imgs = None
        else:
            image_paths, self.poses, near_fars, intrinsics = load_llff_poses(
                datadir, downsample=downsample, split=split, hold_every=hold_every,
                near_scaling=self.near_scaling)
            imgs = load_llff_images(image_paths, intrinsics, split)
            imgs = (imgs * 255).to(torch.uint8)
            if split == 'train':
                imgs = imgs.view(-1, imgs.shape[-1])
            else:
                imgs = imgs.view(-1, intrinsics.height * intrinsics.width, imgs.shape[-1])
        num_images = len(self.poses)
        if contraction:
            bbox = torch.tensor([[-2., -2., -2.], [2., 2., 2.]])
            self.near_fars = near_fars
        else:
            bbox = torch.tensor([[-1.5, -1.67, -1.], [1.5, 1.67, 1.]])
            self.near_fars = torch.tensor([[0.0, self.ndc_far]]).repeat(num_images, 1)

        # These are used when contraction=True
        self.global_translation = torch.tensor([0, 0, 1.5])
        self.global_scale = torch.tensor([0.9, 0.9, 1])

        super().__init__(
            datadir=datadir,
            split=split,
            scene_bbox=bbox,
            batch_size=batch_size,
            imgs=imgs,
            rays_o=None,
            rays_d=None,
            intrinsics=intrinsics,
            is_ndc=ndc,
            is_contracted=contraction,
        )
        log.info(f"LLFFDataset. {contraction=} {ndc=}. Loaded {split} set from {datadir}. "
                 f"{num_images} poses of shape {self.img_h}x{self.img_w}. "
                 f"Images loaded: {imgs is not None}. Near-far[:3]: {self.near_fars[:3]}. "
                 f"Sampling without replacement={self.use_permutation}. {intrinsics}")

    def __getitem__(self, index):
        h = self.intrinsics.height
        w = self.intrinsics.width
        dev = "cpu"
        if self.split == 'train':
            index = self.get_rand_ids(index)
            image_id = torch.div(index, h * w, rounding_mode='floor')
            y = torch.remainder(index, h * w).div(w, rounding_mode='floor')
            x = torch.remainder(index, h * w).remainder(w)
            x = x + 0.5
            y = y + 0.5
        else:
            image_id = [index]
            x, y = create_meshgrid(height=h, width=w, dev=dev, add_half=True, flat=True)
        out = {"near_fars": self.near_fars[image_id, :].view(-1, 2)}
        if self.imgs is not None:
            out["imgs"] = self.imgs[index] / 255.0  # (num_rays, 3)   this converts to f32
        else:
            out["imgs"] = None

        c2w = self.poses[image_id]       # (num_rays, 3, 4)
        camera_dirs = stack_camera_dirs(x, y, self.intrinsics, True)  # [num_rays, 3]
        rays_o, rays_d = get_rays(camera_dirs, c2w, ndc=self.is_ndc, ndc_near=1.0,
                                  intrinsics=self.intrinsics, normalize_rd=True)  # h*w, 3
        out["rays_o"] = rays_o
        out["rays_d"] = rays_d
        out["bg_color"] = torch.tensor([[1.0, 1.0, 1.0]])
        return out


def _split_poses_bounds(poses_bounds: np.ndarray) -> Tuple[np.ndarray, np.ndarray, Intrinsics]:
    poses = poses_bounds[:, :15].reshape(-1, 3, 5)  # (N_images, 3, 5)
    near_fars = poses_bounds[:, -2:]  # (N_images, 2)
    H, W, focal = poses[0, :, -1]  # original intrinsics, same for all images
    intrinsics = Intrinsics(
        width=W, height=H, focal_x=focal, focal_y=focal, center_x=W / 2, center_y=H / 2)
    return poses[:, :, :4], near_fars, intrinsics


def load_llff_poses_helper(datadir: str, downsample: float, near_scaling: float) -> Tuple[np.ndarray, np.ndarray, Intrinsics]:
    poses_bounds = np.load(os.path.join(datadir, 'poses_bounds.npy'))  # (N_images, 17)
    poses, near_fars, intrinsics = _split_poses_bounds(poses_bounds)

    # Step 1: rescale focal length according to training resolution
    intrinsics.scale(1 / downsample)

    # Step 2: correct poses
    # Original poses has rotation in form "down right back", change to "right up back"
    # See https://github.com/bmild/nerf/issues/34
    poses = np.concatenate([poses[..., 1:2], -poses[..., :1], poses[..., 2:4]], -1)
    # (N_images, 3, 4) exclude H, W, focal
    poses, pose_avg = center_poses(poses)

    # Step 3: correct scale so that the nearest depth is at a little more than 1.0
    # See https://github.com/bmild/nerf/issues/34
    near_original = np.min(near_fars)
    scale_factor = near_original * near_scaling  # 0.75 is the default parameter
    # the nearest depth is at 1/0.75=1.33
    near_fars /= scale_factor
    poses[..., 3] /= scale_factor

    return poses, near_fars, intrinsics


def load_llff_poses(datadir: str,
                    downsample: float,
                    split: str,
                    hold_every: int,
                    near_scaling: float = 0.75) -> Tuple[
                        List[str], torch.Tensor, torch.Tensor, Intrinsics]:
    int_dsample = int(downsample)
    if int_dsample != downsample or int_dsample not in {4, 8}:
        raise ValueError(f"Cannot downsample LLFF dataset by {downsample}.")

    poses, near_fars, intrinsics = load_llff_poses_helper(datadir, downsample, near_scaling)

    image_paths = sorted(glob.glob(os.path.join(datadir, f'images_{int_dsample}/*')))
    assert poses.shape[0] == len(image_paths), \
        'Mismatch between number of images and number of poses! Please rerun COLMAP!'

    # Take training or test split
    i_test = np.arange(0, poses.shape[0], hold_every)
    img_list = i_test if split != 'train' else list(set(np.arange(len(poses))) - set(i_test))
    img_list = np.asarray(img_list)

    image_paths = [image_paths[i] for i in img_list]
    poses = torch.from_numpy(poses[img_list]).float()
    near_fars = torch.from_numpy(near_fars[img_list]).float()

    return image_paths, poses, near_fars, intrinsics


def load_llff_images(image_paths: List[str], intrinsics: Intrinsics, split: str):
    all_rgbs: List[torch.Tensor] = parallel_load_images(
        tqdm_title=f'Loading {split} data',
        dset_type='llff',
        data_dir='',
        num_images=len(image_paths),
        paths=image_paths,
        out_h=intrinsics.height,
        out_w=intrinsics.width,
    )
    return torch.stack(all_rgbs, 0)

thirdparty/kplanes/datasets/llff_dataset.py
oppo-us-research/NeuRBF
import glob
import math
import logging as log
import os
from enum import Enum
from pathlib import Path
from typing import Optional, List, Tuple

import numpy as np
import pandas as pd
import torch

from thirdparty.kplanes.datasets.base_dataset import BaseDataset
from thirdparty.kplanes.datasets.intrinsics import Intrinsics
from thirdparty.kplanes.datasets.ray_utils import average_poses
from thirdparty.kplanes.ops.image.io import read_png


class PhototourismScenes(Enum):
    TREVI = "trevi"
    BRANDENBURG = "brandenburg"
    SACRE = "sacre"

    @staticmethod
    def get_scene_from_datadir(datadir: str) -> 'PhototourismScenes':
        if "sacre" in datadir:
            return PhototourismScenes.SACRE
        if "trevi" in datadir:
            return PhototourismScenes.TREVI
        if "brandenburg" in datadir:
            return PhototourismScenes.BRANDENBURG
        raise NotImplementedError(datadir)


class PhotoTourismDataset(BaseDataset):
    def __init__(self,
                 datadir: str,
                 split: str,
                 batch_size: Optional[int] = None,
                 contraction: bool = False,
                 ndc: bool = False,
                 scene_bbox: Optional[List] = None,
                 global_translation: List[float] = None,
                 global_scale: List[float] = None,
                 downsample: float = 1.0):
        if ndc:
            raise NotImplementedError("PhotoTourism only handles contraction and standard.")
        if downsample != 1.0:
            raise NotImplementedError("PhotoTourism does not handle image downsampling.")
        if (global_scale is None or global_translation is None) and contraction:
            raise ValueError("scale and translation must be specified when contraction is used.")
        if not os.path.isdir(datadir):
            raise ValueError(f"Directory {datadir} does not exist.")

        if split == 'train' or split == 'test':
            pt_data_file = os.path.join(datadir, f"cache_{split}.pt")
            if not os.path.isfile(pt_data_file):
                # Populate cache
                cache_data(datadir=datadir, split=split, out_fname=os.path.basename(pt_data_file))
            pt_data = torch.load(pt_data_file)

            intrinsics = [
                Intrinsics(width=img.shape[1], height=img.shape[0],
                           center_x=img.shape[1] / 2, center_y=img.shape[0] / 2,
                           # focals are unused, we reuse intrinsics from Matt's files.
                           focal_y=0, focal_x=0)
                for img in pt_data["images"]
            ]
            if split == 'train':
                near_fars = torch.cat([
                    pt_data["bounds"][i].expand(intrinsics[i].width * intrinsics[i].height, 2)
                    for i in range(len(intrinsics))
                ], dim=0)
                camera_ids = torch.cat([
                    pt_data["camera_ids"][i].expand(intrinsics[i].width * intrinsics[i].height, 1)
                    for i in range(len(intrinsics))
                ])
                images = torch.cat([img.view(-1, 3) for img in pt_data["images"]], 0)
                rays_o = torch.cat([ro.view(-1, 3) for ro in pt_data["rays_o"]], 0)
                rays_d = torch.cat([rd.view(-1, 3) for rd in pt_data["rays_d"]], 0)
            elif split == 'test':
                images = pt_data["images"]
                rays_o = pt_data["rays_o"]
                rays_d = pt_data["rays_d"]
                near_fars = pt_data["bounds"]
                camera_ids = pt_data["camera_ids"]
        elif split == 'render':
            n_frames, frame_size = 150, 800
            rays_o, rays_d, camera_ids, near_fars = pt_render_poses(
                datadir, n_frames=n_frames, frame_h=frame_size, frame_w=frame_size)
            images = None
            intrinsics = [
                Intrinsics(width=frame_size, height=frame_size, focal_x=0, focal_y=0,
                           center_x=frame_size / 2, center_y=frame_size / 2)
                for _ in range(n_frames)
            ]
        else:
            raise NotImplementedError(split)

        # ugly hack: num_images needs to be the number of training images. This is needed to
        # initialize the appearance embedding tensor to the correct size even if we're just
        # reloading a previous model.
        self.num_images = self.get_num_train_images(datadir)
        self.camera_ids = camera_ids  # noqa
        self.near_fars = near_fars  # noqa

        self.global_scale, self.global_translation = None, None
        if contraction:
            self.global_translation = torch.as_tensor(global_translation).float()
            self.global_scale = torch.as_tensor(global_scale).float()

        if scene_bbox is None:
            raise ValueError("Must specify scene_bbox")
        scene_bbox = torch.tensor(scene_bbox)

        super().__init__(
            datadir=datadir,
            split=split,
            batch_size=batch_size,
            is_ndc=ndc,
            is_contracted=contraction,
            scene_bbox=scene_bbox,
            rays_o=rays_o,  # noqa
            rays_d=rays_d,  # noqa
            intrinsics=intrinsics,
            imgs=images,  # noqa
        )
        log.info(f"PhotoTourismDataset contracted={self.is_contracted}, ndc={self.is_ndc}. "
                 f"Loaded {self.split} set from {self.datadir}: "
                 f"{len(intrinsics)} images of sizes between {min(self.img_h)}x{min(self.img_w)} "
                 f"and {max(self.img_h)}x{max(self.img_w)}. "
                 f"Images loaded: {self.imgs is not None}.")
        if self.is_contracted:
            log.info(f"Contraction parameters: global_translation={self.global_translation}, "
                     f"global_scale={self.global_scale}")
        else:
            log.info(f"Bounding box: {self.scene_bbox}")

    def __getitem__(self, index):
        out, index = super().__getitem__(index, return_idxs=True)
        out["bg_color"] = torch.ones((1, 3), dtype=torch.float32)
        out["timestamps"] = self.camera_ids[index]
        out["near_fars"] = self.near_fars[index]
        if self.imgs is not None:
            out["imgs"] = out["imgs"] / 255.0  # this converts to f32

        if self.split != 'train':  # gen left-image and reshape correctly
            intrinsics = self.intrinsics[index]
            img_h, img_w = intrinsics.height, intrinsics.width
            mid = img_w // 2
            if self.imgs is not None:
                out["imgs_left"] = out["imgs"][:, :mid, :].reshape(-1, 3)
                out["rays_o_left"] = out["rays_o"].view(img_h, img_w, 3)[:, :mid, :].reshape(-1, 3)
                out["rays_d_left"] = out["rays_d"].view(img_h, img_w, 3)[:, :mid, :].reshape(-1, 3)
                out["imgs"] = out["imgs"].view(-1, 3)
            out["rays_o"] = out["rays_o"].reshape(-1, 3)
            out["rays_d"] = out["rays_d"].reshape(-1, 3)
            out["timestamps"] = out["timestamps"].expand(out["rays_o"].shape[0], 1)
            out["near_fars"] = out["near_fars"].expand(out["rays_o"].shape[0], 2)
        return out

    @staticmethod
    def get_num_train_images(datadir) -> int:
        scene = PhototourismScenes.get_scene_from_datadir(datadir)
        if scene == PhototourismScenes.TREVI:
            return 1689
        elif scene == PhototourismScenes.BRANDENBURG:
            return 763
        elif scene == PhototourismScenes.SACRE:
            return 830
        else:
            raise ValueError(scene)


def get_rays_tourism(H, W, kinv, pose):
    """
    phototourism camera intrinsics are defined by H, W and kinv.
    Args:
        H: image height
        W: image width
        kinv (3, 3): inverse of camera intrinsic
        pose (4, 4): camera extrinsic
    Returns:
        rays_o (H, W, 3): ray origins
        rays_d (H, W, 3): ray directions
    """
    yy, xx = torch.meshgrid(torch.arange(0., H, device=kinv.device),
                            torch.arange(0., W, device=kinv.device),
                            indexing='ij')
    pixco = torch.stack([xx, yy, torch.ones_like(xx)], dim=-1)

    directions = torch.matmul(pixco, kinv.T)  # (H, W, 3)

    rays_d = torch.matmul(directions, pose[:3, :3].T)
    rays_d = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)  # (H, W, 3)

    rays_o = pose[:3, -1].expand_as(rays_d)  # (H, W, 3)

    return rays_o, rays_d


def load_camera_metadata(datadir: str, idx) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    try:
        poses = np.load(str(Path(datadir) / "c2w_mats.npy"))[idx]
        kinvs = np.load(str(Path(datadir) / "kinv_mats.npy"))[idx]
        bounds = np.load(str(Path(datadir) / "bds.npy"))[idx]
        res = np.load(str(Path(datadir) / "res_mats.npy"))[idx]
    except FileNotFoundError as e:
        error_msg = (
            f"One of the needed Phototourism files does not exist ({e.filename}). "
            f"They can be downloaded from "
            f"https://drive.google.com/drive/folders/1SVHKRQXiRb98q4KHVEbj8eoWxjNS2QLW"
        )
        log.error(error_msg)
        raise e
    return poses, kinvs, bounds, res


def get_ids_for_split(datadir, split):
    # read all files in the tsv first (split to train and test later)
    tsv = glob.glob(os.path.join(datadir, '*.tsv'))[0]
    files = pd.read_csv(tsv, sep='\t')
    files = files[~files['id'].isnull()]  # remove data without id
    files.reset_index(inplace=True, drop=True)
    files = files[files["split"] == split]

    imagepaths = sorted((Path(datadir) / "dense" / "images").glob("*.jpg"))
    imkey = np.array([os.path.basename(im) for im in imagepaths])
    idx = np.in1d(imkey, files["filename"])
    return idx, imagepaths


def scale_cam_metadata(poses: np.ndarray, kinvs: np.ndarray, bounds: np.ndarray, scale: float = 0.05):
    poses[:, :3, 3:4] *= scale
    bounds = bounds * scale * np.array([0.9, 1.2])

    return poses, kinvs, bounds


def cache_data(datadir: str, split: str, out_fname: str):
    log.info(f"Preparing cached rays for dataset at {datadir} - {split=}.")
    scale = 0.05

    idx, imagepaths = get_ids_for_split(datadir, split)

    imagepaths = np.array(imagepaths)[idx]
    poses, kinvs, bounds, res = load_camera_metadata(datadir, idx)
    poses, kinvs, bounds = scale_cam_metadata(poses, kinvs, bounds, scale=scale)
    img_w = res[:, 0]
    img_h = res[:, 1]
    size = int(np.sum(img_w * img_h))
    log.info(f"Loading dataset from {datadir}. Num images={len(imagepaths)}. Total rays={size}.")

    all_images, all_rays_o, all_rays_d, all_bounds, all_camera_ids = [], [], [], [], []
    for idx, impath in enumerate(imagepaths):
        image = read_png(impath)

        pose = torch.from_numpy(poses[idx]).float()
        kinv = torch.from_numpy(kinvs[idx]).float()
        bound = torch.from_numpy(bounds[idx]).float()

        rays_o, rays_d = get_rays_tourism(image.shape[0], image.shape[1], kinv, pose)

        camera_id = torch.tensor(idx)

        all_images.append(image.mul(255).to(torch.uint8))
        all_rays_o.append(rays_o)
        all_rays_d.append(rays_d)
        all_bounds.append(bound)
        all_camera_ids.append(camera_id)

    torch.save({
        "images": all_images,
        "rays_o": all_rays_o,
        "rays_d": all_rays_d,
        "bounds": all_bounds,
        "camera_ids": all_camera_ids,
    }, os.path.join(datadir, out_fname))


def pt_spiral_path(
        scene: PhototourismScenes,
        poses: torch.Tensor,
        n_frames=120,
        n_rots: float = 1.0,
        zrate=.5) -> torch.Tensor:
    if poses.shape[1] > 3:
        poses = poses[:, :3, :]
    c2w = torch.from_numpy(average_poses(poses.numpy()))  # [3, 4]

    # Generate poses for spiral path.
    render_poses = []
    for theta in np.linspace(0., 2. * np.pi * n_rots, n_frames, endpoint=False):
        rotation = c2w[:3, :3]
        # the additive translation vectors have 3 components (x, y, z axes)
        # each with an additive part - which defines a global shift of all poses
        # and a multiplicative part which changes the amplitude of movement
        # of the poses.
        if scene == PhototourismScenes.SACRE:
            translation = c2w[:, 3:4] + torch.tensor([[
                0.01 + 0.03 * np.cos(theta),
                -0.007 * np.sin(theta),
                0.06 + 0.03 * np.sin(theta * zrate)
            ]]).T
        elif scene == PhototourismScenes.BRANDENBURG:
            translation = c2w[:, 3:4] + torch.tensor([[
                0.08 * np.cos(theta),
                -0.07 - 0.01 * np.sin(theta),
                -0.0 + 0.1 * np.sin(theta * zrate)
            ]]).T
        elif scene == PhototourismScenes.TREVI:
            translation = c2w[:, 3:4] + torch.tensor([[
                -0.05 + 0.2 * np.cos(theta),
                -0.07 - 0.07 * np.sin(theta),
                0.02 + 0.05 * np.sin(theta * zrate)
            ]]).T
        else:
            raise NotImplementedError(scene)
        pose = torch.cat([rotation, translation], dim=1)
        render_poses.append(pose)
    return torch.stack(render_poses, dim=0)


def pt_render_poses(datadir: str, n_frames: int, frame_h: int = 800, frame_w: int = 800):
    scene = PhototourismScenes.get_scene_from_datadir(datadir)
    idx, _ = get_ids_for_split(datadir, split='train')
    train_poses, kinvs, bounds, res = load_camera_metadata(datadir, idx)
    train_poses, kinvs, bounds = scale_cam_metadata(train_poses, kinvs, bounds, scale=0.05)

    # build camera intrinsic from appropriate focal distance and cx, cy. Good for trevi
    k = np.array([[780.0, 0, -frame_w / 2], [0, -780, -frame_h / 2], [0, 0, -1]])
    kinv = torch.from_numpy(np.linalg.inv(k)).to(torch.float32)

    bounds = torch.from_numpy(bounds).float()
    train_poses = torch.from_numpy(train_poses).float()

    r_poses = pt_spiral_path(scene, train_poses, n_frames=n_frames, zrate=1, n_rots=1)

    all_rays_o, all_rays_d, camera_ids, near_fars = [], [], [], []
    for pose_id, pose in enumerate(r_poses):
        pose = pose.float()
        rays_o, rays_d = get_rays_tourism(frame_h, frame_w, kinv, pose)
        rays_o = rays_o.view(-1, 3)
        rays_d = rays_d.view(-1, 3)
        all_rays_o.append(rays_o)
        all_rays_d.append(rays_d)

        # Find the closest cam
        closest_cam_idx = torch.linalg.norm(
            train_poses[:, :3, :].view(train_poses.shape[0], -1) - pose.view(-1), dim=1
        ).argmin()

        # For brandenburg and trevi
        if scene == PhototourismScenes.BRANDENBURG or scene == PhototourismScenes.TREVI:
            near_fars.append((
                bounds[closest_cam_idx] + torch.tensor([0.01, 0.0])
            ))
        elif scene == PhototourismScenes.SACRE:
            near_fars.append((
                bounds[closest_cam_idx] + torch.tensor([0.05, 0.0])
            ))
    # camera-IDs. They are floats interpolating between 2 appearance embeddings.
    x = torch.linspace(-1, 1, len(r_poses))
    s = 0.3
    camera_ids = 1/(s * math.sqrt(2 * torch.pi)) * torch.exp(- (x)**2 / (2 * s**2))
    camera_ids = (camera_ids - camera_ids.min()) / camera_ids.max()
    all_rays_o = torch.stack(all_rays_o, dim=0)
    all_rays_d = torch.stack(all_rays_d, dim=0)
    near_fars = torch.stack(near_fars, dim=0)
    return all_rays_o, all_rays_d, camera_ids, near_fars

thirdparty/kplanes/datasets/phototourism_dataset.py
oppo-us-research/NeuRBF
from typing import Tuple, Optional

import numpy as np
import torch

from .intrinsics import Intrinsics

__all__ = (
    "create_meshgrid",
    "get_ray_directions",
    "stack_camera_dirs",
    "get_rays",
    "ndc_rays_blender",
    "center_poses",
    "generate_spiral_path",
    "generate_hemispherical_orbit",
    "generate_spherical_poses",
    "normalize",
    "average_poses",
    "viewmatrix",
)


def create_meshgrid(height: int,
                    width: int,
                    dev: str = 'cpu',
                    add_half: bool = True,
                    flat: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
    xs = torch.arange(width, dtype=torch.float32, device=dev)
    ys = torch.arange(height, dtype=torch.float32, device=dev)
    if add_half:
        xs += 0.5
        ys += 0.5
    # generate grid by stacking coordinates
    yy, xx = torch.meshgrid([ys, xs], indexing="ij")  # both HxW
    if flat:
        return xx.flatten(), yy.flatten()
    return xx, yy


def stack_camera_dirs(x: torch.Tensor, y: torch.Tensor, intrinsics: Intrinsics, opengl_camera: bool):
    # the direction here is without +0.5 pixel centering as calibration is not so accurate
    # see https://github.com/bmild/nerf/issues/24
    x = x.float()
    y = y.float()
    return torch.stack([
        (x - intrinsics.center_x) / intrinsics.focal_x,
        (y - intrinsics.center_y) / intrinsics.focal_y
        * (-1.0 if opengl_camera else 1.0),
        torch.full_like(x, fill_value=-1.0 if opengl_camera else 1.0)
    ], -1)  # (H, W, 3)


def get_ray_directions(intrinsics: Intrinsics, opengl_camera: bool, add_half: bool = True) -> torch.Tensor:
    """
    Get ray directions for all pixels in camera coordinate.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems
    Inputs:
        height, width, focal: image height, width and focal length
    Outputs:
        directions: (height, width, 3), the direction of the rays in camera coordinate
    """
    xx, yy = create_meshgrid(intrinsics.height, intrinsics.width, add_half=add_half)

    return stack_camera_dirs(xx, yy, intrinsics, opengl_camera)


def get_rays(directions: torch.Tensor,
             c2w: torch.Tensor,
             ndc: bool,
             ndc_near: float = 1.0,
             intrinsics: Optional[Intrinsics] = None,
             normalize_rd: bool = True):
    """Get ray origin and normalized directions in world coordinate for all pixels in one image.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Args:
        directions:
        c2w:
        ndc:
        ndc_near:
        intrinsics:
        normalize_rd:

    Returns:

    """
    directions = directions.view(-1, 3)  # [n_rays, 3]
    if len(c2w.shape) == 2:
        c2w = c2w[None, ...]
    rd = (directions[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
    ro = torch.broadcast_to(c2w[:, :3, 3], directions.shape)
    if ndc:
        assert intrinsics is not None, "intrinsics must not be None when NDC active."
        ro, rd = ndc_rays_blender(
            intrinsics=intrinsics, near=ndc_near, rays_o=ro, rays_d=rd)
    if normalize_rd:
        rd /= torch.linalg.norm(rd, dim=-1, keepdim=True)
    return ro, rd


def ndc_rays_blender(intrinsics: Intrinsics, near: float, rays_o: torch.Tensor,
                     rays_d: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    # Shift ray origins to near plane
    t = -(near + rays_o[..., 2]) / rays_d[..., 2]
    rays_o = rays_o + t[..., None] * rays_d

    # Projection
    ndc_coef_x = - (2 * intrinsics.focal_x) / intrinsics.width
    ndc_coef_y = - (2 * intrinsics.focal_y) / intrinsics.height
    o0 = ndc_coef_x * rays_o[..., 0] / rays_o[..., 2]
    o1 = ndc_coef_y * rays_o[..., 1] / rays_o[..., 2]
    o2 = 1. + 2. * near / rays_o[..., 2]

    d0 = ndc_coef_x * (rays_d[..., 0] / rays_d[..., 2] - rays_o[..., 0] / rays_o[..., 2])
    d1 = ndc_coef_y * (rays_d[..., 1] / rays_d[..., 2] - rays_o[..., 1] / rays_o[..., 2])
    d2 = -2. * near / rays_o[..., 2]

    rays_o = torch.stack([o0, o1, o2], -1)
    rays_d = torch.stack([d0, d1, d2], -1)

    return rays_o, rays_d


def normalize(v: np.ndarray) -> np.ndarray:
    """Normalize a vector."""
    return v / np.linalg.norm(v)


def average_poses(poses: np.ndarray) -> np.ndarray:
    """
    Calculate the average pose, which is then used to center all poses
    using @center_poses. Its computation is as follows:
    1. Compute the center: the average of pose centers.
    2. Compute the z axis: the normalized average z axis.
    3. Compute axis y': the average y axis.
    4. Compute x' = y' cross product z, then normalize it as the x axis.
    5. Compute the y axis: z cross product x.

    Note that at step 3, we cannot directly use y' as y axis since it's
    not necessarily orthogonal to z axis. We need to pass from x to y.
    Inputs:
        poses: (N_images, 3, 4)
    Outputs:
        pose_avg: (3, 4) the average pose
    """
    # 1. Compute the center
    center = poses[..., 3].mean(0)  # (3)
    # 2. Compute the z axis
    z = normalize(poses[..., 2].mean(0))  # (3)
    # 3. Compute axis y' (no need to normalize as it's not the final output)
    y_ = poses[..., 1].mean(0)  # (3)
    # 4. Compute the x axis
    x = normalize(np.cross(z, y_))  # (3)
    # 5. Compute the y axis (as z and x are normalized, y is already of norm 1)
    y = np.cross(x, z)  # (3)

    pose_avg = np.stack([x, y, z, center], 1)  # (3, 4)

    return pose_avg


def center_poses(poses):
    """
    Center the poses so that we can use NDC.
    See https://github.com/bmild/nerf/issues/34
    Inputs:
        poses: (N_images, 3, 4)
    Outputs:
        poses_centered: (N_images, 3, 4) the centered poses
        pose_avg: (3, 4) the average pose
    """
    pose_avg = average_poses(poses)  # (3, 4)
    pose_avg_homo = np.eye(4)
    pose_avg_homo[:3] = pose_avg  # convert to homogeneous coordinate for faster computation
    pose_avg_homo = pose_avg_homo
    # by simply adding 0, 0, 0, 1 as the last row
    last_row = np.tile(np.array([0, 0, 0, 1]), (len(poses), 1, 1))  # (N_images, 1, 4)
    poses_homo = \
        np.concatenate([poses, last_row], 1)  # (N_images, 4, 4) homogeneous coordinate

    poses_centered = np.linalg.inv(pose_avg_homo) @ poses_homo  # (N_images, 4, 4)
    poses_centered = poses_centered[:, :3]  # (N_images, 3, 4)

    return poses_centered, pose_avg_homo


def viewmatrix(z, up, pos):
    vec2 = normalize(z)
    vec1_avg = up
    vec0 = normalize(np.cross(vec1_avg, vec2))
    vec1 = normalize(np.cross(vec2, vec0))
    return np.stack([-vec0, vec1, vec2, pos], axis=1)


def generate_spiral_path(poses: np.ndarray,
                         near_fars: np.ndarray,
                         n_frames=120,
                         n_rots=2,
                         zrate=.5,
                         dt=0.75,
                         percentile=70) -> np.ndarray:
    """Calculates a forward facing spiral path for rendering.

    From https://github.com/google-research/google-research/blob/342bfc150ef1155c5254c1e6bd0c912893273e8d/regnerf/internal/datasets.py
    and https://github.com/apchenstu/TensoRF/blob/main/dataLoader/llff.py

    :param poses: [N, 3, 4]
    :param near_fars:
    :param n_frames:
    :param n_rots:
    :param zrate:
    :param dt:
    :return:
    """
    # center pose
    c2w = average_poses(poses)  # [3, 4]

    # Get average pose
    up = normalize(poses[:, :3, 1].sum(0))

    # Find a reasonable "focus depth" for this dataset as a weighted average
    # of near and far bounds in disparity space.
    close_depth, inf_depth = np.min(near_fars) * 1.0, np.max(near_fars) * 5.0
    focal = 1.0 / (((1.0 - dt) / close_depth + dt / inf_depth))

    # Get radii for spiral path using 90th percentile of camera positions.
    positions = poses[:, :3, 3]
    radii = np.percentile(np.abs(positions), percentile, 0)
    radii = np.concatenate([radii, [1.]])

    # Generate poses for spiral path.
    render_poses = []
    for theta in np.linspace(0., 2. * np.pi * n_rots, n_frames, endpoint=False):
        t = radii * [np.cos(theta), -np.sin(theta), -np.sin(theta * zrate), 1.]
        position = c2w @ t
        lookat = c2w @ np.array([0, 0, -focal, 1.0])
        z_axis = normalize(position - lookat)
        render_poses.append(viewmatrix(z_axis, up, position))
    return np.stack(render_poses, axis=0)


def generate_hemispherical_orbit(poses: torch.Tensor, n_frames=120):
    """Calculates a render path which orbits around the z-axis.
    Based on https://github.com/google-research/google-research/blob/342bfc150ef1155c5254c1e6bd0c912893273e8d/regnerf/internal/datasets.py
    """
    origins = poses[:, :3, 3]
    radius = torch.sqrt(torch.mean(torch.sum(origins ** 2, dim=-1)))

    # Assume that z-axis points up towards approximate camera hemisphere
    sin_phi = torch.mean(origins[:, 2], dim=0) / radius
    cos_phi = torch.sqrt(1 - sin_phi ** 2)
    render_poses = []

    up = torch.tensor([0., 0., 1.])
    for theta in np.linspace(0., 2. * np.pi, n_frames, endpoint=False):
        camorigin = radius * torch.tensor(
            [cos_phi * np.cos(theta), cos_phi * np.sin(theta), sin_phi])
        render_poses.append(torch.from_numpy(viewmatrix(camorigin, up, camorigin)))

    render_poses = torch.stack(render_poses, dim=0)
    return render_poses


trans_t = lambda t: torch.Tensor([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, t],
    [0, 0, 0, 1]]).float()

rot_phi = lambda phi: torch.Tensor([
    [1, 0, 0, 0],
    [0, np.cos(phi), -np.sin(phi), 0],
    [0, np.sin(phi), np.cos(phi), 0],
    [0, 0, 0, 1]]).float()

rot_theta = lambda th: torch.Tensor([
    [np.cos(th), 0, -np.sin(th), 0],
    [0, 1, 0, 0],
    [np.sin(th), 0, np.cos(th), 0],
    [0, 0, 0, 1]]).float()


def generate_spherical_poses(theta, phi, radius):
    c2w = trans_t(radius)
    c2w = rot_phi(phi / 180. * np.pi) @ c2w
    c2w = rot_theta(theta / 180. * np.pi) @ c2w
    c2w = torch.Tensor(np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])) @ c2w
    return c2w

thirdparty/kplanes/datasets/ray_utils.py
oppo-us-research/NeuRBF
import json
import logging as log
import os
from typing import Tuple, Optional, Any

import numpy as np
import torch

from .data_loading import parallel_load_images
from .ray_utils import get_ray_directions, generate_hemispherical_orbit, get_rays
from .intrinsics import Intrinsics
from .base_dataset import BaseDataset


class SyntheticNerfDataset(BaseDataset):
    def __init__(self,
                 datadir,
                 split: str,
                 batch_size: Optional[int] = None,
                 downsample: float = 1.0,
                 max_frames: Optional[int] = None):
        self.downsample = downsample
        self.max_frames = max_frames
        self.near_far = [2.0, 6.0]

        if split == 'render':
            frames, transform = load_360_frames(datadir, 'test', self.max_frames)
            imgs, poses = load_360_images(frames, datadir, 'test', self.downsample)
            render_poses = generate_hemispherical_orbit(poses, n_frames=120)
            self.poses = render_poses
            intrinsics = load_360_intrinsics(
                transform, img_h=imgs[0].shape[0], img_w=imgs[0].shape[1],
                downsample=self.downsample)
            imgs = None
        else:
            frames, transform = load_360_frames(datadir, split, self.max_frames)
            imgs, poses = load_360_images(frames, datadir, split, self.downsample)
            intrinsics = load_360_intrinsics(
                transform, img_h=imgs[0].shape[0], img_w=imgs[0].shape[1],
                downsample=self.downsample)
        rays_o, rays_d, imgs = create_360_rays(
            imgs, poses, merge_all=split == 'train', intrinsics=intrinsics)
        super().__init__(
            datadir=datadir,
            split=split,
            scene_bbox=get_360_bbox(datadir, is_contracted=False),
            is_ndc=False,
            is_contracted=False,
            batch_size=batch_size,
            imgs=imgs,
            rays_o=rays_o,
            rays_d=rays_d,
            intrinsics=intrinsics,
        )
        log.info(f"SyntheticNerfDataset. Loaded {split} set from {datadir}."
                 f"{len(poses)} images of shape {self.img_h}x{self.img_w}. "
                 f"Images loaded: {imgs is not None}. "
                 f"Sampling without replacement={self.use_permutation}. {intrinsics}")

    def __getitem__(self, index):
        out = super().__getitem__(index)
        pixels = out["imgs"]
        if self.split == 'train':
            bg_color = torch.rand((1, 3), dtype=pixels.dtype, device=pixels.device)
        else:
            if pixels is None:
                bg_color = torch.ones((1, 3), dtype=torch.float32, device='cuda:0')
            else:
                bg_color = torch.ones((1, 3), dtype=pixels.dtype, device=pixels.device)
        # Alpha compositing
        if pixels is not None:
            pixels = pixels[:, :3] * pixels[:, 3:] + bg_color * (1.0 - pixels[:, 3:])
        out["imgs"] = pixels
        out["bg_color"] = bg_color
        out["near_fars"] = torch.tensor([[2.0, 6.0]])
        return out


def get_360_bbox(datadir, is_contracted=False):
    if is_contracted:
        radius = 2
    elif "ship" in datadir:
        radius = 1.5
    else:
        radius = 1.3
    return torch.tensor([[-radius, -radius, -radius], [radius, radius, radius]])


def create_360_rays(
              imgs: Optional[torch.Tensor],
              poses: torch.Tensor,
              merge_all: bool,
              intrinsics: Intrinsics) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    directions = get_ray_directions(intrinsics, opengl_camera=True)  # [H, W, 3]
    num_frames = poses.shape[0]

    all_rays_o, all_rays_d = [], []
    for i in range(num_frames):
        rays_o, rays_d = get_rays(directions, poses[i], ndc=False, normalize_rd=True)  # h*w, 3
        all_rays_o.append(rays_o)
        all_rays_d.append(rays_d)

    all_rays_o = torch.cat(all_rays_o, 0).to(dtype=torch.float32)  # [n_frames * h * w, 3]
    all_rays_d = torch.cat(all_rays_d, 0).to(dtype=torch.float32)  # [n_frames * h * w, 3]
    if imgs is not None:
        imgs = imgs.view(-1, imgs.shape[-1]).to(dtype=torch.float32)   # [N*H*W, 3/4]
    if not merge_all:
        num_pixels = intrinsics.height * intrinsics.width
        if imgs is not None:
            imgs = imgs.view(num_frames, num_pixels, -1)  # [N, H*W, 3/4]
        all_rays_o = all_rays_o.view(num_frames, num_pixels, 3)  # [N, H*W, 3]
        all_rays_d = all_rays_d.view(num_frames, num_pixels, 3)  # [N, H*W, 3]
    return all_rays_o, all_rays_d, imgs


def load_360_frames(datadir, split, max_frames: int) -> Tuple[Any, Any]:
    with open(os.path.join(datadir, f"transforms_{split}.json"), 'r') as f:
        meta = json.load(f)
        frames = meta['frames']

        # Subsample frames
        tot_frames = len(frames)
        num_frames = min(tot_frames, max_frames or tot_frames)
        if split == 'train' or split == 'test':
            subsample = int(round(tot_frames / num_frames))
            frame_ids = np.arange(tot_frames)[::subsample]
            if subsample > 1:
                log.info(f"Subsampling {split} set to 1 every {subsample} images.")
        else:
            frame_ids = np.arange(num_frames)
        frames = np.take(frames, frame_ids).tolist()
    return frames, meta


def load_360_images(frames, datadir, split, downsample) -> Tuple[torch.Tensor, torch.Tensor]:
    img_poses = parallel_load_images(
        dset_type="synthetic",
        tqdm_title=f'Loading {split} data',
        num_images=len(frames),
        frames=frames,
        data_dir=datadir,
        out_h=None,
        out_w=None,
        downsample=downsample,
    )
    imgs, poses = zip(*img_poses)
    imgs = torch.stack(imgs, 0)  # [N, H, W, 3/4]
    poses = torch.stack(poses, 0)  # [N, ????]
    return imgs, poses


def load_360_intrinsics(transform, img_h, img_w, downsample) -> Intrinsics:
    height = img_h
    width = img_w
    # load intrinsics
    if 'fl_x' in transform or 'fl_y' in transform:
        fl_x = (transform['fl_x'] if 'fl_x' in transform else transform['fl_y']) / downsample
        fl_y = (transform['fl_y'] if 'fl_y' in transform else transform['fl_x']) / downsample
    elif 'camera_angle_x' in transform or 'camera_angle_y' in transform:
        # blender, assert in radians. already downscaled since we use H/W
        fl_x = width / (2 * np.tan(transform['camera_angle_x'] / 2)) if 'camera_angle_x' in transform else None
        fl_y = height / (2 * np.tan(transform['camera_angle_y'] / 2)) if 'camera_angle_y' in transform else None
        if fl_x is None:
            fl_x = fl_y
        if fl_y is None:
            fl_y = fl_x
    else:
        raise RuntimeError('Failed to load focal length, please check the transforms.json!')

    cx = (transform['cx'] / downsample) if 'cx' in transform else (width / 2)
    cy = (transform['cy'] / downsample) if 'cy' in transform else (height / 2)
    return Intrinsics(height=height, width=width, focal_x=fl_x, focal_y=fl_y, center_x=cx, center_y=cy)

thirdparty/kplanes/datasets/synthetic_nerf_dataset.py
oppo-us-research/NeuRBF
import glob
import json
import logging as log
import math
import os
import time
from collections import defaultdict
from typing import Optional, List, Tuple, Any, Dict

import numpy as np
import torch

from .base_dataset import BaseDataset
from .data_loading import parallel_load_images
from .intrinsics import Intrinsics
from .llff_dataset import load_llff_poses_helper
from .ray_utils import (
    generate_spherical_poses, create_meshgrid, stack_camera_dirs, get_rays, generate_spiral_path
)
from .synthetic_nerf_dataset import (
    load_360_images, load_360_intrinsics,
)


class Video360Dataset(BaseDataset):
    len_time: int
    max_cameras: Optional[int]
    max_tsteps: Optional[int]
    timestamps: Optional[torch.Tensor]

    def __init__(self,
                 datadir: str,
                 split: str,
                 batch_size: Optional[int] = None,
                 downsample: float = 1.0,
                 keyframes: bool = False,
                 max_cameras: Optional[int] = None,
                 max_tsteps: Optional[int] = None,
                 isg: bool = False,
                 contraction: bool = False,
                 ndc: bool = False,
                 scene_bbox: Optional[List] = None,
                 near_scaling: float = 0.9,
                 ndc_far: float = 2.6):
        self.keyframes = keyframes
        self.max_cameras = max_cameras
        self.max_tsteps = max_tsteps
        self.downsample = downsample
        self.isg = isg
        self.ist = False
        # self.lookup_time = False
        self.per_cam_near_fars = None
        self.global_translation = torch.tensor([0, 0, 0])
        self.global_scale = torch.tensor([1, 1, 1])
        self.near_scaling = near_scaling
        self.ndc_far = ndc_far
        self.median_imgs = None
        if contraction and ndc:
            raise ValueError("Options 'contraction' and 'ndc' are exclusive.")
        if "lego" in datadir or "dnerf" in datadir:
            dset_type = "synthetic"
        else:
            dset_type = "llff"

        # Note: timestamps are stored normalized between -1, 1.
        if dset_type == "llff":
            if split == "render":
                assert ndc, "Unable to generate render poses without ndc: don't know near-far."
                per_cam_poses, per_cam_near_fars, intrinsics, _ = load_llffvideo_poses(
                    datadir, downsample=self.downsample, split='all', near_scaling=self.near_scaling)
                render_poses = generate_spiral_path(
                    per_cam_poses.numpy(), per_cam_near_fars.numpy(), n_frames=300,
                    n_rots=2, zrate=0.5, dt=self.near_scaling, percentile=60)
                self.poses = torch.from_numpy(render_poses).float()
                self.per_cam_near_fars = torch.tensor([[0.4, self.ndc_far]])
                timestamps = torch.linspace(0, 299, len(self.poses))
                imgs = None
            else:
                per_cam_poses, per_cam_near_fars, intrinsics, videopaths = load_llffvideo_poses(
                    datadir, downsample=self.downsample, split=split, near_scaling=self.near_scaling)
                if split == 'test':
                    keyframes = False
                poses, imgs, timestamps, self.median_imgs = load_llffvideo_data(
                    videopaths=videopaths, cam_poses=per_cam_poses, intrinsics=intrinsics,
                    split=split, keyframes=keyframes, keyframes_take_each=30)
                self.poses = poses.float()
                if contraction:
                    self.per_cam_near_fars = per_cam_near_fars.float()
                else:
                    self.per_cam_near_fars = torch.tensor(
                        [[0.0, self.ndc_far]]).repeat(per_cam_near_fars.shape[0], 1)
            # These values are tuned for the salmon video
            self.global_translation = torch.tensor([0, 0, 2.])
            self.global_scale = torch.tensor([0.5, 0.6, 1])
            # Normalize timestamps between -1, 1
            timestamps = (timestamps.float() / 299) * 2 - 1
        elif dset_type == "synthetic":
            assert not contraction, "Synthetic video dataset does not work with contraction."
            assert not ndc, "Synthetic video dataset does not work with NDC."
            if split == 'render':
                num_tsteps = 120
                dnerf_durations = {'hellwarrior': 100, 'mutant': 150, 'hook': 100, 'bouncingballs': 150, 'lego': 50, 'trex': 200, 'standup': 150, 'jumpingjacks': 200}
                for scene in dnerf_durations.keys():
                    if 'dnerf' in datadir and scene in datadir:
                        num_tsteps = dnerf_durations[scene]
                render_poses = torch.stack([
                    generate_spherical_poses(angle, -30.0, 4.0)
                    for angle in np.linspace(-180, 180, num_tsteps + 1)[:-1]
                ], 0)
                imgs = None
                self.poses = render_poses
                timestamps = torch.linspace(0.0, 1.0, render_poses.shape[0])
                _, transform = load_360video_frames(
                    datadir, 'train', max_cameras=self.max_cameras, max_tsteps=self.max_tsteps)
                img_h, img_w = 800, 800
            else:
                frames, transform = load_360video_frames(
                    datadir, split, max_cameras=self.max_cameras, max_tsteps=self.max_tsteps)
                imgs, self.poses = load_360_images(frames, datadir, split, self.downsample)
                timestamps = torch.tensor(
                    [fetch_360vid_info(f)[0] for f in frames], dtype=torch.float32)
                img_h, img_w = imgs[0].shape[:2]
            if ndc:
                self.per_cam_near_fars = torch.tensor([[0.0, self.ndc_far]])
            else:
                self.per_cam_near_fars = torch.tensor([[2.0, 6.0]])
            if "dnerf" in datadir:
                # dnerf time is between 0, 1. Normalize to -1, 1
                timestamps = timestamps * 2 - 1
            else:
                # lego (our vid) time is like dynerf: between 0, 30.
                timestamps = (timestamps.float() / torch.amax(timestamps)) * 2 - 1
            intrinsics = load_360_intrinsics(
                transform, img_h=img_h, img_w=img_w, downsample=self.downsample)
        else:
            raise ValueError(datadir)

        self.timestamps = timestamps
        if split == 'train':
            self.timestamps = self.timestamps[:, None, None].repeat(
                1, intrinsics.height, intrinsics.width).reshape(-1)  # [n_frames * h * w]
        assert self.timestamps.min() >= -1.0 and self.timestamps.max() <= 1.0, "timestamps out of range."
        if imgs is not None and imgs.dtype != torch.uint8:
            imgs = (imgs * 255).to(torch.uint8)
        if self.median_imgs is not None and self.median_imgs.dtype != torch.uint8:
            self.median_imgs = (self.median_imgs * 255).to(torch.uint8)
        if split == 'train':
            imgs = imgs.view(-1, imgs.shape[-1])
        elif imgs is not None:
            imgs = imgs.view(-1, intrinsics.height * intrinsics.width, imgs.shape[-1])

        # ISG/IST weights are computed on 4x subsampled data.
        weights_subsampled = int(4 / downsample)
        if scene_bbox is not None:
            scene_bbox = torch.tensor(scene_bbox)
        else:
            scene_bbox = get_bbox(datadir, is_contracted=contraction, dset_type=dset_type)
        super().__init__(
            datadir=datadir,
            split=split,
            batch_size=batch_size,
            is_ndc=ndc,
            is_contracted=contraction,
            scene_bbox=scene_bbox,
            rays_o=None,
            rays_d=None,
            intrinsics=intrinsics,
            imgs=imgs,
            sampling_weights=None,  # Start without importance sampling, by default
            weights_subsampled=weights_subsampled,
        )

        self.isg_weights = None
        self.ist_weights = None
        if split == "train" and dset_type == 'llff':  # Only use importance sampling with DyNeRF videos
            if os.path.exists(os.path.join(datadir, f"isg_weights.pt")):
                self.isg_weights = torch.load(os.path.join(datadir, f"isg_weights.pt"))
                log.info(f"Reloaded {self.isg_weights.shape[0]} ISG weights from file.")
            else:
                # Precompute ISG weights
                t_s = time.time()
                gamma = 1e-3 if self.keyframes else 2e-2
                self.isg_weights = dynerf_isg_weight(
                    imgs.view(-1, intrinsics.height, intrinsics.width, imgs.shape[-1]),
                    median_imgs=self.median_imgs, gamma=gamma)
                # Normalize into a probability distribution, to speed up sampling
                self.isg_weights = (self.isg_weights.reshape(-1) / torch.sum(self.isg_weights))
                torch.save(self.isg_weights, os.path.join(datadir, f"isg_weights.pt"))
                t_e = time.time()
                log.info(f"Computed {self.isg_weights.shape[0]} ISG weights in {t_e - t_s:.2f}s.")

            if os.path.exists(os.path.join(datadir, f"ist_weights.pt")):
                self.ist_weights = torch.load(os.path.join(datadir, f"ist_weights.pt"))
                log.info(f"Reloaded {self.ist_weights.shape[0]} IST weights from file.")
            else:
                # Precompute IST weights
                t_s = time.time()
                self.ist_weights = dynerf_ist_weight(
                    imgs.view(-1, self.img_h, self.img_w, imgs.shape[-1]),
                    num_cameras=self.median_imgs.shape[0])
                # Normalize into a probability distribution, to speed up sampling
                self.ist_weights = (self.ist_weights.reshape(-1) / torch.sum(self.ist_weights))
                torch.save(self.ist_weights, os.path.join(datadir, f"ist_weights.pt"))
                t_e = time.time()
                log.info(f"Computed {self.ist_weights.shape[0]} IST weights in {t_e - t_s:.2f}s.")

        if self.isg:
            self.enable_isg()

        log.info(f"VideoDataset contracted={self.is_contracted}, ndc={self.is_ndc}. "
                 f"Loaded {self.split} set from {self.datadir}: "
                 f"{len(self.poses)} images of size {self.img_h}x{self.img_w}. "
                 f"Images loaded: {self.imgs is not None}. "
                 f"{len(torch.unique(timestamps))} timestamps. Near-far: {self.per_cam_near_fars}. "
                 f"ISG={self.isg}, IST={self.ist}, weights_subsampled={self.weights_subsampled}. "
                 f"Sampling without replacement={self.use_permutation}. {intrinsics}")

    def enable_isg(self):
        self.isg = True
        self.ist = False
        self.sampling_weights = self.isg_weights
        log.info(f"Enabled ISG weights.")

    def switch_isg2ist(self):
        self.isg = False
        self.ist = True
        self.sampling_weights = self.ist_weights
        log.info(f"Switched from ISG to IST weights.")

    def __getitem__(self, index):
        h = self.intrinsics.height
        w = self.intrinsics.width
        dev = "cpu"
        if self.split == 'train':
            index = self.get_rand_ids(index)  # [batch_size // (weights_subsampled**2)]
            if self.weights_subsampled == 1 or self.sampling_weights is None:
                # Nothing special to do, either weights_subsampled = 1, or not using weights.
                image_id = torch.div(index, h * w, rounding_mode='floor')
                y = torch.remainder(index, h * w).div(w, rounding_mode='floor')
                x = torch.remainder(index, h * w).remainder(w)
            else:
                # We must deal with the fact that ISG/IST weights are computed on a dataset with
                # different 'downsampling' factor. E.g. if the weights were computed on 4x
                # downsampled data and the current dataset is 2x downsampled, `weights_subsampled`
                # will be 4 / 2 = 2.
                # Split each subsampled index into its 16 components in 2D.
                hsub, wsub = h // self.weights_subsampled, w // self.weights_subsampled
                image_id = torch.div(index, hsub * wsub, rounding_mode='floor')
                ysub = torch.remainder(index, hsub * wsub).div(wsub, rounding_mode='floor')
                xsub = torch.remainder(index, hsub * wsub).remainder(wsub)
                # xsub, ysub is the first point in the 4x4 square of finely sampled points
                x, y = [], []
                for ah in range(self.weights_subsampled):
                    for aw in range(self.weights_subsampled):
                        x.append(xsub * self.weights_subsampled + aw)
                        y.append(ysub * self.weights_subsampled + ah)
                x = torch.cat(x)
                y = torch.cat(y)
                image_id = image_id.repeat(self.weights_subsampled ** 2)
                # Inverse of the process to get x, y from index. image_id stays the same.
                index = x + y * w + image_id * h * w
            x, y = x + 0.5, y + 0.5
        else:
            image_id = [index]
            x, y = create_meshgrid(height=h, width=w, dev=dev, add_half=True, flat=True)

        out = {
            "timestamps": self.timestamps[index],      # (num_rays or 1, )
            "imgs": None,
        }
        if self.split == 'train':
            num_frames_per_camera = len(self.imgs) // (len(self.per_cam_near_fars) * h * w)
            camera_id = torch.div(image_id, num_frames_per_camera, rounding_mode='floor')  # (num_rays)
            out['near_fars'] = self.per_cam_near_fars[camera_id, :]
        else:
            out['near_fars'] = self.per_cam_near_fars  # Only one test camera

        if self.imgs is not None:
            out['imgs'] = (self.imgs[index] / 255.0).view(-1, self.imgs.shape[-1])

        c2w = self.poses[image_id]                                    # [num_rays or 1, 3, 4]
        camera_dirs = stack_camera_dirs(x, y, self.intrinsics, True)  # [num_rays, 3]
        out['rays_o'], out['rays_d'] = get_rays(
            camera_dirs, c2w, ndc=self.is_ndc, ndc_near=1.0, intrinsics=self.intrinsics,
            normalize_rd=True)                                        # [num_rays, 3]

        imgs = out['imgs']
        # Decide BG color
        bg_color = torch.ones((1, 3), dtype=torch.float32, device=dev)
        if self.split == 'train' and imgs.shape[-1] == 4:
            bg_color = torch.rand((1, 3), dtype=torch.float32, device=dev)
        out['bg_color'] = bg_color
        # Alpha compositing
        if imgs is not None and imgs.shape[-1] == 4:
            imgs = imgs[:, :3] * imgs[:, 3:] + bg_color * (1.0 - imgs[:, 3:])
        out['imgs'] = imgs

        return out


def get_bbox(datadir: str, dset_type: str, is_contracted=False) -> torch.Tensor:
    """Returns a default bounding box based on the dataset type, and contraction state.

    Args:
        datadir (str): Directory where data is stored
        dset_type (str): A string defining dataset type (e.g. synthetic, llff)
        is_contracted (bool): Whether the dataset will use contraction

    Returns:
        Tensor: 3x2 bounding box tensor
    """
    if is_contracted:
        radius = 2
    elif dset_type == 'synthetic':
        radius = 1.5
    elif dset_type == 'llff':
        return torch.tensor([[-3.0, -1.67, -1.2], [3.0, 1.67, 1.2]])
    else:
        radius = 1.3
    return torch.tensor([[-radius, -radius, -radius], [radius, radius, radius]])


def fetch_360vid_info(frame: Dict[str, Any]):
    timestamp = None
    fp = frame['file_path']
    if '_r' in fp:
        timestamp = int(fp.split('t')[-1].split('_')[0])
    if 'r_' in fp:
        pose_id = int(fp.split('r_')[-1])
    else:
        pose_id = int(fp.split('r')[-1])
    if timestamp is None:  # will be None for dnerf
        timestamp = frame['time']
    return timestamp, pose_id


def load_360video_frames(datadir, split, max_cameras: int, max_tsteps: Optional[int]) -> Tuple[Any, Any]:
    with open(os.path.join(datadir, f"transforms_{split}.json"), 'r') as fp:
        meta = json.load(fp)
    frames = meta['frames']

    timestamps = set()
    pose_ids = set()
    fpath2poseid = defaultdict(list)
    for frame in frames:
        timestamp, pose_id = fetch_360vid_info(frame)
        timestamps.add(timestamp)
        pose_ids.add(pose_id)
        fpath2poseid[frame['file_path']].append(pose_id)
    timestamps = sorted(timestamps)
    pose_ids = sorted(pose_ids)

    if max_cameras is not None:
        num_poses = min(len(pose_ids), max_cameras or len(pose_ids))
        subsample_poses = int(round(len(pose_ids) / num_poses))
        pose_ids = set(pose_ids[::subsample_poses])
        log.info(f"Selected subset of {len(pose_ids)} camera poses: {pose_ids}.")

    if max_tsteps is not None:
        num_timestamps = min(len(timestamps), max_tsteps or len(timestamps))
        subsample_time = int(math.floor(len(timestamps) / (num_timestamps - 1)))
        timestamps = set(timestamps[::subsample_time])
        log.info(f"Selected subset of timestamps: {sorted(timestamps)} of length {len(timestamps)}")

    sub_frames = []
    for frame in frames:
        timestamp, pose_id = fetch_360vid_info(frame)
        if timestamp in timestamps and pose_id in pose_ids:
            sub_frames.append(frame)
    # We need frames to be sorted by pose_id
    sub_frames = sorted(sub_frames, key=lambda f: fpath2poseid[f['file_path']])
    return sub_frames, meta


def load_llffvideo_poses(datadir: str,
                         downsample: float,
                         split: str,
                         near_scaling: float) -> Tuple[
                            torch.Tensor, torch.Tensor, Intrinsics, List[str]]:
    """Load poses and metadata for LLFF video.

    Args:
        datadir (str): Directory containing the videos and pose information
        downsample (float): How much to downsample videos. The default for LLFF videos is 2.0
        split (str): 'train' or 'test'.
        near_scaling (float): How much to scale the near bound of poses.

    Returns:
        Tensor: A tensor of size [N, 4, 4] containing c2w poses for each camera.
        Tensor: A tensor of size [N, 2] containing near, far bounds for each camera.
        Intrinsics: The camera intrinsics. These are the same for every camera.
        List[str]: List of length N containing the path to each camera's data.
    """
    poses, near_fars, intrinsics = load_llff_poses_helper(datadir, downsample, near_scaling)

    videopaths = np.array(glob.glob(os.path.join(datadir, '*.mp4')))  # [n_cameras]
    assert poses.shape[0] == len(videopaths), \
        'Mismatch between number of cameras and number of poses!'
    videopaths.sort()

    # The first camera is reserved for testing, following https://github.com/facebookresearch/Neural_3D_Video/releases/tag/v1.0
    if split == 'train':
        split_ids = np.arange(1, poses.shape[0])
    elif split == 'test':
        split_ids = np.array([0])
    else:
        split_ids = np.arange(poses.shape[0])
    if 'coffee_martini' in datadir:
        # https://github.com/fengres/mixvoxels/blob/0013e4ad63c80e5f14eb70383e2b073052d07fba/dataLoader/llff_video.py#L323
        log.info(f"Deleting unsynchronized camera from coffee-martini video.")
        split_ids = np.setdiff1d(split_ids, 12)
    poses = torch.from_numpy(poses[split_ids])
    near_fars = torch.from_numpy(near_fars[split_ids])
    videopaths = videopaths[split_ids].tolist()

    return poses, near_fars, intrinsics, videopaths


def load_llffvideo_data(videopaths: List[str],
                        cam_poses: torch.Tensor,
                        intrinsics: Intrinsics,
                        split: str,
                        keyframes: bool,
                        keyframes_take_each: Optional[int] = None,
                        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    if keyframes and (keyframes_take_each is None or keyframes_take_each < 1):
        raise ValueError(f"'keyframes_take_each' must be a positive number, "
                         f"but is {keyframes_take_each}.")

    loaded = parallel_load_images(
        dset_type="video",
        tqdm_title=f"Loading {split} data",
        num_images=len(videopaths),
        paths=videopaths,
        poses=cam_poses,
        out_h=intrinsics.height,
        out_w=intrinsics.width,
        load_every=keyframes_take_each if keyframes else 1,
    )
    imgs, poses, median_imgs, timestamps = zip(*loaded)
    # Stack everything together
    timestamps = torch.cat(timestamps, 0)  # [N]
    poses = torch.cat(poses, 0)            # [N, 3, 4]
    imgs = torch.cat(imgs, 0)              # [N, h, w, 3]
    median_imgs = torch.stack(median_imgs, 0)  # [num_cameras, h, w, 3]

    return poses, imgs, timestamps, median_imgs


@torch.no_grad()
def dynerf_isg_weight(imgs, median_imgs, gamma):
    # imgs is [num_cameras * num_frames, h, w, 3]
    # median_imgs is [num_cameras, h, w, 3]
    assert imgs.dtype == torch.uint8
    assert median_imgs.dtype == torch.uint8
    num_cameras, h, w, c = median_imgs.shape
    squarediff = (
        imgs.view(num_cameras, -1, h, w, c)
            .float()  # creates new tensor, so later operations can be in-place
            .div_(255.0)
            .sub_(
                median_imgs[:, None, ...].float().div_(255.0)
            )
            .square_()  # noqa
    )  # [num_cameras, num_frames, h, w, 3]
    # differences = median_imgs[:, None, ...] - imgs.view(num_cameras, -1, h, w, c)  # [num_cameras, num_frames, h, w, 3]
    # squarediff = torch.square_(differences)
    psidiff = squarediff.div_(squarediff + gamma**2)
    psidiff = (1./3) * torch.sum(psidiff, dim=-1)  # [num_cameras, num_frames, h, w]
    return psidiff  # valid probabilities, each in [0, 1]


@torch.no_grad()
def dynerf_ist_weight(imgs, num_cameras, alpha=0.1, frame_shift=25):  # DyNerf uses alpha=0.1
    assert imgs.dtype == torch.uint8
    N, h, w, c = imgs.shape
    frames = imgs.view(num_cameras, -1, h, w, c).float()  # [num_cameras, num_timesteps, h, w, 3]
    max_diff = None
    shifts = list(range(frame_shift + 1))[1:]
    for shift in shifts:
        shift_left = torch.cat([frames[:, shift:, ...], torch.zeros(num_cameras, shift, h, w, c)], dim=1)
        shift_right = torch.cat([torch.zeros(num_cameras, shift, h, w, c), frames[:, :-shift, ...]], dim=1)
        mymax = torch.maximum(torch.abs_(shift_left - frames), torch.abs_(shift_right - frames))
        if max_diff is None:
            max_diff = mymax
        else:
            max_diff = torch.maximum(max_diff, mymax)  # [num_timesteps, h, w, 3]
    max_diff = torch.mean(max_diff, dim=-1)  # [num_timesteps, h, w]
    max_diff = max_diff.clamp_(min=alpha)
    return max_diff

thirdparty/kplanes/datasets/video_datasets.py
oppo-us-research/NeuRBF
"""
Density proposal field
"""
from typing import Optional, Callable
import logging as log

import torch
import torch.nn as nn
import tinycudann as tcnn

from thirdparty.kplanes.models.kplane_field import interpolate_ms_features, normalize_aabb, init_grid_param
from thirdparty.kplanes.raymarching.spatial_distortions import SpatialDistortion


class KPlaneDensityField(nn.Module):
    def __init__(self,
                 aabb,
                 resolution,
                 num_input_coords,
                 num_output_coords,
                 density_activation: Callable,
                 spatial_distortion: Optional[SpatialDistortion] = None,
                 linear_decoder: bool = True):
        super().__init__()
        self.aabb = nn.Parameter(aabb, requires_grad=False)
        self.spatial_distortion = spatial_distortion
        self.hexplane = num_input_coords == 4
        self.feature_dim = num_output_coords
        self.density_activation = density_activation
        self.linear_decoder = linear_decoder
        activation = "ReLU"
        if self.linear_decoder:
            activation = "None"

        self.grids = init_grid_param(
            grid_nd=2, in_dim=num_input_coords, out_dim=num_output_coords, reso=resolution,
            a=0.1, b=0.15)
        self.sigma_net = tcnn.Network(
            n_input_dims=self.feature_dim,
            n_output_dims=1,
            network_config={
                "otype": "FullyFusedMLP",
                "activation": activation,
                "output_activation": "None",
                "n_neurons": 64,
                "n_hidden_layers": 1,
            },
        )
        log.info(f"Initialized KPlaneDensityField. hexplane={self.hexplane} - "
                 f"resolution={resolution}")
        log.info(f"KPlaneDensityField grids: \n{self.grids}")

    def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
        if self.spatial_distortion is not None:
            pts = self.spatial_distortion(pts)
            pts = pts / 2  # from [-2, 2] to [-1, 1]
        else:
            pts = normalize_aabb(pts, self.aabb)
        n_rays, n_samples = pts.shape[:2]
        if timestamps is not None and self.hexplane:
            timestamps = timestamps[:, None].expand(-1, n_samples)[..., None]  # [n_rays, n_samples, 1]
            pts = torch.cat((pts, timestamps), dim=-1)  # [n_rays, n_samples, 4]

        pts = pts.reshape(-1, pts.shape[-1])
        features = interpolate_ms_features(
            pts, ms_grids=[self.grids], grid_dimensions=2, concat_features=False, num_levels=None)
        density = self.density_activation(
            self.sigma_net(features).to(pts)
        ).view(n_rays, n_samples, 1)
        return density

    def forward(self, pts: torch.Tensor):
        return self.get_density(pts)

    def get_params(self):
        field_params = {k: v for k, v in self.grids.named_parameters(prefix="grids")}
        nn_params = {k: v for k, v in self.sigma_net.named_parameters(prefix="sigma_net")}
        other_params = {k: v for k, v in self.named_parameters() if (
            k not in nn_params.keys() and k not in field_params.keys()
        )}
        return {
            "nn": list(nn_params.values()),
            "field": list(field_params.values()),
            "other": list(other_params.values()),
        }

thirdparty/kplanes/models/density_fields.py
oppo-us-research/NeuRBF
import itertools
import logging as log
from typing import Optional, Union, List, Dict, Sequence, Iterable, Collection, Callable

import torch
import torch.nn as nn
import tinycudann as tcnn

from thirdparty.kplanes.ops.interpolation import grid_sample_wrapper
from thirdparty.kplanes.raymarching.spatial_distortions import SpatialDistortion


def get_normalized_directions(directions):
    """SH encoding must be in the range [0, 1]

    Args:
        directions: batch of directions
    """
    return (directions + 1.0) / 2.0


def normalize_aabb(pts, aabb):
    return (pts - aabb[0]) * (2.0 / (aabb[1] - aabb[0])) - 1.0


def init_grid_param(
        grid_nd: int,
        in_dim: int,
        out_dim: int,
        reso: Sequence[int],
        a: float = 0.1,
        b: float = 0.5):
    assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
    has_time_planes = in_dim == 4
    assert grid_nd <= in_dim
    coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
    grid_coefs = nn.ParameterList()
    for ci, coo_comb in enumerate(coo_combs):
        new_grid_coef = nn.Parameter(torch.empty(
            [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
        ))
        if has_time_planes and 3 in coo_comb:  # Initialize time planes to 1
            nn.init.ones_(new_grid_coef)
        else:
            nn.init.uniform_(new_grid_coef, a=a, b=b)
        grid_coefs.append(new_grid_coef)

    return grid_coefs


def interpolate_ms_features(pts: torch.Tensor,
                            ms_grids: Collection[Iterable[nn.Module]],
                            grid_dimensions: int,
                            concat_features: bool,
                            num_levels: Optional[int],
                            ) -> torch.Tensor:
    coo_combs = list(itertools.combinations(
        range(pts.shape[-1]), grid_dimensions)
    )
    if num_levels is None:
        num_levels = len(ms_grids)
    multi_scale_interp = [] if concat_features else 0.
    grid: nn.ParameterList
    for scale_id, grid in enumerate(ms_grids[:num_levels]):
        interp_space = 1.
        for ci, coo_comb in enumerate(coo_combs):
            # interpolate in plane
            feature_dim = grid[ci].shape[1]  # shape of grid[ci]: 1, out_dim, *reso
            interp_out_plane = (
                grid_sample_wrapper(grid[ci], pts[..., coo_comb])
                .view(-1, feature_dim)
            )
            # compute product over planes
            interp_space = interp_space * interp_out_plane

        # combine over scales
        if concat_features:
            multi_scale_interp.append(interp_space)
        else:
            multi_scale_interp = multi_scale_interp + interp_space

    if concat_features:
        multi_scale_interp = torch.cat(multi_scale_interp, dim=-1)
    return multi_scale_interp


class KPlaneField(nn.Module):
    def __init__(
        self,
        aabb,
        grid_config: Union[str, List[Dict]],
        concat_features_across_scales: bool,
        multiscale_res: Optional[Sequence[int]],
        use_appearance_embedding: bool,
        appearance_embedding_dim: int,
        spatial_distortion: Optional[SpatialDistortion],
        density_activation: Callable,
        linear_decoder: bool,
        linear_decoder_layers: Optional[int],
        num_images: Optional[int],
    ) -> None:
        super().__init__()

        self.aabb = nn.Parameter(aabb, requires_grad=False)
        self.spatial_distortion = spatial_distortion
        self.grid_config = grid_config

        self.multiscale_res_multipliers: List[int] = multiscale_res or [1]
        self.concat_features = concat_features_across_scales
        self.density_activation = density_activation
        self.linear_decoder = linear_decoder

        # 1. Init planes
        self.grids = nn.ModuleList()
        self.feature_dim = 0
        for res in self.multiscale_res_multipliers:
            # initialize coordinate grid
            config = self.grid_config[0].copy()
            # Resolution fix: multi-res only on spatial planes
            config["resolution"] = [
                r * res for r in config["resolution"][:3]
            ] + config["resolution"][3:]
            gp = init_grid_param(
                grid_nd=config["grid_dimensions"],
                in_dim=config["input_coordinate_dim"],
                out_dim=config["output_coordinate_dim"],
                reso=config["resolution"],
            )
            # shape[1] is out-dim - Concatenate over feature len for each scale
            if self.concat_features:
                self.feature_dim += gp[-1].shape[1]
            else:
                self.feature_dim = gp[-1].shape[1]
            self.grids.append(gp)
        log.info(f"Initialized model grids: {self.grids}")

        # 2. Init appearance code-related parameters
        self.use_average_appearance_embedding = True  # for test-time
        self.use_appearance_embedding = use_appearance_embedding
        self.num_images = num_images
        self.appearance_embedding = None
        if use_appearance_embedding:
            assert self.num_images is not None
            self.appearance_embedding_dim = appearance_embedding_dim
            # this will initialize as normal_(0.0, 1.0)
            self.appearance_embedding = nn.Embedding(self.num_images, self.appearance_embedding_dim)
        else:
            self.appearance_embedding_dim = 0

        # 3. Init decoder params
        self.direction_encoder = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "SphericalHarmonics",
                "degree": 4,
            },
        )

        # 3. Init decoder network
        if self.linear_decoder:
            assert linear_decoder_layers is not None
            # The NN learns a basis that is used instead of spherical harmonics
            # Input is an encoded view direction, output is weights for
            # combining the color features into RGB
            # This architecture is based on instant-NGP
            self.color_basis = tcnn.Network(
                n_input_dims=3 + self.appearance_embedding_dim,
                n_output_dims=3 * self.feature_dim,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 128,
                    "n_hidden_layers": linear_decoder_layers,
                },
            )
            # sigma_net just does a linear transformation on the features to get density
            self.sigma_net = tcnn.Network(
                n_input_dims=self.feature_dim,
                n_output_dims=1,
                network_config={
                    "otype": "CutlassMLP",
                    "activation": "None",
                    "output_activation": "None",
                    "n_neurons": 128,
                    "n_hidden_layers": 0,
                },
            )
        else:
            self.geo_feat_dim = 15
            self.sigma_net = tcnn.Network(
                n_input_dims=self.feature_dim,
                n_output_dims=self.geo_feat_dim + 1,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 64,
                    "n_hidden_layers": 1,
                },
            )
            self.in_dim_color = (
                    self.direction_encoder.n_output_dims
                    + self.geo_feat_dim
                    + self.appearance_embedding_dim
            )
            self.color_net = tcnn.Network(
                n_input_dims=self.in_dim_color,
                n_output_dims=3,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "Sigmoid",
                    "n_neurons": 64,
                    "n_hidden_layers": 2,
                },
            )

    def step_after_iter(self, step):
        pass
    
    def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
        """Computes and returns the densities."""
        if self.spatial_distortion is not None:
            pts = self.spatial_distortion(pts)
            pts = pts / 2  # from [-2, 2] to [-1, 1]
        else:
            pts = normalize_aabb(pts, self.aabb)
        n_rays, n_samples = pts.shape[:2]
        if timestamps is not None:
            timestamps = timestamps[:, None].expand(-1, n_samples)[..., None]  # [n_rays, n_samples, 1]
            pts = torch.cat((pts, timestamps), dim=-1)  # [n_rays, n_samples, 4]

        pts = pts.reshape(-1, pts.shape[-1])
        features = interpolate_ms_features(
            pts, ms_grids=self.grids,  # noqa
            grid_dimensions=self.grid_config[0]["grid_dimensions"],
            concat_features=self.concat_features, num_levels=None)
        if len(features) < 1:
            features = torch.zeros((0, 1)).to(features.device)
        if self.linear_decoder:
            density_before_activation = self.sigma_net(features)  # [batch, 1]
        else:
            features = self.sigma_net(features)
            features, density_before_activation = torch.split(
                features, [self.geo_feat_dim, 1], dim=-1)

        density = self.density_activation(
            density_before_activation.to(pts)
        ).view(n_rays, n_samples, 1)
        return density, features

    def forward(self,
                pts: torch.Tensor,
                directions: torch.Tensor,
                timestamps: Optional[torch.Tensor] = None):
        camera_indices = None
        if self.use_appearance_embedding:
            if timestamps is None:
                raise AttributeError("timestamps (appearance-ids) are not provided.")
            camera_indices = timestamps
            timestamps = None
        density, features = self.get_density(pts, timestamps)
        n_rays, n_samples = pts.shape[:2]

        directions = directions.view(-1, 1, 3).expand(pts.shape).reshape(-1, 3)
        if not self.linear_decoder:
            directions = get_normalized_directions(directions)
            encoded_directions = self.direction_encoder(directions)

        if self.linear_decoder:
            color_features = [features]
        else:
            color_features = [encoded_directions, features.view(-1, self.geo_feat_dim)]

        if self.use_appearance_embedding:
            if camera_indices.dtype == torch.float32:
                # Interpolate between two embeddings. Currently they are hardcoded below.
                #emb1_idx, emb2_idx = 100, 121  # trevi
                emb1_idx, emb2_idx = 11, 142  # sacre
                emb_fn = self.appearance_embedding
                emb1 = emb_fn(torch.full_like(camera_indices, emb1_idx, dtype=torch.long))
                emb1 = emb1.view(emb1.shape[0], emb1.shape[2])
                emb2 = emb_fn(torch.full_like(camera_indices, emb2_idx, dtype=torch.long))
                emb2 = emb2.view(emb2.shape[0], emb2.shape[2])
                embedded_appearance = torch.lerp(emb1, emb2, camera_indices)
            elif self.training:
                embedded_appearance = self.appearance_embedding(camera_indices)
            else:
                if hasattr(self, "test_appearance_embedding"):
                    embedded_appearance = self.test_appearance_embedding(camera_indices)
                elif self.use_average_appearance_embedding:
                    embedded_appearance = torch.ones(
                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                    ) * self.appearance_embedding.mean(dim=0)
                else:
                    embedded_appearance = torch.zeros(
                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                    )

            # expand embedded_appearance from n_rays, dim to n_rays*n_samples, dim
            ea_dim = embedded_appearance.shape[-1]
            embedded_appearance = embedded_appearance.view(-1, 1, ea_dim).expand(n_rays, n_samples, -1).reshape(-1, ea_dim)
            if not self.linear_decoder:
                color_features.append(embedded_appearance)

        color_features = torch.cat(color_features, dim=-1)

        if self.linear_decoder:
            if self.use_appearance_embedding:
                basis_values = self.color_basis(torch.cat([directions, embedded_appearance], dim=-1))
            else:
                basis_values = self.color_basis(directions)  # [batch, color_feature_len * 3]
            basis_values = basis_values.view(color_features.shape[0], 3, -1)  # [batch, 3, color_feature_len]
            rgb = torch.sum(color_features[:, None, :] * basis_values, dim=-1)  # [batch, 3]
            rgb = rgb.to(directions)
            rgb = torch.sigmoid(rgb).view(n_rays, n_samples, 3)
        else:
            rgb = self.color_net(color_features).to(directions).view(n_rays, n_samples, 3)

        return {"rgb": rgb, "density": density}

    def get_params(self):
        field_params = {k: v for k, v in self.grids.named_parameters(prefix="grids")}
        nn_params = [
            self.sigma_net.named_parameters(prefix="sigma_net"),
            self.direction_encoder.named_parameters(prefix="direction_encoder"),
        ]
        if self.linear_decoder:
            nn_params.append(self.color_basis.named_parameters(prefix="color_basis"))
        else:
            nn_params.append(self.color_net.named_parameters(prefix="color_net"))
        nn_params = {k: v for plist in nn_params for k, v in plist}
        other_params = {k: v for k, v in self.named_parameters() if (
            k not in nn_params.keys() and k not in field_params.keys()
        )}
        return {
            "nn": list(nn_params.values()),
            "field": list(field_params.values()),
            "other": list(other_params.values()),
        }

thirdparty/kplanes/models/kplane_field.py
oppo-us-research/NeuRBF
from typing import List, Sequence, Optional, Union, Dict, Tuple
import time

import numpy as np
import torch
import torch.nn as nn

from thirdparty.kplanes.models.density_fields import KPlaneDensityField
from thirdparty.kplanes.models.kplane_field import KPlaneField
from nerf_kplanes.network import KPlaneRBFField
from thirdparty.kplanes.ops.activations import init_density_activation
from thirdparty.kplanes.raymarching.ray_samplers import (
    UniformLinDispPiecewiseSampler, UniformSampler,
    ProposalNetworkSampler, RayBundle, RaySamples
)
from thirdparty.kplanes.raymarching.spatial_distortions import SceneContraction, SpatialDistortion
from thirdparty.kplanes.utils.timer import CudaTimer


class LowrankModel(nn.Module):
    def __init__(self,
                 grid_config: Union[str, List[Dict]],
                 # boolean flags
                 is_ndc: bool,
                 is_contracted: bool,
                 aabb: torch.Tensor,
                 # Model arguments
                 multiscale_res: Sequence[int],
                 density_activation: Optional[str] = 'trunc_exp',
                 concat_features_across_scales: bool = False,
                 linear_decoder: bool = True,
                 linear_decoder_layers: Optional[int] = 1,
                 # RBF related
                 field_type: str = 'kplane', 
                 # Spatial distortion
                 global_translation: Optional[torch.Tensor] = None,
                 global_scale: Optional[torch.Tensor] = None,
                 # proposal-sampling arguments
                 num_proposal_iterations: int = 1,
                 use_same_proposal_network: bool = False,
                 proposal_net_args_list: List[Dict] = None,
                 num_proposal_samples: Optional[Tuple[int]] = None,
                 num_samples: Optional[int] = None,
                 single_jitter: bool = False,
                 proposal_warmup: int = 5000,
                 proposal_update_every: int = 5,
                 use_proposal_weight_anneal: bool = True,
                 proposal_weights_anneal_max_num_iters: int = 1000,
                 proposal_weights_anneal_slope: float = 10.0,
                 # appearance embedding (phototourism)
                 use_appearance_embedding: bool = False,
                 appearance_embedding_dim: int = 0,
                 num_images: Optional[int] = None,
                 **kwargs,
                 ):
        super().__init__()
        if isinstance(grid_config, str):
            self.config: List[Dict] = eval(grid_config)
        else:
            self.config: List[Dict] = grid_config
        self.multiscale_res = multiscale_res
        self.is_ndc = is_ndc
        self.is_contracted = is_contracted
        self.concat_features_across_scales = concat_features_across_scales
        self.linear_decoder = linear_decoder
        self.linear_decoder_layers = linear_decoder_layers
        self.density_act = init_density_activation(density_activation)
        self.timer = CudaTimer(enabled=False)

        self.spatial_distortion: Optional[SpatialDistortion] = None
        if self.is_contracted:
            self.spatial_distortion = SceneContraction(
                order=float('inf'), global_scale=global_scale,
                global_translation=global_translation)

        self.field_type = field_type
        t = time.time()
        if self.field_type == 'kplane':
            self.field = KPlaneField(
                aabb,
                grid_config=self.config,
                concat_features_across_scales=self.concat_features_across_scales,
                multiscale_res=self.multiscale_res,
                use_appearance_embedding=use_appearance_embedding,
                appearance_embedding_dim=appearance_embedding_dim,
                spatial_distortion=self.spatial_distortion,
                density_activation=self.density_act,
                linear_decoder=self.linear_decoder,
                linear_decoder_layers=self.linear_decoder_layers,
                num_images=num_images,
            )
        elif self.field_type == 'kplane_rbf':
            self.field = KPlaneRBFField(
                aabb,
                grid_config=self.config,
                concat_features_across_scales=self.concat_features_across_scales,
                multiscale_res=self.multiscale_res,
                use_appearance_embedding=use_appearance_embedding,
                appearance_embedding_dim=appearance_embedding_dim,
                spatial_distortion=self.spatial_distortion,
                density_activation=self.density_act,
                linear_decoder=self.linear_decoder,
                linear_decoder_layers=self.linear_decoder_layers,
                num_images=num_images,
                rbf_config=kwargs['rbf_config'],
                init_data=kwargs['init_data'],
            )
        else:
            raise NotImplementedError()
        self.t_init = time.time() - t

        # Initialize proposal-sampling nets
        self.density_fns = []
        self.num_proposal_iterations = num_proposal_iterations
        self.proposal_net_args_list = proposal_net_args_list
        self.proposal_warmup = proposal_warmup
        self.proposal_update_every = proposal_update_every
        self.use_proposal_weight_anneal = use_proposal_weight_anneal
        self.proposal_weights_anneal_max_num_iters = proposal_weights_anneal_max_num_iters
        self.proposal_weights_anneal_slope = proposal_weights_anneal_slope
        self.proposal_networks = torch.nn.ModuleList()
        if use_same_proposal_network:
            assert len(self.proposal_net_args_list) == 1, "Only one proposal network is allowed."
            prop_net_args = self.proposal_net_args_list[0]
            network = KPlaneDensityField(
                aabb, spatial_distortion=self.spatial_distortion,
                density_activation=self.density_act, linear_decoder=self.linear_decoder, **prop_net_args)
            self.proposal_networks.append(network)
            self.density_fns.extend([network.get_density for _ in range(self.num_proposal_iterations)])
        else:
            for i in range(self.num_proposal_iterations):
                prop_net_args = self.proposal_net_args_list[min(i, len(self.proposal_net_args_list) - 1)]
                network = KPlaneDensityField(
                    aabb, spatial_distortion=self.spatial_distortion,
                    density_activation=self.density_act, linear_decoder=self.linear_decoder, **prop_net_args,
                )
                self.proposal_networks.append(network)
            self.density_fns.extend([network.get_density for network in self.proposal_networks])

        update_schedule = lambda step: np.clip(
            np.interp(step, [0, self.proposal_warmup], [0, self.proposal_update_every]),
            1,
            self.proposal_update_every,
        )
        if self.is_contracted or self.is_ndc:
            initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
        else:
            initial_sampler = UniformSampler(single_jitter=single_jitter)
        self.proposal_sampler = ProposalNetworkSampler(
            num_nerf_samples_per_ray=num_samples,
            num_proposal_samples_per_ray=num_proposal_samples,
            num_proposal_network_iterations=self.num_proposal_iterations,
            single_jitter=single_jitter,
            update_sched=update_schedule,
            initial_sampler=initial_sampler
        )

        # Count parameters
        self.n_params = sum(np.prod(p.shape) for p in self.parameters())
        if self.field_type == 'kplane_rbf' and self.field.rbf_type.endswith('_a'):
            in_dim = self.field.in_dim
            self.n_params -= self.field.n_kernel * int((in_dim-1)*in_dim/2)

    def step_before_iter(self, step):
        if self.use_proposal_weight_anneal:
            # anneal the weights of the proposal network before doing PDF sampling
            N = self.proposal_weights_anneal_max_num_iters
            # https://arxiv.org/pdf/2111.12077.pdf eq. 18
            train_frac = np.clip(step / N, 0, 1)
            bias = lambda x, b: (b * x) / ((b - 1) * x + 1)
            anneal = bias(train_frac, self.proposal_weights_anneal_slope)
            self.proposal_sampler.set_anneal(anneal)

    def step_after_iter(self, step):
        if self.use_proposal_weight_anneal:
            self.proposal_sampler.step_cb(step)
        return self.field.step_after_iter(step)

    @staticmethod
    def render_rgb(rgb: torch.Tensor, weights: torch.Tensor, bg_color: Optional[torch.Tensor]):
        comp_rgb = torch.sum(weights * rgb, dim=-2)
        accumulated_weight = torch.sum(weights, dim=-2)
        if bg_color is None:
            pass
        else:
            comp_rgb = comp_rgb + (1.0 - accumulated_weight) * bg_color
        return comp_rgb

    @staticmethod
    def render_depth(weights: torch.Tensor, ray_samples: RaySamples, rays_d: torch.Tensor):
        steps = (ray_samples.starts + ray_samples.ends) / 2
        one_minus_transmittance = torch.sum(weights, dim=-2)
        depth = torch.sum(weights * steps, dim=-2) + one_minus_transmittance * rays_d[..., -1:]
        return depth

    @staticmethod
    def render_accumulation(weights: torch.Tensor):
        accumulation = torch.sum(weights, dim=-2)
        return accumulation

    def forward(self, rays_o, rays_d, bg_color, near_far: torch.Tensor, timestamps=None):
        """
        rays_o : [batch, 3]
        rays_d : [batch, 3]
        timestamps : [batch]
        near_far : [batch, 2]
        """
        # Fix shape for near-far
        nears, fars = torch.split(near_far, [1, 1], dim=-1)
        if nears.shape[0] != rays_o.shape[0]:
            ones = torch.ones_like(rays_o[..., 0:1])
            nears = ones * nears
            fars = ones * fars

        ray_bundle = RayBundle(origins=rays_o, directions=rays_d, nears=nears, fars=fars)
        # Note: proposal sampler mustn't use timestamps (=camera-IDs) with appearance-embedding,
        #       since the appearance embedding should not affect density. We still pass them in the
        #       call below, but they will not be used as long as density-field resolutions
        #       are be 3D.
        ray_samples, weights_list, ray_samples_list = self.proposal_sampler.generate_ray_samples(
            ray_bundle, timestamps=timestamps, density_fns=self.density_fns)

        field_out = self.field(ray_samples.get_positions(), ray_bundle.directions, timestamps)
        rgb, density = field_out["rgb"], field_out["density"]

        weights = ray_samples.get_weights(density)
        weights_list.append(weights)
        ray_samples_list.append(ray_samples)

        rgb = self.render_rgb(rgb=rgb, weights=weights, bg_color=bg_color)
        depth = self.render_depth(weights=weights, ray_samples=ray_samples, rays_d=ray_bundle.directions)
        accumulation = self.render_accumulation(weights=weights)
        outputs = {
            "rgb": rgb,
            "accumulation": accumulation,
            "depth": depth,
        }

        # These use a lot of GPU memory, so we avoid storing them for eval.
        if self.training:
            outputs["weights_list"] = weights_list
            outputs["ray_samples_list"] = ray_samples_list
        for i in range(self.num_proposal_iterations):
            outputs[f"prop_depth_{i}"] = self.render_depth(
                weights=weights_list[i], ray_samples=ray_samples_list[i], rays_d=ray_bundle.directions)
        return outputs

    def get_params(self, lr_config: Dict):
        model_params = self.field.get_params()
        pn_params = [pn.get_params() for pn in self.proposal_networks]
        field_params = model_params["field"] + [p for pnp in pn_params for p in pnp["field"]]
        nn_params = model_params["nn"] + [p for pnp in pn_params for p in pnp["nn"]]
        other_params = model_params["other"] + [p for pnp in pn_params for p in pnp["other"]]
        out = [
            {"params": field_params, "lr": lr_config['base']},
            {"params": nn_params, "lr": lr_config['base']},
            {"params": other_params, "lr": lr_config['base']},
        ]
        if self.field_type == 'kplane_rbf':
            out.append({"params": model_params['lc0'], "lr": lr_config['lc0']})
            out.append({"params": model_params['kc0'], "lr": lr_config['kc0']})
            out.append({"params": model_params['ks0'], "lr": lr_config['ks0']})
        return out

thirdparty/kplanes/models/lowrank_model.py
oppo-us-research/NeuRBF
import torch
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd

__all__ = (
    "trunc_exp",
    "init_density_activation",
)

from torch.nn import functional as F


class TruncatedExponential(Function):  # pylint: disable=abstract-method
    # Implementation from torch-ngp:
    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, x):  # pylint: disable=arguments-differ
        ctx.save_for_backward(x)
        return torch.exp(x)

    @staticmethod
    @custom_bwd
    def backward(ctx, g):  # pylint: disable=arguments-differ
        x = ctx.saved_tensors[0]
        return g * torch.exp(torch.clamp(x, min=-15, max=15))


trunc_exp = TruncatedExponential.apply


def init_density_activation(activation_type: str):
    if activation_type == 'trunc_exp':
        return lambda x: trunc_exp(x - 1)
    elif activation_type == 'relu':
        return F.relu
    else:
        raise ValueError(activation_type)

thirdparty/kplanes/ops/activations.py
oppo-us-research/NeuRBF
""" Functions for image loading, saving, and color map changes """
#################################################################################
# Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# SPDX-FileCopyrightText: Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES
# SPDX-License-Identifier: BSD-3-Clause
#################################################################################

# Visualizing and Communicating Errors in Rendered Images
# Ray Tracing Gems II, 2021,
# by Pontus Andersson, Jim Nilsson, and Tomas Akenine-Moller.
# Pointer to the chapter: https://research.nvidia.com/publication/2021-08_Visualizing-and-Communicating.

# Visualizing Errors in Rendered High Dynamic Range Images
# Eurographics 2021,
# by Pontus Andersson, Jim Nilsson, Peter Shirley, and Tomas Akenine-Moller.
# Pointer to the paper: https://research.nvidia.com/publication/2021-05_HDR-FLIP.

# FLIP: A Difference Evaluator for Alternating Images
# High Performance Graphics 2020,
# by Pontus Andersson, Jim Nilsson, Tomas Akenine-Moller,
# Magnus Oskarsson, Kalle Astrom, and Mark D. Fairchild.
# Pointer to the paper: https://research.nvidia.com/publication/2020-07_FLIP.

# Code by Pontus Andersson, Jim Nilsson, and Tomas Akenine-Moller.

import numpy as np
from PIL import Image


def HWCtoCHW(x):
    """
    Transforms an image from HxWxC layout to CxHxW

    :param x: image with HxWxC layout
    :return: image with CxHxW layout
    """
    return np.rollaxis(x, 2)


def CHWtoHWC(x):
    """
    Transforms an image from CxHxW layout to HxWxC

    :param x: image with CxHxW layout
    :return: image with HxWxC layout
    """
    return np.swapaxes(np.swapaxes(x, 0, 1), 1, 2)


def save_image(img_file, img):
    """
    Saves image as png

    :param img_file: image's filename
    :param img: float image (in the [0,1] range) to save
    """
    img_array = (np.round(np.clip(img, 0.0, 1.0) * 255.0)).astype(np.uint8)
    im = Image.fromarray(img_array)
    im.save(img_file)


def load_image_array(img_file):
    """
    Loads an image (in the [0, 255] range) and transforms it into a numpy array and into the [0, 1] range

    :param img_file: image's filename
    :return: float image (in the [0,1] range) on CxHxW layout
    """
    img = Image.open(img_file, 'r').convert('RGB')
    img = np.asarray(img).astype(np.float32)
    img = HWCtoCHW(img)
    img = img / 255.0
    return img


def read_exr(filename):
    """
    Read color data from EXR image file. Set negative values and nans to 0.

    :param filename: string describing file path
    :return: RGB image in float32 format (with HxWxC layout)
    """
    import OpenEXR as exr
    import Imath

    exrfile = exr.InputFile(filename)
    header = exrfile.header()

    dw = header['dataWindow']
    isize = (dw.max.y - dw.min.y + 1, dw.max.x - dw.min.x + 1)

    channelData = dict()

    # Convert all channels in the image to numpy arrays
    for c in header['channels']:
        C = exrfile.channel(c, Imath.PixelType(Imath.PixelType.FLOAT))
        C = np.frombuffer(C, dtype=np.float32)
        C = np.reshape(C, isize)

        channelData[c] = C

    if len(channelData) == 1:
        channelData['R'] = channelData['G'] = channelData['B'] = channelData[
            next(iter(channelData))]

    colorChannels = ['R', 'G', 'B']
    img = np.concatenate([channelData[c][..., np.newaxis] for c in colorChannels], axis=2)

    return np.maximum(np.nan_to_num(np.array(img)),
                      0.0)  # added maximum to avoid negative values in images


def index2color(index_map, color_map):
    """
    Transforms grayscale index map to colors in a color map

    :param index_map: integer matrix containing integer versions of the FLIP errors, scaled to [0,255]
    :param color_map: matrix of size 256x3 where each row contains a color
    :return: error map in color map colors
    """
    dim = index_map.shape
    index_map = index_map.flatten().astype(int)
    column_stacked_colors = color_map[index_map, :]
    heat_map = np.reshape(column_stacked_colors.transpose(), (3, dim[0], dim[1]))
    return heat_map


def get_magma_map():
    """
    Return the magma map, described here: https://bids.github.io/colormap/

    :return: matrix of size 256x3 where each row contains a color in the magma map
    """
    color_map = [[0.001462, 0.000466, 0.013866],
                 [0.002258, 0.001295, 0.018331],
                 [0.003279, 0.002305, 0.023708],
                 [0.004512, 0.003490, 0.029965],
                 [0.005950, 0.004843, 0.037130],
                 [0.007588, 0.006356, 0.044973],
                 [0.009426, 0.008022, 0.052844],
                 [0.011465, 0.009828, 0.060750],
                 [0.013708, 0.011771, 0.068667],
                 [0.016156, 0.013840, 0.076603],
                 [0.018815, 0.016026, 0.084584],
                 [0.021692, 0.018320, 0.092610],
                 [0.024792, 0.020715, 0.100676],
                 [0.028123, 0.023201, 0.108787],
                 [0.031696, 0.025765, 0.116965],
                 [0.035520, 0.028397, 0.125209],
                 [0.039608, 0.031090, 0.133515],
                 [0.043830, 0.033830, 0.141886],
                 [0.048062, 0.036607, 0.150327],
                 [0.052320, 0.039407, 0.158841],
                 [0.056615, 0.042160, 0.167446],
                 [0.060949, 0.044794, 0.176129],
                 [0.065330, 0.047318, 0.184892],
                 [0.069764, 0.049726, 0.193735],
                 [0.074257, 0.052017, 0.202660],
                 [0.078815, 0.054184, 0.211667],
                 [0.083446, 0.056225, 0.220755],
                 [0.088155, 0.058133, 0.229922],
                 [0.092949, 0.059904, 0.239164],
                 [0.097833, 0.061531, 0.248477],
                 [0.102815, 0.063010, 0.257854],
                 [0.107899, 0.064335, 0.267289],
                 [0.113094, 0.065492, 0.276784],
                 [0.118405, 0.066479, 0.286321],
                 [0.123833, 0.067295, 0.295879],
                 [0.129380, 0.067935, 0.305443],
                 [0.135053, 0.068391, 0.315000],
                 [0.140858, 0.068654, 0.324538],
                 [0.146785, 0.068738, 0.334011],
                 [0.152839, 0.068637, 0.343404],
                 [0.159018, 0.068354, 0.352688],
                 [0.165308, 0.067911, 0.361816],
                 [0.171713, 0.067305, 0.370771],
                 [0.178212, 0.066576, 0.379497],
                 [0.184801, 0.065732, 0.387973],
                 [0.191460, 0.064818, 0.396152],
                 [0.198177, 0.063862, 0.404009],
                 [0.204935, 0.062907, 0.411514],
                 [0.211718, 0.061992, 0.418647],
                 [0.218512, 0.061158, 0.425392],
                 [0.225302, 0.060445, 0.431742],
                 [0.232077, 0.059889, 0.437695],
                 [0.238826, 0.059517, 0.443256],
                 [0.245543, 0.059352, 0.448436],
                 [0.252220, 0.059415, 0.453248],
                 [0.258857, 0.059706, 0.457710],
                 [0.265447, 0.060237, 0.461840],
                 [0.271994, 0.060994, 0.465660],
                 [0.278493, 0.061978, 0.469190],
                 [0.284951, 0.063168, 0.472451],
                 [0.291366, 0.064553, 0.475462],
                 [0.297740, 0.066117, 0.478243],
                 [0.304081, 0.067835, 0.480812],
                 [0.310382, 0.069702, 0.483186],
                 [0.316654, 0.071690, 0.485380],
                 [0.322899, 0.073782, 0.487408],
                 [0.329114, 0.075972, 0.489287],
                 [0.335308, 0.078236, 0.491024],
                 [0.341482, 0.080564, 0.492631],
                 [0.347636, 0.082946, 0.494121],
                 [0.353773, 0.085373, 0.495501],
                 [0.359898, 0.087831, 0.496778],
                 [0.366012, 0.090314, 0.497960],
                 [0.372116, 0.092816, 0.499053],
                 [0.378211, 0.095332, 0.500067],
                 [0.384299, 0.097855, 0.501002],
                 [0.390384, 0.100379, 0.501864],
                 [0.396467, 0.102902, 0.502658],
                 [0.402548, 0.105420, 0.503386],
                 [0.408629, 0.107930, 0.504052],
                 [0.414709, 0.110431, 0.504662],
                 [0.420791, 0.112920, 0.505215],
                 [0.426877, 0.115395, 0.505714],
                 [0.432967, 0.117855, 0.506160],
                 [0.439062, 0.120298, 0.506555],
                 [0.445163, 0.122724, 0.506901],
                 [0.451271, 0.125132, 0.507198],
                 [0.457386, 0.127522, 0.507448],
                 [0.463508, 0.129893, 0.507652],
                 [0.469640, 0.132245, 0.507809],
                 [0.475780, 0.134577, 0.507921],
                 [0.481929, 0.136891, 0.507989],
                 [0.488088, 0.139186, 0.508011],
                 [0.494258, 0.141462, 0.507988],
                 [0.500438, 0.143719, 0.507920],
                 [0.506629, 0.145958, 0.507806],
                 [0.512831, 0.148179, 0.507648],
                 [0.519045, 0.150383, 0.507443],
                 [0.525270, 0.152569, 0.507192],
                 [0.531507, 0.154739, 0.506895],
                 [0.537755, 0.156894, 0.506551],
                 [0.544015, 0.159033, 0.506159],
                 [0.550287, 0.161158, 0.505719],
                 [0.556571, 0.163269, 0.505230],
                 [0.562866, 0.165368, 0.504692],
                 [0.569172, 0.167454, 0.504105],
                 [0.575490, 0.169530, 0.503466],
                 [0.581819, 0.171596, 0.502777],
                 [0.588158, 0.173652, 0.502035],
                 [0.594508, 0.175701, 0.501241],
                 [0.600868, 0.177743, 0.500394],
                 [0.607238, 0.179779, 0.499492],
                 [0.613617, 0.181811, 0.498536],
                 [0.620005, 0.183840, 0.497524],
                 [0.626401, 0.185867, 0.496456],
                 [0.632805, 0.187893, 0.495332],
                 [0.639216, 0.189921, 0.494150],
                 [0.645633, 0.191952, 0.492910],
                 [0.652056, 0.193986, 0.491611],
                 [0.658483, 0.196027, 0.490253],
                 [0.664915, 0.198075, 0.488836],
                 [0.671349, 0.200133, 0.487358],
                 [0.677786, 0.202203, 0.485819],
                 [0.684224, 0.204286, 0.484219],
                 [0.690661, 0.206384, 0.482558],
                 [0.697098, 0.208501, 0.480835],
                 [0.703532, 0.210638, 0.479049],
                 [0.709962, 0.212797, 0.477201],
                 [0.716387, 0.214982, 0.475290],
                 [0.722805, 0.217194, 0.473316],
                 [0.729216, 0.219437, 0.471279],
                 [0.735616, 0.221713, 0.469180],
                 [0.742004, 0.224025, 0.467018],
                 [0.748378, 0.226377, 0.464794],
                 [0.754737, 0.228772, 0.462509],
                 [0.761077, 0.231214, 0.460162],
                 [0.767398, 0.233705, 0.457755],
                 [0.773695, 0.236249, 0.455289],
                 [0.779968, 0.238851, 0.452765],
                 [0.786212, 0.241514, 0.450184],
                 [0.792427, 0.244242, 0.447543],
                 [0.798608, 0.247040, 0.444848],
                 [0.804752, 0.249911, 0.442102],
                 [0.810855, 0.252861, 0.439305],
                 [0.816914, 0.255895, 0.436461],
                 [0.822926, 0.259016, 0.433573],
                 [0.828886, 0.262229, 0.430644],
                 [0.834791, 0.265540, 0.427671],
                 [0.840636, 0.268953, 0.424666],
                 [0.846416, 0.272473, 0.421631],
                 [0.852126, 0.276106, 0.418573],
                 [0.857763, 0.279857, 0.415496],
                 [0.863320, 0.283729, 0.412403],
                 [0.868793, 0.287728, 0.409303],
                 [0.874176, 0.291859, 0.406205],
                 [0.879464, 0.296125, 0.403118],
                 [0.884651, 0.300530, 0.400047],
                 [0.889731, 0.305079, 0.397002],
                 [0.894700, 0.309773, 0.393995],
                 [0.899552, 0.314616, 0.391037],
                 [0.904281, 0.319610, 0.388137],
                 [0.908884, 0.324755, 0.385308],
                 [0.913354, 0.330052, 0.382563],
                 [0.917689, 0.335500, 0.379915],
                 [0.921884, 0.341098, 0.377376],
                 [0.925937, 0.346844, 0.374959],
                 [0.929845, 0.352734, 0.372677],
                 [0.933606, 0.358764, 0.370541],
                 [0.937221, 0.364929, 0.368567],
                 [0.940687, 0.371224, 0.366762],
                 [0.944006, 0.377643, 0.365136],
                 [0.947180, 0.384178, 0.363701],
                 [0.950210, 0.390820, 0.362468],
                 [0.953099, 0.397563, 0.361438],
                 [0.955849, 0.404400, 0.360619],
                 [0.958464, 0.411324, 0.360014],
                 [0.960949, 0.418323, 0.359630],
                 [0.963310, 0.425390, 0.359469],
                 [0.965549, 0.432519, 0.359529],
                 [0.967671, 0.439703, 0.359810],
                 [0.969680, 0.446936, 0.360311],
                 [0.971582, 0.454210, 0.361030],
                 [0.973381, 0.461520, 0.361965],
                 [0.975082, 0.468861, 0.363111],
                 [0.976690, 0.476226, 0.364466],
                 [0.978210, 0.483612, 0.366025],
                 [0.979645, 0.491014, 0.367783],
                 [0.981000, 0.498428, 0.369734],
                 [0.982279, 0.505851, 0.371874],
                 [0.983485, 0.513280, 0.374198],
                 [0.984622, 0.520713, 0.376698],
                 [0.985693, 0.528148, 0.379371],
                 [0.986700, 0.535582, 0.382210],
                 [0.987646, 0.543015, 0.385210],
                 [0.988533, 0.550446, 0.388365],
                 [0.989363, 0.557873, 0.391671],
                 [0.990138, 0.565296, 0.395122],
                 [0.990871, 0.572706, 0.398714],
                 [0.991558, 0.580107, 0.402441],
                 [0.992196, 0.587502, 0.406299],
                 [0.992785, 0.594891, 0.410283],
                 [0.993326, 0.602275, 0.414390],
                 [0.993834, 0.609644, 0.418613],
                 [0.994309, 0.616999, 0.422950],
                 [0.994738, 0.624350, 0.427397],
                 [0.995122, 0.631696, 0.431951],
                 [0.995480, 0.639027, 0.436607],
                 [0.995810, 0.646344, 0.441361],
                 [0.996096, 0.653659, 0.446213],
                 [0.996341, 0.660969, 0.451160],
                 [0.996580, 0.668256, 0.456192],
                 [0.996775, 0.675541, 0.461314],
                 [0.996925, 0.682828, 0.466526],
                 [0.997077, 0.690088, 0.471811],
                 [0.997186, 0.697349, 0.477182],
                 [0.997254, 0.704611, 0.482635],
                 [0.997325, 0.711848, 0.488154],
                 [0.997351, 0.719089, 0.493755],
                 [0.997351, 0.726324, 0.499428],
                 [0.997341, 0.733545, 0.505167],
                 [0.997285, 0.740772, 0.510983],
                 [0.997228, 0.747981, 0.516859],
                 [0.997138, 0.755190, 0.522806],
                 [0.997019, 0.762398, 0.528821],
                 [0.996898, 0.769591, 0.534892],
                 [0.996727, 0.776795, 0.541039],
                 [0.996571, 0.783977, 0.547233],
                 [0.996369, 0.791167, 0.553499],
                 [0.996162, 0.798348, 0.559820],
                 [0.995932, 0.805527, 0.566202],
                 [0.995680, 0.812706, 0.572645],
                 [0.995424, 0.819875, 0.579140],
                 [0.995131, 0.827052, 0.585701],
                 [0.994851, 0.834213, 0.592307],
                 [0.994524, 0.841387, 0.598983],
                 [0.994222, 0.848540, 0.605696],
                 [0.993866, 0.855711, 0.612482],
                 [0.993545, 0.862859, 0.619299],
                 [0.993170, 0.870024, 0.626189],
                 [0.992831, 0.877168, 0.633109],
                 [0.992440, 0.884330, 0.640099],
                 [0.992089, 0.891470, 0.647116],
                 [0.991688, 0.898627, 0.654202],
                 [0.991332, 0.905763, 0.661309],
                 [0.990930, 0.912915, 0.668481],
                 [0.990570, 0.920049, 0.675675],
                 [0.990175, 0.927196, 0.682926],
                 [0.989815, 0.934329, 0.690198],
                 [0.989434, 0.941470, 0.697519],
                 [0.989077, 0.948604, 0.704863],
                 [0.988717, 0.955742, 0.712242],
                 [0.988367, 0.962878, 0.719649],
                 [0.988033, 0.970012, 0.727077],
                 [0.987691, 0.977154, 0.734536],
                 [0.987387, 0.984288, 0.742002],
                 [0.987053, 0.991438, 0.749504]]
    return np.asarray(color_map)


def get_viridis_map():
    """
    Return the viridis map, described here: https://bids.github.io/colormap/

    :return: matrix of size 256x3 where each row contains a color in the virdis map
    """
    color_map = [[0.267004, 0.004874, 0.329415],
                 [0.268510, 0.009605, 0.335427],
                 [0.269944, 0.014625, 0.341379],
                 [0.271305, 0.019942, 0.347269],
                 [0.272594, 0.025563, 0.353093],
                 [0.273809, 0.031497, 0.358853],
                 [0.274952, 0.037752, 0.364543],
                 [0.276022, 0.044167, 0.370164],
                 [0.277018, 0.050344, 0.375715],
                 [0.277941, 0.056324, 0.381191],
                 [0.278791, 0.062145, 0.386592],
                 [0.279566, 0.067836, 0.391917],
                 [0.280267, 0.073417, 0.397163],
                 [0.280894, 0.078907, 0.402329],
                 [0.281446, 0.084320, 0.407414],
                 [0.281924, 0.089666, 0.412415],
                 [0.282327, 0.094955, 0.417331],
                 [0.282656, 0.100196, 0.422160],
                 [0.282910, 0.105393, 0.426902],
                 [0.283091, 0.110553, 0.431554],
                 [0.283197, 0.115680, 0.436115],
                 [0.283229, 0.120777, 0.440584],
                 [0.283187, 0.125848, 0.444960],
                 [0.283072, 0.130895, 0.449241],
                 [0.282884, 0.135920, 0.453427],
                 [0.282623, 0.140926, 0.457517],
                 [0.282290, 0.145912, 0.461510],
                 [0.281887, 0.150881, 0.465405],
                 [0.281412, 0.155834, 0.469201],
                 [0.280868, 0.160771, 0.472899],
                 [0.280255, 0.165693, 0.476498],
                 [0.279574, 0.170599, 0.479997],
                 [0.278826, 0.175490, 0.483397],
                 [0.278012, 0.180367, 0.486697],
                 [0.277134, 0.185228, 0.489898],
                 [0.276194, 0.190074, 0.493001],
                 [0.275191, 0.194905, 0.496005],
                 [0.274128, 0.199721, 0.498911],
                 [0.273006, 0.204520, 0.501721],
                 [0.271828, 0.209303, 0.504434],
                 [0.270595, 0.214069, 0.507052],
                 [0.269308, 0.218818, 0.509577],
                 [0.267968, 0.223549, 0.512008],
                 [0.266580, 0.228262, 0.514349],
                 [0.265145, 0.232956, 0.516599],
                 [0.263663, 0.237631, 0.518762],
                 [0.262138, 0.242286, 0.520837],
                 [0.260571, 0.246922, 0.522828],
                 [0.258965, 0.251537, 0.524736],
                 [0.257322, 0.256130, 0.526563],
                 [0.255645, 0.260703, 0.528312],
                 [0.253935, 0.265254, 0.529983],
                 [0.252194, 0.269783, 0.531579],
                 [0.250425, 0.274290, 0.533103],
                 [0.248629, 0.278775, 0.534556],
                 [0.246811, 0.283237, 0.535941],
                 [0.244972, 0.287675, 0.537260],
                 [0.243113, 0.292092, 0.538516],
                 [0.241237, 0.296485, 0.539709],
                 [0.239346, 0.300855, 0.540844],
                 [0.237441, 0.305202, 0.541921],
                 [0.235526, 0.309527, 0.542944],
                 [0.233603, 0.313828, 0.543914],
                 [0.231674, 0.318106, 0.544834],
                 [0.229739, 0.322361, 0.545706],
                 [0.227802, 0.326594, 0.546532],
                 [0.225863, 0.330805, 0.547314],
                 [0.223925, 0.334994, 0.548053],
                 [0.221989, 0.339161, 0.548752],
                 [0.220057, 0.343307, 0.549413],
                 [0.218130, 0.347432, 0.550038],
                 [0.216210, 0.351535, 0.550627],
                 [0.214298, 0.355619, 0.551184],
                 [0.212395, 0.359683, 0.551710],
                 [0.210503, 0.363727, 0.552206],
                 [0.208623, 0.367752, 0.552675],
                 [0.206756, 0.371758, 0.553117],
                 [0.204903, 0.375746, 0.553533],
                 [0.203063, 0.379716, 0.553925],
                 [0.201239, 0.383670, 0.554294],
                 [0.199430, 0.387607, 0.554642],
                 [0.197636, 0.391528, 0.554969],
                 [0.195860, 0.395433, 0.555276],
                 [0.194100, 0.399323, 0.555565],
                 [0.192357, 0.403199, 0.555836],
                 [0.190631, 0.407061, 0.556089],
                 [0.188923, 0.410910, 0.556326],
                 [0.187231, 0.414746, 0.556547],
                 [0.185556, 0.418570, 0.556753],
                 [0.183898, 0.422383, 0.556944],
                 [0.182256, 0.426184, 0.557120],
                 [0.180629, 0.429975, 0.557282],
                 [0.179019, 0.433756, 0.557430],
                 [0.177423, 0.437527, 0.557565],
                 [0.175841, 0.441290, 0.557685],
                 [0.174274, 0.445044, 0.557792],
                 [0.172719, 0.448791, 0.557885],
                 [0.171176, 0.452530, 0.557965],
                 [0.169646, 0.456262, 0.558030],
                 [0.168126, 0.459988, 0.558082],
                 [0.166617, 0.463708, 0.558119],
                 [0.165117, 0.467423, 0.558141],
                 [0.163625, 0.471133, 0.558148],
                 [0.162142, 0.474838, 0.558140],
                 [0.160665, 0.478540, 0.558115],
                 [0.159194, 0.482237, 0.558073],
                 [0.157729, 0.485932, 0.558013],
                 [0.156270, 0.489624, 0.557936],
                 [0.154815, 0.493313, 0.557840],
                 [0.153364, 0.497000, 0.557724],
                 [0.151918, 0.500685, 0.557587],
                 [0.150476, 0.504369, 0.557430],
                 [0.149039, 0.508051, 0.557250],
                 [0.147607, 0.511733, 0.557049],
                 [0.146180, 0.515413, 0.556823],
                 [0.144759, 0.519093, 0.556572],
                 [0.143343, 0.522773, 0.556295],
                 [0.141935, 0.526453, 0.555991],
                 [0.140536, 0.530132, 0.555659],
                 [0.139147, 0.533812, 0.555298],
                 [0.137770, 0.537492, 0.554906],
                 [0.136408, 0.541173, 0.554483],
                 [0.135066, 0.544853, 0.554029],
                 [0.133743, 0.548535, 0.553541],
                 [0.132444, 0.552216, 0.553018],
                 [0.131172, 0.555899, 0.552459],
                 [0.129933, 0.559582, 0.551864],
                 [0.128729, 0.563265, 0.551229],
                 [0.127568, 0.566949, 0.550556],
                 [0.126453, 0.570633, 0.549841],
                 [0.125394, 0.574318, 0.549086],
                 [0.124395, 0.578002, 0.548287],
                 [0.123463, 0.581687, 0.547445],
                 [0.122606, 0.585371, 0.546557],
                 [0.121831, 0.589055, 0.545623],
                 [0.121148, 0.592739, 0.544641],
                 [0.120565, 0.596422, 0.543611],
                 [0.120092, 0.600104, 0.542530],
                 [0.119738, 0.603785, 0.541400],
                 [0.119512, 0.607464, 0.540218],
                 [0.119423, 0.611141, 0.538982],
                 [0.119483, 0.614817, 0.537692],
                 [0.119699, 0.618490, 0.536347],
                 [0.120081, 0.622161, 0.534946],
                 [0.120638, 0.625828, 0.533488],
                 [0.121380, 0.629492, 0.531973],
                 [0.122312, 0.633153, 0.530398],
                 [0.123444, 0.636809, 0.528763],
                 [0.124780, 0.640461, 0.527068],
                 [0.126326, 0.644107, 0.525311],
                 [0.128087, 0.647749, 0.523491],
                 [0.130067, 0.651384, 0.521608],
                 [0.132268, 0.655014, 0.519661],
                 [0.134692, 0.658636, 0.517649],
                 [0.137339, 0.662252, 0.515571],
                 [0.140210, 0.665859, 0.513427],
                 [0.143303, 0.669459, 0.511215],
                 [0.146616, 0.673050, 0.508936],
                 [0.150148, 0.676631, 0.506589],
                 [0.153894, 0.680203, 0.504172],
                 [0.157851, 0.683765, 0.501686],
                 [0.162016, 0.687316, 0.499129],
                 [0.166383, 0.690856, 0.496502],
                 [0.170948, 0.694384, 0.493803],
                 [0.175707, 0.697900, 0.491033],
                 [0.180653, 0.701402, 0.488189],
                 [0.185783, 0.704891, 0.485273],
                 [0.191090, 0.708366, 0.482284],
                 [0.196571, 0.711827, 0.479221],
                 [0.202219, 0.715272, 0.476084],
                 [0.208030, 0.718701, 0.472873],
                 [0.214000, 0.722114, 0.469588],
                 [0.220124, 0.725509, 0.466226],
                 [0.226397, 0.728888, 0.462789],
                 [0.232815, 0.732247, 0.459277],
                 [0.239374, 0.735588, 0.455688],
                 [0.246070, 0.738910, 0.452024],
                 [0.252899, 0.742211, 0.448284],
                 [0.259857, 0.745492, 0.444467],
                 [0.266941, 0.748751, 0.440573],
                 [0.274149, 0.751988, 0.436601],
                 [0.281477, 0.755203, 0.432552],
                 [0.288921, 0.758394, 0.428426],
                 [0.296479, 0.761561, 0.424223],
                 [0.304148, 0.764704, 0.419943],
                 [0.311925, 0.767822, 0.415586],
                 [0.319809, 0.770914, 0.411152],
                 [0.327796, 0.773980, 0.406640],
                 [0.335885, 0.777018, 0.402049],
                 [0.344074, 0.780029, 0.397381],
                 [0.352360, 0.783011, 0.392636],
                 [0.360741, 0.785964, 0.387814],
                 [0.369214, 0.788888, 0.382914],
                 [0.377779, 0.791781, 0.377939],
                 [0.386433, 0.794644, 0.372886],
                 [0.395174, 0.797475, 0.367757],
                 [0.404001, 0.800275, 0.362552],
                 [0.412913, 0.803041, 0.357269],
                 [0.421908, 0.805774, 0.351910],
                 [0.430983, 0.808473, 0.346476],
                 [0.440137, 0.811138, 0.340967],
                 [0.449368, 0.813768, 0.335384],
                 [0.458674, 0.816363, 0.329727],
                 [0.468053, 0.818921, 0.323998],
                 [0.477504, 0.821444, 0.318195],
                 [0.487026, 0.823929, 0.312321],
                 [0.496615, 0.826376, 0.306377],
                 [0.506271, 0.828786, 0.300362],
                 [0.515992, 0.831158, 0.294279],
                 [0.525776, 0.833491, 0.288127],
                 [0.535621, 0.835785, 0.281908],
                 [0.545524, 0.838039, 0.275626],
                 [0.555484, 0.840254, 0.269281],
                 [0.565498, 0.842430, 0.262877],
                 [0.575563, 0.844566, 0.256415],
                 [0.585678, 0.846661, 0.249897],
                 [0.595839, 0.848717, 0.243329],
                 [0.606045, 0.850733, 0.236712],
                 [0.616293, 0.852709, 0.230052],
                 [0.626579, 0.854645, 0.223353],
                 [0.636902, 0.856542, 0.216620],
                 [0.647257, 0.858400, 0.209861],
                 [0.657642, 0.860219, 0.203082],
                 [0.668054, 0.861999, 0.196293],
                 [0.678489, 0.863742, 0.189503],
                 [0.688944, 0.865448, 0.182725],
                 [0.699415, 0.867117, 0.175971],
                 [0.709898, 0.868751, 0.169257],
                 [0.720391, 0.870350, 0.162603],
                 [0.730889, 0.871916, 0.156029],
                 [0.741388, 0.873449, 0.149561],
                 [0.751884, 0.874951, 0.143228],
                 [0.762373, 0.876424, 0.137064],
                 [0.772852, 0.877868, 0.131109],
                 [0.783315, 0.879285, 0.125405],
                 [0.793760, 0.880678, 0.120005],
                 [0.804182, 0.882046, 0.114965],
                 [0.814576, 0.883393, 0.110347],
                 [0.824940, 0.884720, 0.106217],
                 [0.835270, 0.886029, 0.102646],
                 [0.845561, 0.887322, 0.099702],
                 [0.855810, 0.888601, 0.097452],
                 [0.866013, 0.889868, 0.095953],
                 [0.876168, 0.891125, 0.095250],
                 [0.886271, 0.892374, 0.095374],
                 [0.896320, 0.893616, 0.096335],
                 [0.906311, 0.894855, 0.098125],
                 [0.916242, 0.896091, 0.100717],
                 [0.926106, 0.897330, 0.104071],
                 [0.935904, 0.898570, 0.108131],
                 [0.945636, 0.899815, 0.112838],
                 [0.955300, 0.901065, 0.118128],
                 [0.964894, 0.902323, 0.123941],
                 [0.974417, 0.903590, 0.130215],
                 [0.983868, 0.904867, 0.136897],
                 [0.993248, 0.906157, 0.143936]]
    return np.asarray(color_map)

thirdparty/kplanes/ops/flip/data.py
oppo-us-research/NeuRBF
import glob
import os
from typing import List, Optional

import cv2
import torch
import torchvision.transforms.functional as TF
from PIL import Image
import logging as log
import numpy as np
import imageio.v3 as iio
import imageio


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.
    """
    Image.fromarray(data).save(path)


def read_png(file_name: str, resize_h: Optional[int] = None, resize_w: Optional[int] = None) -> torch.Tensor:
    """Reads a PNG image from path, potentially resizing it.
    """
    img = Image.open(file_name).convert('RGB')  # PIL outputs BGR by default
    if resize_h is not None and resize_w is not None:
        img.resize((resize_w, resize_h), Image.LANCZOS)
    img = TF.to_tensor(img)  # TF converts to C, H, W
    img = img.permute(1, 2, 0).contiguous()  # H, W, C
    return img


def glob_imgs(path, exts=None):
    """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.
    """
    if exts is None:
        exts = ['*.png', '*.PNG', '*.jpg', '*.jpeg', '*.JPG', '*.JPEG']
    imgs = []
    for ext in exts:
        imgs.extend(glob.glob(os.path.join(path, ext)))
    return imgs


def write_video_to_file(file_name, frames: List[np.ndarray]):
    log.info(f"Saving video ({len(frames)} frames) to {file_name}")
    # Photo tourism image sizes differ
    sizes = np.array([frame.shape[:2] for frame in frames])
    same_size_frames = np.unique(sizes, axis=0).shape[0] == 1
    if same_size_frames:
        imageio.mimwrite(file_name, frames, fps=30, quality=10, macro_block_size=None)
    else:
        height = sizes[:, 0].max()
        width = sizes[:, 1].max()
        video = cv2.VideoWriter(
            file_name, cv2.VideoWriter_fourcc(*'mp4v'), 5, (width, height))
        for img in frames:
            image = np.zeros((height, width, 3), dtype=np.uint8)
            h, w = img.shape[:2]
            if isinstance(img, torch.Tensor):
                img = img.numpy()
            image[(height-h)//2:(height-h)//2+h, (width-w)//2:(width-w)//2+w, :] = img
            video.write(image[:, :, ::-1])  # opencv uses BGR instead of RGB
        cv2.destroyAllWindows()
        video.release()


def read_mp4(file_name: str) -> List[torch.Tensor]:
    all_frames = iio.imread(
        file_name, plugin='pyav', format='rgb24', constant_framerate=True, thread_count=2
    )
    out_frames = [torch.from_numpy(f) for f in all_frames]
    return out_frames

thirdparty/kplanes/ops/image/io.py
oppo-us-research/NeuRBF
import math
import tempfile
import subprocess
import os
import re
from typing import List

import numpy as np
import skimage.metrics
import torch
from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure
import lpips

from .io import write_video_to_file, write_png


ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(data_range=1.0)

""" 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.Tensor): Image tensor of shape [H,W3]

    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 * math.log10(1.0 / mse)


def ssim(rgb, gts):
    """
    Modified from https://github.com/google/mipnerf/blob/16e73dfdb52044dcceb47cda5243a686391a6e0f/internal/math.py#L58
    """
    filter_size = 11
    filter_sigma = 1.5
    k1 = 0.01
    k2 = 0.03
    max_val = 1.0
    rgb = rgb.cpu().numpy()
    gts = gts.cpu().numpy()
    assert len(rgb.shape) == 3
    assert rgb.shape[-1] == 3
    assert rgb.shape == gts.shape
    import scipy.signal

    # Construct a 1D Gaussian blur filter.
    hw = filter_size // 2
    shift = (2 * hw - filter_size + 1) / 2
    f_i = ((np.arange(filter_size) - hw + shift) / filter_sigma)**2
    filt = np.exp(-0.5 * f_i)
    filt /= np.sum(filt)

    # Blur in x and y (faster than the 2D convolution).
    def convolve2d(z, f):
        return scipy.signal.convolve2d(z, f, mode='valid')

    filt_fn = lambda z: np.stack([
        convolve2d(convolve2d(z[..., i], filt[:, None]), filt[None, :])
        for i in range(z.shape[-1])], -1)
    mu0 = filt_fn(rgb)
    mu1 = filt_fn(gts)
    mu00 = mu0 * mu0
    mu11 = mu1 * mu1
    mu01 = mu0 * mu1
    sigma00 = filt_fn(rgb**2) - mu00
    sigma11 = filt_fn(gts**2) - mu11
    sigma01 = filt_fn(rgb * gts) - mu01

    # Clip the variances and covariances to valid values.
    # Variance must be non-negative:
    sigma00 = np.maximum(0., sigma00)
    sigma11 = np.maximum(0., sigma11)
    sigma01 = np.sign(sigma01) * np.minimum(
        np.sqrt(sigma00 * sigma11), np.abs(sigma01))
    c1 = (k1 * max_val)**2
    c2 = (k2 * max_val)**2
    numer = (2 * mu01 + c1) * (2 * sigma01 + c2)
    denom = (mu00 + mu11 + c1) * (sigma00 + sigma11 + c2)
    ssim_map = numer / denom
    return np.mean(ssim_map)


def ssim_old(rgb, gts):
    """Calculate the SSIM metric.

    Assumes the RGB image is in [0,1]

    Args:
        rgb (torch.Tensor): Image tensor of shape [H,W,3]
        gts (torch.Tensor): 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(),
        channel_axis=2,
        gaussian_weights=True,
        sigma=1.5,
        use_sample_covariance=False)


def msssim(rgb, gts):
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    return ms_ssim(torch.permute(rgb[None, ...], (0, 3, 1, 2)),
                   torch.permute(gts[None, ...], (0, 3, 1, 2))).item()


__LPIPS__ = {}


def init_lpips(net_name, device):
    return lpips.LPIPS(net=net_name, version='0.1').eval().to(device)


def rgb_lpips(rgb, gts, net_name='alex', device='cpu'):
    if net_name not in __LPIPS__:
        __LPIPS__[net_name] = init_lpips(net_name, device)
    gts = gts.permute([2, 0, 1]).contiguous().to(device)
    rgb = rgb.permute([2, 0, 1]).contiguous().to(device)
    return __LPIPS__[net_name](gts, rgb, normalize=True).item()


def jod(pred_frames: List[np.ndarray], gt_frames: List[np.ndarray]) -> float:
    with tempfile.TemporaryDirectory() as tmpdir:
        file_pred = os.path.join(tmpdir, "pred.mp4")
        write_video_to_file(file_pred, pred_frames)
        file_gt = os.path.join(tmpdir, "gt.mp4")
        write_video_to_file(file_gt, gt_frames)
        result = subprocess.check_output([
            'fvvdp', '--test', file_pred, '--ref', file_gt, '--gpu', '0',
            '--display', 'standard_fhd'])
        result = float(result.decode().strip().split('=')[1])
    return result


def flip(pred_frames: List[np.ndarray], gt_frames: List[np.ndarray], interval: int = 10) -> float:
    def extract_from_result(text: str, prompt: str):
        m = re.search(prompt, text)
        return float(m.group(1))

    all_results = []
    with tempfile.TemporaryDirectory() as tmpdir:
        pred_fname = os.path.join(tmpdir, "pred.png")
        gt_fname = os.path.join(tmpdir, "gt.png")
        for i in range(len(pred_frames)):
            if (i % interval) != 0:
                continue
            write_png(pred_fname, pred_frames[i])
            write_png(gt_fname, gt_frames[i])
            result = subprocess.check_output(
                ['python', 'thirdparty/kplanes/ops/flip/flip.py', '--reference', gt_fname, '--test', pred_fname]
            ).decode()
            all_results.append(extract_from_result(result, r'Mean: (\d+\.\d+)'))
    return sum(all_results) / len(all_results)

thirdparty/kplanes/ops/image/metrics.py
oppo-us-research/NeuRBF
import torch
from torch.nn import functional as F


def grid_sample_wrapper(grid: torch.Tensor, coords: torch.Tensor, align_corners: bool = True) -> torch.Tensor:
    grid_dim = coords.shape[-1]

    if grid.dim() == grid_dim + 1:
        # no batch dimension present, need to add it
        grid = grid.unsqueeze(0)
    if coords.dim() == 2:
        coords = coords.unsqueeze(0)

    if grid_dim == 2 or grid_dim == 3:
        grid_sampler = F.grid_sample
    else:
        raise NotImplementedError(f"Grid-sample was called with {grid_dim}D data but is only "
                                  f"implemented for 2 and 3D data.")

    coords = coords.view([coords.shape[0]] + [1] * (grid_dim - 1) + list(coords.shape[1:]))
    B, feature_dim = grid.shape[:2]
    n = coords.shape[-2]
    interp = grid_sampler(
        grid,  # [B, feature_dim, reso, ...]
        coords,  # [B, 1, ..., n, grid_dim]
        align_corners=align_corners,
        mode='bilinear', padding_mode='border')
    interp = interp.view(B, feature_dim, n).transpose(-1, -2)  # [B, n, feature_dim]
    interp = interp.squeeze()  # [B?, n, feature_dim?]
    return interp

thirdparty/kplanes/ops/interpolation.py
oppo-us-research/NeuRBF
import torch


def ray_samples_to_sdist(ray_samples):
    """Convert ray samples to s space"""
    starts = ray_samples.spacing_starts
    ends = ray_samples.spacing_ends
    sdist = torch.cat([starts[..., 0], ends[..., -1:, 0]], dim=-1)  # (num_rays, num_samples + 1)
    return sdist


def lossfun_distortion(t, w):
    """
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L142
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L266
    """
    ut = (t[..., 1:] + t[..., :-1]) / 2
    dut = torch.abs(ut[..., :, None] - ut[..., None, :])
    loss_inter = torch.sum(w * torch.sum(w[..., None, :] * dut, dim=-1), dim=-1)

    loss_intra = torch.sum(w**2 * (t[..., 1:] - t[..., :-1]), dim=-1) / 3

    return loss_inter + loss_intra


def distortion_loss(weights_list, ray_samples_list):
    """From mipnerf360"""
    # takes samples and weights as proposed by the 'nerf' network
    c = ray_samples_to_sdist(ray_samples_list[-1])
    w = weights_list[-1][..., 0]
    loss = torch.mean(lossfun_distortion(c, w))
    return loss

thirdparty/kplanes/ops/losses/distortion_loss.py
oppo-us-research/NeuRBF
"""
Implements the histogram loss from mipnerf360
"""
import torch


EPS = 1.0e-7


def outer(
    t0_starts: torch.Tensor,
    t0_ends: torch.Tensor,
    t1_starts: torch.Tensor,
    t1_ends: torch.Tensor,
    y1: torch.Tensor,
) -> torch.Tensor:
    """Faster version of
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L117
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L64
    Args:
        t0_starts: start of the interval edges
        t0_ends: end of the interval edges
        t1_starts: start of the interval edges
        t1_ends: end of the interval edges
        y1: weights
    """
    cy1 = torch.cat([torch.zeros_like(y1[..., :1]), torch.cumsum(y1, dim=-1)], dim=-1)

    idx_lo = torch.searchsorted(t1_starts.contiguous(), t0_starts.contiguous(), right=True) - 1
    idx_lo = torch.clamp(idx_lo, min=0, max=y1.shape[-1] - 1)
    idx_hi = torch.searchsorted(t1_ends.contiguous(), t0_ends.contiguous(), right=True)
    idx_hi = torch.clamp(idx_hi, min=0, max=y1.shape[-1] - 1)
    cy1_lo = torch.take_along_dim(cy1[..., :-1], idx_lo, dim=-1)
    cy1_hi = torch.take_along_dim(cy1[..., 1:], idx_hi, dim=-1)
    y0_outer = cy1_hi - cy1_lo

    return y0_outer


def lossfun_outer(
    t: torch.Tensor,  # [..., "num_samples+1"],
    w: torch.Tensor,  # [..., "num_samples"],
    t_env: torch.Tensor,  # [..., "num_samples+1"],
    w_env: torch.Tensor,  # [..., "num_samples"],
):
    """
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L136
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L80
    Args:
        t: interval edges
        w: weights
        t_env: interval edges of the upper bound enveloping histogram (from proposal model)
        w_env: weights that should upper bound the inner (t,w) histogram (from proposal model)
    """
    w_outer = outer(t[..., :-1], t[..., 1:], t_env[..., :-1], t_env[..., 1:], w_env)
    return torch.clip(w - w_outer, min=0) ** 2 / (w + EPS)


def ray_samples_to_sdist(ray_samples):
    """Convert ray samples to s space"""
    starts = ray_samples.spacing_starts
    ends = ray_samples.spacing_ends
    sdist = torch.cat([starts[..., 0], ends[..., -1:, 0]], dim=-1)  # (num_rays, num_samples + 1)
    return sdist


def interlevel_loss(weights_list, ray_samples_list):
    """Calculates the proposal loss in the MipNeRF-360 paper.
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/model.py#L515
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/train_utils.py#L133
    """
    c = ray_samples_to_sdist(ray_samples_list[-1]).detach()
    w = weights_list[-1][..., 0].detach()
    loss_interlevel = 0.0
    for ray_samples, weights in zip(ray_samples_list[:-1], weights_list[:-1]):
        sdist = ray_samples_to_sdist(ray_samples)
        cp = sdist  # (num_rays, num_samples + 1)
        wp = weights[..., 0]  # (num_rays, num_samples)
        loss_interlevel += torch.mean(lossfun_outer(c, w, cp, wp))
    return loss_interlevel

thirdparty/kplanes/ops/losses/histogram_loss.py
oppo-us-research/NeuRBF
import math

import torch
import torch.optim


def get_cosine_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    eta_min: float = 0.0,
    num_cycles: float = 0.999,
    last_epoch: int = -1,
):
    """
    https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/optimization.py#L129
    """
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(eta_min, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)


def get_log_linear_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    eta_min: float = 2e-5,
    eta_max: float = 1e-2,
    num_cycles: float = 0.999,
    last_epoch: int = -1,
):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))  # in [0,1]
        return math.exp(progress * math.log(eta_min) + (1 - progress) * math.log(eta_max)) / eta_max

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)


def get_step_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    milestones,
    gamma: float,
    num_warmup_steps: int,
    last_epoch: int = -1,
):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        out = 1.0
        for m in milestones:
            if current_step < m:
                break
            out *= gamma
        return out
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)

thirdparty/kplanes/ops/lr_scheduling.py
oppo-us-research/NeuRBF
"""The ray samplers are almost completely copied from NeRF-studio

https://github.com/nerfstudio-project/nerfstudio/blob/628e4fe1a638e7fb3b7ad33d4d91a4b1d63a9b68/nerfstudio/model_components/ray_samplers.py

Copyright 2022 The Nerfstudio Team. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from abc import abstractmethod
from dataclasses import dataclass
from typing import Optional, Callable, Tuple, List

import torch
from torch import nn


@dataclass
class RaySamples:
    """xyz coordinate for ray origin."""
    origins: torch.Tensor  # [bs:..., 3]
    """Direction of ray."""
    directions: torch.Tensor  # [bs:..., 3]
    """Where the frustum starts along a ray."""
    starts: torch.Tensor  # [bs:..., 1]
    """Where the frustum ends along a ray."""
    ends: torch.Tensor  # [bs:..., 1]
    """"width" of each sample."""
    deltas: Optional[torch.Tensor] = None  # [bs, ...?, 1]
    """Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling."""
    spacing_starts: Optional[torch.Tensor] = None  # [bs, ...?, num_samples, 1]
    """End of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling."""
    spacing_ends: Optional[torch.Tensor] = None  # [bs, ...?, num_samples, 1]
    """Function to convert bins to euclidean distance."""
    spacing_to_euclidean_fn: Optional[Callable] = None

    def get_positions(self) -> torch.Tensor:
        """Calulates "center" position of frustum. Not weighted by mass.
        Returns:
            xyz positions (..., 3).
        """
        return self.origins + self.directions * (self.starts + self.ends) / 2  # world space

    def get_weights2(self, densities: torch.Tensor) -> torch.Tensor:
        densities = densities.squeeze(2)
        deltas = self.deltas.squeeze(2)
        delta_mask = deltas > 0
        deltas = deltas[delta_mask]

        delta_density = torch.zeros_like(densities)
        delta_density[delta_mask] = deltas * densities[delta_mask]
        alphas = 1 - torch.exp(-delta_density)

        transmittance = torch.cat(
            (
                torch.ones(alphas.shape[0], 1, device=alphas.device),
                torch.cumprod(1.0 - alphas, dim=-1)
            ), dim=-1
        )
        weights = alphas * transmittance[:, :-1]
        return weights[..., None]

    def get_weights(self, densities: torch.Tensor) -> torch.Tensor:
        """Return weights based on predicted densities
        Args:
            densities: Predicted densities for samples along ray (..., num_samples, 1)
        Returns:
            Weights for each sample  (..., num_samples, 1)
        """
        delta_mask = self.deltas > 0
        deltas = self.deltas[delta_mask]

        delta_density = torch.zeros_like(densities)
        delta_density[delta_mask] = deltas * densities[delta_mask]
        alphas = 1 - torch.exp(-delta_density)

        transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)
        transmittance = torch.cat(
            [torch.zeros((*transmittance.shape[:1], 1, 1), device=densities.device), transmittance], dim=-2
        )
        transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
        weights = alphas * transmittance  # [..., "num_samples"]
        return weights


@dataclass
class RayBundle:
    """A bundle of ray parameters."""

    """Ray origins (XYZ)"""
    origins: torch.Tensor  # [..., 3]
    """Unit ray direction vector"""
    directions: torch.Tensor  # [..., 3]
    """Distance along ray to start sampling"""
    nears: Optional[torch.Tensor] = None  # [..., 1]
    """Rays Distance along ray to stop sampling"""
    fars: Optional[torch.Tensor] = None  # [..., 1]

    def __len__(self):
        num_rays = torch.numel(self.origins) // self.origins.shape[-1]
        return num_rays

    def get_ray_samples(
        self,
        bin_starts: torch.Tensor,
        bin_ends: torch.Tensor,
        spacing_starts: Optional[torch.Tensor] = None,
        spacing_ends: Optional[torch.Tensor] = None,
        spacing_to_euclidean_fn: Optional[Callable] = None,
    ) -> RaySamples:
        """Produces samples for each ray by projection points along the ray direction. Currently samples uniformly.
        Args:
            bin_starts: Distance from origin to start of bin.
                TensorType["bs":..., "num_samples", 1]
            bin_ends: Distance from origin to end of bin.
        Returns:
            Samples projected along ray.
        """
        deltas = bin_ends - bin_starts
        return RaySamples(
            origins=self.origins[..., None, :],  # [..., 1, 3]
            directions=self.directions[..., None, :],  # [..., 1, 3]
            starts=bin_starts,  # [..., num_samples, 1]  world
            ends=bin_ends,  # [..., num_samples, 1]      world
            deltas=deltas,  # [..., num_samples, 1]  world coo
            spacing_starts=spacing_starts,  # [..., num_samples, 1]
            spacing_ends=spacing_ends,  # [..., num_samples, 1]
            spacing_to_euclidean_fn=spacing_to_euclidean_fn,
        )


class Sampler(nn.Module):
    """Generate Samples
    Args:
        num_samples: number of samples to take
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
    ) -> None:
        super().__init__()
        self.num_samples = num_samples

    @abstractmethod
    def generate_ray_samples(self) -> RaySamples:
        """Generate Ray Samples"""

    def forward(self, *args, **kwargs) -> RaySamples:
        """Generate ray samples"""
        return self.generate_ray_samples(*args, **kwargs)


class SpacedSampler(Sampler):
    """Sample points according to a function.
    Args:
        num_samples: Number of samples per ray
        spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).
        spacing_fn_inv: The inverse of spacing_fn.
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        spacing_fn: Callable,
        spacing_fn_inv: Callable,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(num_samples=num_samples)
        self.train_stratified = train_stratified
        self.single_jitter = single_jitter
        self.spacing_fn = spacing_fn
        self.spacing_fn_inv = spacing_fn_inv

    # noinspection PyMethodOverriding
    def generate_ray_samples(
        self,
        ray_bundle: RayBundle,
        num_samples: Optional[int] = None,
    ) -> RaySamples:
        """Generates position samples accoring to spacing function.
        Args:
            ray_bundle: Ray-origins, directions, etc.
            num_samples: Number of samples per ray
        Returns:
            Positions and deltas for samples along a ray
        """
        num_samples = num_samples or self.num_samples
        assert num_samples is not None
        num_rays = ray_bundle.origins.shape[0]

        bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)[None, ...]  # [1, num_samples+1]

        # TODO More complicated than it needs to be.
        if self.train_stratified and self.training:
            if self.single_jitter:
                t_rand = torch.rand((num_rays, 1), dtype=bins.dtype, device=bins.device)
            else:
                t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device)
            bin_centers = (bins[..., 1:] + bins[..., :-1]) / 2.0
            bin_upper = torch.cat([bin_centers, bins[..., -1:]], -1)
            bin_lower = torch.cat([bins[..., :1], bin_centers], -1)
            bins = bin_lower + (bin_upper - bin_lower) * t_rand
        else:
            bins = bins.repeat(num_rays, 1)

        # s_near, s_far in [0, 1]
        s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars))
        spacing_to_euclidean_fn = lambda x: self.spacing_fn_inv(x * s_far + (1 - x) * s_near)
        # euclidean = world
        euclidean_bins = spacing_to_euclidean_fn(bins)  # [num_rays, num_samples+1]

        return ray_bundle.get_ray_samples(
            bin_starts=euclidean_bins[..., :-1, None],  # world [near, far]
            bin_ends=euclidean_bins[..., 1:, None],     # world [near, far]
            spacing_starts=bins[..., :-1, None],        # [0, 1]
            spacing_ends=bins[..., 1:, None],           # [0, 1]
            spacing_to_euclidean_fn=spacing_to_euclidean_fn,
        )


class UniformSampler(SpacedSampler):
    """Sample uniformly along a ray
    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: x,
            spacing_fn_inv=lambda x: x,
            train_stratified=train_stratified,
            single_jitter=single_jitter,
        )


class LinearDisparitySampler(SpacedSampler):
    """Sample linearly in disparity along a ray
    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: 1 / x,
            spacing_fn_inv=lambda x: 1 / x,
            train_stratified=train_stratified,
            single_jitter=single_jitter,
        )


class UniformLinDispPiecewiseSampler(SpacedSampler):
    """Piecewise sampler along a ray that allocates the first half of the samples uniformly and the second half
    using linearly in disparity spacing.
    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: torch.where(x < 1, x / 2, 1 - 1 / (2 * x)),
            spacing_fn_inv=lambda x: torch.where(x < 0.5, 2 * x, 1 / (2 - 2 * x)),
            train_stratified=train_stratified,
            single_jitter=single_jitter,
        )


class PDFSampler(Sampler):
    """Sample based on probability distribution
    Args:
        num_samples: Number of samples per ray
        train_stratified: Randomize location within each bin during training.
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
        include_original: Add original samples to ray.
        histogram_padding: Amount to weights prior to computing PDF.
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified: bool = True,
        single_jitter: bool = False,
        include_original: bool = True,
        histogram_padding: float = 0.01,
    ) -> None:
        super().__init__(num_samples=num_samples)
        self.train_stratified = train_stratified
        self.include_original = include_original
        self.histogram_padding = histogram_padding
        self.single_jitter = single_jitter

    # noinspection PyMethodOverriding
    def generate_ray_samples(
        self,
        ray_bundle: RayBundle,
        ray_samples: Optional[RaySamples] = None,
        weights: Optional[torch.Tensor] = None,
        num_samples: Optional[int] = None,
        eps: float = 1e-5,
    ) -> RaySamples:
        """Generates position samples given a distribution.
        Args:
            ray_bundle: Ray-origins, directions, etc.
            ray_samples: Existing ray samples
            weights: Weights for each bin  [..., "num_samples", 1]
            num_samples: Number of samples per ray
            eps: Small value to prevent numerical issues.
        Returns:
            Positions and deltas for samples along a ray
        """
        if ray_samples is None or ray_bundle is None:
            raise ValueError("ray_samples must be provided")

        num_samples = num_samples or self.num_samples
        assert num_samples is not None
        num_bins = num_samples + 1

        weights = weights[..., 0] + self.histogram_padding

        # Add small offset to rays with zero weight to prevent NaNs
        weights_sum = torch.sum(weights, dim=-1, keepdim=True)
        padding = torch.relu(eps - weights_sum)
        weights = weights + padding / weights.shape[-1]
        weights_sum += padding

        pdf = weights / weights_sum
        cdf = torch.min(torch.ones_like(pdf), torch.cumsum(pdf, dim=-1))
        cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)

        if self.train_stratified and self.training:
            # Stratified samples between 0 and 1
            u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
            u = u.expand((*cdf.shape[:-1], num_bins))
            if self.single_jitter:
                rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins
            else:
                rand = torch.rand((*cdf.shape[:-1], num_samples + 1), device=cdf.device) / num_bins
            u = u + rand
        else:
            # Uniform samples between 0 and 1
            u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
            u = u + 1.0 / (2 * num_bins)
            u = u.expand(size=(*cdf.shape[:-1], num_bins))
        u = u.contiguous()

        assert (
            ray_samples.spacing_starts is not None and ray_samples.spacing_ends is not None
        ), "ray_sample spacing_starts and spacing_ends must be provided"
        assert ray_samples.spacing_to_euclidean_fn is not None, "ray_samples.spacing_to_euclidean_fn must be provided"
        existing_bins = torch.cat(
            [
                ray_samples.spacing_starts[..., 0],
                ray_samples.spacing_ends[..., -1:, 0],
            ],
            dim=-1,
        )  # [0, 1]

        inds = torch.searchsorted(cdf, u, right=True)
        below = torch.clamp(inds - 1, 0, existing_bins.shape[-1] - 1)
        above = torch.clamp(inds, 0, existing_bins.shape[-1] - 1)
        cdf_g0 = torch.gather(cdf, -1, below)
        bins_g0 = torch.gather(existing_bins, -1, below)
        cdf_g1 = torch.gather(cdf, -1, above)
        bins_g1 = torch.gather(existing_bins, -1, above)

        t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1)
        bins = bins_g0 + t * (bins_g1 - bins_g0)

        if self.include_original:
            bins, _ = torch.sort(torch.cat([existing_bins, bins], -1), -1)

        # Stop gradients
        bins = bins.detach()

        euclidean_bins = ray_samples.spacing_to_euclidean_fn(bins)

        return ray_bundle.get_ray_samples(
            bin_starts=euclidean_bins[..., :-1, None],
            bin_ends=euclidean_bins[..., 1:, None],
            spacing_starts=bins[..., :-1, None],
            spacing_ends=bins[..., 1:, None],
            spacing_to_euclidean_fn=ray_samples.spacing_to_euclidean_fn,
        )


class ProposalNetworkSampler(Sampler):
    """Sampler that uses a proposal network to generate samples."""
    def __init__(
        self,
        num_proposal_samples_per_ray: Tuple[int] = (64,),
        num_nerf_samples_per_ray: int = 32,
        num_proposal_network_iterations: int = 2,
        single_jitter: bool = False,
        update_sched: Callable = lambda x: 1,
        initial_sampler: Optional[Sampler] = None,
    ) -> None:
        super().__init__()
        self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
        self.num_nerf_samples_per_ray = num_nerf_samples_per_ray
        self.num_proposal_network_iterations = num_proposal_network_iterations
        self.update_sched = update_sched
        if self.num_proposal_network_iterations < 1:
            raise ValueError("num_proposal_network_iterations must be >= 1")

        # samplers
        if initial_sampler is None:
            initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
        self.initial_sampler = initial_sampler
        self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)

        self._anneal = 1.0
        self._steps_since_update = 0
        self._step = 0

    def set_anneal(self, anneal: float) -> None:
        """Set the anneal value for the proposal network."""
        self._anneal = anneal

    def step_cb(self, step):
        """Callback to register a training step has passed. This is used to keep track of the sampling schedule"""
        self._step = step
        self._steps_since_update += 1

    def generate_ray_samples(
        self,
        ray_bundle: Optional[RayBundle] = None,
        timestamps: Optional[float] = None,
        density_fns: Optional[List[Callable]] = None,
    ) -> Tuple[RaySamples, List, List]:
        assert ray_bundle is not None
        assert density_fns is not None
        assert len(density_fns) == self.num_proposal_network_iterations

        weights_list = []
        ray_samples_list = []

        n = self.num_proposal_network_iterations
        weights = None
        ray_samples = None
        updated = self._steps_since_update > self.update_sched(self._step) or self._step < 10
        for i_level in range(n + 1):
            is_prop = i_level < n
            num_samples = self.num_proposal_samples_per_ray[i_level] if is_prop else self.num_nerf_samples_per_ray
            if i_level == 0:
                # Uniform sampling because we need to start with some samples
                ray_samples = self.initial_sampler(ray_bundle, num_samples=num_samples)
            else:
                # PDF sampling based on the last samples and their weights
                # Perform annealing to the weights. This will be a no-op if self._anneal is 1.0.
                assert weights is not None
                annealed_weights = torch.pow(weights, self._anneal)
                ray_samples = self.pdf_sampler(ray_bundle, ray_samples, annealed_weights, num_samples=num_samples)
            if is_prop:
                if updated:
                    # always update on the first step or the inf check in grad scaling crashes
                    density = density_fns[i_level](ray_samples.get_positions(), timestamps)  # world space
                else:
                    with torch.no_grad():
                        density = density_fns[i_level](ray_samples.get_positions(), timestamps)
                weights = ray_samples.get_weights(density)
                weights_list.append(weights)  # (num_rays, num_samples)
                ray_samples_list.append(ray_samples)
        if updated:
            self._steps_since_update = 0

        assert ray_samples is not None
        return ray_samples, weights_list, ray_samples_list

    def __str__(self):
        return (f"ProposalNetworkSampler("
                f"num_proposal_samples_per_ray={self.num_proposal_samples_per_ray}, "
                f"num_nerf_samples_per_ray={self.num_nerf_samples_per_ray}, "
                f"num_proposal_network_iterations={self.num_proposal_network_iterations})")

thirdparty/kplanes/raymarching/ray_samplers.py
oppo-us-research/NeuRBF
from typing import Optional, Union

import torch
import torch.nn as nn


class SpatialDistortion(nn.Module):
    """Apply spatial distortions"""

    def forward(
        self, positions: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            positions: Sample to distort (shape: batch-size, ..., 3)
        Returns:
            distorted sample - same shape
        """


class SceneContraction(SpatialDistortion):
    """Contract unbounded space using the contraction was proposed in MipNeRF-360.
        We use the following contraction equation:
        .. math::
            f(x) = \\begin{cases}
                x & ||x|| \\leq 1 \\\\
                (2 - \\frac{1}{||x||})(\\frac{x}{||x||}) & ||x|| > 1
            \\end{cases}
        If the order is not specified, we use the Frobenius norm, this will contract the space to a sphere of
        radius 1. If the order is L_inf (order=float("inf")), we will contract the space to a cube of side length 2.
        If using voxel based encodings such as the Hash encoder, we recommend using the L_inf norm.
        Args:
            order: Order of the norm. Default to the Frobenius norm. Must be set to None for Gaussians.
    """

    def __init__(self,
                 order: Optional[Union[float, int]] = None,
                 global_translation: Optional[torch.Tensor] = None,
                 global_scale: Optional[torch.Tensor] = None,
                 ) -> None:
        super().__init__()
        self.order = order
        if global_translation is None:
            global_translation = torch.tensor([0.0, 0.0, 0.0])
        self.global_translation = nn.Parameter(global_translation, requires_grad=False)
        if global_scale is None:
            global_scale = torch.tensor([1.0, 1.0, 1.0])
            
        self.global_scale = nn.Parameter(global_scale, requires_grad=False)

    def forward(self, positions):
        # Apply global scale and translation
        positions = (
            positions * self.global_scale[None, None, :]
            + self.global_translation[None, None, :]
        )

        mag = torch.linalg.norm(positions, ord=self.order, dim=-1)
        mask = mag >= 1
        x_new = positions.clone()
        x_new[mask] = (2 - (1 / mag[mask][..., None])) * (positions[mask] / mag[mask][..., None])

        return x_new

    def __str__(self):
        return (f"SceneContraction(global_translation={self.global_translation}, "
                f"global_scale={self.global_scale})")

thirdparty/kplanes/raymarching/spatial_distortions.py
oppo-us-research/NeuRBF
import abc
import random
import logging as log
import math
import os
from copy import copy
from typing import Iterable, Optional, Union, Dict, Tuple, Sequence, MutableMapping

import numpy as np
import torch
import torch.utils.data
from torch.utils.tensorboard import SummaryWriter

from thirdparty.kplanes.utils.timer import CudaTimer
from thirdparty.kplanes.utils.ema import EMA
from thirdparty.kplanes.models.lowrank_model import LowrankModel
from thirdparty.kplanes.utils.my_tqdm import tqdm
from thirdparty.kplanes.ops.image import metrics
from thirdparty.kplanes.ops.image.io import write_png
from thirdparty.kplanes.runners.regularization import Regularizer
from thirdparty.kplanes.ops.lr_scheduling import (
    get_cosine_schedule_with_warmup, get_step_schedule_with_warmup
)


class BaseTrainer(abc.ABC):
    def __init__(self,
                 train_data_loader: Iterable,
                 num_steps: int,
                 logdir: str,
                 expname: str,
                 train_fp16: bool,
                 save_every: int,
                 valid_every: int,
                 save_outputs: bool,
                 device: Union[str, torch.device],
                 **kwargs):
        self.train_data_loader = train_data_loader
        self.num_steps = num_steps
        self.train_fp16 = train_fp16
        self.save_every = save_every
        self.valid_every = valid_every
        self.save_outputs = save_outputs
        self.device = device
        self.eval_batch_size = kwargs.get('eval_batch_size', 8129)
        self.extra_args = kwargs
        self.timer = CudaTimer(enabled=False)

        self.log_dir = os.path.join(logdir, expname)
        os.makedirs(self.log_dir, exist_ok=True)
        self.writer = SummaryWriter(log_dir=self.log_dir)

        self.global_step: Optional[int] = None
        self.loss_info: Optional[Dict[str, EMA]] = None

        self.model = self.init_model(**self.extra_args)
        self.optimizer = self.init_optim(**self.extra_args)
        self.scheduler = self.init_lr_scheduler(**self.extra_args)
        self.criterion = torch.nn.MSELoss(reduction='mean')
        self.regularizers = self.init_regularizers(**self.extra_args)
        self.gscaler = torch.cuda.amp.GradScaler(enabled=self.train_fp16)

        self.model.to(self.device)

    @abc.abstractmethod
    def eval_step(self, data, **kwargs) -> MutableMapping[str, torch.Tensor]:
        self.model.eval()
        return None  # noqa

    def train_step(self, data, **kwargs) -> bool:
        self.model.train()
        data = self._move_data_to_device(data)
        if "timestamps" not in data:
            data["timestamps"] = None
        self.timer.check("move-to-device")

        with torch.cuda.amp.autocast(enabled=self.train_fp16):
            fwd_out = self.model(
                data['rays_o'], data['rays_d'], bg_color=data['bg_color'],
                near_far=data['near_fars'], timestamps=data['timestamps'])
            self.timer.check("model-forward")
            # Reconstruction loss
            recon_loss = self.criterion(fwd_out['rgb'], data['imgs'])
            # Regularization
            loss = recon_loss
            for r in self.regularizers:
                reg_loss = r.regularize(self.model, model_out=fwd_out)
                loss = loss + reg_loss
            self.timer.check("regularizaion-forward")
        # Update weights
        self.optimizer.zero_grad(set_to_none=True)
        self.gscaler.scale(loss).backward()
        self.timer.check("backward")
        self.gscaler.step(self.optimizer)
        scale = self.gscaler.get_scale()
        self.gscaler.update()
        self.timer.check("scaler-step")

        # Report on losses
        if self.global_step % self.calc_metrics_every == 0:
            with torch.no_grad():
                recon_loss_val = recon_loss.item()
                self.loss_info[f"mse"].update(recon_loss_val)
                self.loss_info[f"psnr"].update(-10 * math.log10(recon_loss_val))
                for r in self.regularizers:
                    r.report(self.loss_info)

        return scale <= self.gscaler.get_scale()

    def post_step(self, progress_bar):
        init_data = self.model.step_after_iter(self.global_step)
        if self.global_step % self.calc_metrics_every == 0:
            progress_bar.set_postfix_str(
                losses_to_postfix(self.loss_info, lr=self.lr), refresh=False)
            for loss_name, loss_val in self.loss_info.items():
                self.writer.add_scalar(f"train/loss/{loss_name}", loss_val.value, self.global_step)
                if self.timer.enabled:
                    tsum = 0.
                    tstr = "Timings: "
                    for tname, tval in self.timer.timings.items():
                        tstr += f"{tname}={tval:.1f}ms  "
                        tsum += tval
                    tstr += f"tot={tsum:.1f}ms"
                    log.info(tstr)
        progress_bar.update(1)
        if self.valid_every > -1 and self.global_step % self.valid_every == 0:
            print()
            self.validate()
        if self.save_every > -1 and self.global_step % self.save_every == 0:
            print()
            self.save_model()
        return init_data

    def pre_epoch(self):
        self.loss_info = self.init_epoch_info()

    def train(self):
        """Override this if some very specific training procedure is needed."""
        if self.global_step is None:
            self.global_step = 0
        log.info(f"Starting training from step {self.global_step + 1}")
        pb = tqdm(initial=self.global_step, total=self.num_steps)
        try:
            self.pre_epoch()
            batch_iter = iter(self.train_data_loader)
            while self.global_step < self.num_steps:
                self.timer.reset()
                self.model.step_before_iter(self.global_step)
                self.global_step += 1
                self.timer.check("step-before-iter")
                try:
                    data = next(batch_iter)
                    self.timer.check("dloader-next")
                except StopIteration:
                    self.pre_epoch()
                    batch_iter = iter(self.train_data_loader)
                    data = next(batch_iter)
                    log.info("Reset data-iterator")

                try:
                    step_successful = self.train_step(data)
                except StopIteration:
                    self.pre_epoch()
                    batch_iter = iter(self.train_data_loader)
                    log.info("Reset data-iterator")
                    step_successful = True

                if step_successful and self.scheduler is not None:
                    self.scheduler.step()
                for r in self.regularizers:
                    r.step(self.global_step)
                init_data = self.post_step(progress_bar=pb)
                self.timer.check("after-step")
                if init_data is not None:
                    break
                
        finally:
            pb.close()
            self.writer.close()
        return init_data

    def _move_data_to_device(self, data):
        data["rays_o"] = data["rays_o"].to(self.device)
        data["rays_d"] = data["rays_d"].to(self.device)
        data["imgs"] = data["imgs"].to(self.device)
        data["near_fars"] = data["near_fars"].to(self.device)
        if "timestamps" in data:
            data["timestamps"] = data["timestamps"].to(self.device)
        bg_color = data["bg_color"]
        if isinstance(bg_color, torch.Tensor):
            bg_color = bg_color.to(self.device)
        data["bg_color"] = bg_color
        return data

    def _normalize_err(self, preds: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
        err = torch.abs(preds - gt)
        err = err.mean(-1, keepdim=True)  # mean over channels
        # normalize between 0, 1 where 1 corresponds to the 90th percentile
        # err = err.clamp_max(torch.quantile(err, 0.9))
        err = self._normalize_01(err)
        return err.repeat(1, 1, 3)

    @staticmethod
    def _normalize_01(t: torch.Tensor) -> torch.Tensor:
        return (t - t.min()) / t.max()

    def _normalize_depth(self, depth: torch.Tensor, img_h: int, img_w: int) -> torch.Tensor:
        return (
            self._normalize_01(depth)
        ).cpu().reshape(img_h, img_w)[..., None]

    def calc_metrics(self, preds: torch.Tensor, gt: torch.Tensor):
        if gt.shape[-1] == 4:
            gt = gt[..., :3] * gt[..., 3:] + (1.0 - gt[..., 3:])

        err = (gt - preds) ** 2
        return {
            "mse": torch.mean(err),
            "psnr": metrics.psnr(preds, gt),
            "ssim": metrics.ssim(preds, gt),
            "ms-ssim": metrics.msssim(preds, gt),
            "lpips_alex": metrics.rgb_lpips(preds, gt, net_name='alex', device=err.device),
            "lpips_vgg": metrics.rgb_lpips(preds, gt, net_name='vgg', device=err.device)
        }

    def evaluate_metrics(self,
                         gt: Optional[torch.Tensor],
                         preds: MutableMapping[str, torch.Tensor],
                         dset,
                         img_idx: int,
                         name: Optional[str] = None,
                         save_outputs: bool = True) -> Tuple[dict, np.ndarray, Optional[np.ndarray]]:
        if isinstance(dset.img_h, int):
            img_h, img_w = dset.img_h, dset.img_w
        else:
            img_h, img_w = dset.img_h[img_idx], dset.img_w[img_idx]
        preds_rgb = (
            preds["rgb"]
            .reshape(img_h, img_w, 3)
            .cpu()
            .clamp(0, 1)
        )
        if not torch.isfinite(preds_rgb).all():
            log.warning(f"Predictions have {torch.isnan(preds_rgb).sum()} NaNs, "
                        f"{torch.isinf(preds_rgb).sum()} infs.")
            preds_rgb = torch.nan_to_num(preds_rgb, nan=0.0)
        out_img = preds_rgb
        summary = dict()

        out_depth = None
        if "depth" in preds:
            out_depth = preds["depth"].cpu().reshape(img_h, img_w)[..., None]
        preds.pop("depth")

        # This is used for proposal-depth keys
        for k in preds.keys():
            if "depth" in k:
                prop_depth = preds[k].cpu().reshape(img_h, img_w)[..., None]
                out_depth = torch.cat((out_depth, prop_depth), dim=1) if out_depth is not None else prop_depth

        if gt is not None:
            gt = gt.reshape(img_h, img_w, -1).cpu()
            if gt.shape[-1] == 4:
                gt = gt[..., :3] * gt[..., 3:] + (1.0 - gt[..., 3:])
            summary.update(self.calc_metrics(preds_rgb, gt))
            out_img_with_gt = torch.cat((out_img, gt), dim=1)
            # out_img = torch.cat((out_img, self._normalize_err(preds_rgb, gt)), dim=0)
            out_img_with_gt = torch.cat((out_img_with_gt, 
                (preds_rgb - gt).abs().mean(-1, keepdim=True).expand(-1, -1, out_img.shape[-1])*10), dim=1).clamp(0, 1)

        out_img_np: np.ndarray = (out_img * 255.0).byte().numpy()
        out_img_with_gt_np: np.ndarray = (out_img_with_gt * 255.0).byte().numpy()
        out_depth_np: Optional[np.ndarray] = None
        if out_depth is not None:
            out_depth = self._normalize_01(out_depth)
            out_depth_np = (out_depth * 255.0).repeat(1, 1, 3).byte().numpy()

        if save_outputs:
            out_name = f"rgb-step{self.global_step}-{img_idx}"
            out_with_gt_name = f"rgb_with_gt-step{self.global_step}-{img_idx}"
            if name is not None and name != "":
                out_name += "-" + name
                out_with_gt_name += "-" + name
            write_png(os.path.join(self.log_dir, out_name + ".png"), out_img_np)
            write_png(os.path.join(self.log_dir, out_with_gt_name + ".png"), out_img_with_gt_np)
            if out_depth is not None:
                depth_name = f"depth-step{self.global_step}-{img_idx}"
                write_png(os.path.join(self.log_dir, depth_name + ".png"), out_depth_np)

        return summary, out_img_with_gt_np, out_depth_np

    @abc.abstractmethod
    def validate(self):
        pass

    def report_test_metrics(self, scene_metrics: Dict[str, Sequence[float]], extra_name: Optional[str]):
        log_text = f"step {self.global_step}/{self.num_steps}"
        if extra_name is not None:
            log_text += f" | {extra_name}"
        scene_metrics_agg: Dict[str, float] = {}
        for k in scene_metrics:
            ak = f"{k}_{extra_name}"
            scene_metrics_agg[ak] = np.mean(np.asarray(scene_metrics[k])).item()
            log_text += f" | {k}: {scene_metrics_agg[ak]:.4f}"
            self.writer.add_scalar(f"test/{ak}", scene_metrics_agg[ak], self.global_step)

        log.info(log_text)
        return scene_metrics_agg

    def get_save_dict(self):
        return {
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "lr_scheduler": self.scheduler.state_dict() if self.scheduler is not None else None,
            "global_step": self.global_step
        }

    def save_model(self):
        model_fname = os.path.join(self.log_dir, f'model.pth')
        log.info(f'Saving model checkpoint to: {model_fname}')
        torch.save(self.get_save_dict(), model_fname)

    def load_model(self, checkpoint_data, training_needed: bool = True):
        self.model.load_state_dict(checkpoint_data["model"], strict=False)
        log.info("=> Loaded model state from checkpoint")

        if training_needed:
            self.optimizer.load_state_dict(checkpoint_data["optimizer"])
            log.info("=> Loaded optimizer state from checkpoint")

        if training_needed and self.scheduler is not None:
            self.scheduler.load_state_dict(checkpoint_data['lr_scheduler'])
            log.info("=> Loaded scheduler state from checkpoint")

        self.global_step = checkpoint_data["global_step"]
        log.info(f"=> Loaded step {self.global_step} from checkpoints")

    @abc.abstractmethod
    def init_epoch_info(self) -> Dict[str, EMA]:
        pass

    # noinspection PyUnresolvedReferences,PyProtectedMember
    def init_lr_scheduler(self, **kwargs) -> Optional[torch.optim.lr_scheduler._LRScheduler]:
        eta_min = 0
        lr_sched = None
        max_steps = self.num_steps
        scheduler_type = kwargs['scheduler_type']
        log.info(f"Initializing LR Scheduler of type {scheduler_type} with "
                 f"{max_steps} maximum steps.")
        if scheduler_type == "cosine":
            lr_sched = torch.optim.lr_scheduler.CosineAnnealingLR(
                self.optimizer,
                T_max=max_steps,
                eta_min=eta_min)
        elif scheduler_type == "warmup_cosine":
            lr_sched = get_cosine_schedule_with_warmup(
                self.optimizer, num_warmup_steps=512, num_training_steps=max_steps)
        elif scheduler_type == "step":
            lr_sched = torch.optim.lr_scheduler.MultiStepLR(
                self.optimizer,
                milestones=[
                    max_steps // 2,
                    max_steps * 3 // 4,
                    max_steps * 5 // 6,
                    max_steps * 9 // 10,
                ],
                gamma=0.33)
        elif scheduler_type == "warmup_step":
            lr_sched = get_step_schedule_with_warmup(
                self.optimizer, milestones=[
                    max_steps // 2,
                    max_steps * 3 // 4,
                    max_steps * 5 // 6,
                    max_steps * 9 // 10,
                ],
                gamma=0.33,
                num_warmup_steps=512)
        return lr_sched

    def init_optim(self, **kwargs) -> torch.optim.Optimizer:
        optim_type = kwargs['optim_type']
        if optim_type == 'adam':
            optim = torch.optim.Adam(params=self.model.get_params(kwargs['lr_config']), eps=1e-15)
        else:
            raise NotImplementedError()
        return optim

    @abc.abstractmethod
    def init_model(self, **kwargs) -> torch.nn.Module:
        pass

    def get_regularizers(self, **kwargs) -> Sequence[Regularizer]:
        return ()

    def init_regularizers(self, **kwargs):
        # Keep only the regularizers with a positive weight
        regularizers = [r for r in self.get_regularizers(**kwargs) if r.weight > 0]
        return regularizers

    @property
    def lr(self):
        return self.optimizer.param_groups[0]['lr']

    @property
    def calc_metrics_every(self):
        return 1


def losses_to_postfix(loss_dict: Dict[str, EMA], lr: Optional[float]) -> str:
    pfix = [f"{lname}={lval}" for lname, lval in loss_dict.items()]
    if lr is not None:
        pfix.append(f"lr={lr:.2e}")
    return "  ".join(pfix)


def init_dloader_random(_):
    seed = torch.initial_seed() % 2**32  # worker-specific seed initialized by pytorch
    np.random.seed(seed)
    random.seed(seed)


def initialize_model(
        runner: Union['StaticTrainer', 'PhototourismTrainer', 'VideoTrainer'],
        **kwargs) -> LowrankModel:
    """Initialize a `LowrankModel` according to the **kwargs parameters.

    Args:
        runner: The runner object which will hold the model.
                Needed here to fetch dataset parameters.
        **kwargs: Extra parameters to pass to the model

    Returns:
        Initialized LowrankModel.
    """
    from .phototourism_trainer import PhototourismTrainer
    extra_args = copy(kwargs)
    extra_args.pop('global_scale', None)
    extra_args.pop('global_translation', None)

    dset = runner.test_dataset
    try:
        global_translation = dset.global_translation
    except AttributeError:
        global_translation = None
    try:
        global_scale = dset.global_scale
    except AttributeError:
        global_scale = None

    num_images = None
    if runner.train_dataset is not None:
        try:
            num_images = runner.train_dataset.num_images
        except AttributeError:
            num_images = None
    else:
        try:
            num_images = runner.test_dataset.num_images
        except AttributeError:
            num_images = None
    model = LowrankModel(
        grid_config=extra_args.pop("grid_config"),
        aabb=dset.scene_bbox,
        is_ndc=dset.is_ndc,
        is_contracted=dset.is_contracted,
        global_scale=global_scale,
        global_translation=global_translation,
        use_appearance_embedding=isinstance(runner, PhototourismTrainer),
        num_images=num_images,
        **extra_args)
    log.info(f"Initialized {model.__class__} model with "
             f"{sum(np.prod(p.shape) for p in model.parameters()):,} parameters, "
             f"using ndc {model.is_ndc} and contraction {model.is_contracted}. "
             f"Linear decoder: {model.linear_decoder}.")
    return model

thirdparty/kplanes/runners/base_trainer.py
oppo-us-research/NeuRBF
import logging as log
import math
import os
from collections import defaultdict
from typing import Dict, MutableMapping, Union, Any

import pandas as pd
import torch
import torch.utils.data

from thirdparty.kplanes.datasets.phototourism_dataset import PhotoTourismDataset
from thirdparty.kplanes.utils.ema import EMA
from thirdparty.kplanes.models.lowrank_model import LowrankModel
from thirdparty.kplanes.utils.my_tqdm import tqdm
from thirdparty.kplanes.ops.image import metrics
from .base_trainer import BaseTrainer, init_dloader_random, initialize_model
from .regularization import (
    PlaneTV, TimeSmoothness, HistogramLoss, L1TimePlanes, DistortionLoss
)


class PhototourismTrainer(BaseTrainer):
    def __init__(self,
                 tr_loader: torch.utils.data.DataLoader,
                 tr_dset: torch.utils.data.TensorDataset,
                 ts_dset: torch.utils.data.TensorDataset,
                 num_steps: int,
                 logdir: str,
                 expname: str,
                 train_fp16: bool,
                 save_every: int,
                 valid_every: int,
                 save_outputs: bool,
                 device: Union[str, torch.device],
                 **kwargs
                 ):
        self.train_dataset = tr_dset
        self.test_dataset = ts_dset
        super().__init__(
            train_data_loader=tr_loader,
            num_steps=num_steps,
            logdir=logdir,
            expname=expname,
            train_fp16=train_fp16,
            save_every=save_every,
            valid_every=valid_every,
            save_outputs=save_outputs,
            device=device,
            **kwargs)

    def eval_step(self, data, **kwargs) -> MutableMapping[str, torch.Tensor]:
        """
        Note that here `data` contains a whole image. we need to split it up before tracing
        for memory constraints.
        """
        super().eval_step(data, **kwargs)
        batch_size = self.eval_batch_size
        with torch.cuda.amp.autocast(enabled=self.train_fp16), torch.no_grad():
            rays_o = data["rays_o"]
            rays_d = data["rays_d"]
            timestamp = data["timestamps"]
            near_far = data["near_fars"]
            bg_color = data["bg_color"]
            if isinstance(bg_color, torch.Tensor):
                bg_color = bg_color.to(self.device)
            preds = defaultdict(list)
            for b in range(math.ceil(rays_o.shape[0] / batch_size)):
                rays_o_b = rays_o[b * batch_size: (b + 1) * batch_size].to(self.device)
                rays_d_b = rays_d[b * batch_size: (b + 1) * batch_size].to(self.device)
                timestamps_b = timestamp[b * batch_size: (b + 1) * batch_size].to(self.device)
                near_far_b = near_far[b * batch_size: (b + 1) * batch_size].to(self.device)
                outputs = self.model(
                    rays_o_b, rays_d_b, timestamps=timestamps_b, bg_color=bg_color,
                    near_far=near_far_b)
                for k, v in outputs.items():
                    if "rgb" in k or "depth" in k:
                        preds[k].append(v.cpu())
        return {k: torch.cat(v, 0) for k, v in preds.items()}

    def train_step(self, data: Dict[str, Union[int, torch.Tensor]], **kwargs):
        return super().train_step(data, **kwargs)

    def post_step(self, progress_bar):
        return super().post_step(progress_bar)

    def pre_epoch(self):
        super().pre_epoch()
        # Reset randomness in train-dataset
        self.train_dataset.reset_iter()

    @torch.no_grad()
    def validate(self):
        with torch.autograd.enable_grad():
            self.optimize_appearance_codes()

        dataset = self.test_dataset
        per_scene_metrics = defaultdict(list)
        pred_frames, out_depths = [], []
        pb = tqdm(total=len(dataset), desc=f"Test scene ({dataset.name})")
        for img_idx, data in enumerate(dataset):
            preds = self.eval_step(data)
            out_metrics, out_img, out_depth = self.evaluate_metrics(
                data["imgs"], preds, dset=dataset, img_idx=img_idx, name=None,
                save_outputs=self.save_outputs)
            pred_frames.append(out_img)
            if out_depth is not None:
                out_depths.append(out_depth)
            for k, v in out_metrics.items():
                per_scene_metrics[k].append(v)
            pb.set_postfix_str(f"PSNR={out_metrics['psnr']:.2f}", refresh=False)
            pb.update(1)
        pb.close()
        val_metrics = [
            self.report_test_metrics(per_scene_metrics, extra_name=None),
        ]
        df = pd.DataFrame.from_records(val_metrics)
        df.to_csv(os.path.join(self.log_dir, f"test_metrics_step{self.global_step}.csv"))

    def calc_metrics(self, preds: torch.Tensor, gt: torch.Tensor):
        """
        Compute error metrics. This function gets called by `evaluate_metrics` in the base
        trainer class.
        :param preds:
        :param gt:
        :return:
        """
        mid = gt.shape[1] // 2
        gt_right = gt[:, mid:]
        preds_rgb_right = preds[:, mid:]

        err = (gt_right - preds_rgb_right) ** 2
        return {
            "mse": torch.mean(err),
            "psnr": metrics.psnr(preds_rgb_right, gt_right),
            "ssim": metrics.ssim(preds_rgb_right, gt_right),
            "ms-ssim": metrics.msssim(preds_rgb_right, gt_right),
        }

    def init_epoch_info(self):
        ema_weight = 0.9
        loss_info = defaultdict(lambda: EMA(ema_weight))
        return loss_info

    def init_model(self, **kwargs) -> LowrankModel:
        return initialize_model(self, **kwargs)

    def get_regularizers(self, **kwargs):
        return [
            PlaneTV(kwargs.get('plane_tv_weight', 0.0), what='field'),
            PlaneTV(kwargs.get('plane_tv_weight_proposal_net', 0.0), what='proposal_network'),
            L1TimePlanes(kwargs.get('l1_time_planes', 0.0), what='field'),
            L1TimePlanes(kwargs.get('l1_time_planes_proposal_net', 0.0), what='proposal_network'),
            TimeSmoothness(kwargs.get('time_smoothness_weight', 0.0), what='field'),
            TimeSmoothness(kwargs.get('time_smoothness_weight_proposal_net', 0.0), what='proposal_network'),
            HistogramLoss(kwargs.get('histogram_loss_weight', 0.0)),
            DistortionLoss(kwargs.get('distortion_loss_weight', 0.0)),
        ]

    @property
    def calc_metrics_every(self):
        return 5

    def optimize_appearance_step(self, data, im_id):
        rays_o = data["rays_o_left"]
        rays_d = data["rays_d_left"]
        imgs = data["imgs_left"]
        near_far = data["near_fars"]
        bg_color = data["bg_color"]
        if isinstance(bg_color, torch.Tensor):
            bg_color = bg_color.to(self.device)

        epochs = self.extra_args['app_optim_n_epochs']
        batch_size = self.eval_batch_size
        n_steps = math.ceil(rays_o.shape[0] / batch_size)

        camera_id = torch.full((batch_size, ), fill_value=im_id, dtype=torch.int32, device=self.device)

        app_optim = torch.optim.Adam(params=self.model.field.test_appearance_embedding.parameters(), lr=self.extra_args['app_optim_lr'])
        lr_sched = torch.optim.lr_scheduler.StepLR(app_optim, step_size=3 * n_steps, gamma=0.1)
        lowest_loss, lowest_loss_count = 100_000_000, 0
        grad_scaler = torch.cuda.amp.GradScaler(enabled=self.train_fp16)
        for n in range(epochs):
            idx = torch.randperm(rays_o.shape[0])
            for b in range(n_steps):
                batch_ids = idx[b * batch_size: (b + 1) * batch_size]
                rays_o_b = rays_o[batch_ids].to(self.device)
                rays_d_b = rays_d[batch_ids].to(self.device)
                imgs_b = imgs[batch_ids].to(self.device)
                near_far_b = near_far[batch_ids].to(self.device)
                camera_id_b = camera_id[:len(batch_ids)]

                with torch.cuda.amp.autocast(enabled=self.train_fp16):
                    fwd_out = self.model(
                        rays_o_b, rays_d_b, timestamps=camera_id_b, bg_color=bg_color,
                        near_far=near_far_b)
                    recon_loss = self.criterion(fwd_out['rgb'], imgs_b)
                app_optim.zero_grad(set_to_none=True)
                grad_scaler.scale(recon_loss).backward()
                grad_scaler.step(app_optim)
                grad_scaler.update()
                lr_sched.step()

                self.writer.add_scalar(
                    f"appearance_loss_{self.global_step}/recon_loss_{im_id}", recon_loss.item(),
                    b + n * n_steps)

                if recon_loss.item() < lowest_loss:
                    lowest_loss = recon_loss.item()
                    lowest_loss_count = 0
                lowest_loss_count += 1
            # 1 epoch without improvement -> stop
            if lowest_loss_count > 1 * n_steps:
                break

    def optimize_appearance_codes(self):
        """Optimize the appearance embedding of all test poses.

        At test time, since the correct embedding for a test pose is unknown, half of the test
        ground-truth image is used to optimize the appearance embedding corresponding to that test
        image. Evaluation is then performed on the other half of the image.
        """
        dset = self.test_dataset
        num_test_imgs = len(dset)

        # 1. Initialize test appearance code to average code.
        if not hasattr(self.model.field, "test_appearance_embedding"):
            tst_embedding = torch.nn.Embedding(
                num_test_imgs, self.model.field.appearance_embedding_dim
            ).to(self.device)
            with torch.autograd.no_grad():
                tst_embedding.weight.copy_(
                    self.model.field.appearance_embedding.weight
                        .detach()
                        .mean(dim=0, keepdim=True)
                        .expand(num_test_imgs, -1)
                )
            self.model.field.test_appearance_embedding = tst_embedding

        # 2. Setup parameter trainability
        self.model.eval()
        param_trainable = {}
        for pn, p in self.model.named_parameters():
            param_trainable[pn] = p.requires_grad
            p.requires_grad_(False)
        self.model.field.test_appearance_embedding.requires_grad_(True)

        # 3. Optimize
        pb = tqdm(total=len(dset), desc=f"Test-time appearance-code optimization")
        for img_idx, data in enumerate(dset):
            self.optimize_appearance_step(data, img_idx)
            pb.update(1)
        pb.close()

        # 4. Reset parameter trainability
        for pn, p in self.model.named_parameters():
            p.requires_grad_(param_trainable[pn])
        self.model.field.test_appearance_embedding.requires_grad_(False)


def init_tr_data(data_downsample, data_dir, **kwargs):
    batch_size = kwargs['batch_size']
    log.info(f"Loading PhotoTourismDataset with downsample={data_downsample}")
    tr_dset = PhotoTourismDataset(
        data_dir, split='train', batch_size=batch_size,
        contraction=kwargs['contract'], ndc=kwargs['ndc'],
        scene_bbox=kwargs['scene_bbox'], global_scale=kwargs.get('global_scale', None),
        global_translation=kwargs.get('global_translation', None),
    )
    tr_loader = torch.utils.data.DataLoader(
        tr_dset, batch_size=None, num_workers=4,  prefetch_factor=4, pin_memory=True,
        worker_init_fn=init_dloader_random)
    return {"tr_loader": tr_loader, "tr_dset": tr_dset}


def init_ts_data(data_dir, split, **kwargs):
    ts_dset = PhotoTourismDataset(
        data_dir, split=split, batch_size=None,
        contraction=kwargs['contract'], ndc=kwargs['ndc'],
        scene_bbox=kwargs['scene_bbox'], global_scale=kwargs.get('global_scale', None),
        global_translation=kwargs.get('global_translation', None),
    )
    return {"ts_dset": ts_dset}


def load_data(data_downsample, data_dirs, validate_only, render_only, **kwargs):
    assert len(data_dirs) == 1
    od: Dict[str, Any] = {}
    if not validate_only and not render_only:
        od.update(init_tr_data(data_downsample, data_dirs[0], **kwargs))
    else:
        od.update(tr_loader=None, tr_dset=None)
    test_split = 'render' if render_only else 'test'
    od.update(init_ts_data(data_dirs[0], split=test_split, **kwargs))
    return od

thirdparty/kplanes/runners/phototourism_trainer.py
oppo-us-research/NeuRBF
import abc
import os
from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.optim.lr_scheduler
from torch import nn

from thirdparty.kplanes.models.lowrank_model import LowrankModel
from thirdparty.kplanes.ops.losses.histogram_loss import interlevel_loss
from thirdparty.kplanes.raymarching.ray_samplers import RaySamples


def compute_plane_tv(t):
    batch_size, c, h, w = t.shape
    count_h = batch_size * c * (h - 1) * w
    count_w = batch_size * c * h * (w - 1)
    h_tv = torch.square(t[..., 1:, :] - t[..., :h-1, :]).sum()
    w_tv = torch.square(t[..., :, 1:] - t[..., :, :w-1]).sum()
    return 2 * (h_tv / count_h + w_tv / count_w)  # This is summing over batch and c instead of avg


def compute_plane_smoothness(t):
    batch_size, c, h, w = t.shape
    # Convolve with a second derivative filter, in the time dimension which is dimension 2
    first_difference = t[..., 1:, :] - t[..., :h-1, :]  # [batch, c, h-1, w]
    second_difference = first_difference[..., 1:, :] - first_difference[..., :h-2, :]  # [batch, c, h-2, w]
    # Take the L2 norm of the result
    return torch.square(second_difference).mean()


class Regularizer():
    def __init__(self, reg_type, initialization):
        self.reg_type = reg_type
        self.initialization = initialization
        self.weight = float(self.initialization)
        self.last_reg = None

    def step(self, global_step):
        pass

    def report(self, d):
        if self.last_reg is not None:
            d[self.reg_type].update(self.last_reg.item())

    def regularize(self, *args, **kwargs) -> torch.Tensor:
        out = self._regularize(*args, **kwargs) * self.weight
        self.last_reg = out.detach()
        return out

    @abc.abstractmethod
    def _regularize(self, *args, **kwargs) -> torch.Tensor:
        raise NotImplementedError()

    def __str__(self):
        return f"Regularizer({self.reg_type}, weight={self.weight})"


class PlaneTV(Regularizer):
    def __init__(self, initial_value, what: str = 'field'):
        if what not in {'field', 'proposal_network'}:
            raise ValueError(f'what must be one of "field" or "proposal_network" '
                             f'but {what} was passed.')
        name = f'planeTV-{what[:2]}'
        super().__init__(name, initial_value)
        self.what = what

    def step(self, global_step):
        pass

    def _regularize(self, model: LowrankModel, **kwargs):
        multi_res_grids: Sequence[nn.ParameterList]
        if self.what == 'field':
            multi_res_grids = model.field.grids
        elif self.what == 'proposal_network':
            multi_res_grids = [p.grids for p in model.proposal_networks]
        else:
            raise NotImplementedError(self.what)
        total = 0
        # Note: input to compute_plane_tv should be of shape [batch_size, c, h, w]
        for grids in multi_res_grids:
            if len(grids) == 3:
                spatial_grids = [0, 1, 2]
            else:
                spatial_grids = [0, 1, 3]  # These are the spatial grids; the others are spatiotemporal
            for grid_id in spatial_grids:
                total += compute_plane_tv(grids[grid_id])
            for grid in grids:
                # grid: [1, c, h, w]
                total += compute_plane_tv(grid)
        return total


class TimeSmoothness(Regularizer):
    def __init__(self, initial_value, what: str = 'field'):
        if what not in {'field', 'proposal_network'}:
            raise ValueError(f'what must be one of "field" or "proposal_network" '
                             f'but {what} was passed.')
        name = f'time-smooth-{what[:2]}'
        super().__init__(name, initial_value)
        self.what = what

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        multi_res_grids: Sequence[nn.ParameterList]
        if self.what == 'field':
            multi_res_grids = model.field.grids
        elif self.what == 'proposal_network':
            multi_res_grids = [p.grids for p in model.proposal_networks]
        else:
            raise NotImplementedError(self.what)
        total = 0
        # model.grids is 6 x [1, rank * F_dim, reso, reso]
        for grids in multi_res_grids:
            if len(grids) == 3:
                time_grids = []
            else:
                time_grids = [2, 4, 5]
            for grid_id in time_grids:
                total += compute_plane_smoothness(grids[grid_id])
        return torch.as_tensor(total)


class HistogramLoss(Regularizer):
    def __init__(self, initial_value):
        super().__init__('histogram-loss', initial_value)

        self.visualize = False
        self.count = 0

    def _regularize(self, model: LowrankModel, model_out, **kwargs) -> torch.Tensor:
        if self.visualize:
            if self.count % 500 == 0:
                prop_idx = 0
                fine_idx = 1
                # proposal info
                weights_proposal = model_out["weights_list"][prop_idx].detach().cpu().numpy()
                spacing_starts_proposal = model_out["ray_samples_list"][prop_idx].spacing_starts
                spacing_ends_proposal = model_out["ray_samples_list"][prop_idx].spacing_ends
                sdist_proposal = torch.cat([
                    spacing_starts_proposal[..., 0],
                    spacing_ends_proposal[..., -1:, 0]
                ], dim=-1).detach().cpu().numpy()

                # fine info
                weights_fine = model_out["weights_list"][fine_idx].detach().cpu().numpy()
                spacing_starts_fine = model_out["ray_samples_list"][fine_idx].spacing_starts
                spacing_ends_fine = model_out["ray_samples_list"][fine_idx].spacing_ends
                sdist_fine = torch.cat([
                    spacing_starts_fine[..., 0],
                    spacing_ends_fine[..., -1:, 0]
                ], dim=-1).detach().cpu().numpy()

                for i in range(10):  # plot 10 rays
                    fix, ax1 = plt.subplots()

                    delta = np.diff(sdist_proposal[i], axis=-1)
                    ax1.bar(sdist_proposal[i, :-1], weights_proposal[i].squeeze() / delta, width=delta, align="edge", label='proposal', alpha=0.7, color="b")
                    ax1.legend()
                    ax2 = ax1.twinx()

                    delta = np.diff(sdist_fine[i], axis=-1)
                    ax2.bar(sdist_fine[i, :-1], weights_fine[i].squeeze() / delta, width=delta, align="edge", label='fine', alpha=0.3, color='r')
                    ax2.legend()
                    os.makedirs(f'histogram_loss/{self.count}', exist_ok=True)
                    plt.savefig(f'./histogram_loss/{self.count}/batch_{i}.png')
                    plt.close()
                    plt.cla()
                    plt.clf()
            self.count += 1
        return interlevel_loss(model_out['weights_list'], model_out['ray_samples_list'])


class L1ProposalNetwork(Regularizer):
    def __init__(self, initial_value):
        super().__init__('l1-proposal-network', initial_value)

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        grids = [p.grids for p in model.proposal_networks]
        total = 0.0
        for pn_grids in grids:
            for grid in pn_grids:
                total += torch.abs(grid).mean()
        return torch.as_tensor(total)


class DepthTV(Regularizer):
    def __init__(self, initial_value):
        super().__init__('tv-depth', initial_value)

    def _regularize(self, model: LowrankModel, model_out, **kwargs) -> torch.Tensor:
        depth = model_out['depth']
        tv = compute_plane_tv(
            depth.reshape(64, 64)[None, None, :, :]
        )
        return tv


class L1TimePlanes(Regularizer):
    def __init__(self, initial_value, what='field'):
        if what not in {'field', 'proposal_network'}:
            raise ValueError(f'what must be one of "field" or "proposal_network" '
                             f'but {what} was passed.')
        super().__init__(f'l1-time-{what[:2]}', initial_value)
        self.what = what

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        # model.grids is 6 x [1, rank * F_dim, reso, reso]
        multi_res_grids: Sequence[nn.ParameterList]
        if self.what == 'field':
            multi_res_grids = model.field.grids
        elif self.what == 'proposal_network':
            multi_res_grids = [p.grids for p in model.proposal_networks]
        else:
            raise NotImplementedError(self.what)

        total = 0.0
        for grids in multi_res_grids:
            if len(grids) == 3:
                continue
            else:
                # These are the spatiotemporal grids
                spatiotemporal_grids = [2, 4, 5]
            for grid_id in spatiotemporal_grids:
                total += torch.abs(1 - grids[grid_id]).mean()
        return torch.as_tensor(total)


class DistortionLoss(Regularizer):
    def __init__(self, initial_value):
        super().__init__('distortion-loss', initial_value)

    def _regularize(self, model: LowrankModel, model_out, **kwargs) -> torch.Tensor:
        """
        Efficient O(N) realization of distortion loss.
        from https://github.com/sunset1995/torch_efficient_distloss/blob/main/torch_efficient_distloss/eff_distloss.py
        There are B rays each with N sampled points.
        """
        w = model_out['weights_list'][-1]
        rs: RaySamples = model_out['ray_samples_list'][-1]
        m = (rs.starts + rs.ends) / 2
        interval = rs.deltas

        loss_uni = (1/3) * (interval * w.pow(2)).sum(dim=-1).mean()
        wm = (w * m)
        w_cumsum = w.cumsum(dim=-1)
        wm_cumsum = wm.cumsum(dim=-1)
        loss_bi_0 = wm[..., 1:] * w_cumsum[..., :-1]
        loss_bi_1 = w[..., 1:] * wm_cumsum[..., :-1]
        loss_bi = 2 * (loss_bi_0 - loss_bi_1).sum(dim=-1).mean()
        return loss_bi + loss_uni

thirdparty/kplanes/runners/regularization.py
oppo-us-research/NeuRBF
import math
import os
from collections import defaultdict
from typing import Dict, MutableMapping, Union, Sequence, Any

import pandas as pd
import torch
import torch.utils.data

from thirdparty.kplanes.datasets import SyntheticNerfDataset, LLFFDataset
from thirdparty.kplanes.models.lowrank_model import LowrankModel
from thirdparty.kplanes.utils.ema import EMA
from thirdparty.kplanes.utils.my_tqdm import tqdm
from thirdparty.kplanes.utils.parse_args import parse_optint
from .base_trainer import BaseTrainer, init_dloader_random, initialize_model
from .regularization import (
    PlaneTV, HistogramLoss, L1ProposalNetwork, DepthTV, DistortionLoss,
)


class StaticTrainer(BaseTrainer):
    def __init__(self,
                 tr_loader: torch.utils.data.DataLoader,
                 ts_dset: torch.utils.data.TensorDataset,
                 tr_dset: torch.utils.data.TensorDataset,
                 num_steps: int,
                 logdir: str,
                 expname: str,
                 train_fp16: bool,
                 save_every: int,
                 valid_every: int,
                 save_outputs: bool,
                 device: Union[str, torch.device],
                 **kwargs
                 ):
        self.test_dataset = ts_dset
        self.train_dataset = tr_dset
        self.is_ndc = self.test_dataset.is_ndc
        self.is_contracted = self.test_dataset.is_contracted

        super().__init__(
            train_data_loader=tr_loader,
            num_steps=num_steps,
            logdir=logdir,
            expname=expname,
            train_fp16=train_fp16,
            save_every=save_every,
            valid_every=valid_every,
            save_outputs=save_outputs,
            device=device,
            **kwargs
        )

    def eval_step(self, data, **kwargs) -> MutableMapping[str, torch.Tensor]:
        """
        Note that here `data` contains a whole image. we need to split it up before tracing
        for memory constraints.
        """
        super().eval_step(data, **kwargs)
        batch_size = self.eval_batch_size
        channels = {"rgb", "depth", "proposal_depth"}
        with torch.cuda.amp.autocast(enabled=self.train_fp16), torch.no_grad():
            rays_o = data["rays_o"]
            rays_d = data["rays_d"]
            # near_far and bg_color are constant over mini-batches
            near_far = data["near_fars"].to(self.device)
            bg_color = data["bg_color"]
            if isinstance(bg_color, torch.Tensor):
                bg_color = bg_color.to(self.device)
            preds = defaultdict(list)
            for b in range(math.ceil(rays_o.shape[0] / batch_size)):
                rays_o_b = rays_o[b * batch_size: (b + 1) * batch_size].to(self.device)
                rays_d_b = rays_d[b * batch_size: (b + 1) * batch_size].to(self.device)
                outputs = self.model(rays_o_b, rays_d_b, near_far=near_far,
                                     bg_color=bg_color)
                for k, v in outputs.items():
                    if k in channels or "depth" in k:
                        preds[k].append(v.cpu())
        return {k: torch.cat(v, 0) for k, v in preds.items()}

    def train_step(self, data: Dict[str, Union[int, torch.Tensor]], **kwargs):
        return super().train_step(data, **kwargs)

    def post_step(self, progress_bar):
        return super().post_step(progress_bar)

    def pre_epoch(self):
        super().pre_epoch()
        self.train_dataset.reset_iter()

    @torch.no_grad()
    def validate(self):
        dataset = self.test_dataset
        per_scene_metrics = defaultdict(list)
        pb = tqdm(total=len(dataset), desc=f"Test scene {dataset.name}")
        for img_idx, data in enumerate(dataset):
            ts_render = self.eval_step(data)
            out_metrics, _, _ = self.evaluate_metrics(
                data["imgs"], ts_render, dset=dataset, img_idx=img_idx,
                name=None, save_outputs=self.save_outputs)
            for k, v in out_metrics.items():
                per_scene_metrics[k].append(v)
            pb.set_postfix_str(f"PSNR={out_metrics['psnr']:.2f}", refresh=False)
            pb.update(1)
        pb.close()
        val_metrics = [
            self.report_test_metrics(per_scene_metrics, extra_name="")
        ]
        df = pd.DataFrame.from_records(val_metrics)
        df.to_csv(os.path.join(self.log_dir, f"test_metrics_step{self.global_step}.csv"))
        self.val_metrics = val_metrics

    def get_save_dict(self):
        base_save_dict = super().get_save_dict()
        return base_save_dict

    def load_model(self, checkpoint_data, training_needed: bool = True):
        super().load_model(checkpoint_data, training_needed)

    def init_epoch_info(self):
        ema_weight = 0.9  # higher places higher weight to new observations
        loss_info = defaultdict(lambda: EMA(ema_weight))
        return loss_info

    def init_model(self, **kwargs) -> LowrankModel:
        return initialize_model(self, **kwargs)

    def get_regularizers(self, **kwargs):
        return [
            PlaneTV(kwargs.get('plane_tv_weight', 0.0), what='field'),
            PlaneTV(kwargs.get('plane_tv_weight_proposal_net', 0.0), what='proposal_network'),
            HistogramLoss(kwargs.get('histogram_loss_weight', 0.0)),
            L1ProposalNetwork(kwargs.get('l1_proposal_net_weight', 0.0)),
            DepthTV(kwargs.get('depth_tv_weight', 0.0)),
            DistortionLoss(kwargs.get('distortion_loss_weight', 0.0)),
        ]

    @property
    def calc_metrics_every(self):
        return 5


def decide_dset_type(dd) -> str:
    if ("chair" in dd or "drums" in dd or "ficus" in dd or "hotdog" in dd
            or "lego" in dd or "materials" in dd or "mic" in dd
            or "ship" in dd):
        return "synthetic"
    elif ("fern" in dd or "flower" in dd or "fortress" in dd
          or "horns" in dd or "leaves" in dd or "orchids" in dd
          or "room" in dd or "trex" in dd):
        return "llff"
    else:
        raise RuntimeError(f"data_dir {dd} not recognized as LLFF or Synthetic dataset.")


def init_tr_data(data_downsample: float, data_dirs: Sequence[str], **kwargs):
    batch_size = int(kwargs['batch_size'])
    assert len(data_dirs) == 1
    data_dir = data_dirs[0]

    dset_type = decide_dset_type(data_dir)
    if dset_type == "synthetic":
        max_tr_frames = parse_optint(kwargs.get('max_tr_frames'))
        dset = SyntheticNerfDataset(
            data_dir, split='train', downsample=data_downsample,
            max_frames=max_tr_frames, batch_size=batch_size)
    elif dset_type == "llff":
        hold_every = parse_optint(kwargs.get('hold_every'))
        dset = LLFFDataset(
            data_dir, split='train', downsample=int(data_downsample), hold_every=hold_every,
            batch_size=batch_size, contraction=kwargs['contract'], ndc=kwargs['ndc'],
            ndc_far=float(kwargs['ndc_far']), near_scaling=float(kwargs['near_scaling']))
    else:
        raise ValueError(f"Dataset type {dset_type} invalid.")
    dset.reset_iter()

    tr_loader = torch.utils.data.DataLoader(
        dset, num_workers=4, prefetch_factor=4, pin_memory=True,
        batch_size=None, worker_init_fn=init_dloader_random)

    return {
        "tr_dset": dset,
        "tr_loader": tr_loader,
    }


def init_ts_data(data_dirs: Sequence[str], split: str, **kwargs):
    assert len(data_dirs) == 1
    data_dir = data_dirs[0]
    dset_type = decide_dset_type(data_dir)
    if dset_type == "synthetic":
        max_ts_frames = parse_optint(kwargs.get('max_ts_frames'))
        dset = SyntheticNerfDataset(
            data_dir, split=split, downsample=1, max_frames=max_ts_frames)
    elif dset_type == "llff":
        hold_every = parse_optint(kwargs.get('hold_every'))
        dset = LLFFDataset(
            data_dir, split=split, downsample=4, hold_every=hold_every,
            contraction=kwargs['contract'], ndc=kwargs['ndc'],
            ndc_far=float(kwargs['ndc_far']), near_scaling=float(kwargs['near_scaling']))
    else:
        raise ValueError(f"Dataset type {dset_type} invalid.")
    return {"ts_dset": dset}


def load_data(data_downsample, data_dirs, validate_only, render_only, **kwargs):
    od: Dict[str, Any] = {}
    if not validate_only:
        od.update(init_tr_data(data_downsample, data_dirs, **kwargs))
    else:
        od.update(tr_loader=None, tr_dset=None)
    test_split = 'render' if render_only else 'test'
    od.update(init_ts_data(data_dirs, split=test_split, **kwargs))
    return od

thirdparty/kplanes/runners/static_trainer.py
oppo-us-research/NeuRBF
import logging as log
import math
import os
from collections import defaultdict
from typing import Dict, MutableMapping, Union, Any, List

import pandas as pd
import torch
import torch.utils.data

from thirdparty.kplanes.datasets.video_datasets import Video360Dataset
from thirdparty.kplanes.utils.ema import EMA
from thirdparty.kplanes.utils.my_tqdm import tqdm
from thirdparty.kplanes.ops.image import metrics
from thirdparty.kplanes.ops.image.io import write_video_to_file
from thirdparty.kplanes.models.lowrank_model import LowrankModel
from .base_trainer import BaseTrainer, init_dloader_random, initialize_model
from .regularization import (
    PlaneTV, TimeSmoothness, HistogramLoss, L1TimePlanes, DistortionLoss
)


class VideoTrainer(BaseTrainer):
    def __init__(self,
                 tr_loader: torch.utils.data.DataLoader,
                 tr_dset: torch.utils.data.TensorDataset,
                 ts_dset: torch.utils.data.TensorDataset,
                 num_steps: int,
                 logdir: str,
                 expname: str,
                 train_fp16: bool,
                 save_every: int,
                 valid_every: int,
                 save_outputs: bool,
                 isg_step: int,
                 ist_step: int,
                 device: Union[str, torch.device],
                 **kwargs
                 ):
        self.train_dataset = tr_dset
        self.test_dataset = ts_dset
        self.ist_step = ist_step
        self.isg_step = isg_step
        self.save_video = save_outputs
        # Switch to compute extra video metrics (FLIP, JOD)
        self.compute_video_metrics = False
        super().__init__(
            train_data_loader=tr_loader,
            num_steps=num_steps,
            logdir=logdir,
            expname=expname,
            train_fp16=train_fp16,
            save_every=save_every,
            valid_every=valid_every,
            save_outputs=False,  # False since we're saving video
            device=device,
            **kwargs)

    def eval_step(self, data, **kwargs) -> MutableMapping[str, torch.Tensor]:
        """
        Note that here `data` contains a whole image. we need to split it up before tracing
        for memory constraints.
        """
        super().eval_step(data, **kwargs)
        batch_size = self.eval_batch_size
        with torch.cuda.amp.autocast(enabled=self.train_fp16), torch.no_grad():
            rays_o = data["rays_o"]
            rays_d = data["rays_d"]
            timestamp = data["timestamps"]
            near_far = data["near_fars"].to(self.device)
            bg_color = data["bg_color"]
            if isinstance(bg_color, torch.Tensor):
                bg_color = bg_color.to(self.device)
            preds = defaultdict(list)
            for b in range(math.ceil(rays_o.shape[0] / batch_size)):
                rays_o_b = rays_o[b * batch_size: (b + 1) * batch_size].to(self.device)
                rays_d_b = rays_d[b * batch_size: (b + 1) * batch_size].to(self.device)
                timestamps_d_b = timestamp.expand(rays_o_b.shape[0]).to(self.device)
                outputs = self.model(
                    rays_o_b, rays_d_b, timestamps=timestamps_d_b, bg_color=bg_color,
                    near_far=near_far)
                for k, v in outputs.items():
                    if "rgb" in k or "depth" in k:
                        preds[k].append(v.cpu())
        return {k: torch.cat(v, 0) for k, v in preds.items()}

    def train_step(self, data: Dict[str, Union[int, torch.Tensor]], **kwargs):
        scale_ok = super().train_step(data, **kwargs)

        if self.global_step == self.isg_step:
            self.train_dataset.enable_isg()
            raise StopIteration  # Whenever we change the dataset
        if self.global_step == self.ist_step:
            self.train_dataset.switch_isg2ist()
            raise StopIteration  # Whenever we change the dataset

        return scale_ok

    def post_step(self, progress_bar):
        super().post_step(progress_bar)

    def pre_epoch(self):
        super().pre_epoch()
        # Reset randomness in train-dataset
        self.train_dataset.reset_iter()

    @torch.no_grad()
    def validate(self):
        dataset = self.test_dataset
        per_scene_metrics: Dict[str, Union[float, List]] = defaultdict(list)
        pred_frames, out_depths = [], []
        pb = tqdm(total=len(dataset), desc=f"Test scene ({dataset.name})")
        for img_idx, data in enumerate(dataset):
            preds = self.eval_step(data)
            out_metrics, out_img, out_depth = self.evaluate_metrics(
                data["imgs"], preds, dset=dataset, img_idx=img_idx, name=None,
                save_outputs=self.save_outputs)
            pred_frames.append(out_img)
            if out_depth is not None:
                out_depths.append(out_depth)
            for k, v in out_metrics.items():
                per_scene_metrics[k].append(v)
            pb.set_postfix_str(f"PSNR={out_metrics['psnr']:.2f}", refresh=False)
            pb.update(1)
        pb.close()
        if self.save_video:
            write_video_to_file(
                os.path.join(self.log_dir, f"step{self.global_step}.mp4"),
                pred_frames
            )
            if len(out_depths) > 0:
                write_video_to_file(
                    os.path.join(self.log_dir, f"step{self.global_step}-depth.mp4"),
                    out_depths
                )
        # Calculate JOD (on whole video)
        if self.compute_video_metrics:
            per_scene_metrics["JOD"] = metrics.jod(
                [f[:dataset.img_h, :, :] for f in pred_frames],
                [f[dataset.img_h: 2*dataset.img_h, :, :] for f in pred_frames],
            )
            per_scene_metrics["FLIP"] = metrics.flip(
                [f[:dataset.img_h, :, :] for f in pred_frames],
                [f[dataset.img_h: 2*dataset.img_h, :, :] for f in pred_frames],
            )

        val_metrics = [
            self.report_test_metrics(per_scene_metrics, extra_name=None),
        ]
        df = pd.DataFrame.from_records(val_metrics)
        df.to_csv(os.path.join(self.log_dir, f"test_metrics_step{self.global_step}.csv"))

    def get_save_dict(self):
        base_save_dict = super().get_save_dict()
        return base_save_dict

    def load_model(self, checkpoint_data, training_needed: bool = True):
        super().load_model(checkpoint_data, training_needed)
        if self.train_dataset is not None:
            if -1 < self.isg_step < self.global_step < self.ist_step:
                self.train_dataset.enable_isg()
            elif -1 < self.ist_step < self.global_step:
                self.train_dataset.switch_isg2ist()

    def init_epoch_info(self):
        ema_weight = 0.9
        loss_info = defaultdict(lambda: EMA(ema_weight))
        return loss_info

    def init_model(self, **kwargs) -> LowrankModel:
        return initialize_model(self, **kwargs)

    def get_regularizers(self, **kwargs):
        return [
            PlaneTV(kwargs.get('plane_tv_weight', 0.0), what='field'),
            PlaneTV(kwargs.get('plane_tv_weight_proposal_net', 0.0), what='proposal_network'),
            L1TimePlanes(kwargs.get('l1_time_planes', 0.0), what='field'),
            L1TimePlanes(kwargs.get('l1_time_planes_proposal_net', 0.0), what='proposal_network'),
            TimeSmoothness(kwargs.get('time_smoothness_weight', 0.0), what='field'),
            TimeSmoothness(kwargs.get('time_smoothness_weight_proposal_net', 0.0), what='proposal_network'),
            HistogramLoss(kwargs.get('histogram_loss_weight', 0.0)),
            DistortionLoss(kwargs.get('distortion_loss_weight', 0.0)),
        ]

    @property
    def calc_metrics_every(self):
        return 5


def init_tr_data(data_downsample, data_dir, **kwargs):
    isg = kwargs.get('isg', False)
    ist = kwargs.get('ist', False)
    keyframes = kwargs.get('keyframes', False)
    batch_size = kwargs['batch_size']
    log.info(f"Loading Video360Dataset with downsample={data_downsample}")
    tr_dset = Video360Dataset(
        data_dir, split='train', downsample=data_downsample,
        batch_size=batch_size,
        max_cameras=kwargs.get('max_train_cameras', None),
        max_tsteps=kwargs['max_train_tsteps'] if keyframes else None,
        isg=isg, keyframes=keyframes, contraction=kwargs['contract'], ndc=kwargs['ndc'],
        near_scaling=float(kwargs.get('near_scaling', 0)), ndc_far=float(kwargs.get('ndc_far', 0)),
        scene_bbox=kwargs['scene_bbox'],
    )
    if ist:
        tr_dset.switch_isg2ist()  # this should only happen in case we're reloading

    g = torch.Generator()
    g.manual_seed(0)
    tr_loader = torch.utils.data.DataLoader(
        tr_dset, batch_size=None, num_workers=4,  prefetch_factor=4, pin_memory=True,
        worker_init_fn=init_dloader_random, generator=g)
    return {"tr_loader": tr_loader, "tr_dset": tr_dset}


def init_ts_data(data_dir, split, **kwargs):
    if 'dnerf' in data_dir:
        downsample = 1.0
    else:
        downsample = 2.0
    ts_dset = Video360Dataset(
        data_dir, split=split, downsample=downsample,
        max_cameras=kwargs.get('max_test_cameras', None), max_tsteps=kwargs.get('max_test_tsteps', None),
        contraction=kwargs['contract'], ndc=kwargs['ndc'],
        near_scaling=float(kwargs.get('near_scaling', 0)), ndc_far=float(kwargs.get('ndc_far', 0)),
        scene_bbox=kwargs['scene_bbox'],
    )
    return {"ts_dset": ts_dset}


def load_data(data_downsample, data_dirs, validate_only, render_only, **kwargs):
    assert len(data_dirs) == 1
    od: Dict[str, Any] = {}
    if not validate_only and not render_only:
        od.update(init_tr_data(data_downsample, data_dirs[0], **kwargs))
    else:
        od.update(tr_loader=None, tr_dset=None)
    test_split = 'render' if render_only else 'test'
    od.update(init_ts_data(data_dirs[0], split=test_split, **kwargs))
    return od

thirdparty/kplanes/runners/video_trainer.py
oppo-us-research/NeuRBF
"""Entry point for simple renderings, given a trainer and some poses."""
import os
import logging as log
from typing import Union

import torch

from thirdparty.kplanes.models.lowrank_model import LowrankModel
from thirdparty.kplanes.utils.my_tqdm import tqdm
from thirdparty.kplanes.ops.image.io import write_video_to_file
from thirdparty.kplanes.runners.static_trainer import StaticTrainer
from thirdparty.kplanes.runners.video_trainer import VideoTrainer


@torch.no_grad()
def render_to_path(trainer: Union[VideoTrainer, StaticTrainer], extra_name: str = "") -> None:
    """Render all poses in the `test_dataset`, saving them to file
    Args:
        trainer: The trainer object which is used for rendering
        extra_name: String to append to the saved file-name
    """
    dataset = trainer.test_dataset

    pb = tqdm(total=100, desc=f"Rendering scene")
    frames = []
    for img_idx, data in enumerate(dataset):
        ts_render = trainer.eval_step(data)

        if isinstance(dataset.img_h, int):
            img_h, img_w = dataset.img_h, dataset.img_w
        else:
            img_h, img_w = dataset.img_h[img_idx], dataset.img_w[img_idx]
        preds_rgb = (
            ts_render["rgb"]
            .reshape(img_h, img_w, 3)
            .cpu()
            .clamp(0, 1)
            .mul(255.0)
            .byte()
            .numpy()
        )
        frames.append(preds_rgb)
        pb.update(1)
    pb.close()

    out_fname = os.path.join(trainer.log_dir, f"rendering_path_{extra_name}.mp4")
    write_video_to_file(out_fname, frames)
    log.info(f"Saved rendering path with {len(frames)} frames to {out_fname}")


def normalize_for_disp(img):
    img = img - torch.min(img)
    img = img / torch.max(img)
    return img


@torch.no_grad()
def decompose_space_time(trainer: StaticTrainer, extra_name: str = "") -> None:
    """Render space-time decomposition videos for poses in the `test_dataset`.

    The space-only part of the decomposition is obtained by setting the time-planes to 1.
    The time-only part is obtained by simple subtraction of the space-only part from the full
    rendering.

    Args:
        trainer: The trainer object which is used for rendering
        extra_name: String to append to the saved file-name
    """
    chosen_cam_idx = 15
    model: LowrankModel = trainer.model
    dataset = trainer.test_dataset

    # Store original parameters from main field and proposal-network field
    parameters = []
    for multires_grids in model.field.grids:
        parameters.append([grid.data for grid in multires_grids])
    pn_parameters = []
    for pn in model.proposal_networks:
        pn_parameters.append([grid_plane.data for grid_plane in pn.grids])

    camdata = None
    for img_idx, data in enumerate(dataset):
        if img_idx == chosen_cam_idx:
            camdata = data
    if camdata is None:
        raise ValueError(f"Cam idx {chosen_cam_idx} invalid.")

    num_frames = img_idx + 1
    frames = []
    for img_idx in tqdm(range(num_frames), desc="Rendering scene with separate space and time components"):
        # Linearly interpolated timestamp, normalized between -1, 1
        camdata["timestamps"] = torch.Tensor([img_idx / num_frames]) * 2 - 1

        if isinstance(dataset.img_h, int):
            img_h, img_w = dataset.img_h, dataset.img_w
        else:
            img_h, img_w = dataset.img_h[img_idx], dataset.img_w[img_idx]

        # Full model: turn on time-planes
        for i in range(len(model.field.grids)):
            for plane_idx in [2, 4, 5]:
                model.field.grids[i][plane_idx].data = parameters[i][plane_idx]
        for i in range(len(model.proposal_networks)):
            for plane_idx in [2, 4, 5]:
                model.proposal_networks[i].grids[plane_idx].data = pn_parameters[i][plane_idx]
        preds = trainer.eval_step(camdata)
        full_out = preds["rgb"].reshape(img_h, img_w, 3).cpu()

        # Space-only model: turn off time-planes
        for i in range(len(model.field.grids)):
            for plane_idx in [2, 4, 5]:  # time-grids off
                model.field.grids[i][plane_idx].data = torch.ones_like(parameters[i][plane_idx])
        for i in range(len(model.proposal_networks)):
            for plane_idx in [2, 4, 5]:
                model.proposal_networks[i].grids[plane_idx].data = torch.ones_like(pn_parameters[i][plane_idx])
        preds = trainer.eval_step(camdata)
        spatial_out = preds["rgb"].reshape(img_h, img_w, 3).cpu()

        # Temporal model: full - space
        temporal_out = normalize_for_disp(full_out - spatial_out)

        frames.append(
            torch.cat([full_out, spatial_out, temporal_out], dim=1)
                 .clamp(0, 1)
                 .mul(255.0)
                 .byte()
                 .numpy()
        )

    out_fname = os.path.join(trainer.log_dir, f"spacetime_{extra_name}.mp4")
    write_video_to_file(out_fname, frames)
    log.info(f"Saved rendering path with {len(frames)} frames to {out_fname}")

thirdparty/kplanes/utils/create_rendering.py
oppo-us-research/NeuRBF

class EMA():
    def __init__(self, weighting=0.9):
        self.weighting = weighting
        self.val = None

    def update(self, val):
        if self.val is None:
            self.val = val
        else:
            self.val = self.weighting * val + (1 - self.weighting) * self.val

    @property
    def value(self):
        return self.val

    def __str__(self):
        return f"{self.val:.2e}"

thirdparty/kplanes/utils/ema.py
oppo-us-research/NeuRBF
"""Evaluate several metrics for a pretrained model. Handles video and static."""
import re
import glob
import os
from collections import defaultdict

import numpy as np
import torch

from thirdparty.kplanes.ops.image import metrics
from thirdparty.kplanes.ops.image.io import read_mp4, read_png


def eval_static_metrics(static_dir):
    all_file_names = glob.glob(os.path.join(static_dir, r"*.png"))
    # Collect all GT+Pred files per step
    files_per_step = defaultdict(list)
    for f in all_file_names:
        if "depth" in f:
            continue
        match = re.match(r".*step([0-9]+)-([0-9])+\.png", f)
        if match is None:
            continue
        step = int(match.group(1))
        files_per_step[step].append(f)
    steps = list(files_per_step.keys())
    max_step = max(steps)
    print(f"Evaluating static metrics for {static_dir} at step "
          f"{max_step} with {len(files_per_step[max_step])} files.")
    frames = [read_png(f) for f in files_per_step[max_step]]
    h1 = frames[0].shape[0] // 3
    pred_frames = [f[:h1] for f in frames]
    gt_frames = [f[h1:2*h1] for f in frames]

    psnrs, ssims, msssims, lpipss = [], [], [], []
    for pred, gt in zip(pred_frames, gt_frames):
        psnrs.append(metrics.psnr(pred, gt))
        ssims.append(metrics.ssim(pred, gt))
        msssims.append(metrics.msssim(pred, gt))
        lpipss.append(metrics.rgb_lpips(pred, gt, net_name="alex"))
    psnr = np.mean(psnrs)
    ssim = np.mean(ssims)
    msssim = np.mean(msssims)
    lpips = np.mean(lpipss)

    print()
    print(f"Images at {static_dir} - step {max_step}. Metrics:")
    print(f"PSNR = {psnr}")
    print(f"SSIM = {ssim}")
    print(f"MS-SSIM = {msssim}")
    print(f"Alex-LPIPS= {lpips}")
    print()


def eval_video_metrics(video_path):
    frames = read_mp4(video_path)
    h1 = frames[0].shape[0] // 3
    pred_frames = [f[:h1].numpy() for f in frames]
    gt_frames = [f[h1:2*h1].numpy() for f in frames]

    psnrs, ssims, msssims, lpipss = [], [], [], []
    for pred, gt in zip(pred_frames, gt_frames):
        pred = torch.from_numpy(pred).float().div(255)
        gt = torch.from_numpy(gt).float().div(255)
        psnrs.append(metrics.psnr(pred, gt))
        ssims.append(metrics.ssim(pred, gt))
        msssims.append(metrics.msssim(pred, gt))
        lpipss.append(metrics.rgb_lpips(pred, gt, net_name="alex"))
    psnr = np.mean(psnrs)
    ssim = np.mean(ssims)
    msssim = np.mean(msssims)
    lpips = np.mean(lpipss)

    flip = metrics.flip(pred_frames=pred_frames, gt_frames=gt_frames, interval=10)
    # jod = metrics.jod(pred_frames=pred_frames, gt_frames=gt_frames)

    print()
    print(f"Video at {video_path} metrics:")
    print(f"PSNR = {psnr}")
    print(f"SSIM = {ssim}")
    # print(f"MS-SSIM = {msssim}")
    # print(f"Alex-LPIPS = {lpips}")
    # print(f"FLIP = {flip}")
    # print(f"JOD = {jod}")
    print()
    print()


dnerf_scenes = ['hellwarrior', 'mutant', 'hook', 'bouncingballs', 'lego', 'trex', 'standup', 'jumpingjacks']
types = ['linear', 'mlp']

if __name__ == "__main__":
    for modeltype in types:
        for scene in dnerf_scenes:
            eval_video_metrics(f"logs/dnerf_{modeltype}_refactor1/{scene}_concat32_lr0.01_time0.1_tv0.0001_proptime0.001_proptv0.0001_distort0/step30000.mp4")

thirdparty/kplanes/utils/eval_model_metrics.py
oppo-us-research/NeuRBF
import os
from tqdm.auto import tqdm as original_tqdm
import functools
import logging as log


tqdm_disabled = os.environ.get('DISABLE_TQDM', False)
if tqdm_disabled:
    log.info("Disabling tqdm")
tqdm = original_tqdm
tqdm.__init__ = functools.partialmethod(
    tqdm.__init__, mininterval=0.5, disable=tqdm_disabled)

thirdparty/kplanes/utils/my_tqdm.py
oppo-us-research/NeuRBF
from typing import Optional


def parse_optfloat(val, default_val=None) -> Optional[float]:
    if val == "None" or val is None:
        return default_val
    return float(val)


def parse_optint(val, default_val=None) -> Optional[int]:
    if val == "None" or val is None:
        return default_val
    return int(val)

thirdparty/kplanes/utils/parse_args.py
oppo-us-research/NeuRBF
"""Accurate timer for CUDA code"""
import torch.cuda


class CudaTimer:
    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self.start, self.end = None, None
        self.timings = {}
        if self.enabled:
            self.start = torch.cuda.Event(enable_timing=True)
            self.end = torch.cuda.Event(enable_timing=True)
            self.prev_time_gpu = self.start.record()

    def reset(self):
        if self.enabled:
            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.enabled:
            self.end.record()
            torch.cuda.synchronize()
            gpu_time = self.start.elapsed_time(self.end)
            self.timings[name] = gpu_time

            self.prev_time_gpu = self.start.record()

thirdparty/kplanes/utils/timer.py
oppo-us-research/NeuRBF
from ..tensorf.models.tensorBase import *
from .quaternion_utils import *
import numpy as np


class RenderingEquationEncoding(torch.nn.Module):
    def __init__(self, num_theta, num_phi, device):
        super(RenderingEquationEncoding, self).__init__()

        self.num_theta = num_theta
        self.num_phi = num_phi

        omega, omega_la, omega_mu = init_predefined_omega(num_theta, num_phi)
        self.omega = omega.view(1, num_theta, num_phi, 3).to(device)
        self.omega_la = omega_la.view(1, num_theta, num_phi, 3).to(device)
        self.omega_mu = omega_mu.view(1, num_theta, num_phi, 3).to(device)

    def forward(self, omega_o, a, la, mu):
        Smooth = F.relu((omega_o[:, None, None] * self.omega).sum(dim=-1, keepdim=True)) # N, num_theta, num_phi, 1

        la = F.softplus(la - 1)
        mu = F.softplus(mu - 1)
        exp_input = -la * (self.omega_la * omega_o[:, None, None]).sum(dim=-1, keepdim=True).pow(2) -mu * (self.omega_mu * omega_o[:, None, None]).sum(dim=-1, keepdim=True).pow(2)
        out = a * Smooth * torch.exp(exp_input)

        return out

class RenderingNet(torch.nn.Module):
    def __init__(self, num_theta = 8, num_phi=16, data_dim_color=192, featureC=256, device='cpu', 
                 view_pe=-1, fea_pe=-1, btn_freq=[]):
        super(RenderingNet, self).__init__()

        self.ch_cd = 3
        self.ch_s = 3
        self.ch_normal = 3
        self.ch_bottleneck = 128

        self.num_theta = 8
        self.num_phi = 16
        self.num_asg = self.num_theta * self.num_phi

        self.ch_asg_feature = 128
        self.ch_per_theta = self.ch_asg_feature // self.num_theta

        self.ch_a = 2
        self.ch_la = 1
        self.ch_mu = 1
        self.ch_per_asg = self.ch_a + self.ch_la + self.ch_mu

        self.ch_normal_dot_viewdir = 1

        self.view_pe = view_pe
        self.fea_pe = fea_pe

        if len(btn_freq) >= 2:
            self.btn_freq = torch.linspace(np.log2(btn_freq[0]), np.log2(btn_freq[1]), 
                                            self.ch_bottleneck, device=0)
            self.btn_freq = torch.exp2(self.btn_freq)
        else:
            self.btn_freq = None

        self.ree_function = RenderingEquationEncoding(num_theta, num_phi, device)

        self.spatial_mlp = torch.nn.Sequential(
                torch.nn.Linear(data_dim_color, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, self.ch_cd + self.ch_s + self.ch_bottleneck + self.ch_normal + self.ch_asg_feature)).to(device)

        self.asg_mlp = torch.nn.Sequential(torch.nn.Linear(self.ch_per_theta, self.num_phi * self.ch_per_asg)).to(device)

        in_dim = self.ch_bottleneck + self.num_asg * self.ch_a + self.ch_normal_dot_viewdir
        if self.view_pe > -1:
           in_dim += 3
        if self.view_pe > 0:
           in_dim += 3 * self.view_pe * 2
        self.directional_mlp = torch.nn.Sequential(
                torch.nn.Linear(in_dim, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, featureC),
                torch.nn.GELU(),
                torch.nn.Linear(featureC, 3)).to(device)


    def spatial_mlp_forward(self, x):
        out = self.spatial_mlp(x)
        sections = [self.ch_cd, self.ch_s, self.ch_normal, self.ch_bottleneck, self.ch_asg_feature]
        diffuse_color, tint, normals, bottleneck, asg_features = torch.split(out, sections, dim=-1)
        normals = -F.normalize(normals, dim=1)
        return diffuse_color, tint, normals, bottleneck, asg_features

    def asg_mlp_forward(self, asg_feature):
        N = asg_feature.size(0)
        asg_feature = asg_feature.view(N, self.num_theta, -1)
        asg_params = self.asg_mlp(asg_feature)
        asg_params = asg_params.view(N, self.num_theta, self.num_phi, -1)

        a, la, mu = torch.split(asg_params, [self.ch_a, self.ch_la, self.ch_mu], dim=-1)
        return a, la, mu

    def directional_mlp_forward(self, x):
        out = self.directional_mlp(x)
        return out

    def reflect(self, viewdir, normal):
        out = 2 * (viewdir * normal).sum(dim=-1, keepdim=True) * normal - viewdir
        return out

    def forward(self, pts, viewdir, feature):
        diffuse_color, tint, normal, bottleneck, asg_feature = self.spatial_mlp_forward(feature)
        refdir = self.reflect(-viewdir, normal)

        a, la, mu = self.asg_mlp_forward(asg_feature)
        ree = self.ree_function(refdir, a, la, mu) # N, num_theta, num_phi, ch_per_asg
        ree = ree.view(ree.size(0), -1)

        if self.btn_freq is not None:
            bottleneck = bottleneck + torch.sin(bottleneck * self.btn_freq[None])

        normal_dot_viewdir = ((-viewdir) * normal).sum(dim=-1, keepdim=True)
        dir_mlp_input = [bottleneck, ree, normal_dot_viewdir]
        if self.view_pe > -1:
            dir_mlp_input += [viewdir]
        if self.view_pe > 0:
            dir_mlp_input += [positional_encoding(viewdir, self.view_pe)]
        dir_mlp_input = torch.cat(dir_mlp_input, dim=-1)
        specular_color = self.directional_mlp_forward(dir_mlp_input)

        raw_rgb = diffuse_color + tint * specular_color
        rgb = torch.sigmoid(raw_rgb)

        return rgb, normal

thirdparty/nrff/nrff.py
oppo-us-research/NeuRBF
import torch
import numpy as np



def quaternion_product(p, q):
    p_r = p[..., [0]]
    p_i = p[..., 1:]
    q_r = q[..., [0]]
    q_i = q[..., 1:]

    out_r = p_r * q_r - (p_i * q_i).sum(dim=-1)
    out_i = p_r * q_i + q_r * p_i + torch.linalg.cross(p_i, q_i, dim=-1)

    return torch.cat([out_r, out_i], dim=-1)

def quaternion_inverse(p):
    p_r = p[..., [0]]
    p_i = -p[..., 1:]

    return torch.cat([p_r, p_i], dim=-1)

def quaternion_rotate(p, q):
    q_inv = quaternion_inverse(q)

    qp = quaternion_product(q, p)
    out = quaternion_product(qp, q_inv)
    return out

def build_q(vec, angle):
    out_r = torch.cos(angle / 2)
    out_i = torch.sin(angle / 2) * vec

    return torch.cat([out_r, out_i], dim=-1)


def cartesian2quaternion(x):
    zeros_ = x.new_zeros([*x.shape[:-1], 1])
    return torch.cat([zeros_, x], dim=-1)


def spherical2cartesian(theta, phi):
    x = torch.cos(phi) * torch.sin(theta)
    y = torch.sin(phi) * torch.sin(theta)
    z = torch.cos(theta)

    return [x, y, z]

def init_predefined_omega(n_theta, n_phi):
    theta_list = torch.linspace(0, np.pi, n_theta)
    phi_list = torch.linspace(0, np.pi*2, n_phi)

    out_omega = []
    out_omega_lambda = []
    out_omega_mu = []

    for i in range(n_theta):
        theta = theta_list[i].view(1, 1)

        for j in range(n_phi):
            phi = phi_list[j].view(1, 1)

            omega = spherical2cartesian(theta, phi)
            omega = torch.stack(omega, dim=-1).view(1, 3)

            omega_lambda = spherical2cartesian(theta+np.pi/2, phi)
            omega_lambda = torch.stack(omega_lambda, dim=-1).view(1, 3)

            p = cartesian2quaternion(omega_lambda)
            q = build_q(omega, torch.tensor(np.pi/2).view(1, 1))
            omega_mu = quaternion_rotate(p, q)[..., 1:]

            out_omega.append(omega)
            out_omega_lambda.append(omega_lambda)
            out_omega_mu.append(omega_mu)


    out_omega = torch.stack(out_omega, dim=0)
    out_omega_lambda = torch.stack(out_omega_lambda, dim=0)
    out_omega_mu = torch.stack(out_omega_mu, dim=0)

    return out_omega, out_omega_lambda, out_omega_mu

thirdparty/nrff/quaternion_utils.py
oppo-us-research/NeuRBF
import torch,cv2
from torch.utils.data import Dataset
import json
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T


from .ray_utils import *


class BlenderDataset(Dataset):
    def __init__(self, datadir, split='train', downsample=1.0, is_stack=False, N_vis=-1):

        self.N_vis = N_vis
        self.root_dir = datadir
        self.split = split
        self.is_stack = is_stack
        self.img_wh = (int(800/downsample),int(800/downsample))
        self.define_transforms()

        self.scene_bbox = torch.tensor([[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]])
        self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
        self.read_meta()
        self.define_proj_mat()

        self.white_bg = True
        self.near_far = [2.0,6.0]
        
        self.center = torch.mean(self.scene_bbox, axis=0).float().view(1, 1, 3)
        self.radius = (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
        self.downsample=downsample

    def read_depth(self, filename):
        depth = np.array(read_pfm(filename)[0], dtype=np.float32)  # (800, 800)
        return depth
    
    def read_meta(self):

        with open(os.path.join(self.root_dir, f"transforms_{self.split}.json"), 'r') as f:
            self.meta = json.load(f)

        w, h = self.img_wh
        self.focal = 0.5 * 800 / np.tan(0.5 * self.meta['camera_angle_x'])  # original focal length
        self.focal *= self.img_wh[0] / 800  # modify focal length to match size self.img_wh


        # ray directions for all pixels, same for all images (same H, W, focal)
        self.directions = get_ray_directions(h, w, [self.focal,self.focal])  # (h, w, 3)
        self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
        self.intrinsics = torch.tensor([[self.focal,0,w/2],[0,self.focal,h/2],[0,0,1]]).float()

        self.image_paths = []
        self.poses = []
        self.all_rays = []
        self.all_rgbs = []
        self.all_masks = []
        self.all_depth = []
        self.downsample=1.0

        img_eval_interval = 1 if self.N_vis < 0 else len(self.meta['frames']) // self.N_vis
        idxs = list(range(0, len(self.meta['frames']), img_eval_interval))
        for i in tqdm(idxs, desc=f'Loading data {self.split} ({len(idxs)})'):#img_list:#

            frame = self.meta['frames'][i]
            pose = np.array(frame['transform_matrix']) @ self.blender2opencv
            c2w = torch.FloatTensor(pose)
            self.poses += [c2w]

            image_path = os.path.join(self.root_dir, f"{frame['file_path']}.png")
            self.image_paths += [image_path]
            img = Image.open(image_path)
            
            if self.downsample!=1.0:
                img = img.resize(self.img_wh, Image.LANCZOS)
            img = self.transform(img)  # (c, h, w)
            img = img.view(img.shape[0], -1).permute(1, 0)  # (h*w, c) RGBA
            img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) if img.shape[-1] == 4 else img  # blend A to RGB
            self.all_rgbs += [img]


            rays_o, rays_d = get_rays(self.directions, c2w)  # both (h*w, 3)
            self.all_rays += [torch.cat([rays_o, rays_d], 1)]  # (h*w, 6)


        self.poses = torch.stack(self.poses)
        if not self.is_stack:
            self.all_rays = torch.cat(self.all_rays, 0)  # (len(self.meta['frames])*h*w, 3)
            self.all_rgbs = torch.cat(self.all_rgbs, 0)  # (len(self.meta['frames])*h*w, 3)

#             self.all_depth = torch.cat(self.all_depth, 0)  # (len(self.meta['frames])*h*w, 3)
        else:
            self.all_rays = torch.stack(self.all_rays, 0)  # (len(self.meta['frames]),h*w, 3)
            self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1,*self.img_wh[::-1], 3)  # (len(self.meta['frames]),h,w,3)
            # self.all_masks = torch.stack(self.all_masks, 0).reshape(-1,*self.img_wh[::-1])  # (len(self.meta['frames]),h,w,3)


    def define_transforms(self):
        self.transform = T.ToTensor()
        
    def define_proj_mat(self):
        self.proj_mat = self.intrinsics.unsqueeze(0) @ torch.inverse(self.poses)[:,:3]

    def world2ndc(self,points,lindisp=None):
        device = points.device
        return (points - self.center.to(device)) / self.radius.to(device)
        
    def __len__(self):
        return len(self.all_rgbs)

    def __getitem__(self, idx):

        if self.split == 'train':  # use data in the buffers
            sample = {'rays': self.all_rays[idx],
                      'rgbs': self.all_rgbs[idx]}

        else:  # create data for each image separately

            img = self.all_rgbs[idx]
            rays = self.all_rays[idx]
            mask = self.all_masks[idx] # for quantity evaluation

            sample = {'rays': rays,
                      'rgbs': img,
                      'mask': mask}
        return sample

thirdparty/tensorf/dataLoader/blender.py
oppo-us-research/NeuRBF
from .tensorBase import *


class TensorVM(TensorBase):
    def __init__(self, aabb, gridSize, device, **kargs):
        super(TensorVM, self).__init__(aabb, gridSize, device, **kargs)
        

    def init_svd_volume(self, res, device):
        self.plane_coef = torch.nn.Parameter(
            0.1 * torch.randn((3, self.app_n_comp + self.density_n_comp, res, res), device=device))
        self.line_coef = torch.nn.Parameter(
            0.1 * torch.randn((3, self.app_n_comp + self.density_n_comp, res, 1), device=device))
        self.basis_mat = torch.nn.Linear(self.app_n_comp * 3, self.app_dim, bias=False, device=device)

    
    def get_optparam_groups(self, lr_init_spatialxyz = 0.02, lr_init_network = 0.001):
        grad_vars = [{'params': self.line_coef, 'lr': lr_init_spatialxyz}, {'params': self.plane_coef, 'lr': lr_init_spatialxyz},
                         {'params': self.basis_mat.parameters(), 'lr':lr_init_network}]
        if isinstance(self.renderModule, torch.nn.Module):
            grad_vars += [{'params':self.renderModule.parameters(), 'lr':lr_init_network}]
        return grad_vars

    def compute_features(self, xyz_sampled):

        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach()
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach()

        plane_feats = F.grid_sample(self.plane_coef[:, -self.density_n_comp:], coordinate_plane, align_corners=True).view(
                                        -1, *xyz_sampled.shape[:1])
        line_feats = F.grid_sample(self.line_coef[:, -self.density_n_comp:], coordinate_line, align_corners=True).view(
                                        -1, *xyz_sampled.shape[:1])
        
        sigma_feature = torch.sum(plane_feats * line_feats, dim=0)
        
        
        plane_feats = F.grid_sample(self.plane_coef[:, :self.app_n_comp], coordinate_plane, align_corners=True).view(3 * self.app_n_comp, -1)
        line_feats = F.grid_sample(self.line_coef[:, :self.app_n_comp], coordinate_line, align_corners=True).view(3 * self.app_n_comp, -1)
        
        
        app_features = self.basis_mat((plane_feats * line_feats).T)
        
        return sigma_feature, app_features

    def compute_densityfeature(self, xyz_sampled):
        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach().view(3, -1, 1, 2)
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)

        plane_feats = F.grid_sample(self.plane_coef[:, -self.density_n_comp:], coordinate_plane, align_corners=True).view(
                                        -1, *xyz_sampled.shape[:1])
        line_feats = F.grid_sample(self.line_coef[:, -self.density_n_comp:], coordinate_line, align_corners=True).view(
                                        -1, *xyz_sampled.shape[:1])
        
        sigma_feature = torch.sum(plane_feats * line_feats, dim=0)
        
        
        return sigma_feature
    
    def compute_appfeature(self, xyz_sampled):
        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach().view(3, -1, 1, 2)
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)
        
        plane_feats = F.grid_sample(self.plane_coef[:, :self.app_n_comp], coordinate_plane, align_corners=True).view(3 * self.app_n_comp, -1)
        line_feats = F.grid_sample(self.line_coef[:, :self.app_n_comp], coordinate_line, align_corners=True).view(3 * self.app_n_comp, -1)
        
        
        app_features = self.basis_mat((plane_feats * line_feats).T)
        
        
        return app_features
    

    def vectorDiffs(self, vector_comps):
        total = 0
        
        for idx in range(len(vector_comps)):
            n_comp, n_size = vector_comps[idx].shape[:-1]
            
            dotp = torch.matmul(vector_comps[idx].view(n_comp,n_size), vector_comps[idx].view(n_comp,n_size).transpose(-1,-2))
            non_diagonal = dotp.view(-1)[1:].view(n_comp-1, n_comp+1)[...,:-1]
            total = total + torch.mean(torch.abs(non_diagonal))
        return total

    def vector_comp_diffs(self):
        
        return self.vectorDiffs(self.line_coef[:,-self.density_n_comp:]) + self.vectorDiffs(self.line_coef[:,:self.app_n_comp])
    
    
    @torch.no_grad()
    def up_sampling_VM(self, plane_coef, line_coef, res_target):

        for i in range(len(self.vecMode)):
            vec_id = self.vecMode[i]
            mat_id_0, mat_id_1 = self.matMode[i]

            plane_coef[i] = torch.nn.Parameter(
                F.interpolate(plane_coef[i].data, size=(res_target[mat_id_1], res_target[mat_id_0]), mode='bilinear',
                              align_corners=True))
            line_coef[i] = torch.nn.Parameter(
                F.interpolate(line_coef[i].data, size=(res_target[vec_id], 1), mode='bilinear', align_corners=True))

        return plane_coef, line_coef

    @torch.no_grad()
    def upsample_volume_grid(self, res_target):
        # self.app_plane, self.app_line = self.up_sampling_VM(self.app_plane, self.app_line, res_target)
        # self.density_plane, self.density_line = self.up_sampling_VM(self.density_plane, self.density_line, res_target)

        scale = res_target[0]/self.line_coef.shape[2] #assuming xyz have the same scale
        plane_coef = F.interpolate(self.plane_coef.detach().data, scale_factor=scale, mode='bilinear',align_corners=True)
        line_coef  = F.interpolate(self.line_coef.detach().data, size=(res_target[0],1), mode='bilinear',align_corners=True)
        self.plane_coef, self.line_coef = torch.nn.Parameter(plane_coef), torch.nn.Parameter(line_coef)
        self.compute_stepSize(res_target)
        print(f'upsamping to {res_target}')


class TensorVMSplit(TensorBase):
    def __init__(self, aabb, gridSize, device, **kargs):
        super(TensorVMSplit, self).__init__(aabb, gridSize, device, **kargs)


    def init_svd_volume(self, res, device):
        self.density_plane, self.density_line = self.init_one_svd(self.density_n_comp, self.gridSize, 0.1, device)
        self.app_plane, self.app_line = self.init_one_svd(self.app_n_comp, self.gridSize, 0.1, device)
        if self.app_dim is not None:
            self.basis_mat = torch.nn.Linear(sum(self.app_n_comp), self.app_dim, bias=False).to(device)
        else:
            self.app_dim = sum(self.app_n_comp)


    def init_one_svd(self, n_component, gridSize, scale, device):
        plane_coef, line_coef = [], []
        for i in range(len(self.vecMode)):
            vec_id = self.vecMode[i]
            mat_id_0, mat_id_1 = self.matMode[i]
            plane_coef.append(torch.nn.Parameter(
                scale * torch.randn((1, n_component[i], gridSize[mat_id_1], gridSize[mat_id_0]))))  #
            line_coef.append(
                torch.nn.Parameter(scale * torch.randn((1, n_component[i], gridSize[vec_id], 1))))

        return torch.nn.ParameterList(plane_coef).to(device), torch.nn.ParameterList(line_coef).to(device)
    

    def get_optparam_groups(self, lr_init_spatialxyz = 0.02, lr_init_network = 0.001):
        grad_vars = [{'params': self.density_line, 'lr': lr_init_spatialxyz}, 
                     {'params': self.density_plane, 'lr': lr_init_spatialxyz},
                     {'params': self.app_line, 'lr': lr_init_spatialxyz}, 
                     {'params': self.app_plane, 'lr': lr_init_spatialxyz},]
        if hasattr(self, 'basis_mat'):
            grad_vars += [{'params': self.basis_mat.parameters(), 'lr':lr_init_network}]
        if isinstance(self.renderModule, torch.nn.Module):
            grad_vars += [{'params':self.renderModule.parameters(), 'lr':lr_init_network}]
        return grad_vars


    def vectorDiffs(self, vector_comps):
        total = 0
        
        for idx in range(len(vector_comps)):
            n_comp, n_size = vector_comps[idx].shape[1:-1]
            
            dotp = torch.matmul(vector_comps[idx].view(n_comp,n_size), vector_comps[idx].view(n_comp,n_size).transpose(-1,-2))
            non_diagonal = dotp.view(-1)[1:].view(n_comp-1, n_comp+1)[...,:-1]
            total = total + torch.mean(torch.abs(non_diagonal))
        return total

    def vector_comp_diffs(self):
        return self.vectorDiffs(self.density_line) + self.vectorDiffs(self.app_line)
    
    def density_L1(self):
        total = 0
        for idx in range(len(self.density_plane)):
            total = total + torch.mean(torch.abs(self.density_plane[idx])) + torch.mean(torch.abs(self.density_line[idx]))# + torch.mean(torch.abs(self.app_plane[idx])) + torch.mean(torch.abs(self.density_plane[idx]))
        return total
    
    def TV_loss_density(self, reg):
        total = 0
        for idx in range(len(self.density_plane)):
            total = total + reg(self.density_plane[idx]) * 1e-2 #+ reg(self.density_line[idx]) * 1e-3
        return total
        
    def TV_loss_app(self, reg):
        total = 0
        for idx in range(len(self.app_plane)):
            total = total + reg(self.app_plane[idx]) * 1e-2 #+ reg(self.app_line[idx]) * 1e-3
        return total

    def compute_densityfeature(self, xyz_sampled, others):
        # xyz_sampled: [n_pts, 3]

        # plane + line basis
        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach().view(3, -1, 1, 2)  # [3, n_pts, 1, 2]
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))  # [3, n_pts]
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)  # [3, n_pts, 1, 2]

        sigma_feature = torch.zeros((xyz_sampled.shape[0],), device=xyz_sampled.device)
        for idx_plane in range(len(self.density_plane)):
            # density_plane: [1, c, h, w]
            plane_coef_point = F.grid_sample(self.density_plane[idx_plane], coordinate_plane[[idx_plane]],
                                                align_corners=True).view(-1, xyz_sampled.shape[0])  # [c, n_pts]
            # density_line: [1, c, d, 1]
            line_coef_point = F.grid_sample(self.density_line[idx_plane], coordinate_line[[idx_plane]],
                                            align_corners=True).view(-1, xyz_sampled.shape[0])  # [c, n_pts]
            sigma_feature = sigma_feature + torch.sum(plane_coef_point * line_coef_point, dim=0)

        return sigma_feature


    def compute_appfeature(self, xyz_sampled, others):

        # plane + line basis
        coordinate_plane = torch.stack((xyz_sampled[..., self.matMode[0]], xyz_sampled[..., self.matMode[1]], xyz_sampled[..., self.matMode[2]])).detach().view(3, -1, 1, 2)
        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)

        plane_coef_point,line_coef_point = [],[]
        for idx_plane in range(len(self.app_plane)):
            plane_coef_point.append(F.grid_sample(self.app_plane[idx_plane], coordinate_plane[[idx_plane]],
                                                align_corners=True).view(-1, *xyz_sampled.shape[:1]))
            line_coef_point.append(F.grid_sample(self.app_line[idx_plane], coordinate_line[[idx_plane]],
                                            align_corners=True).view(-1, *xyz_sampled.shape[:1]))
        plane_coef_point, line_coef_point = torch.cat(plane_coef_point), torch.cat(line_coef_point)

        if hasattr(self, 'basis_mat'):
            return self.basis_mat((plane_coef_point * line_coef_point).T)
        else:
            return (plane_coef_point * line_coef_point).T


    @torch.no_grad()
    def up_sampling_VM(self, plane_coef, line_coef, res_target):

        for i in range(len(self.vecMode)):
            vec_id = self.vecMode[i]
            mat_id_0, mat_id_1 = self.matMode[i]
            plane_coef[i] = torch.nn.Parameter(
                F.interpolate(plane_coef[i].data, size=(res_target[mat_id_1], res_target[mat_id_0]), mode='bilinear',
                              align_corners=True))
            line_coef[i] = torch.nn.Parameter(
                F.interpolate(line_coef[i].data, size=(res_target[vec_id], 1), mode='bilinear', align_corners=True))

        return plane_coef, line_coef

    @torch.no_grad()
    def upsample_volume_grid(self, res_target, res_target_real):
        self.app_plane, self.app_line = self.up_sampling_VM(self.app_plane, self.app_line, res_target)
        self.density_plane, self.density_line = self.up_sampling_VM(self.density_plane, self.density_line, res_target)

        self.update_stepSize(res_target)
        print(f'upsamping to {res_target}')

    @torch.no_grad()
    def shrink(self, new_aabb):
        print("====> shrinking ...")
        xyz_min, xyz_max = new_aabb
        t_l, b_r = (xyz_min - self.aabb[0]) / self.units, (xyz_max - self.aabb[0]) / self.units
        t_l, b_r = torch.round(torch.round(t_l)).long(), torch.round(b_r).long() + 1
        b_r = torch.stack([b_r, self.gridSize]).amin(0)

        for i in range(len(self.vecMode)):
            mode0 = self.vecMode[i]
            self.density_line[i] = torch.nn.Parameter(
                self.density_line[i].data[...,t_l[mode0]:b_r[mode0],:]
            )
            self.app_line[i] = torch.nn.Parameter(
                self.app_line[i].data[...,t_l[mode0]:b_r[mode0],:]
            )
            mode0, mode1 = self.matMode[i]
            self.density_plane[i] = torch.nn.Parameter(
                self.density_plane[i].data[...,t_l[mode1]:b_r[mode1],t_l[mode0]:b_r[mode0]]
            )
            self.app_plane[i] = torch.nn.Parameter(
                self.app_plane[i].data[...,t_l[mode1]:b_r[mode1],t_l[mode0]:b_r[mode0]]
            )

        if not torch.all(self.alphaMask.gridSize == self.gridSize):
            t_l_r, b_r_r = t_l / (self.gridSize-1), (b_r-1) / (self.gridSize-1)
            correct_aabb = torch.zeros_like(new_aabb)
            correct_aabb[0] = (1-t_l_r)*self.aabb[0] + t_l_r*self.aabb[1]
            correct_aabb[1] = (1-b_r_r)*self.aabb[0] + b_r_r*self.aabb[1]
            print("aabb", new_aabb, "\ncorrect aabb", correct_aabb)
            new_aabb = correct_aabb

        newSize = b_r - t_l
        self.aabb = new_aabb
        self.update_stepSize((newSize[0], newSize[1], newSize[2]))


class TensorCP(TensorBase):
    def __init__(self, aabb, gridSize, device, **kargs):
        super(TensorCP, self).__init__(aabb, gridSize, device, **kargs)


    def init_svd_volume(self, res, device):
        self.density_line = self.init_one_svd(self.density_n_comp[0], self.gridSize, 0.2, device)
        self.app_line = self.init_one_svd(self.app_n_comp[0], self.gridSize, 0.2, device)
        self.basis_mat = torch.nn.Linear(self.app_n_comp[0], self.app_dim, bias=False).to(device)


    def init_one_svd(self, n_component, gridSize, scale, device):
        line_coef = []
        for i in range(len(self.vecMode)):
            vec_id = self.vecMode[i]
            line_coef.append(
                torch.nn.Parameter(scale * torch.randn((1, n_component, gridSize[vec_id], 1))))
        return torch.nn.ParameterList(line_coef).to(device)

    
    def get_optparam_groups(self, lr_init_spatialxyz = 0.02, lr_init_network = 0.001):
        grad_vars = [{'params': self.density_line, 'lr': lr_init_spatialxyz},
                     {'params': self.app_line, 'lr': lr_init_spatialxyz},
                     {'params': self.basis_mat.parameters(), 'lr':lr_init_network}]
        if isinstance(self.renderModule, torch.nn.Module):
            grad_vars += [{'params':self.renderModule.parameters(), 'lr':lr_init_network}]
        return grad_vars

    def compute_densityfeature(self, xyz_sampled):

        coordinate_line = torch.stack((xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)


        line_coef_point = F.grid_sample(self.density_line[0], coordinate_line[[0]],
                                            align_corners=True).view(-1, *xyz_sampled.shape[:1])
        line_coef_point = line_coef_point * F.grid_sample(self.density_line[1], coordinate_line[[1]],
                                        align_corners=True).view(-1, *xyz_sampled.shape[:1])
        line_coef_point = line_coef_point * F.grid_sample(self.density_line[2], coordinate_line[[2]],
                                        align_corners=True).view(-1, *xyz_sampled.shape[:1])
        sigma_feature = torch.sum(line_coef_point, dim=0)
        
        
        return sigma_feature
    
    def compute_appfeature(self, xyz_sampled):

        coordinate_line = torch.stack(
            (xyz_sampled[..., self.vecMode[0]], xyz_sampled[..., self.vecMode[1]], xyz_sampled[..., self.vecMode[2]]))
        coordinate_line = torch.stack((torch.zeros_like(coordinate_line), coordinate_line), dim=-1).detach().view(3, -1, 1, 2)


        line_coef_point = F.grid_sample(self.app_line[0], coordinate_line[[0]],
                                            align_corners=True).view(-1, *xyz_sampled.shape[:1])
        line_coef_point = line_coef_point * F.grid_sample(self.app_line[1], coordinate_line[[1]],
                                                          align_corners=True).view(-1, *xyz_sampled.shape[:1])
        line_coef_point = line_coef_point * F.grid_sample(self.app_line[2], coordinate_line[[2]],
                                                          align_corners=True).view(-1, *xyz_sampled.shape[:1])

        return self.basis_mat(line_coef_point.T)
    

    @torch.no_grad()
    def up_sampling_Vector(self, density_line_coef, app_line_coef, res_target):

        for i in range(len(self.vecMode)):
            vec_id = self.vecMode[i]
            density_line_coef[i] = torch.nn.Parameter(
                F.interpolate(density_line_coef[i].data, size=(res_target[vec_id], 1), mode='bilinear', align_corners=True))
            app_line_coef[i] = torch.nn.Parameter(
                F.interpolate(app_line_coef[i].data, size=(res_target[vec_id], 1), mode='bilinear', align_corners=True))

        return density_line_coef, app_line_coef

    @torch.no_grad()
    def upsample_volume_grid(self, res_target):
        self.density_line, self.app_line = self.up_sampling_Vector(self.density_line, self.app_line, res_target)

        self.update_stepSize(res_target)
        print(f'upsamping to {res_target}')

    @torch.no_grad()
    def shrink(self, new_aabb):
        print("====> shrinking ...")
        xyz_min, xyz_max = new_aabb
        t_l, b_r = (xyz_min - self.aabb[0]) / self.units, (xyz_max - self.aabb[0]) / self.units

        t_l, b_r = torch.round(torch.round(t_l)).long(), torch.round(b_r).long() + 1
        b_r = torch.stack([b_r, self.gridSize]).amin(0)


        for i in range(len(self.vecMode)):
            mode0 = self.vecMode[i]
            self.density_line[i] = torch.nn.Parameter(
                self.density_line[i].data[...,t_l[mode0]:b_r[mode0],:]
            )
            self.app_line[i] = torch.nn.Parameter(
                self.app_line[i].data[...,t_l[mode0]:b_r[mode0],:]
            )

        if not torch.all(self.alphaMask.gridSize == self.gridSize):
            t_l_r, b_r_r = t_l / (self.gridSize-1), (b_r-1) / (self.gridSize-1)
            correct_aabb = torch.zeros_like(new_aabb)
            correct_aabb[0] = (1-t_l_r)*self.aabb[0] + t_l_r*self.aabb[1]
            correct_aabb[1] = (1-b_r_r)*self.aabb[0] + b_r_r*self.aabb[1]
            print("aabb", new_aabb, "\ncorrect aabb", correct_aabb)
            new_aabb = correct_aabb

        newSize = b_r - t_l
        self.aabb = new_aabb
        self.update_stepSize((newSize[0], newSize[1], newSize[2]))

    def density_L1(self):
        total = 0
        for idx in range(len(self.density_line)):
            total = total + torch.mean(torch.abs(self.density_line[idx]))
        return total

    def TV_loss_density(self, reg):
        total = 0
        for idx in range(len(self.density_line)):
            total = total + reg(self.density_line[idx]) * 1e-3
        return total

    def TV_loss_app(self, reg):
        total = 0
        for idx in range(len(self.app_line)):
            total = total + reg(self.app_line[idx]) * 1e-3
        return total
thirdparty/tensorf/models/tensoRF.py
oppo-us-research/NeuRBF
import torch
import torch.nn
import torch.nn.functional as F
from .sh import eval_sh_bases
import numpy as np
import time


def positional_encoding(positions, freqs):

        freq_bands = (2**torch.arange(freqs).float()).to(positions.device)  # (F,)
        pts = (positions[..., None] * freq_bands).reshape(
            positions.shape[:-1] + (freqs * positions.shape[-1], ))  # (..., DF)
        pts = torch.cat([torch.sin(pts), torch.cos(pts)], dim=-1)
        return pts

def raw2alpha(sigma, dist):
    # sigma, dist  [N_rays, N_samples]
    alpha = 1. - torch.exp(-sigma*dist)

    T = torch.cumprod(torch.cat([torch.ones(alpha.shape[0], 1).to(alpha.device), 1. - alpha + 1e-10], -1), -1)

    weights = alpha * T[:, :-1]  # [N_rays, N_samples]
    return alpha, weights, T[:,-1:]


def SHRender(xyz_sampled, viewdirs, features):
    sh_mult = eval_sh_bases(2, viewdirs)[:, None]
    rgb_sh = features.view(-1, 3, sh_mult.shape[-1])
    rgb = torch.relu(torch.sum(sh_mult * rgb_sh, dim=-1) + 0.5)
    return rgb


def RGBRender(xyz_sampled, viewdirs, features):

    rgb = features
    return rgb

class AlphaGridMask(torch.nn.Module):
    def __init__(self, device, aabb, alpha_volume):
        super(AlphaGridMask, self).__init__()
        self.device = device

        self.aabb=aabb.to(self.device)
        self.aabbSize = self.aabb[1] - self.aabb[0]
        self.invgridSize = 1.0/self.aabbSize * 2
        self.alpha_volume = alpha_volume.view(1,1,*alpha_volume.shape[-3:])
        self.gridSize = torch.LongTensor([alpha_volume.shape[-1],alpha_volume.shape[-2],alpha_volume.shape[-3]]).to(self.device)

    def sample_alpha(self, xyz_sampled, normalize=True):
        if normalize:
            xyz_sampled = self.normalize_coord(xyz_sampled)
        alpha_vals = F.grid_sample(self.alpha_volume, xyz_sampled.view(1,-1,1,1,3), align_corners=True).view(-1)
        return alpha_vals

    def normalize_coord(self, xyz_sampled):
        return (xyz_sampled-self.aabb[0]) * self.invgridSize - 1


class MLPRender_Fea(torch.nn.Module):
    def __init__(self, inChanel, viewpe=6, feape=6, featureC=128):
        super(MLPRender_Fea, self).__init__()

        self.in_mlpC = 2*viewpe*3 + 2*feape*inChanel + 3 + inChanel
        self.viewpe = viewpe
        self.feape = feape
        layer1 = torch.nn.Linear(self.in_mlpC, featureC)
        layer2 = torch.nn.Linear(featureC, featureC)
        layer3 = torch.nn.Linear(featureC,3)

        self.mlp = torch.nn.Sequential(layer1, torch.nn.ReLU(inplace=True), layer2, torch.nn.ReLU(inplace=True), layer3)
        torch.nn.init.constant_(self.mlp[-1].bias, 0)

    def forward(self, pts, viewdirs, features):
        indata = [features, viewdirs]
        if self.feape > 0:
            indata += [positional_encoding(features, self.feape)]
        if self.viewpe > 0:
            indata += [positional_encoding(viewdirs, self.viewpe)]
        mlp_in = torch.cat(indata, dim=-1)
        rgb = self.mlp(mlp_in)
        rgb = torch.sigmoid(rgb)

        return rgb, None

class MLPRender_PE(torch.nn.Module):
    def __init__(self,inChanel, viewpe=6, pospe=6, featureC=128):
        super(MLPRender_PE, self).__init__()

        self.in_mlpC = (3+2*viewpe*3)+ (3+2*pospe*3)  + inChanel #
        self.viewpe = viewpe
        self.pospe = pospe
        layer1 = torch.nn.Linear(self.in_mlpC, featureC)
        layer2 = torch.nn.Linear(featureC, featureC)
        layer3 = torch.nn.Linear(featureC,3)

        self.mlp = torch.nn.Sequential(layer1, torch.nn.ReLU(inplace=True), layer2, torch.nn.ReLU(inplace=True), layer3)
        torch.nn.init.constant_(self.mlp[-1].bias, 0)

    def forward(self, pts, viewdirs, features):
        indata = [features, viewdirs]
        if self.pospe > 0:
            indata += [positional_encoding(pts, self.pospe)]
        if self.viewpe > 0:
            indata += [positional_encoding(viewdirs, self.viewpe)]
        mlp_in = torch.cat(indata, dim=-1)
        rgb = self.mlp(mlp_in)
        rgb = torch.sigmoid(rgb)

        return rgb

class MLPRender(torch.nn.Module):
    def __init__(self,inChanel, viewpe=6, featureC=128):
        super(MLPRender, self).__init__()

        self.in_mlpC = (3+2*viewpe*3) + inChanel
        self.viewpe = viewpe
        
        layer1 = torch.nn.Linear(self.in_mlpC, featureC)
        layer2 = torch.nn.Linear(featureC, featureC)
        layer3 = torch.nn.Linear(featureC,3)

        self.mlp = torch.nn.Sequential(layer1, torch.nn.ReLU(inplace=True), layer2, torch.nn.ReLU(inplace=True), layer3)
        torch.nn.init.constant_(self.mlp[-1].bias, 0)

    def forward(self, pts, viewdirs, features):
        indata = [features, viewdirs]
        if self.viewpe > 0:
            indata += [positional_encoding(viewdirs, self.viewpe)]
        mlp_in = torch.cat(indata, dim=-1)
        rgb = self.mlp(mlp_in)
        rgb = torch.sigmoid(rgb)

        return rgb



class TensorBase(torch.nn.Module):
    def __init__(self, aabb, gridSize, device, 
                 density_n_comp = 8, appearance_n_comp = 24, app_dim = 27,
                    shadingMode = 'MLP_PE', alphaMask = None, near_far=[2.0,6.0],
                    density_shift = -10, alphaMask_thres=0.001, distance_scale=25, rayMarch_weight_thres=0.0001,
                    pos_pe = 6, view_pe = 6, fea_pe = 6, featureC=128, step_ratio=2.0,
                    fea2denseAct = 'softplus', gridSize_real=None, **kwargs):
        super(TensorBase, self).__init__()
        self.args = kwargs['args']

        if not hasattr(self, 'density_n_comp'):
            self.density_n_comp = density_n_comp
        if not hasattr(self, 'app_n_comp'):
            self.app_n_comp = appearance_n_comp
        self.app_dim = app_dim
        self.aabb = aabb
        self.alphaMask = alphaMask
        self.device=device

        self.density_shift = density_shift
        self.alphaMask_thres = alphaMask_thres
        self.distance_scale = distance_scale
        self.rayMarch_weight_thres = rayMarch_weight_thres
        self.fea2denseAct = fea2denseAct

        self.near_far = near_far
        self.step_ratio = step_ratio


        self.update_stepSize(gridSize)

        self.matMode = [[0,1], [0,2], [1,2]]
        self.vecMode =  [2, 1, 0]
        self.comp_w = [1,1,1]

        self.init_svd_volume(gridSize_real, device)

        self.shadingMode, self.pos_pe, self.view_pe, self.fea_pe, self.featureC = shadingMode, pos_pe, view_pe, fea_pe, featureC
        if self.args.model_name.startswith('models.tensoRF.'):
            self.init_render_func(shadingMode, self.app_dim, pos_pe, view_pe, fea_pe, featureC, device, args=kwargs['args'])

    def init_render_func(self, shadingMode, in_channel, pos_pe, view_pe, fea_pe, featureC, device, args):
        if shadingMode == 'MLP_PE':
            self.renderModule = MLPRender_PE(in_channel, view_pe, pos_pe, featureC).to(device)
        elif shadingMode == 'MLP_Fea':
            self.renderModule = MLPRender_Fea(in_channel, view_pe, fea_pe, featureC).to(device)
        elif shadingMode == 'MLP':
            self.renderModule = MLPRender(in_channel, view_pe, featureC).to(device)
        elif shadingMode == 'SH':
            self.renderModule = SHRender
        elif shadingMode == 'RGB':
            assert in_channel == 3
            self.renderModule = RGBRender
        elif self.shadingMode == 'ASG_Fea':
            from ...nrff.nrff import RenderingNet
            self.renderModule = RenderingNet(
                8, 16, in_channel, featureC, device=device, view_pe=view_pe, 
                btn_freq=args['btn_freq'] if 'btn_freq' in args else [],
                )
        else:
            print("Unrecognized shading module")
            exit()
        print("pos_pe", pos_pe, "view_pe", view_pe, "fea_pe", fea_pe)
        print(self.renderModule)

    def update_stepSize(self, gridSize):
        print("aabb", self.aabb.view(-1))
        print("grid size", gridSize)
        self.aabbSize = self.aabb[1] - self.aabb[0]
        self.invaabbSize = 2.0/self.aabbSize
        self.gridSize= torch.LongTensor(gridSize).to(self.device)
        self.units=self.aabbSize / (self.gridSize-1)
        self.stepSize=torch.mean(self.units)*self.step_ratio
        self.aabbDiag = torch.sqrt(torch.sum(torch.square(self.aabbSize)))
        self.nSamples=int((self.aabbDiag / self.stepSize).item()) + 1
        print("sampling step size: ", self.stepSize)
        print("sampling number: ", self.nSamples)

    def init_svd_volume(self, res, device):
        pass

    def compute_features(self, xyz_sampled):
        pass
    
    def compute_rbf(self, xyz_sampled):
        return {}

    def compute_densityfeature(self, xyz_sampled):
        pass
    
    def compute_appfeature(self, xyz_sampled):
        pass
    
    def normalize_coord(self, xyz_sampled):
        return (xyz_sampled-self.aabb[0]) * self.invaabbSize - 1

    def get_optparam_groups(self, lr_init_spatial = 0.02, lr_init_network = 0.001):
        pass

    def get_kwargs(self):
        return {
            'aabb': self.aabb,
            'gridSize':self.gridSize.tolist(),
            'density_n_comp': self.density_n_comp,
            'appearance_n_comp': self.app_n_comp,
            'app_dim': self.app_dim,

            'density_shift': self.density_shift,
            'alphaMask_thres': self.alphaMask_thres,
            'distance_scale': self.distance_scale,
            'rayMarch_weight_thres': self.rayMarch_weight_thres,
            'fea2denseAct': self.fea2denseAct,

            'near_far': self.near_far,
            'step_ratio': self.step_ratio,

            'shadingMode': self.shadingMode,
            'pos_pe': self.pos_pe,
            'view_pe': self.view_pe,
            'fea_pe': self.fea_pe,
            'featureC': self.featureC
        }

    def save(self, path):
        kwargs = self.get_kwargs()
        ckpt = {'kwargs': kwargs, 'state_dict': self.state_dict()}
        if self.alphaMask is not None:
            alpha_volume = self.alphaMask.alpha_volume.bool().cpu().numpy()
            ckpt.update({'alphaMask.shape':alpha_volume.shape})
            ckpt.update({'alphaMask.mask':np.packbits(alpha_volume.reshape(-1))})
            ckpt.update({'alphaMask.aabb': self.alphaMask.aabb.cpu()})
        torch.save(ckpt, path)

    def load(self, ckpt):
        if 'alphaMask.aabb' in ckpt.keys():
            length = np.prod(ckpt['alphaMask.shape'])
            alpha_volume = torch.from_numpy(np.unpackbits(ckpt['alphaMask.mask'])[:length].reshape(ckpt['alphaMask.shape']))
            self.alphaMask = AlphaGridMask(self.device, ckpt['alphaMask.aabb'].to(self.device), alpha_volume.float().to(self.device))
        self.load_state_dict(ckpt['state_dict'])


    def sample_ray_ndc(self, rays_o, rays_d, is_train=True, N_samples=-1):
        N_samples = N_samples if N_samples > 0 else self.nSamples
        near, far = self.near_far
        interpx = torch.linspace(near, far, N_samples).unsqueeze(0).to(rays_o)
        if is_train:
            interpx += torch.rand_like(interpx).to(rays_o) * ((far - near) / N_samples)

        rays_pts = rays_o[..., None, :] + rays_d[..., None, :] * interpx[..., None]
        mask_outbbox = ((self.aabb[0] > rays_pts) | (rays_pts > self.aabb[1])).any(dim=-1)
        return rays_pts, interpx, ~mask_outbbox

    def sample_ray(self, rays_o, rays_d, is_train=True, N_samples=-1):
        N_samples = N_samples if N_samples>0 else self.nSamples
        stepsize = self.stepSize
        near, far = self.near_far
        vec = torch.where(rays_d==0, torch.full_like(rays_d, 1e-6), rays_d)
        rate_a = (self.aabb[1] - rays_o) / vec
        rate_b = (self.aabb[0] - rays_o) / vec
        t_min = torch.minimum(rate_a, rate_b).amax(-1).clamp(min=near, max=far)

        rng = torch.arange(N_samples)[None].float()
        if is_train:
            rng = rng.repeat(rays_d.shape[-2],1)
            rng += torch.rand_like(rng[:,[0]])
        step = stepsize * rng.to(rays_o.device)
        interpx = (t_min[...,None] + step)

        rays_pts = rays_o[...,None,:] + rays_d[...,None,:] * interpx[...,None]
        mask_outbbox = ((self.aabb[0]>rays_pts) | (rays_pts>self.aabb[1])).any(dim=-1)

        return rays_pts, interpx, ~mask_outbbox


    def shrink(self, new_aabb, voxel_size):
        pass

    @torch.no_grad()
    def getDenseAlpha(self,gridSize=None):
        gridSize = self.gridSize if gridSize is None else gridSize

        samples = torch.stack(torch.meshgrid(
            torch.linspace(0, 1, gridSize[0]),
            torch.linspace(0, 1, gridSize[1]),
            torch.linspace(0, 1, gridSize[2]),
        ), -1).to(self.device)
        dense_xyz = self.aabb[0] * (1-samples) + self.aabb[1] * samples

        alpha = torch.zeros_like(dense_xyz[...,0])
        for i in range(gridSize[0]):
            alpha[i] = self.compute_alpha(dense_xyz[i].view(-1,3), self.stepSize).view((gridSize[1], gridSize[2]))
        return alpha, dense_xyz

    @torch.no_grad()
    def updateAlphaMask(self, gridSize=(200,200,200)):

        alpha, dense_xyz = self.getDenseAlpha(gridSize)
        dense_xyz = dense_xyz.transpose(0,2).contiguous()
        alpha = alpha.clamp(0,1).transpose(0,2).contiguous()[None,None]
        total_voxels = gridSize[0] * gridSize[1] * gridSize[2]

        ks = 3
        alpha = F.max_pool3d(alpha, kernel_size=ks, padding=ks // 2, stride=1).view(gridSize[::-1])
        alpha[alpha>=self.alphaMask_thres] = 1
        alpha[alpha<self.alphaMask_thres] = 0

        self.alphaMask = AlphaGridMask(self.device, self.aabb, alpha)

        valid_xyz = dense_xyz[alpha>0.5]

        xyz_min = valid_xyz.amin(0)
        xyz_max = valid_xyz.amax(0)

        new_aabb = torch.stack((xyz_min, xyz_max))

        total = torch.sum(alpha)
        print(f"bbox: {xyz_min, xyz_max} alpha rest %%%f"%(total/total_voxels*100))
        return new_aabb

    @torch.no_grad()
    def filtering_rays(self, all_rays, all_rgbs, N_samples=256, chunk=10240*5, bbox_only=False):
        print('========> filtering rays ...')
        tt = time.time()

        N = torch.tensor(all_rays.shape[:-1]).prod()

        mask_filtered = []
        idx_chunks = torch.split(torch.arange(N), chunk)
        for idx_chunk in idx_chunks:
            rays_chunk = all_rays[idx_chunk].to(self.device)

            rays_o, rays_d = rays_chunk[..., :3], rays_chunk[..., 3:6]
            if bbox_only:
                vec = torch.where(rays_d == 0, torch.full_like(rays_d, 1e-6), rays_d)
                rate_a = (self.aabb[1] - rays_o) / vec
                rate_b = (self.aabb[0] - rays_o) / vec
                t_min = torch.minimum(rate_a, rate_b).amax(-1)#.clamp(min=near, max=far)
                t_max = torch.maximum(rate_a, rate_b).amin(-1)#.clamp(min=near, max=far)
                mask_inbbox = t_max > t_min

            else:
                xyz_sampled, _,_ = self.sample_ray(rays_o, rays_d, N_samples=N_samples, is_train=False)
                mask_inbbox= (self.alphaMask.sample_alpha(xyz_sampled).view(xyz_sampled.shape[:-1]) > 0).any(-1)

            mask_filtered.append(mask_inbbox.cpu())

        mask_filtered = torch.cat(mask_filtered).view(all_rgbs.shape[:-1])

        print(f'Ray filtering done! takes {time.time()-tt} s. ray mask ratio: {torch.sum(mask_filtered) / N}')
        return all_rays[mask_filtered], all_rgbs[mask_filtered]


    def feature2density(self, density_features):
        if self.fea2denseAct == "softplus":
            return F.softplus(density_features+self.density_shift)
        elif self.fea2denseAct == "relu":
            return F.relu(density_features)


    def compute_alpha(self, xyz_locs, length=1):

        if self.alphaMask is not None:
            alphas = self.alphaMask.sample_alpha(xyz_locs)
            alpha_mask = alphas > 0
        else:
            alpha_mask = torch.ones_like(xyz_locs[:,0], dtype=bool)


        sigma = torch.zeros(xyz_locs.shape[:-1], device=xyz_locs.device)

        if alpha_mask.any():
            xyz_sampled = self.normalize_coord(xyz_locs[alpha_mask])
            others = self.compute_rbf(xyz_sampled) if hasattr(self, 'lcd0') else {}
            sigma_feature = self.compute_densityfeature(xyz_sampled, others)
            validsigma = self.feature2density(sigma_feature)
            sigma[alpha_mask] = validsigma


        alpha = 1 - torch.exp(-sigma*length).view(xyz_locs.shape[:-1])

        return alpha


    def forward(self, rays_chunk, white_bg=True, is_train=False, ndc_ray=False, N_samples=-1):

        # sample points
        viewdirs = rays_chunk[:, 3:6]
        if ndc_ray:
            xyz_sampled, z_vals, ray_valid = self.sample_ray_ndc(rays_chunk[:, :3], viewdirs, is_train=is_train,N_samples=N_samples)
            dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])), dim=-1)
            rays_norm = torch.norm(viewdirs, dim=-1, keepdim=True)
            dists = dists * rays_norm
            viewdirs = viewdirs / rays_norm
        else:
            xyz_sampled, z_vals, ray_valid = self.sample_ray(rays_chunk[:, :3], viewdirs, is_train=is_train,N_samples=N_samples)
            dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])), dim=-1)
        viewdirs = viewdirs.view(-1, 1, 3).expand(xyz_sampled.shape)
        
        # xyz_sampled: [n_ray, n_pts_per_ray, 3]
        # ray_valid: [n_ray, n_pts_per_ray]
        
        if self.alphaMask is not None:
            alphas = self.alphaMask.sample_alpha(xyz_sampled[ray_valid])  # [n_pts_valid]
            alpha_mask = alphas > 0
            ray_invalid = ~ray_valid
            ray_invalid[ray_valid] |= (~alpha_mask)
            ray_valid = ~ray_invalid  # [n_ray, n_pts_per_ray]

        sigma = torch.zeros(xyz_sampled.shape[:-1], device=xyz_sampled.device)
        rgb = torch.zeros((*xyz_sampled.shape[:2], 3), device=xyz_sampled.device)

        if ray_valid.any():
            xyz_sampled = self.normalize_coord(xyz_sampled)
            others = self.compute_rbf(xyz_sampled[ray_valid]) if hasattr(self, 'lcd0') else {}
            sigma_feature = self.compute_densityfeature(xyz_sampled[ray_valid], others)

            validsigma = self.feature2density(sigma_feature)
            sigma[ray_valid] = validsigma

        alpha, weight, bg_weight = raw2alpha(sigma, dists * self.distance_scale)
        app_mask = weight > self.rayMarch_weight_thres  # [n_ray, n_pts_per_ray]

        normals = None
        valid_viewdirs = None
        valid_weights = None
        if app_mask.any():
            valid_viewdirs = viewdirs[app_mask]
            valid_weights = weight[app_mask]
            valid_xyz = xyz_sampled[app_mask]

            if hasattr(self, 'lcd0'):
                for k, v in others.items():
                    temp = torch.zeros([*xyz_sampled.shape[:2], *v.shape[1:]], device=v.device, dtype=v.dtype)
                    temp[ray_valid] = v
                    others[k] = temp[app_mask]
            else:
                others = self.compute_rbf(valid_xyz)
            app_features = self.compute_appfeature(valid_xyz, others)
            valid_rgbs, normals = self.renderModule(valid_xyz, valid_viewdirs, app_features)
            if rgb.dtype is not valid_rgbs.dtype:
                rgb = torch.zeros((*xyz_sampled.shape[:2], 3), device=xyz_sampled.device, dtype=valid_rgbs.dtype)
            rgb[app_mask] = valid_rgbs

        acc_map = torch.sum(weight, -1)
        rgb_map = torch.sum(weight[..., None] * rgb, -2)

        if white_bg or (is_train and torch.rand((1,))<0.5):
            rgb_map = rgb_map + (1. - acc_map[..., None])

        rgb_map = rgb_map.clamp(0,1)

        with torch.no_grad():
            depth_map = torch.sum(weight * z_vals, -1)
            depth_map = depth_map + (1. - acc_map) * rays_chunk[..., -1]

        return {'rgb_map':rgb_map,
                'depth_map':depth_map,
                'normals':normals,
                'valid_viewdirs':valid_viewdirs,
                'valid_weights':valid_weights
            }


    @torch.no_grad()
    def get_pts_data(self, pts, normalize_pts=False):
        if normalize_pts:
            pts = (pts - self.aabb_rbf[0]) * 2.0 / (self.aabb_rbf[1] - self.aabb_rbf[0]) - 1

        chunk_size = 2**21
        density_all = []
        features_all = []
        for i in range(0, pts.shape[0], chunk_size):
            pts_i = pts[i:i+chunk_size]

            pts_valid = torch.ones(pts_i.shape[0], dtype=torch.bool, device=pts_i.device)
            if self.alphaMask is not None:
                alphas = self.alphaMask.sample_alpha(pts_i, normalize=False)
                alpha_mask = alphas > 0
                pts_invalid = ~pts_valid
                pts_invalid[pts_valid] |= (~alpha_mask)
                pts_valid = ~pts_invalid  # [n_pts]

            sigma = torch.zeros(pts_i.shape[:-1], device=pts_i.device)
            app_features = torch.zeros((*pts_i.shape[:-1], self.app_dim), device=pts_i.device)

            # Get density
            if pts_valid.any():
                others = self.compute_rbf(pts_i[pts_valid])
                sigma_feature = self.compute_densityfeature(pts_i[pts_valid], others)
                validsigma = self.feature2density(sigma_feature)
                sigma[pts_valid] = validsigma

            # Get appearance features
            if pts_valid.any():
                app_features[pts_valid] = self.compute_appfeature(pts_i[pts_valid], others)
                
            density_all.append(sigma.cpu())
            features_all.append(app_features.cpu())
        density_all = torch.cat(density_all, 0)
        features_all = torch.cat(features_all, 0)
        del sigma, app_features
        torch.cuda.empty_cache()

        return density_all, features_all


    def step_after_iter(self, step):
        pass


def count_params(tensorf, verbose=True):
    if hasattr(tensorf, 'density_field'):
        n_params_density = sum([torch.numel(p) for p in tensorf.density_field.parameters() if p.requires_grad])
    elif hasattr(tensorf, 'density_plane'):
        n_params_density = sum([torch.numel(p) for p in tensorf.density_plane.parameters() if p.requires_grad]) + \
            sum([torch.numel(p) for p in tensorf.density_line.parameters() if p.requires_grad])
    if hasattr(tensorf, 'appearance_field'):
        n_params_app = sum([torch.numel(p) for p in tensorf.appearance_field.parameters() if p.requires_grad])
    elif hasattr(tensorf, 'app_plane'):
        n_params_app = sum([torch.numel(p) for p in tensorf.app_plane.parameters() if p.requires_grad]) + \
            sum([torch.numel(p) for p in tensorf.app_line.parameters() if p.requires_grad])

    if hasattr(tensorf, 'lc0'):
        n_params_app += sum([torch.numel(p) for p in tensorf.lc0.parameters() if p.requires_grad])
    if hasattr(tensorf, 'lcd0'):
        n_params_density += sum([torch.numel(p) for p in tensorf.lcd0.parameters() if p.requires_grad])
    if hasattr(tensorf, 'lcb0'):
        n_params_app += sum([torch.numel(p) for p in [tensorf.lcb0] if p.requires_grad])
    if hasattr(tensorf, 'kc0'):
        n_params_app += sum([torch.numel(p) for p in tensorf.kc0.parameters()])
    if hasattr(tensorf, 'ks0'):
        n_params_app += sum([torch.numel(p) for p in tensorf.ks0.parameters()]) // 9 * 6
    
    n_params_mlp = 0
    if hasattr(tensorf, 'renderModule'):
        n_params_mlp += sum([torch.numel(p) for p in tensorf.renderModule.parameters() if p.requires_grad])
    if hasattr(tensorf, 'basis_mat'):
        n_params_mlp += sum([torch.numel(p) for p in tensorf.basis_mat.parameters() if p.requires_grad])
    if hasattr(tensorf, 'basis_mat_ds'):
        n_params_mlp += sum([torch.numel(p) for p in tensorf.basis_mat_ds.parameters() if p.requires_grad])

    n_params = {'mlp': n_params_mlp, 'density': n_params_density, 'app': n_params_app}
    n_params['feat'] = n_params['density'] + n_params['app']
    n_params['total'] = n_params['mlp'] + n_params['feat']
    if verbose:
        print('Num feature parameters:', n_params['feat'])
        print('Num MLP parameters:', n_params['mlp'])
    return n_params

thirdparty/tensorf/models/tensorBase.py
oppo-us-research/NeuRBF
import configargparse
import importlib.util
import os
from typing import List, Dict, Any
from omegaconf import OmegaConf
import pprint
import glob


def config_parser(cmd=None, mode='config'):
    parser = configargparse.ArgumentParser()
    parser.add_argument('--config_init', type=str, help='config init file path')
    parser.add_argument('--config', type=str, help='config file path')
    parser.add_argument("--expname", type=str,
                        help='experiment name')
    parser.add_argument("--basedir", type=str, default='./log',
                        help='where to store ckpts and logs')
    parser.add_argument("--add_timestamp", type=int, default=0,
                        help='add timestamp to dir')
    parser.add_argument("--datadir", type=str, default='./data/nerf_synthetic/lego',
                        help='input data directory')
    parser.add_argument("--data_name", type=str, default=None)
    parser.add_argument("--progress_refresh_rate", type=int, default=10,
                        help='how many iterations to show psnrs or iters')

    parser.add_argument('--with_depth', action='store_true')
    parser.add_argument('--downsample_train', type=float, default=1.0)
    parser.add_argument('--downsample_test', type=float, default=1.0)

    parser.add_argument('--model_name', type=str, default='TensorVMSplitRBF')

    # loader options
    parser.add_argument("--batch_size", type=int, default=4096)
    parser.add_argument("--batch_size_init", type=int, default=None)
    parser.add_argument("--n_iters", type=int, default=30000)

    parser.add_argument('--dataset_name', type=str, default='blender')


    # training options
    # learning rate
    parser.add_argument("--lr_init", type=float, default=0.02,
                        help='learning rate')
    parser.add_argument("--lr_basis", type=float, default=1e-3,
                        help='learning rate')
    parser.add_argument("--lr_decay_iters", type=int, default=-1,
                        help = 'number of iterations the lr will decay to the target ratio; -1 will set it to n_iters')
    parser.add_argument("--lr_decay_target_ratio", type=float, default=0.1,
                        help='the target decay ratio; after decay_iters inital lr decays to lr*ratio')
    parser.add_argument("--lr_upsample_reset", type=int, default=1,
                        help='reset lr to inital after upsampling')

    # loss
    parser.add_argument("--L1_weight_inital", type=float, default=0.0,
                        help='loss weight')
    parser.add_argument("--L1_weight_rest", type=float, default=0,
                        help='loss weight')
    parser.add_argument("--Ortho_weight", type=float, default=0.0,
                        help='loss weight')
    parser.add_argument("--TV_weight_density", type=float, default=0.0,
                        help='loss weight')
    parser.add_argument("--TV_weight_app", type=float, default=0.0,
                        help='loss weight')

    # model
    # volume options
    parser.add_argument("--n_lamb_sigma", type=int, action="append")
    parser.add_argument("--n_lamb_sh", type=int, action="append")
    parser.add_argument("--data_dim_color", type=int, default=27)

    parser.add_argument("--rm_weight_mask_thre", type=float, default=0.0001,
                        help='mask points in ray marching')
    parser.add_argument("--alpha_mask_thre", type=float, default=0.0001,
                        help='threshold for creating alpha mask volume')
    parser.add_argument("--distance_scale", type=float, default=25,
                        help='scaling sampling distance for computation')
    parser.add_argument("--density_shift", type=float, default=-10,
                        help='shift density in softplus; making density = 0  when feature == 0')

    # network decoder
    parser.add_argument("--shadingMode", type=str, default="MLP_PE",
                        help='which shading mode to use')
    parser.add_argument("--pos_pe", type=int, default=6,
                        help='number of pe for pos')
    parser.add_argument("--view_pe", type=int, default=6,
                        help='number of pe for view')
    parser.add_argument("--fea_pe", type=int, default=6,
                        help='number of pe for features')
    parser.add_argument("--featureC", type=int, default=128,
                        help='hidden feature channel in MLP')



    parser.add_argument("--ckpt", type=str, default=None,
                        help='specific weights npy file to reload for coarse network')
    parser.add_argument("--render_only", type=int, default=0)
    parser.add_argument("--render_test", type=int, default=0)
    parser.add_argument("--render_train", type=int, default=0)
    parser.add_argument("--render_path", type=int, default=0)
    parser.add_argument("--export_mesh", type=int, default=0)

    # rendering options
    parser.add_argument('--lindisp', default=False, action="store_true",
                        help='use disparity depth sampling')
    parser.add_argument("--perturb", type=float, default=1.,
                        help='set to 0. for no jitter, 1. for jitter')
    parser.add_argument("--accumulate_decay", type=float, default=0.998)
    parser.add_argument("--fea2denseAct", type=str, default='softplus')
    parser.add_argument('--ndc_ray', type=int, default=0)
    parser.add_argument('--nSamples', type=int, default=1e6,
                        help='sample point each ray, pass 1e6 if automatic adjust')
    parser.add_argument('--step_ratio',type=float,default=0.5)


    ## blender flags
    parser.add_argument("--white_bkgd", action='store_true',
                        help='set to render synthetic data on a white bkgd (always use for dvoxels)')



    parser.add_argument('--N_voxel_init',
                        type=int,
                        default=100**3)
    parser.add_argument('--N_voxel_final',
                        type=int,
                        default=300**3)
    parser.add_argument("--upsamp_list", type=int, action="append")
    parser.add_argument("--update_AlphaMask_list", type=int, action="append")

    parser.add_argument('--idx_view',
                        type=int,
                        default=0)
    # logging/saving options
    parser.add_argument("--N_vis", type=int, default=5,
                        help='N images to vis')
    parser.add_argument("--vis_every", type=int, default=10000,
                        help='frequency of visualize the image')
    
    if cmd is not None:
        args = parser.parse_args(cmd)
    else:
        args = parser.parse_args()
    args = OmegaConf.create(vars(args))

    # Import config
    if mode == 'config_init':
        args.config = args.config_init
    if args.config is not None:
        spec = importlib.util.spec_from_file_location(os.path.basename(args.config), args.config)
        cfg = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(cfg)
        config = OmegaConf.create(cfg.config)

        # Override from config into argparse
        for k, v in config.items():
            # Keep certain args if they are set in cli
            if k in ['data_name', 'batch_size_init'] and getattr(args, k) is not None:
                continue
            setattr(args, k, v)

        # Replace placeholders with actual value
        args['datadir'] = args['datadir'].replace('data_name', args['data_name'])
        args['expname'] = args['expname'].replace('data_name', args['data_name'])
        if 'rbf_config' in args:
            if 'init_data_fp' in args['rbf_config']:
                args['rbf_config']['init_data_fp'] = args['rbf_config']['init_data_fp'].replace('data_name', args['data_name'])

    # Add version id
    if args.add_version_id:
        exp_list = glob.glob(f'{args["basedir"]}/{args["expname"]}-v*'.replace('[', '[[]'))
        if len(exp_list) == 0:
            vid = 0
        else:
            vid = max([int(v.split(f'{args["expname"]}-v')[-1]) for v in exp_list]) + 1
        args["expname"] += f'-v{vid}'

    return args

thirdparty/tensorf/opt.py
oppo-us-research/NeuRBF
import torch, os, imageio, sys
from tqdm.auto import tqdm
from .dataLoader.ray_utils import get_rays
from .utils import *
from .dataLoader.ray_utils import ndc_rays_blender


def OctreeRender_trilinear_fast(rays, tensorf, chunk=4096, N_samples=-1, ndc_ray=False, white_bg=True, is_train=False, device='cuda'):

    rgbs, alphas, depth_maps, weights, uncertainties = [], [], [], [], []
    N_rays_all = rays.shape[0]
    for chunk_idx in range(N_rays_all // chunk + int(N_rays_all % chunk > 0)):
        rays_chunk = rays[chunk_idx * chunk:(chunk_idx + 1) * chunk].to(device)

        out = tensorf(rays_chunk, is_train=is_train, white_bg=white_bg, ndc_ray=ndc_ray, N_samples=N_samples)

        rgbs.append(out['rgb_map'])
        depth_maps.append(out['depth_map'])

    return torch.cat(rgbs), None, torch.cat(depth_maps), None, out

@torch.no_grad()
def evaluation(test_dataset,tensorf, args, renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
               white_bg=False, ndc_ray=False, compute_extra_metrics=True, device='cuda', 
               save_img=True, save_video=True):
    PSNRs, rgb_maps, depth_maps = [], [], []
    ssims,l_alex,l_vgg=[],[],[]
    os.makedirs(savePath, exist_ok=True)
    os.makedirs(savePath+"/rgbd", exist_ok=True)

    try:
        tqdm._instances.clear()
    except Exception:
        pass

    near_far = test_dataset.near_far
    img_eval_interval = 1 if N_vis < 0 else max(test_dataset.all_rays.shape[0] // N_vis,1)
    idxs = list(range(0, test_dataset.all_rays.shape[0], img_eval_interval))
    for idx, samples in tqdm(enumerate(test_dataset.all_rays[0::img_eval_interval]), file=sys.stdout):

        W, H = test_dataset.img_wh
        rays = samples.view(-1,samples.shape[-1])

        rgb_map, _, depth_map, _, _ = renderer(rays, tensorf, chunk=4096, N_samples=N_samples,
                                        ndc_ray=ndc_ray, white_bg = white_bg, device=device)
        rgb_map = rgb_map.clamp(0.0, 1.0)

        rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()

        depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
        if len(test_dataset.all_rgbs):
            gt_rgb = test_dataset.all_rgbs[idxs[idx]].view(H, W, 3)
            loss = torch.mean((rgb_map - gt_rgb) ** 2)
            PSNRs.append(-10.0 * np.log(loss.item()) / np.log(10.0))

            if compute_extra_metrics:
                ssim = rgb_ssim(rgb_map, gt_rgb, 1)
                l_a = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'alex', tensorf.device)
                l_v = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'vgg', tensorf.device)
                ssims.append(ssim)
                l_alex.append(l_a)
                l_vgg.append(l_v)

        rgb_map = (rgb_map.numpy() * 255).astype('uint8')
        rgb_maps.append(rgb_map)
        depth_maps.append(depth_map)
        if savePath is not None and save_img:
            imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
            rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
            imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)

    if save_video:
        imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
        imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)

    if PSNRs:
        psnr = np.mean(np.asarray(PSNRs))
        if compute_extra_metrics:
            ssim = np.mean(np.asarray(ssims))
            l_a = np.mean(np.asarray(l_alex))
            l_v = np.mean(np.asarray(l_vgg))
            np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
        else:
            np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))


    return PSNRs, ssims, l_vgg, l_alex

@torch.no_grad()
def evaluation_path(test_dataset,tensorf, c2ws, renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
                    white_bg=False, ndc_ray=False, compute_extra_metrics=True, device='cuda'):
    PSNRs, rgb_maps, depth_maps = [], [], []
    ssims,l_alex,l_vgg=[],[],[]
    os.makedirs(savePath, exist_ok=True)
    os.makedirs(savePath+"/rgbd", exist_ok=True)

    try:
        tqdm._instances.clear()
    except Exception:
        pass

    near_far = test_dataset.near_far
    for idx, c2w in tqdm(enumerate(c2ws)):

        W, H = test_dataset.img_wh

        c2w = torch.FloatTensor(c2w)
        rays_o, rays_d = get_rays(test_dataset.directions, c2w)  # both (h*w, 3)
        if ndc_ray:
            rays_o, rays_d = ndc_rays_blender(H, W, test_dataset.focal[0], 1.0, rays_o, rays_d)
        rays = torch.cat([rays_o, rays_d], 1)  # (h*w, 6)

        rgb_map, _, depth_map, _, _ = renderer(rays, tensorf, chunk=8192, N_samples=N_samples,
                                        ndc_ray=ndc_ray, white_bg = white_bg, device=device)
        rgb_map = rgb_map.clamp(0.0, 1.0)

        rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()

        depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)

        rgb_map = (rgb_map.numpy() * 255).astype('uint8')
        rgb_maps.append(rgb_map)
        depth_maps.append(depth_map)
        if savePath is not None:
            imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
            rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
            imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)

    imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=8)
    imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=8)

    if PSNRs:
        psnr = np.mean(np.asarray(PSNRs))
        if compute_extra_metrics:
            ssim = np.mean(np.asarray(ssims))
            l_a = np.mean(np.asarray(l_alex))
            l_v = np.mean(np.asarray(l_vgg))
            np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
        else:
            np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))


    return PSNRs

thirdparty/tensorf/renderer.py
oppo-us-research/NeuRBF
import cv2, torch
import numpy as np
from PIL import Image
import torchvision.transforms as T
import torch.nn.functional as F
import scipy.signal

mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]))


def visualize_depth_numpy(depth, minmax=None, cmap=cv2.COLORMAP_JET):
    """
    depth: (H, W)
    """

    x = np.nan_to_num(depth) # change nan to 0
    if minmax is None:
        mi = np.min(x[x>0]) # get minimum positive depth (ignore background)
        ma = np.max(x)
    else:
        mi,ma = minmax

    x = (x-mi)/(ma-mi+1e-8) # normalize to 0~1
    x = (255*x).astype(np.uint8)
    x_ = cv2.applyColorMap(x, cmap)
    return x_, [mi,ma]

def init_log(log, keys):
    for key in keys:
        log[key] = torch.tensor([0.0], dtype=float)
    return log

def visualize_depth(depth, minmax=None, cmap=cv2.COLORMAP_JET):
    """
    depth: (H, W)
    """
    if type(depth) is not np.ndarray:
        depth = depth.cpu().numpy()

    x = np.nan_to_num(depth) # change nan to 0
    if minmax is None:
        mi = np.min(x[x>0]) # get minimum positive depth (ignore background)
        ma = np.max(x)
    else:
        mi,ma = minmax

    x = (x-mi)/(ma-mi+1e-8) # normalize to 0~1
    x = (255*x).astype(np.uint8)
    x_ = Image.fromarray(cv2.applyColorMap(x, cmap))
    x_ = T.ToTensor()(x_)  # (3, H, W)
    return x_, [mi,ma]

def N_to_reso(n_voxels, bbox):
    xyz_min, xyz_max = bbox
    dim = len(xyz_min)
    voxel_size = ((xyz_max - xyz_min).prod() / n_voxels).pow(1 / dim)
    return ((xyz_max - xyz_min) / voxel_size).long().tolist()


def scale_reso(res, voxels):
    res_ref = [voxels**(1/3)] * 3
    params_ref = res_ref[0] * res_ref[1] + res_ref[0] * res_ref[2] + res_ref[1] * res_ref[2] + res_ref[0] + res_ref[1] + res_ref[2]
    params = res[0] * res[1] + res[0] * res[2] + res[1] * res[2] + res[0] + res[1] + res[2]
    res = (torch.tensor(res) * (params_ref / params)**0.5).long().tolist()
    return res


def cal_n_samples(reso, step_ratio=0.5):
    return int(np.linalg.norm(reso)/step_ratio)



__LPIPS__ = {}
def init_lpips(net_name, device):
    assert net_name in ['alex', 'vgg']
    import lpips
    print(f'init_lpips: lpips_{net_name}')
    return lpips.LPIPS(net=net_name, version='0.1').eval().to(device)

def rgb_lpips(np_gt, np_im, net_name, device):
    if net_name not in __LPIPS__:
        __LPIPS__[net_name] = init_lpips(net_name, device)
    gt = torch.from_numpy(np_gt).permute([2, 0, 1]).contiguous().to(device)
    im = torch.from_numpy(np_im).permute([2, 0, 1]).contiguous().to(device)
    return __LPIPS__[net_name](gt, im, normalize=True).item()


def findItem(items, target):
    for one in items:
        if one[:len(target)]==target:
            return one
    return None


''' Evaluation metrics (ssim, lpips)
'''
def rgb_ssim(img0, img1, max_val,
             filter_size=11,
             filter_sigma=1.5,
             k1=0.01,
             k2=0.03,
             return_map=False):
    # Modified from https://github.com/google/mipnerf/blob/16e73dfdb52044dcceb47cda5243a686391a6e0f/internal/math.py#L58
    assert len(img0.shape) == 3
    assert img0.shape[-1] == 3
    assert img0.shape == img1.shape

    # Construct a 1D Gaussian blur filter.
    hw = filter_size // 2
    shift = (2 * hw - filter_size + 1) / 2
    f_i = ((np.arange(filter_size) - hw + shift) / filter_sigma)**2
    filt = np.exp(-0.5 * f_i)
    filt /= np.sum(filt)

    # Blur in x and y (faster than the 2D convolution).
    def convolve2d(z, f):
        return scipy.signal.convolve2d(z, f, mode='valid')

    filt_fn = lambda z: np.stack([
        convolve2d(convolve2d(z[...,i], filt[:, None]), filt[None, :])
        for i in range(z.shape[-1])], -1)
    mu0 = filt_fn(img0)
    mu1 = filt_fn(img1)
    mu00 = mu0 * mu0
    mu11 = mu1 * mu1
    mu01 = mu0 * mu1
    sigma00 = filt_fn(img0**2) - mu00
    sigma11 = filt_fn(img1**2) - mu11
    sigma01 = filt_fn(img0 * img1) - mu01

    # Clip the variances and covariances to valid values.
    # Variance must be non-negative:
    sigma00 = np.maximum(0., sigma00)
    sigma11 = np.maximum(0., sigma11)
    sigma01 = np.sign(sigma01) * np.minimum(
        np.sqrt(sigma00 * sigma11), np.abs(sigma01))
    c1 = (k1 * max_val)**2
    c2 = (k2 * max_val)**2
    numer = (2 * mu01 + c1) * (2 * sigma01 + c2)
    denom = (mu00 + mu11 + c1) * (sigma00 + sigma11 + c2)
    ssim_map = numer / denom
    ssim = np.mean(ssim_map)
    return ssim_map if return_map else ssim


import torch.nn as nn
class TVLoss(nn.Module):
    def __init__(self,TVLoss_weight=1):
        super(TVLoss,self).__init__()
        self.TVLoss_weight = TVLoss_weight

    def forward(self,x):
        batch_size = x.size()[0]
        if len(x.shape) == 4:
            h_x = x.size()[2]
            w_x = x.size()[3]
            count_h = self._tensor_size(x[:,:,1:,:])
            count_w = self._tensor_size(x[:,:,:,1:])
            h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum()
            w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum()
            return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size
        elif len(x.shape) == 5:
            d_x = x.size()[2]
            h_x = x.size()[3]
            w_x = x.size()[4]
            count_d = self._tensor_size(x[:,:,1:,:,:])
            count_h = self._tensor_size(x[:,:,:,1:,:])
            count_w = self._tensor_size(x[:,:,:,:,1:])
            d_tv = torch.pow((x[:,:,1:,:,:]-x[:,:,:d_x-1,:,:]),2).sum()
            h_tv = torch.pow((x[:,:,:,1:,:]-x[:,:,:,:h_x-1,:]),2).sum()
            w_tv = torch.pow((x[:,:,:,:,1:]-x[:,:,:,:,:w_x-1]),2).sum()
            return self.TVLoss_weight*2*(d_tv/count_d+h_tv/count_h+w_tv/count_w)/batch_size
        else:
            raise ValueError

    def _tensor_size(self,t):
        return t.shape[1:].numel()



import plyfile
import skimage.measure
def convert_sdf_samples_to_ply(
    pytorch_3d_sdf_tensor,
    ply_filename_out,
    bbox,
    level=0.5,
    offset=None,
    scale=None,
):
    """
    Convert sdf samples to .ply

    :param pytorch_3d_sdf_tensor: a torch.FloatTensor of shape (n,n,n)
    :voxel_grid_origin: a list of three floats: the bottom, left, down origin of the voxel grid
    :voxel_size: float, the size of the voxels
    :ply_filename_out: string, path of the filename to save to

    This function adapted from: https://github.com/RobotLocomotion/spartan
    """

    numpy_3d_sdf_tensor = pytorch_3d_sdf_tensor.numpy()
    voxel_size = list((bbox[1]-bbox[0]) / np.array(pytorch_3d_sdf_tensor.shape))

    verts, faces, normals, values = skimage.measure.marching_cubes(
        numpy_3d_sdf_tensor, level=level, spacing=voxel_size
    )
    faces = faces[...,::-1] # inverse face orientation

    # transform from voxel coordinates to camera coordinates
    # note x and y are flipped in the output of marching_cubes
    mesh_points = np.zeros_like(verts)
    mesh_points[:, 0] = bbox[0,0] + verts[:, 0]
    mesh_points[:, 1] = bbox[0,1] + verts[:, 1]
    mesh_points[:, 2] = bbox[0,2] + verts[:, 2]

    # apply additional offset and scale
    if scale is not None:
        mesh_points = mesh_points / scale
    if offset is not None:
        mesh_points = mesh_points - offset

    # try writing to the ply file

    num_verts = verts.shape[0]
    num_faces = faces.shape[0]

    verts_tuple = np.zeros((num_verts,), dtype=[("x", "f4"), ("y", "f4"), ("z", "f4")])

    for i in range(0, num_verts):
        verts_tuple[i] = tuple(mesh_points[i, :])

    faces_building = []
    for i in range(0, num_faces):
        faces_building.append(((faces[i, :].tolist(),)))
    faces_tuple = np.array(faces_building, dtype=[("vertex_indices", "i4", (3,))])

    el_verts = plyfile.PlyElement.describe(verts_tuple, "vertex")
    el_faces = plyfile.PlyElement.describe(faces_tuple, "face")

    ply_data = plyfile.PlyData([el_verts, el_faces])
    print("saving mesh to %s" % (ply_filename_out))
    ply_data.write(ply_filename_out)

thirdparty/tensorf/utils.py
oppo-us-research/NeuRBF
import torch
import torch.nn as nn
import torch.nn.functional as F

class FreqEncoder(nn.Module):
    def __init__(self, input_dim, max_freq_log2, N_freqs,
                 log_sampling=True, include_input=True,
                 periodic_fns=(torch.sin, torch.cos)):
    
        super().__init__()

        self.input_dim = input_dim
        self.include_input = include_input
        self.periodic_fns = periodic_fns

        self.output_dim = 0
        if self.include_input:
            self.output_dim += self.input_dim

        self.output_dim += self.input_dim * N_freqs * len(self.periodic_fns)

        if log_sampling:
            self.freq_bands = 2. ** torch.linspace(0., max_freq_log2, N_freqs)
        else:
            self.freq_bands = torch.linspace(2. ** 0., 2. ** max_freq_log2, N_freqs)

        self.freq_bands = self.freq_bands.numpy().tolist()

    def forward(self, input, **kwargs):

        out = []
        if self.include_input:
            out.append(input)

        for i in range(len(self.freq_bands)):
            freq = self.freq_bands[i]
            for p_fn in self.periodic_fns:
                out.append(p_fn(input * freq))

        out = torch.cat(out, dim=-1)


        return out

def get_encoder(encoding, input_dim=3, 
                multires=6, 
                degree=4,
                num_levels=16, level_dim=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=2048, align_corners=True,  # False
                allocate_params=True,
                levels_omit=[],
                **kwargs):

    if encoding == 'None':
        return lambda x, **kwargs: x, input_dim
    
    elif encoding == 'frequency':
        from .freqencoder import FreqEncoder
        encoder = FreqEncoder(input_dim=input_dim, degree=multires)

    elif encoding == 'sphere_harmonics':
        from .shencoder import SHEncoder
        encoder = SHEncoder(input_dim=input_dim, degree=degree)

    elif encoding == 'hashgrid':
        from .gridencoder import GridEncoder
        encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='hash', align_corners=align_corners, allocate_params=allocate_params, levels_omit=levels_omit)
    
    elif encoding == 'tiledgrid':
        from .gridencoder import GridEncoder
        encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='tiled', align_corners=align_corners)
    
    elif encoding == 'ash':
        from .ashencoder import AshEncoder
        encoder = AshEncoder(input_dim=input_dim, output_dim=16, log2_hashmap_size=log2_hashmap_size, resolution=desired_resolution)

    else:
        raise NotImplementedError('Unknown encoding mode, choose from [None, frequency, sphere_harmonics, hashgrid, tiledgrid]')

    return encoder, encoder.output_dim
thirdparty/torch_ngp/encoding.py
oppo-us-research/NeuRBF
import numpy as np

import torch
import torch.nn as nn
from torch.autograd import Function
from torch.autograd.function import once_differentiable
from torch.cuda.amp import custom_bwd, custom_fwd 

try:
    import _gridencoder as _backend
except ImportError:
    from .backend import _backend

_gridtype_to_id = {
    'hash': 0,
    'tiled': 1,
}

class _grid_encode(Function):
    @staticmethod
    @custom_fwd
    def forward(ctx, inputs, embeddings, offsets, per_level_scale, base_resolution, calc_grad_inputs=False, gridtype=0, align_corners=False):
        # inputs: [B, D], float in [0, 1]
        # embeddings: [sO, C], float
        # offsets: [L + 1], int
        # RETURN: [B, F], float

        inputs = inputs.contiguous()

        B, D = inputs.shape # batch size, coord dim
        L = offsets.shape[0] - 1 # level
        C = embeddings.shape[1] # embedding dim for each level
        S = np.log2(per_level_scale) # resolution multiplier at each level, apply log2 for later CUDA exp2f
        H = base_resolution # base resolution

        # manually handle autocast (only use half precision embeddings, inputs must be float for enough precision)
        # if C % 2 != 0, force float, since half for atomicAdd is very slow.
        if torch.is_autocast_enabled() and C % 2 == 0:
            embeddings = embeddings.to(torch.half)

        # L first, optimize cache for cuda kernel, but needs an extra permute later
        outputs = torch.empty(L, B, C, device=inputs.device, dtype=embeddings.dtype)

        if calc_grad_inputs:
            dy_dx = torch.empty(B, L * D * C, device=inputs.device, dtype=embeddings.dtype)
        else:
            dy_dx = None

        _backend.grid_encode_forward(inputs, embeddings, offsets, outputs, B, D, C, L, S, H, dy_dx, gridtype, align_corners)

        # permute back to [B, L * C]
        outputs = outputs.permute(1, 0, 2).reshape(B, L * C)

        ctx.save_for_backward(inputs, embeddings, offsets, dy_dx)
        ctx.dims = [B, D, C, L, S, H, gridtype]
        ctx.align_corners = align_corners

        return outputs
    
    @staticmethod
    #@once_differentiable
    @custom_bwd
    def backward(ctx, grad):

        inputs, embeddings, offsets, dy_dx = ctx.saved_tensors
        B, D, C, L, S, H, gridtype = ctx.dims
        align_corners = ctx.align_corners

        # grad: [B, L * C] --> [L, B, C]
        grad = grad.view(B, L, C).permute(1, 0, 2).contiguous()

        grad_embeddings = torch.zeros_like(embeddings)

        if dy_dx is not None:
            grad_inputs = torch.zeros_like(inputs, dtype=embeddings.dtype)
        else:
            grad_inputs = None

        _backend.grid_encode_backward(grad, inputs, embeddings, offsets, grad_embeddings, B, D, C, L, S, H, dy_dx, grad_inputs, gridtype, align_corners)

        if dy_dx is not None:
            grad_inputs = grad_inputs.to(inputs.dtype)

        return grad_inputs, grad_embeddings, None, None, None, None, None, None
        


grid_encode = _grid_encode.apply


class GridEncoder(nn.Module):
    def __init__(self, input_dim=3, num_levels=16, level_dim=2, per_level_scale=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=None, gridtype='hash', align_corners=False, allocate_params=True,
    levels_omit=[]):
        super().__init__()

        # the finest resolution desired at the last level, if provided, overridee per_level_scale
        if desired_resolution is not None and num_levels > 1:
            per_level_scale = np.exp2(np.log2(desired_resolution / base_resolution) / (num_levels - 1))

        self.input_dim = input_dim # coord dims, 2 or 3
        self.num_levels = num_levels # num levels, each level multiply resolution by 2
        self.level_dim = level_dim # encode channels per level
        self.per_level_scale = per_level_scale # multiply resolution by this scale at each level.
        self.log2_hashmap_size = log2_hashmap_size
        self.base_resolution = base_resolution
        self.gridtype = gridtype
        self.gridtype_id = _gridtype_to_id[gridtype] # "tiled" or "hash"
        self.align_corners = align_corners

        channels_keep = []
        for i in range(num_levels):
            if i not in levels_omit:
                channels_keep += [i*level_dim+x for x in range(level_dim)]
        channels_keep = torch.tensor(channels_keep, dtype=torch.long)
        self.register_buffer('channels_keep', channels_keep)
        self.output_dim_orig = num_levels * level_dim
        self.output_dim = channels_keep.shape[0]

        # allocate parameters
        offsets = []
        offset = 0
        self.max_params = 2 ** log2_hashmap_size
        for i in range(num_levels):
            resolution = int(np.ceil(base_resolution * per_level_scale ** i))
            params_in_level = min(self.max_params, (resolution if align_corners else resolution + 1) ** input_dim) # limit max number
            params_in_level = int(np.ceil(params_in_level / 8) * 8) # make divisible
            offsets.append(offset)
            offset += params_in_level
        offsets.append(offset)
        offsets = torch.from_numpy(np.array(offsets, dtype=np.int32))
        self.register_buffer('offsets', offsets)
        
        self.n_params = offsets[-1] * level_dim

        # parameters
        if allocate_params:
            self.embeddings = nn.Parameter(torch.empty(offset, level_dim))
            self.reset_parameters()
    
    def reset_parameters(self):
        std = 1e-4
        self.embeddings.data.uniform_(-std, std)

    def __repr__(self):
        return f"GridEncoder: input_dim={self.input_dim} num_levels={self.num_levels} level_dim={self.level_dim} resolution={self.base_resolution} -> {int(round(self.base_resolution * self.per_level_scale ** (self.num_levels - 1)))} per_level_scale={self.per_level_scale:.4f} params={tuple(self.embeddings.shape)} gridtype={self.gridtype} align_corners={self.align_corners}"
    
    def forward(self, inputs, bound=1):
        # inputs: [..., input_dim], normalized real world positions in [-bound, bound]
        # return: [..., num_levels * level_dim]

        inputs = (inputs + bound) / (2 * bound) # map to [0, 1]

        prefix_shape = list(inputs.shape[:-1])
        inputs = inputs.view(-1, self.input_dim)

        outputs = grid_encode(inputs, self.embeddings, self.offsets, self.per_level_scale, self.base_resolution, inputs.requires_grad, self.gridtype_id, self.align_corners)
        outputs = outputs.view(prefix_shape + [self.output_dim_orig])

        return outputs[..., self.channels_keep]

thirdparty/torch_ngp/gridencoder/grid.py
oppo-us-research/NeuRBF
import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np

def mape_loss(pred, target, reduction='mean'):
    # pred, target: [B, 1], torch tenspr
    difference = (pred - target).abs()
    scale = 1 / (target.abs() + 1e-2)
    loss = difference * scale

    if reduction == 'mean':
        loss = loss.mean()
    
    return loss

def huber_loss(pred, target, delta=0.1, reduction='mean'):
    rel = (pred - target).abs()
    sqr = 0.5 / delta * rel * rel
    loss = torch.where(rel > delta, rel - 0.5 * delta, sqr)

    if reduction == 'mean':
        loss = loss.mean()

    return loss

thirdparty/torch_ngp/loss.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import numpy as np
import cupy as cp
from pykdtree.kdtree import KDTree


class KMeans():
    def __init__(self, n_clusters, max_iter=100):
        self.n_clusters = n_clusters
        self.max_iter = max_iter

    def init_centers(self, X, n_clusters, sample_weight=None):
        if sample_weight is not None:
            idx = np.random.choice(X.shape[0], size=n_clusters, replace=False, p=sample_weight)
        else:
            idx = np.random.choice(X.shape[0], size=n_clusters, replace=False)
        return X[idx]

    def init_centers_cp(self, X, n_clusters, sample_weight=None):
        if sample_weight is not None:
            idx = cp.random.choice(X.shape[0], size=n_clusters, replace=True, p=sample_weight)
        else:
            idx = cp.random.choice(X.shape[0], size=n_clusters, replace=False)
        return X[idx]

    def compute_centers_loop_np(self, X, labels, sample_weight=None):
        X = X.astype(np.float64)
        sample_weight = sample_weight.astype(np.float64)
        centers = np.zeros((self.n_clusters, X.shape[1]), dtype=np.float64)
        if sample_weight is not None:
            X = X * sample_weight[:, None]
        for k in range(self.n_clusters):
            idx = labels == k
            if idx.sum() == 0:
                centers[k, :] = 0
            else:
                if sample_weight is None:
                    count = idx.sum()
                else:
                    count = sample_weight[idx].sum()
                centers[k, :] = X[idx, :].sum(axis=0) / count
        return centers

    def compute_centers_np(self, X, labels, sample_weight=None):
        ix = np.argsort(labels)
        labels = labels[ix]
        X = X[ix].astype(np.float64)
        if sample_weight is not None:
            sample_weight = sample_weight[ix].astype(np.float64)
            X = X * sample_weight[:, None]
        
        d = np.diff(labels, prepend=0)
        pos = np.flatnonzero(d)
        pos = np.repeat(pos, d[pos])
        pos = np.append(np.insert(pos, 0, 0), len(X))

        X = np.concatenate((np.zeros_like(X[0:1]), X), axis=0)
        X = np.cumsum(X, axis=0)
        if sample_weight is not None:
            sample_weight = np.concatenate((np.zeros_like(sample_weight[0:1]), sample_weight), axis=0)
            sample_weight = np.cumsum(sample_weight, axis=0)

        X = np.diff(X[pos], axis=0)
        if sample_weight is None:
            count = np.diff(pos).clip(min=1)
        else:
            count = np.diff(sample_weight[pos], axis=0)
            count[count==0]=1
        centers = X / count[:, None]

        return centers

    def compute_centers_cupy(self, X, labels, sample_weight=None):
        '''
        X: [p d], cupy float
        labels: [p], cupy int
        sample_weight: [p], cupy float
        '''
        ix = cp.argsort(labels)
        labels = labels[ix]
        X = X[ix]
        if sample_weight is not None:
            sample_weight = sample_weight[ix]
            X = X * sample_weight[:, None]
        
        d = cp.diff(labels, prepend=0)
        pos = cp.flatnonzero(d)
        pos = cp.asarray(np.repeat(pos.get(), d[pos].get()))
        pos = cp.append(cp.concatenate((cp.zeros_like(pos[0:1]), pos)), len(X))

        X = cp.concatenate((cp.zeros_like(X[0:1]), X), axis=0)
        X = cp.cumsum(X, axis=0)
        if sample_weight is not None:
            sample_weight = cp.concatenate((cp.zeros_like(sample_weight[0:1]), sample_weight), axis=0)
            sample_weight = cp.cumsum(sample_weight, axis=0)

        X = cp.diff(X[pos], axis=0)
        if sample_weight is None:
            count = cp.diff(pos)
        else:
            count = cp.diff(sample_weight[pos], axis=0)
        centers = X / count[:, None]

        return centers, count

    def fit(self, X, sample_weight=None, backend=0, gpu=0):
        if backend==0:  # numpy
            self.fit_np(X.cpu().numpy(), sample_weight.cpu().numpy())
        elif backend==1:  # cupy
            with cp.cuda.Device(gpu):
                self.fit_cupy(X, sample_weight)
                cp.get_default_memory_pool().free_all_blocks()
                cp.get_default_pinned_memory_pool().free_all_blocks()
        else:
            raise NotImplementedError

    def fit_np(self, X, sample_weight=None, **kwargs):
        self.centers = self.init_centers(X, sample_weight)
        for i in range(self.max_iter):
            centers_old = self.centers
            self.kdtree = KDTree(self.centers)
            _, self.labels = query_chunked(self.kdtree, X, k=1, sqr_dists=True, chunk_size=int(2e8), return_dist=False)
            self.centers = self.compute_centers_np(X, self.labels, sample_weight)
            if np.all(centers_old == self.centers):
                break
        self.kdtree = KDTree(self.centers)
        _, self.labels = query_chunked(self.kdtree, X, k=1, sqr_dists=True, chunk_size=int(2e8), return_dist=False)

    def fit_cupy(self, X, sample_weight=None, **kwargs):
        X_cp = cp.asarray(X, dtype=cp.float64)
        X = X.cpu().numpy().astype(np.float32, copy=False)
        if sample_weight is not None:
            sample_weight_cp = cp.asarray(sample_weight, dtype=cp.float64)
            sample_weight = sample_weight.cpu().numpy().astype(np.float32, copy=False)
            sample_weight_normalized = sample_weight / sample_weight.sum()
        else:
            sample_weight_cp = None
            sample_weight_normalized = None
        centers_cp = self.init_centers(X_cp, self.n_clusters, sample_weight_normalized)
        self.centers = cp.asnumpy(centers_cp).astype(np.float32)
        for i in range(self.max_iter):
            centers_old_cp = centers_cp
            self.kdtree = KDTree(self.centers)
            _, self.labels = query_chunked(self.kdtree, X, k=1, sqr_dists=True, chunk_size=int(2e8), return_dist=False)
            centers_cp, count = reduce_within_clusters_chunked(X_cp, self.n_clusters, cp.asarray(self.labels), 
                sample_weight_cp, chunk_size=int(1e8 / X_cp.shape[-1]))
            centers_cp[count==0] = centers_old_cp[count==0]
            self.centers = cp.asnumpy(centers_cp).astype(np.float32)
            # if cp.all(centers_old_cp == centers_cp):
            #     break
        self.kdtree = KDTree(self.centers)
        _, self.labels = query_chunked(self.kdtree, X, k=1, sqr_dists=True, chunk_size=int(2e8), return_dist=False)

    def predict(self, X, sample_weight=None):
        _, labels = query_chunked(self.kdtree, X.astype(np.float32, copy=False), k=1, sqr_dists=True, 
            chunk_size=int(2e8), return_dist=False)

        return labels


def query_chunked(kd_tree, x, k, sqr_dists, chunk_size=int(2e8), return_dist=False):
    if chunk_size is None: chunk_size = x.shape[0]
    if chunk_size >= x.shape[0]: return kd_tree.query(x, k=k, sqr_dists=sqr_dists)

    dist = np.zeros([x.shape[0], k], dtype=np.float32) if return_dist else None
    idx = np.zeros([x.shape[0], k], dtype=np.uint32)
    if k == 1:
        if return_dist: dist = dist[:, 0]
        idx = idx[:, 0]
    for i in range(0, x.shape[0], chunk_size):
        dist_i, idx[i:i+chunk_size] = kd_tree.query(x[i:i+chunk_size], k=k, sqr_dists=sqr_dists)
        if return_dist: dist[i:i+chunk_size] = dist_i
    return dist, idx


def reduce_within_clusters(X, n_clusters, labels, sample_weight=None, reduce_weight=True):
    '''
    X: [p ...], cupy float
    labels: [p], cupy int
    sample_weight: [p], cupy float
    '''
    X = X.astype(dtype=cp.float64, copy=False)
    sample_weight = sample_weight.astype(dtype=cp.float64, copy=False)

    ix = cp.argsort(labels)
    labels = labels[ix]
    X = X[ix]
    if sample_weight is not None:
        sample_weight = sample_weight[ix]
        X = X * sample_weight.reshape([X.shape[0], *([1]*(X.ndim-1))])
    
    d = cp.diff(labels, prepend=0)
    pos = cp.flatnonzero(d)
    pos = cp.asarray(np.repeat(pos.get(), d[pos].get()))
    pos = cp.append(cp.concatenate((cp.zeros_like(pos[0:1]), pos)), len(X))

    X = cp.concatenate((cp.zeros_like(X[0:1]), X), axis=0)
    X = cp.cumsum(X, axis=0)
    if sample_weight is not None:
        sample_weight = cp.concatenate((cp.zeros_like(sample_weight[0:1]), sample_weight), axis=0)
        sample_weight = cp.cumsum(sample_weight, axis=0)

    X = cp.diff(X[pos], axis=0)
    if sample_weight is None:
        count = cp.diff(pos)
    else:
        count = cp.diff(sample_weight[pos], axis=0)
    if reduce_weight:
        out = X / count.reshape([X.shape[0], *([1]*(X.ndim-1))])
    else:
        out = X

    if out.shape[0] < n_clusters:
        n_fill = n_clusters - out.shape[0]
        out = cp.concatenate([out, cp.zeros([n_fill, *out.shape[1:]])])
        count = cp.concatenate([count, cp.zeros([n_fill])])

    return out, count


def reduce_within_clusters_chunked(X, n_clusters, labels, sample_weight=None, chunk_size=None):
    if chunk_size is None: chunk_size = X.shape[0]
    for i in range(0, X.shape[0], chunk_size):
        out_i, count_i = reduce_within_clusters(X[i:i+chunk_size], n_clusters, labels[i:i+chunk_size], 
            sample_weight[i:i+chunk_size] if sample_weight is not None else None, reduce_weight=False)
        if i == 0:
            out = out_i
            count = count_i
        else:
            out += out_i
            count += count_i
    out /= count.reshape([out.shape[0], *([1]*(out.ndim-1))])
    return out, count

util_clustering.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Ellipse, Rectangle
from matplotlib.collections import PatchCollection, EllipseCollection
import matplotlib.transforms as transforms
import numpy as np


def plot_ellipses_v1(ax, means, covs, xlim, ylim, alpha=0.5, flip_angle=True):
    '''
    mean: [k d], numpy array
    cov: [k] or [k d] or [k d d], numpy array
    '''
    if covs.ndim == 1:
        covs = covs[..., None, None] * np.eye(2)[None, ...]
    elif covs.ndim == 2:
        covs = covs[..., None] * np.eye(2)[None, ...]
    elif covs.ndim > 3:
        raise NotImplementedError

    a = covs[:, 0, 0]
    b = covs[:, 0, 1]
    c = covs[:, 1, 1]
    temp1 = (a + c) / 2
    temp2 = np.sqrt(((a - c) / 2) ** 2 + b ** 2)
    lambda1 = temp1 + temp2
    lambda2 = temp1 - temp2
    lambda2[lambda2 < 0] = 0

    theta = np.degrees(np.arctan2(lambda1 - a, b))
    theta[(b == 0) * (a >= c)] = 0
    theta[(b == 0) * (a < c)] = 90
    if flip_angle: theta = -theta

    col = EllipseCollection(np.sqrt(lambda1)*2, np.sqrt(lambda2)*2, theta, units='x', offsets=means, 
        offset_transform=ax.transData, alpha=alpha, facecolor='w', edgecolor='y')
    ax.add_collection(col)

    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.tick_params(axis='both', which='major', labelsize=10)

    return ax


def plot_ellipse(ax, mean, cov, **kwargs):
    pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])
    ell_radius_x = np.sqrt(1 + pearson)
    ell_radius_y = np.sqrt(1 - pearson) if pearson <= 1 else 0
    ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2, **kwargs)

    scale_x = np.sqrt(cov[0, 0])
    scale_y = np.sqrt(cov[1, 1])
    transf = transforms.Affine2D().rotate_deg(-45).scale(scale_x, scale_y).translate(mean[0], mean[1])
    ellipse.set_transform(transf + ax.transData)
    
    ax.add_patch(ellipse)

    return ax


def plot_ellipses(ax, means, covs, xlim, ylim, alpha=0.5):
    if covs.ndim == 1:
        covs = covs[..., None, None] * np.eye(2)[None, ...]
    elif covs.ndim == 2:
        covs = covs[..., None] * np.eye(2)[None, ...]
    elif covs.ndim > 3:
        raise NotImplementedError

    for i in range(0, means.shape[0]):
        ax = plot_ellipse(ax, means[i, ...], covs[i, ...], alpha=alpha, facecolor='w', edgecolor='y')

    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    ax.tick_params(axis='both', which='major', labelsize=10)

    return ax

util_figure.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import cupy as cp
from pykdtree.kdtree import KDTree
import torch

import util_misc
import util_clustering


def kc_init_v3(n_kernel, points, points_weight, cmin, cmax, s_dims, device, points_sampling=1, reg_sampling=2, weight_exp=1, weight_thres=0, n_iter=300, **kwargs):
    '''Weighted Kmeans using util_clustering.KMeans'''
    p, d = points.shape

    # Subsample points
    if points_sampling > 1:
        points = points.view(*s_dims.tolist(), d)
        points_weight = points_weight.view(*s_dims.tolist(), 1)
        if s_dims.shape[0] == 2:
            points = points[::points_sampling, ::points_sampling]
            points_weight = points_weight[::points_sampling, ::points_sampling]
        elif s_dims.shape[0] == 3:
            points = points[::points_sampling, ::points_sampling, ::points_sampling]
            points_weight = points_weight[::points_sampling, ::points_sampling, ::points_sampling]
        else:
            raise NotImplementedError
        points = points.reshape(-1, d)
        points_weight = points_weight.reshape(-1, 1)

    points = points.view(-1, d)
    points_weight = points_weight.view(-1)
    points_weight = points_weight / points_weight.mean()
    
    if weight_thres != 0:
        points_weight[points_weight <= weight_thres] = 0
    if weight_exp != 1:
        points_weight = points_weight ** weight_exp

    kc = []

    # Regularly put some kernels
    if reg_sampling > 0:
        kc_i, _, _ = util_misc.get_grid_points(((s_dims/reg_sampling).ceil()+1).to(torch.long).tolist(), 
            align_corners=True, vmin=cmin.tolist(), vmax=cmax.tolist())
        kc_i = kc_i.flip(-1)  # x y ...
        kc.append(kc_i)
        n_kernel -= kc_i.shape[0]

    if n_kernel > 0:
        mask = points_weight > 0
        estimator = util_clustering.KMeans(n_clusters=n_kernel, max_iter=n_iter)
        estimator.fit(points[mask], sample_weight=points_weight[mask], backend=1, gpu=device)
        kc.append(torch.tensor(estimator.centers, device=points.device))

    kc = torch.cat(kc)
    return kc, {}


def kw_init_v3(kc, points, points_weight, s_dims, device, alpha, points_sampling=1, weight_exp=1, weight_thres=0, rbf_type='ivq_a', **kwargs):
    '''
    For each kernel, use weighted mean covariance with points in its cluster.
    kc: [k, d], torch tensor
    '''
    p, d = points.shape

    # Subsample points
    if points_sampling > 1:
        points = points.view(*s_dims.tolist(), d)
        points = points[::points_sampling, ::points_sampling]
        points = points.reshape(-1, d)
        points_weight = points_weight.view(*s_dims.tolist(), 1)
        points_weight = points_weight[::points_sampling, ::points_sampling]
        points_weight = points_weight.reshape(-1, 1)

    points = points.view(-1, d)
    points_weight = points_weight.view(-1)
    points_weight = points_weight / points_weight.mean()
    
    if weight_thres != 0:
        points_weight[points_weight <= weight_thres] = 0
    if weight_exp != 1:
        points_weight = points_weight ** weight_exp

    mask = points_weight > 0
    points = points[mask]
    points_weight = points_weight[mask]

    kc = kc.cpu().numpy()
    kdtree = KDTree(kc)
    _, labels = util_clustering.query_chunked(kdtree, points.cpu().numpy(), k=1, sqr_dists=True, 
        chunk_size=int(2e8), return_dist=False)  # [p]

    points_offset = points.cpu() - kc[labels]  # [p d]
    if rbf_type.endswith('_a') or rbf_type.endswith('_f'):
        points_data = points_offset[..., None] * points_offset[:, None, :]  # [p d d], covariance
    elif rbf_type.endswith('_d'):
        points_data = (points_offset**2)[..., None]  # [p d 1]
    elif rbf_type.endswith('_s'):
        points_data = (points_offset**2).sum(-1, keepdims=True)[..., None]  # [p 1 1]
    else:
        raise NotImplementedError
    del points_offset

    with cp.cuda.Device(device):
        labels = cp.asarray(labels)
        points_weight = cp.asarray(points_weight)
        kw_sq = torch.zeros([kc.shape[0], *points_data.shape[1:]], dtype=torch.float32, device=points.device)
        for i in range(points_data.shape[1]):
            for j in range(points_data.shape[2]):
                kw_sq_ij, count = util_clustering.reduce_within_clusters_chunked(cp.asarray(points_data[:, i, j]), 
                    kc.shape[0], labels, points_weight, chunk_size=int(1e8))
                kw_sq_ij /= alpha
                if rbf_type.endswith('_a') or rbf_type.endswith('_f'):
                    if i == j:
                        kw_sq_ij[count==0] = 1e-10
                    else:
                        kw_sq_ij[count==0] = 0
                else:
                    kw_sq_ij[count==0] = 1e-10
                kw_sq[:, i, j] = torch.tensor(kw_sq_ij, dtype=torch.float32, device=points.device)
                del kw_sq_ij, count

        del labels, points_weight
        cp.get_default_memory_pool().free_all_blocks()
        cp.get_default_pinned_memory_pool().free_all_blocks()

    if rbf_type == 'nlin_s' or rbf_type == 'nlin_d':
        kw_sq[kw_sq < 0] = 1e-10

    return kw_sq

util_init.py
oppo-us-research/NeuRBF
import torch
import numpy as np
import skimage
from scipy.spatial import cKDTree as KDTree
from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure
import lpips


def mse(image_pred, image_gt, valid_mask=None, reduction='mean'):
    """
    Modified from https://github.com/kwea123/MINER_pl/blob/84c089f097890a13b59d5d4ca17ca79f39d707e0/metrics.py#L5
    """
    value = (image_pred-image_gt)**2
    if valid_mask is not None:
        value = value[valid_mask]
    if reduction == 'mean':
        return value.mean()
    return value


def psnr(image_pred, image_gt, valid_mask=None, reduction='mean', vmin=0, vmax=1):
    """
    Modified from https://github.com/kwea123/MINER_pl/blob/84c089f097890a13b59d5d4ca17ca79f39d707e0/metrics.py#L14
    """
    if torch.is_tensor(image_pred):
        return -10*torch.log10(mse(image_pred, image_gt, valid_mask, reduction))
    elif isinstance(image_pred, np.ndarray):
        return -10*np.log10(mse(image_pred, image_gt, valid_mask, reduction))


def mse2psnr(x, vmax=1): return 20*np.log10(vmax) - 10*np.log10(x)


def mae(pred, gt):
    return (pred - gt).abs().mean()


def ssim_ski_func(pred, gt, vmin=0, vmax=1):
    if type(gt) is torch.Tensor:
        gt = gt.cpu().numpy()
    if type(pred) is torch.Tensor:
        pred = pred.cpu().numpy()
    return skimage.metrics.structural_similarity(gt, np.clip(pred, a_min=vmin, a_max=vmax), data_range=vmax - vmin, 
        channel_axis=2)


def ssim_func(rgb, gts):
    """
    Modified from https://github.com/google/mipnerf/blob/16e73dfdb52044dcceb47cda5243a686391a6e0f/internal/math.py#L58
    """
    filter_size = 11
    filter_sigma = 1.5
    k1 = 0.01
    k2 = 0.03
    max_val = 1.0
    rgb = rgb.cpu().numpy()
    gts = gts.cpu().numpy()
    assert len(rgb.shape) == 3
    assert rgb.shape[-1] == 3
    assert rgb.shape == gts.shape
    import scipy.signal

    # Construct a 1D Gaussian blur filter.
    hw = filter_size // 2
    shift = (2 * hw - filter_size + 1) / 2
    f_i = ((np.arange(filter_size) - hw + shift) / filter_sigma)**2
    filt = np.exp(-0.5 * f_i)
    filt /= np.sum(filt)

    # Blur in x and y (faster than the 2D convolution).
    def convolve2d(z, f):
        return scipy.signal.convolve2d(z, f, mode='valid')

    filt_fn = lambda z: np.stack([
        convolve2d(convolve2d(z[..., i], filt[:, None]), filt[None, :])
        for i in range(z.shape[-1])], -1)
    mu0 = filt_fn(rgb)
    mu1 = filt_fn(gts)
    mu00 = mu0 * mu0
    mu11 = mu1 * mu1
    mu01 = mu0 * mu1
    sigma00 = filt_fn(rgb**2) - mu00
    sigma11 = filt_fn(gts**2) - mu11
    sigma01 = filt_fn(rgb * gts) - mu01

    # Clip the variances and covariances to valid values.
    # Variance must be non-negative:
    sigma00 = np.maximum(0., sigma00)
    sigma11 = np.maximum(0., sigma11)
    sigma01 = np.sign(sigma01) * np.minimum(
        np.sqrt(sigma00 * sigma11), np.abs(sigma01))
    c1 = (k1 * max_val)**2
    c2 = (k2 * max_val)**2
    numer = (2 * mu01 + c1) * (2 * sigma01 + c2)
    denom = (mu00 + mu11 + c1) * (sigma00 + sigma11 + c2)
    ssim_map = numer / denom
    return np.mean(ssim_map)


ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(data_range=1.0)


def msssim(rgb, gts):
    """
    Modified from https://github.com/sarafridov/K-Planes/blob/7e3a82dbdda31eddbe2a160bc9ef89e734d9fc54/plenoxels/ops/image/metrics.py#L119
    """
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    return ms_ssim(torch.permute(rgb[None, ...], (0, 3, 1, 2)),
                   torch.permute(gts[None, ...], (0, 3, 1, 2))).item()


__LPIPS__ = {}


def init_lpips(net_name, device):
    """
    Modified from https://github.com/sarafridov/K-Planes/blob/7e3a82dbdda31eddbe2a160bc9ef89e734d9fc54/plenoxels/ops/image/metrics.py#L128
    """
    return lpips.LPIPS(net=net_name, version='0.1').eval().to(device)


def rgb_lpips(rgb, gts, net_name='alex', device='cpu'):
    """
    Modified from https://github.com/sarafridov/K-Planes/blob/7e3a82dbdda31eddbe2a160bc9ef89e734d9fc54/plenoxels/ops/image/metrics.py#L132
    """
    if net_name not in __LPIPS__:
        __LPIPS__[net_name] = init_lpips(net_name, device)
    gts = gts.permute([2, 0, 1]).contiguous().to(device)
    rgb = rgb.permute([2, 0, 1]).contiguous().to(device)
    return __LPIPS__[net_name](gts, rgb, normalize=True).item()


def iou(occ1, occ2, thres):
    """
    Modified from https://github.com/kwea123/MINER_pl/blob/84c089f097890a13b59d5d4ca17ca79f39d707e0/metrics.py#L21
    """
    occ1 = occ1>=thres
    occ2 = occ2>=thres
    area_union = (occ1 | occ2).sum()
    area_intersect = (occ1 & occ2).sum()
    return area_intersect/(area_union+1e-8)


def mesh_metrics(pred_mesh, gt_mesh, n_surface_samples, fscore_tau):
    if pred_mesh.faces.shape[0] == 0:
        return {}

    pred_points, pred_indices = pred_mesh.sample(n_surface_samples, return_index=True)
    pred_points = pred_points.astype(np.float32)
    pred_normals = pred_mesh.face_normals[pred_indices]
    gt_points, gt_indices = gt_mesh.sample(n_surface_samples, return_index=True)
    gt_points = gt_points.astype(np.float32)
    gt_normals = gt_mesh.face_normals[gt_indices]

    kdtree = KDTree(gt_points)
    dist_p2g, indices_p2g = kdtree.query(pred_points)
    kdtree = KDTree(pred_points)
    dist_g2p, indices_g2p = kdtree.query(gt_points)

    out = {}
    out['cd_l1'] = (np.mean(dist_p2g) + np.mean(dist_g2p))
    out['cd_l2'] = (np.mean(dist_p2g**2) + np.mean(dist_g2p**2))

    for tau in fscore_tau:
        precision = np.mean((dist_p2g <= tau).astype(np.float32)) * 100.0
        recall = np.mean((dist_g2p <= tau).astype(np.float32)) * 100.0
        fs = (2 * precision * recall) / (precision + recall + 1e-9)
        out[f'fs_{tau:.0e}'] = fs

    normals_p2g = gt_normals[indices_p2g]
    nc_p2g = np.abs(np.sum(normals_p2g * pred_normals, axis=1))
    normals_g2p = pred_normals[indices_g2p]
    nc_g2p = np.abs(np.sum(normals_g2p * gt_normals, axis=1))
    out['nc'] = 0.5 * (np.mean(nc_p2g) + np.mean(nc_g2p))

    nc_p2g = np.degrees(np.arccos(nc_p2g))
    nc_p2g = np.degrees(np.arccos(nc_g2p))
    out['nae'] = 0.5 * (np.mean(nc_p2g) + np.mean(nc_g2p))

    return out

util_metric.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import GPUtil
import numpy as np
import torch
import math
import trimesh
import mcubes
from einops import rearrange


comb2d = [[0,0],[1,0],[0,1],[1,1]]
comb3d = [[0,0,0],[1,0,0],[0,1,0],[1,1,0],[0,0,1],[1,0,1],[0,1,1],[1,1,1]]


def get_grid_points(dims, align_corners, vmin=-1, vmax=1, device='cpu'):
    num_dim = len(dims)
    if not (isinstance(vmin, tuple) or isinstance(vmin, list)):
        vmin = [vmin] * num_dim
    assert len(vmin) == num_dim
    if not (isinstance(vmax, tuple) or isinstance(vmax, list)):
        vmax = [vmax] * num_dim
    assert len(vmax) == num_dim

    if align_corners:
        side_interval = [(vmax[i] - vmin[i]) / (dim_i - 1) if dim_i > 1 
            else vmax[i] - vmin[i] for i, dim_i in enumerate(dims)]
        points = torch.meshgrid([torch.linspace(vmin[i], vmax[i], steps=dim_i, device=device) if dim_i > 1 
            else torch.tensor([(vmin[i] + vmax[i]) / 2], device=device) for i, dim_i in enumerate(dims)], indexing='ij')
    else:
        side_interval = [(vmax[i] - vmin[i]) / dim_i for i, dim_i in enumerate(dims)]
        points = torch.meshgrid([torch.linspace(vmin[i], vmax[i], 
            steps=dim_i + 1, device=device)[:-1] + side_interval[i] / 2
            for i, dim_i in enumerate(dims)], indexing='ij')
    points_grid = torch.stack(points, dim=-1)
    points = points_grid.reshape(-1, len(dims))
    side_interval = torch.tensor(side_interval, device=device)

    return points, points_grid, side_interval
    # points: (prod(dims), num_dim), order in the last dim: [dim0, dim1, ...].
    # points_grid: (*dims, num_dim), order in the last dim: [dim0, dim1, ...].
    # side_interval: (num_dim), order: [dim0, dim1, ...].


def get_lower_int_power(x, d):
    out = int(x**(1/d))
    if (out+1)**d == x:
        out+=1
    return out**d


def get_rbf_params_per_kernel(rbf_type, in_dim, kc_mult=2):
    if rbf_type.endswith('_a'):
        return int(in_dim*kc_mult + (in_dim+1)*in_dim/2)
    elif rbf_type.endswith('_f'):
        return int(in_dim*kc_mult + in_dim**2)
    elif rbf_type.endswith('_d'):
        return int(in_dim*kc_mult + in_dim)
    elif rbf_type.endswith('_s'):
        return int(in_dim*kc_mult + 1)
    else:
        raise NotImplementedError


def kw_to_kw_sq(kw, kernel_type):
    return kw**2


def kw_to_ks(kw, kernel_type, is_scalar=False):
    if kernel_type.startswith('nlin_') and \
        (is_scalar or (kernel_type[-2:] != '_a' and kernel_type[-2:] != '_f')):
        # Used to avoid power and then sqrt
        ks = 1 / kw
    else:
        kw_sq = kw_to_kw_sq(kw, kernel_type)
        ks = kw_sq_to_ks(kw_sq, kernel_type, is_scalar)
    return ks


def kw_sq_to_ks(kw_sq, kernel_type, is_scalar=False):
    if kernel_type[-2:] == '_a' and not is_scalar:
        if isinstance(kw_sq, np.ndarray):
            ks = np.linalg.inv(kw_sq)
        else:
            ks = torch.linalg.inv(kw_sq).contiguous()
    elif kernel_type[-2:] == '_f' and not is_scalar:
        d = kw_sq.shape[-1]
        if isinstance(kw_sq, np.ndarray):
            cov_inv = np.linalg.inv(kw_sq)  # [... d d]
            w, v = np.linalg.eig(cov_inv)  # [... d], [... d d]
            w = w[..., None]**0.5 * np.eye(d).reshape(*([1]*(len(w.shape)-1)), d, d)  # [... d d]
            ks = (v @ w).swapaxes(-1, -2)
        else:
            cov_inv = torch.linalg.inv(kw_sq).contiguous()  # [... d d]
            w, v = torch.linalg.eig(cov_inv)  # [... d], [... d d]
            w, v = w.real, v.real
            w = torch.diag_embed(w**0.5)  # [... d d]
            ks = (v @ w).swapaxes(-1, -2)  # [... d d]
    else:
        if kernel_type.startswith('nlin_'):
            ks = 1 / kw_sq**0.5
        else:
            ks = 1 / kw_sq
    return ks


def ks_to_kw_sq(ks, kernel_type, is_scalar=False):
    if kernel_type.endswith('_a') and not is_scalar:
        if isinstance(ks, np.ndarray):
            kw_sq = np.linalg.inv(ks)
        else:
            kw_sq = torch.linalg.inv(ks)
    elif kernel_type.endswith('_f') and not is_scalar:
        if isinstance(ks, np.ndarray):
            kw_sq = np.linalg.inv(ks.swapaxes(-1, -2) @ ks)
        else:
            kw_sq = torch.linalg.inv(ks.swapaxes(-1, -2) @ ks)
    else:
        if kernel_type.startswith('nlin_'):
            kw_sq = 1 / ks**2
        else:
            kw_sq = 1 / ks
    return kw_sq


def normalize_kw_sq(kw_sq, rbf_type):
    # Normalize kw so that rbf has value 0.5 when dist=kw
    if rbf_type.startswith('nlin_'):
        kw_sq = kw_sq * 4
    elif rbf_type.startswith('ivmq_'):
        kw_sq = kw_sq / 3
    return kw_sq


def clip_kw_sq(data, kernel_type, cmin, cmax, dims, is_ks, is_flat, vmin_factor=100, vmax_factor=10, vmin_min=1e-3):
    '''
    cmin: (n_dim), tensor, last dim order: x y ...
    cmax: (n_dim), tensor, last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    '''
    with torch.no_grad():
        vmin = ((cmax - cmin) / dims / 2 / vmin_factor).to(data.device)  # [d]
        vmax = ((cmax - cmin) / vmax_factor).to(data.device)  # [d]
        if vmin_min is not None:
            vmin.clip_(min=vmin_min)
        if not is_ks:  # data is kw_sq
            vmin, vmax = kw_to_kw_sq(vmin, kernel_type), kw_to_kw_sq(vmax, kernel_type)
        else:  # data is ks
            vmin, vmax = kw_to_ks(vmax, kernel_type, True), kw_to_ks(vmin, kernel_type, True)

        if kernel_type[-2:] == '_a' or (kernel_type[-2:] == '_f' and not is_ks):
            if not is_flat:
                # Diagonal elements
                for i in range(dims.shape[0]):
                    data[..., i, i].clip_(vmin[i], vmax[i])
                # Off-diagonal elements
                for i in range(dims.shape[0]):
                    for j in range(dims.shape[0]):
                        if i == j:
                            continue
                        val = (data[..., i, i] * data[..., j, j]).sqrt() * 0.999
                        data[..., i, j].clip_(min=-val, max=val)
            else:
                # Diagonal elements
                for i in range(dims.shape[0]):
                    data[..., i*dims.shape[0]+i].clip_(vmin[i], vmax[i])
                # Off-diagonal elements
                for i in range(dims.shape[0]):
                    for j in range(dims.shape[0]):
                        if i == j:
                            continue
                        val = (data[..., i*dims.shape[0]+i] * data[..., j*dims.shape[0]+j]).sqrt() * 0.999
                        data[..., i*dims.shape[0]+j].clip_(min=-val, max=val)
        elif kernel_type[-2:] == '_f' and is_ks:
            if not is_flat:
                for i in range(dims.shape[0]):
                    for j in range(dims.shape[0]):
                        data[..., i, j].clip_(-vmax[j], vmax[j])
            else:
                for i in range(dims.shape[0]):
                    for j in range(dims.shape[0]):
                        data[..., i*dims.shape[0]+j].clip_(-vmax[j], vmax[j])
        elif kernel_type[-2:] == '_d':
            data.clip_(vmin[None, :, None], vmax[None, :, None])
        elif kernel_type[-2:] == '_s':
            data.clip_(vmin.max(), vmax.min())
        else:
            raise NotImplementedError


def rescale_kc(kc, block_centers, patch_size, input_size):
    return kc * (np.array(patch_size) - 1)[None, None, :] / np.array(input_size[:-1])[None, None, :] + block_centers


def rescale_kw_sq(kw_sq, patch_size, input_size):
    if kw_sq.ndim < 4:
        if kw_sq.ndim == 2:
            kw_sq = np.tile(kw_sq[..., None], (1, 1, len(patch_size)))
        return kw_sq * ((np.array(patch_size) - 1)[None, None, :] / np.array(input_size[:-1])[None, None, :]) ** 2
    else:
        raise NotImplementedError


def get_multi_index_nearest(coords, dims, interval, vmin, vmax, align_corners=True):
    """
    coords: (..., n_dim), last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    interval: (n_dim), tensor, last dim order: x y ...
    vmin: (n_dim), tensor, last dim order: x y ...
    vmax: (n_dim), tensor, last dim order: x y ...
    return: (..., n_dim), last dim order: x y ...
    """
    n_dim = coords.shape[-1]
    interval = interval.expand([*([1]*(coords.dim()-1)), -1])  # (..., n_dim)
    if align_corners:
        sub = (coords - vmin.expand([*([1]*(coords.dim()-1)), -1]) + interval/2) / interval  # (..., n_dim)
    else:
        sub = (coords - vmin.expand([*([1]*(coords.dim()-1)), -1])) / interval  # (..., n_dim)
    mask = (sub - sub.round()).abs() < 1e-3
    sub[mask] = sub[mask].round()
    sub[~mask] = sub[~mask].floor()
    sub = sub.clip(min=0).clip(max=dims.expand([*([1]*(sub.dim()-1)), n_dim])-1)  # (..., n_dim)
    return sub


def get_multi_index_multi_nearest(coords, dims, interval, vmin, vmax, align_corners=True):
    """
    coords: (..., n_dim), last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    interval: (n_dim), tensor, last dim order: x y ...
    vmin: (n_dim), tensor, last dim order: x y ...
    vmax: (n_dim), tensor, last dim order: x y ...
    return: (..., n_dim), last dim order: x y ...
    """
    n_dim = coords.shape[-1]
    interval = interval.expand([*([1]*(coords.dim()-1)), -1])  # (..., n_dim)
    if align_corners:
        sub = (coords - vmin.expand([*([1]*(coords.dim()-1)), -1]) + interval/2) / interval  # (..., n_dim)
    else:
        sub = (coords - vmin.expand([*([1]*(coords.dim()-1)), -1])) / interval  # (..., n_dim)
    mask = (sub - sub.round()).abs() < 1e-3
    sub[mask] = sub[mask].round()
    sub[~mask] = sub[~mask].floor()
    sub = sub.clip(min=0).clip(max=dims.expand([*([1]*(sub.dim()-1)), n_dim])-1)  # (..., n_dim)
    sub = sub[..., None, :]

    # coords on block edges are assigned to all adjacent blocks
    for i in range(n_dim):
        temp = coords[..., i] / interval[..., i]
        temp = (temp - temp.round()).abs()
        mask = (temp < 1e-3) * (coords[..., i] > vmin[i]) * (coords[..., i] < vmax[i])
        sub_i = torch.ones_like(sub) * dims.max()
        sub_i[mask] = sub[mask]
        sub_i[mask, ..., i] -= 1
        sub = torch.cat([sub, sub_i], dim=-2)

    return sub


def get_multi_index_nearest_lower(coords, dims, interval, vmin, vmax, align_corners=True):
    """
    Get the multi index of nearest lower neighbor.
    coords: (..., n_dim), last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    interval: (n_dim), tensor, last dim order: x y ...
    vmin: (n_dim), tensor, last dim order: x y ...
    vmax: (n_dim), tensor, last dim order: x y ...
    return: (..., n_dim), last dim order: x y ...
    """
    n_dim = coords.shape[-1]
    interval = interval.expand([*([1]*(coords.dim()-1)), -1])  # (..., n_dim)
    vmin = vmin.expand([*([1]*(coords.dim()-1)), -1])  # (..., n_dim)
    dims = dims.expand([*([1]*(coords.dim()-1)), n_dim])  # (..., n_dim)

    if align_corners:
        sub = torch.div(coords - vmin, interval, rounding_mode='floor')  # (..., n_dim)
    else:
        sub = torch.div(coords - vmin - interval/2, interval, rounding_mode='floor')  # (..., n_dim)
    sub = sub.clip(min=0).clip(max=dims-2)  # (..., n_dim)
    return sub.to(torch.long)


def get_multi_index_first_ring(coords, dims, interval, vmin, vmax, align_corners=True):
    """
    coords: (..., n_dim), last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    interval: (n_dim), last dim order: x y ...
    vmin: (n_dim), tensor, last dim order: x y ...
    vmax: (n_dim), tensor, last dim order: x y ...
    return: (..., n_dim), last dim order: x y ...
    """
    sub = get_multi_index_nearest_lower(coords, dims, interval, vmin, vmax, align_corners=True)

    n_dim = coords.shape[-1]
    if n_dim == 2:
        comb = comb2d
    elif n_dim == 3:
        comb = comb3d
    else:
        raise NotImplementedError
    sub = sub[..., None, :] + \
        torch.tensor(comb, device=sub.device, dtype=torch.long).expand([*([1]*(sub.dim()-1)), -1, n_dim])  # (..., n_ngb, n_dim)
    return sub


def ravel_multi_index(sub, dims):
    """
    sub: (..., n_dim), tensor, last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    """
    idx = (sub * (dims.cumprod(-1) / dims).to(torch.long).expand(*([1]*(sub.dim()-1)), -1)).sum(-1)
    return idx  # (...)


def ravel_multi_index_loop(sub, dims):
    """
    sub: (..., n_dim), tensor, last dim order: x y ...
    dims: (n_dim), tensor, last dim order: x y ...
    """
    mult_factor = (dims.cumprod(-1) / dims).to(torch.long)
    idx = sub[..., 0] * mult_factor[0]
    for i in range(1, len(dims)):
        idx += sub[..., i] * mult_factor[i]
    idx = idx.to(torch.long)
    return idx  # (...)


def select_devices(devices, force_reselect=True, excludeID=[1]):
    if isinstance(devices, str):
        if not devices.endswith('#'):
            if not force_reselect:
                return devices
            else:
                n_device = len(devices.split(','))
        else:
            n_device = int(devices[:-1])
    elif isinstance(devices, list) or isinstance(devices, tuple):
        if not force_reselect:
            return devices
        else:
            n_device = len(devices)

    devices = GPUtil.getAvailable(order='memory', limit=n_device, maxLoad=1.0, maxMemory=0.75, 
        includeNan=False, excludeID=excludeID, excludeUUID=[])  # Find available GPUs by sorting memory usage
    if len(devices) < n_device:
        raise ValueError('Not enough available GPUs.')
    return devices


def count_parameters(model):
    if isinstance(model, torch.nn.Module) or isinstance(model, torch.nn.ModuleList):
        out = np.array([sum(p.numel() for _, p in model.named_parameters() if p.requires_grad), 
            sum(p.numel() for _, p in model.named_parameters() if not p.requires_grad), 0])
    elif isinstance(model, torch.nn.Parameter):
        out = np.array([model.numel() if model.requires_grad else 0, 
            model.numel() if not model.requires_grad else 0, 0])
    else:
        raise NotImplementedError
    out[-1] = out[0] + out[1]
    return out


def extract_geometry(sdf, thres, cmin, cmax):
    if torch.is_tensor(sdf):
        sdf = sdf.squeeze().cpu().numpy()
    assert(len(sdf.shape)==3)
    assert(sdf.shape[0]==sdf.shape[1] and sdf.shape[0]==sdf.shape[2])
    sdf = -sdf.transpose((2, 1, 0))  # w h d, - for inside, + for outside
    resolution = sdf.shape[0]
    vertices, faces = mcubes.marching_cubes(sdf, thres)
    vertices = vertices / (resolution - 1.0) * (cmax - cmin)[None, :] + cmin[None, :]
    return vertices, faces


def extract_mesh(sdf, thres, cmin, cmax):
    vertices, faces = extract_geometry(sdf, thres, cmin, cmax)
    mesh = trimesh.Trimesh(vertices, faces, process=False)  # important, process=True leads to seg fault...
    return mesh


def random_shuffle_chunked(n, chunk_size=10000*10000, device=0):
    device_cuda = 0 if device == 'cpu' else device
    if n <= chunk_size:
        out = torch.randperm(n, device=device_cuda)
    else:
        n_chunk = math.ceil(n / chunk_size)
        out = torch.multinomial(torch.ones([chunk_size, n_chunk], device=device_cuda), n_chunk) * chunk_size
        out += torch.arange(chunk_size, device=device_cuda)[:, None].expand(-1, n_chunk)
        for i in range(n_chunk):
            out[:, i] = out[torch.randperm(chunk_size, device=device_cuda), i]
        out = out.reshape(-1)
        out = out[out < n]
    return out.to(device_cuda)


def stratified_shuffle_chunked(s_dims, b_dims, chunk_size=10000*5000, device=0):
    '''
    s_dims: tensor, ... h w
    b_dims: tensor, ... h w
    '''
    device_cuda = 0 if device == 'cpu' else device
    d = s_dims.shape[0]
    s_dims_new = ((s_dims / b_dims).ceil() * b_dims).to(torch.long)

    _, points_grid, _ = get_grid_points(s_dims_new.tolist(), align_corners=True, vmin=0, vmax=(s_dims_new-1).tolist(), 
        device=device)  # ... h w d
    points_grid = rearrange(points_grid, '(b0 p0) (b1 p1) d -> (b0 b1) (p0 p1) d', b0=b_dims[0], b1=b_dims[1])
    b, p = points_grid.shape[:-1]

    b_chunk = math.floor(chunk_size / p)
    if b_chunk >= b:
        order = torch.multinomial(torch.ones([b, p], device=device_cuda), p).to(device)  # b p
    else:
        order = torch.zeros([b, p], dtype=torch.long, device=device)
        prob = torch.ones([b_chunk, p], device=device_cuda)
        for i in range(0, b, b_chunk):
            b_chunk_i = min(b_chunk, b - i)
            if b_chunk_i != b_chunk:
                prob = torch.ones([b_chunk_i, p], device=device_cuda)
            order[i:i+b_chunk_i, :] = torch.multinomial(prob, p).to(order.device)

    points_grid = points_grid.gather(1, order[..., None].expand(-1, -1, d)).to(torch.long)  # b p d
    s_dims = s_dims.to(device)
    mask = (points_grid < s_dims[None, None]).all(dim=-1)  # b p
    points_grid = ravel_multi_index_loop(points_grid.flip(-1), s_dims.flip(-1))  # b p
    return points_grid.to(dtype=torch.int, device=device_cuda), mask.to(device_cuda)

util_misc.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import torch
from functools import partial
import time
import trimesh
import kornia

import util_misc
import util_init


class PE(torch.nn.Module):
    """
    positional encoding
    """
    def __init__(self, P):
        """
        P: (d, F) encoding matrix
        """
        super().__init__()
        self.register_buffer("P", P)

    @property
    def out_dim(self):
        return self.P.shape[1]*2

    def forward(self, x):
        """
        x: (p, d)
        """
        x_ = (x[..., None] * self.P[None]).sum(-2) # (p, F)
        return torch.cat([torch.sin(x_), torch.cos(x_)], -1) # (p, 2*F)


@torch.jit.script
def gaussian_activation(x, a):
    return torch.exp(-x**2/(2*a**2))


@torch.jit.script
def scaledsin_activation(x, a):
    return torch.sin(x*a)


@torch.jit.script
def rbf_nlin_s_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, 1)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc
    out = (1 - offset.abs() * ks).clamp(0, 1).prod(-1)
    return out, offset


@torch.jit.script
def rbf_nlin_d_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # [..., k, d]
    out = (1 - offset.abs() * ks).clamp(0, 1).prod(-1)
    return out, offset


@torch.jit.script
def rbf_nlin_f_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # [..., k, d]
    out = (1 - (offset[..., None, :] * ks).sum(-1).abs()).clamp(0, 1).prod(-1)
    return out, offset


@torch.jit.script
def rbf_ivq_s_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, 1)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc
    out = (1 / (1 + offset.pow(2).sum(-1) * ks[..., 0]))
    return out, offset


@torch.jit.script
def rbf_ivq_d_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = (1 / (1 + (offset.pow(2) * ks).sum(-1)))
    return out, offset


@torch.jit.script
def rbf_ivq_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = 1 / (1 + ((offset[..., None] * ks).sum(-2) * offset).sum(-1))
    return out, offset


def rbf_ivmq_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = 1 / torch.sqrt(1 + ((offset[..., None] * ks).sum(-2) * offset).sum(-1))
    return out, offset


def rbf_ivd_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = 1 / (torch.sqrt(((offset[..., None] * ks).sum(-2) * offset).sum(-1)) + 1e-10)
    return out, offset


def rbf_ivc_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = 1 / (((offset[..., None] * ks).sum(-2) * offset).sum(-1).sqrt()**3 + 1e-30)
    return out, offset


def rbf_gauss_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    d = x.shape[-1]
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.exp(-0.5*((offset[..., None] * ks).sum(-2) * offset).sum(-1))
    return out, offset


def rbf_sgauss2_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    d = x.shape[-1]
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.exp(-(0.5*((offset[..., None] * ks).sum(-2) * offset).sum(-1))**2)
    return out, offset


def rbf_nsgauss2_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    d = x.shape[-1]
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.exp(-0.5*((offset[..., None] * ks).sum(-2) * offset).sum(-1))**2
    return out, offset


def rbf_markov_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    d = x.shape[-1]
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.exp(-((offset[..., None] * ks).sum(-2) * offset).sum(-1).sqrt())
    return out, offset


def rbf_expsin_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.exp(-torch.sin(((offset[..., None] * ks).sum(-2) * offset).sum(-1).sqrt()))
    return out, offset


def rbf_qd_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = 1 + ((offset[..., None] * ks).sum(-2) * offset).sum(-1)
    return out, offset


def rbf_mqd_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.sqrt(1 + ((offset[..., None] * ks).sum(-2) * offset).sum(-1))
    return out, offset


def rbf_phs1_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = ((offset[..., None] * ks).sum(-2) * offset).sum(-1).sqrt()
    return out, offset


def rbf_phs1_s_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, 1)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.sqrt((offset**2).sum(-1))
    return out, offset


def rbf_phs3_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = ((offset[..., None] * ks).sum(-2) * offset).sum(-1).sqrt()**3
    return out, offset


def rbf_vep_a_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = 1 - ((offset[..., None] * ks).sum(-2) * offset).sum(-1)
    return out, offset


def rbf_cos_s_fb(x, kc, ks):
    """
        Inputs:
            x: (..., d)
            kc: (..., k, d)
            ks: (..., k, d, d)
        Outputs:
            (..., k)
    """
    offset = x[..., None, :] - kc  # (..., k, d)
    out = torch.nn.functional.cosine_similarity(x[..., None, :], kc, dim=-1)
    return out, offset


def fix_params(model, fix_params):
    for name, p in model.named_parameters():
        if name.split('.')[0] in fix_params:
            p.requires_grad = False


def get_param_groups(model):
    param_groups = {}
    for name, p in model.named_parameters():
        if p.requires_grad:
            if name.startswith('backbone'):
                k = 'dec'
            elif name.startswith('encoder'):
                k = 'hg'
            elif name.startswith('hg0'):
                k = 'hg0'
            elif name.startswith('lc0'):
                k = 'lc0'
            elif name.startswith('lcb0'):
                k = 'lcb0'
            elif name.startswith('kc0'):
                k = 'kc0'
            elif name.startswith('ks0'):
                k = 'ks0'
            else:
                raise NotImplementedError
            if k not in param_groups:
                param_groups[k] = []
            param_groups[k].append(p)
            print(k, name, p.shape)
    return param_groups


def parse_optim_type(optim_type):
    if optim_type == 'SparseAdam':
        return torch.optim.SparseAdam
    elif optim_type == 'Adam':
        return torch.optim.Adam
    elif optim_type == 'RAdam':
        return torch.optim.RAdam
    elif optim_type == 'SGD':
        return torch.optim.SGD
    else:
        raise NotImplementedError


def configure_optimizers(param_groups, hparams):
    optims = {}
    for k in param_groups.keys():
        hparams_k = hparams[k]
        optim_type = hparams_k['type']
        optim_fn = parse_optim_type(optim_type)
        if optim_type == 'SGD':
            optims[k] = optim_fn(param_groups[k], lr=hparams_k['lr'], weight_decay=hparams_k['wd'])
        elif optim_type == 'SparseAdam':
            optims[k] = optim_fn(param_groups[k], lr=hparams_k['lr'], betas=hparams_k['betas'], eps=hparams_k['eps'])
        else:
            optims[k] = optim_fn(param_groups[k], lr=hparams_k['lr'], betas=hparams_k['betas'], eps=hparams_k['eps'], weight_decay=hparams_k['wd'])
    return optims


def parse_lr_sch_type(lr_sch_type):
    if lr_sch_type == 'ngp':
        return lambda optimizer: torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 0.33**max(0, step//10000-1))
    elif lr_sch_type == 'cosine':
        return torch.optim.lr_scheduler.CosineAnnealingLR
    else:
        raise NotImplementedError


def configure_lr_schedulers(optims, hparams):
    lr_schs = {}
    for k in optims.keys():
        hparams_k = hparams[k]
        lr_sch_type = hparams_k['type']
        lr_sch_fn = parse_lr_sch_type(lr_sch_type)
        if lr_sch_type == 'ngp':
            lr_schs[k] = lr_sch_fn(optims[k])
        else:
            lr_schs[k] = lr_sch_fn(optims[k], T_max=hparams_k['T_max'], eta_min=optims[k].param_groups[0]['lr']/hparams_k['gamma'])
    return lr_schs


def init_kc_i(n_kernel, rbf_type, dataset, kc_init_config, device):
    with torch.no_grad():
        init_type = kc_init_config['type']
        kc_init_fn = None if init_type == 'none' else \
            partial(eval('util_init.kc_init_' + init_type), **kc_init_config)

        kc = None
        out_other = {}
        if kc_init_fn is not None:
            points = dataset.points.cpu().reshape(-1, dataset.points.shape[-1])
            if kc_init_config['weight_src'] == 'gt_grad':
                points_weight = dataset.gt_grad.cpu()
            elif kc_init_config['weight_src'] == 'gt':
                points_weight = dataset.gt.cpu().reshape(-1, dataset.gt.shape[-1])
                points_weight = (points_weight.abs() <= kc_init_config['weight_src_thres']).to(torch.float)
                points_weight[kc_init_config['weight_src_npts']:] = 0
            elif kc_init_config['weight_src'] == 'gt_inv':
                points_weight = dataset.gt.cpu().reshape(-1, dataset.gt.shape[-1])
                points_weight = 1 / (points_weight.abs() + kc_init_config['weight_src_thres'])
                points_weight[kc_init_config['weight_src_npts']:] = 0
            else:
                raise NotImplementedError
            kc, out_other = kc_init_fn(n_kernel, points=points, 
                points_weight=points_weight, cmin=dataset.cmin, cmax=dataset.cmax, 
                s_dims=dataset.s_dims, device=device, rbf_type=rbf_type)

    return kc, out_other


def init_ks_i(rbf_type, dataset, kw_init_config, kc, out_other, device):
    with torch.no_grad():
        init_type = kw_init_config['type']
        kw_init_fn = None if init_type == 'none' else \
            partial(eval('util_init.kw_init_' + init_type), **kw_init_config)

        ks = None
        if kw_init_fn is not None:
            points = dataset.points.cpu().reshape(-1, dataset.points.shape[-1])
            if kw_init_config['weight_src'] == 'gt_grad':
                points_weight = dataset.gt_grad.cpu()
            elif kw_init_config['weight_src'] == 'gt':
                points_weight = dataset.gt.cpu().reshape(-1, dataset.gt.shape[-1])
                points_weight = (points_weight.abs() <= kw_init_config['weight_src_thres']).to(torch.float)
                points_weight[kw_init_config['weight_src_npts']:] = 0
            elif kw_init_config['weight_src'] == 'gt_inv':
                points_weight = dataset.gt.cpu().reshape(-1, dataset.gt.shape[-1])
                points_weight = 1 / (points_weight.abs() + kw_init_config['weight_src_thres'])
                points_weight[kw_init_config['weight_src_npts']:] = 0
            else:
                raise NotImplementedError
            kw_sq = kw_init_fn(kc, in_other=out_other, points=points, 
                points_weight=points_weight, cmin=dataset.cmin, cmax=dataset.cmax, 
                s_dims=dataset.s_dims, device=device, rbf_type=rbf_type)
            kw_sq = util_misc.normalize_kw_sq(kw_sq, rbf_type)
            util_misc.clip_kw_sq(kw_sq, rbf_type, dataset.cmin.flip(-1), dataset.cmax.flip(-1), 
                dataset.s_dims.flip(-1), is_ks=False, is_flat=False, vmin_min=None)
            ks = util_misc.kw_sq_to_ks(kw_sq, rbf_type)

    return ks


def init_rbf_params(model, dataset, kc_init_config, kw_init_config, device):
    with torch.no_grad():
        for k in model.rbf_suffixes:
            n_kernel = getattr(model, f'kc{k}').weight.shape[0]

            print(f'Initializing kc{k}...')
            t = time.time()
            kc, out_other = init_kc_i(n_kernel, model.rbf_type, dataset, kc_init_config[k], device)
            if kc is not None:
                getattr(model, f'kc{k}').weight.data = kc
                print(f'Init kc{k}: {time.time() - t}')
            else:
                kc = getattr(model, f'kc{k}').weight.data

            print(f'Initializing ks{k}...')
            t = time.time()
            ks = init_ks_i(model.rbf_type, dataset, kw_init_config[k], kc, out_other, device)
            if ks is not None:
                getattr(model, f'ks{k}').weight.data = ks.reshape(n_kernel, -1)
                print(f'Init ks{k}: {time.time() - t}')


def init_nerf_kc_i(n_kernel, rbf_type, points, density, features, cmin, cmax, s_dims, kc_init_config, device):
    with torch.no_grad():
        init_type = kc_init_config['type']
        kc_init_fn = None if init_type == 'none' else \
            partial(eval('util_init.kc_init_' + init_type), **kc_init_config)

        kc = None
        out_other = {}
        if kc_init_fn is not None:
            points = points.cpu().reshape(-1, points.shape[-1])

            if features is None:
                points_weight = density
            else:
                alphas = 1 - torch.exp(-density)

                features = (features - features.min()) / (features.max() - features.min())
                features = features.reshape(*s_dims.tolist(), -1)
                features_grad = kornia.filters.SpatialGradient3d(mode='diff', order=1)(features.movedim(-1, 0)[None])[0]  # [c 3 d h w]
                features_grad = features_grad.movedim((0, 1), (-2, -1)).pow(2).sum(dim=[-2, -1]).sqrt()[..., None]  # [d h w 1]

                points_weight = alphas * features_grad.reshape(-1, 1)

            kc, out_other = kc_init_fn(n_kernel, points=points, points_weight=points_weight, 
                cmin=cmin, cmax=cmax, s_dims=s_dims, device=device, rbf_type=rbf_type)

    return kc, out_other


def init_nerf_ks_i(rbf_type, points, density, features, cmin, cmax, s_dims, kw_init_config, kc, out_other, device):
    with torch.no_grad():
        init_type = kw_init_config['type']
        kw_init_fn = None if init_type == 'none' else \
            partial(eval('util_init.kw_init_' + init_type), **kw_init_config)

        ks = None
        if kw_init_fn is not None:
            points = points.cpu().reshape(-1, points.shape[-1])

            if features is None:
                points_weight = density
            else:
                alphas = 1 - torch.exp(-density)

                features = (features - features.min()) / (features.max() - features.min())
                features = features.reshape(*s_dims.tolist(), -1)
                features_grad = kornia.filters.SpatialGradient3d(mode='diff', order=1)(features.movedim(-1, 0)[None])[0]  # [c 3 d h w]
                features_grad = features_grad.movedim((0, 1), (-2, -1)).pow(2).sum(dim=[-2, -1]).sqrt()[..., None]  # [d h w 1]

                points_weight = alphas * features_grad.reshape(-1, 1)

            kw_sq = kw_init_fn(kc, in_other=out_other, points=points, points_weight=points_weight, 
                cmin=cmin, cmax=cmax, s_dims=s_dims, device=device, rbf_type=rbf_type)
            kw_sq = util_misc.normalize_kw_sq(kw_sq, rbf_type)
            util_misc.clip_kw_sq(kw_sq, rbf_type, cmin.flip(-1), cmax.flip(-1), s_dims.flip(-1), 
                is_ks=False, is_flat=False, vmin_factor=100, vmax_factor=10)
            ks = util_misc.kw_sq_to_ks(kw_sq, rbf_type)

    return ks


def init_nerf_rbf_params(model, points, density, features, kc_init_config, kw_init_config, device):
    with torch.no_grad():
        for k in model.rbf_suffixes:
            n_kernel = getattr(model, f'kc{k}').weight.shape[0]

            print(f'Initializing kc{k}...')
            t = time.time()
            kc, out_other = init_nerf_kc_i(n_kernel, model.rbf_type, points, density, features, model.cmin, 
                                           model.cmax, model.s_dims, kc_init_config[k], device)
            if kc is not None:
                getattr(model, f'kc{k}').weight.data = kc
                print(f'Init kc{k}: {time.time() - t}')
            else:
                kc = getattr(model, f'kc{k}').weight.data

            print(f'Initializing ks{k}...')
            t = time.time()
            ks = init_nerf_ks_i(model.rbf_type, points, density, features, model.cmin, 
                                model.cmax, model.s_dims, kw_init_config[k], kc, out_other, device)
            if ks is not None:
                getattr(model, f'ks{k}').weight.data = ks.reshape(n_kernel, -1)
                print(f'Init ks{k}: {time.time() - t}')

util_network.py
oppo-us-research/NeuRBF
# MIT License

# Copyright (c) 2023 OPPO

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import math

import torch
import torch.nn.functional as F

import util_figure


def write_kernel_summary_2d(writer, step, data_gt, data_gt_grad, data_pred, kc, kw_sq, cmin=-1, cmax=1, suffix='', sidelength=None):
    '''
    data_gt: [1 c h w]
    data_gt_grad: [1 c h w]
    data_pred: [1 c h w]
    kc: [p d]
    kw_sq: [p] or [p d] or [p d d]
    '''
    size_scale_factor = (512**2 / torch.tensor(data_gt.shape[2:]).prod()).sqrt()
    data_gt = F.interpolate(data_gt, scale_factor=size_scale_factor, mode='bilinear', align_corners=False, 
        recompute_scale_factor=True)[0, ...].detach().cpu().numpy().transpose([1, 2, 0])
    data_gt_grad = F.interpolate(data_gt_grad, scale_factor=size_scale_factor, mode='bilinear', align_corners=False,
        recompute_scale_factor=True)[0, ...].detach().cpu().numpy().transpose([1, 2, 0])
    data_pred = F.interpolate(data_pred, scale_factor=size_scale_factor, mode='bilinear', align_corners=False,
        recompute_scale_factor=True)[0, ...].detach().cpu().numpy().transpose([1, 2, 0])

    vmin, vmax = 0, 1
    cmin, cmax = cmin[None, :].astype(np.float32), cmax[None, :].astype(np.float32)
    dims = data_gt.shape[:2][::-1]  # w h

    scale_factor = (np.array(dims, dtype=np.float32)[None, ...] / (cmax - cmin))  # [1 d]
    kc = (kc - cmin) * scale_factor  # Scale centers to [0, dims]
    kc -= 0.5  # The center of top-left pixel is [0, 0]
    if kw_sq.ndim == 1:
        kw_sq = kw_sq[:, None] * scale_factor ** 2
    elif kw_sq.ndim == 2:
        kw_sq = kw_sq * scale_factor ** 2
    elif kw_sq.ndim == 3:
        scale_factor_mat = scale_factor.T * scale_factor
        kw_sq = kw_sq * scale_factor_mat[None, ...]
    else:
        raise NotImplementedError

    xlim_scaled = (-dims[0] * 0.1, dims[0] * 1.1)
    ylim_scaled = (-dims[1] * 0.1, dims[1] * 1.1)

    fig, axes = plt.subplots(1, 1, figsize=[5*1, 5])
    axes = [axes]
    axes[0].imshow(data_gt_grad)
    for i in [0]:
        axes[i] = util_figure.plot_ellipses_v1(axes[i], kc, kw_sq, xlim_scaled, ylim_scaled, alpha=0.2, 
            flip_angle=True)
        axes[i].invert_yaxis()
    if fig is not None:
        writer.add_figure(f'kernels{suffix}', fig, global_step=step)


def write_image_summary_new(writer, step, data_gt, data_gt_grad, data_pred, cmin=-1, cmax=1, suffix=''):
    '''
    data_gt: [1 c h w]
    data_gt_grad: [1 c h w]
    data_pred: [1 c h w]
    '''
    size_scale_factor = (512**2 / torch.tensor(data_gt.shape[2:]).prod()).sqrt()
    data_gt = F.interpolate(data_gt, scale_factor=size_scale_factor, mode='bilinear', align_corners=False, 
        recompute_scale_factor=True)[0, ...].detach().cpu()
    data_pred = F.interpolate(data_pred, scale_factor=size_scale_factor, mode='bilinear', align_corners=False,
        recompute_scale_factor=True)[0, ...].detach().cpu()

    vmin, vmax = 0, 1

    writer.add_images(f'out{suffix}',
                    torch.cat([
                        data_gt.clip(vmin, vmax),
                        data_pred.clip(vmin, vmax),
                        ((data_gt - data_pred).abs()*50).clip(vmin, vmax),
                        ((data_gt - data_pred).pow(2).sum(0, keepdim=True).sqrt().expand(data_gt.shape[0], -1, -1)*50).clip(vmin, vmax),
                    ], -1),
                    global_step=step, dataformats='CHW')


def write_sdf_slice_summary(writer, step, data_gt, data_pred, suffix=''):
    '''
    data_gt: [1 c d h w]
    data_pred: [1 c d h w]
    '''
    size_scale_factor = (512**2 / torch.tensor(data_gt.shape[-2:]).prod()).sqrt()
    slice_id = math.ceil(data_gt.shape[-1]*1/2) - 1
    data_gt = torch.cat([data_gt[..., slice_id, :, :], data_gt[..., slice_id, :], data_gt[..., slice_id]], -2)  # 1 c 3h w
    data_pred = torch.cat([data_pred[..., slice_id, :, :], data_pred[..., slice_id, :], data_pred[..., slice_id]], -2)  # 1 c 3h w

    data_gt = F.interpolate(data_gt, scale_factor=size_scale_factor, mode='bilinear', align_corners=False, 
        recompute_scale_factor=True)[0, ...].detach().cpu().permute(1, 2, 0)  # 3h w c
    data_pred = F.interpolate(data_pred, scale_factor=size_scale_factor, mode='bilinear', align_corners=False,
        recompute_scale_factor=True)[0, ...].detach().cpu().permute(1, 2, 0)  # 3h w c

    vmin, vmax = -0.1, 0.1
    fig, axes = plt.subplots(1, 4, figsize=[5*4, 5*3])
    axes[0].imshow(data_gt, vmin=vmin, vmax=vmax)
    axes[0].contour(data_gt[..., 0], levels=[0], colors='k', linewidths=0.5)
    axes[1].imshow(data_pred, vmin=vmin, vmax=vmax)
    axes[1].contour(data_pred[..., 0], levels=[0], colors='k', linewidths=0.5)
    axes[2].imshow((data_gt - data_pred).abs()*10, vmin=vmin, vmax=vmax)
    axes[2].contour(data_gt[..., 0], levels=[0], colors='r', linewidths=0.5, alpha=1)
    axes[2].contour(data_pred[..., 0], levels=[0], colors='k', linewidths=0.5, alpha=1)
    axes[3].imshow((data_gt - data_pred).abs()*10, vmin=vmin, vmax=vmax)
    for ax in axes:
        ax.invert_yaxis()
        ax.axis('off')

    writer.add_figure(f'out{suffix}', fig, global_step=step)


def write_tensor3d_slice_summary(writer, step, data_gt, data, suffix=''):
    '''
    data: [1 c d h w]
    '''
    size_scale_factor = (512**2 / torch.tensor(data.shape[-2:]).prod()).sqrt()
    slice_id = math.ceil(data.shape[-1]*1/2) - 1
    nrow, ncol = 4, 8+1
    fig, axes = plt.subplots(nrow, ncol, figsize=[2.25*ncol, 2.25*nrow])
    
    data_gt = data_gt[..., slice_id, :, :]  # 1 c h w
    data_gt = F.interpolate(data_gt, scale_factor=size_scale_factor, mode='bilinear', align_corners=False, 
        recompute_scale_factor=True)[0, ...].detach().cpu().permute(1, 2, 0)  # 3h w c
    for i in range(nrow):
        axes[i, 0].imshow(data_gt)
        axes[i, 0].invert_yaxis()
        axes[i, 0].axis('off')

    for i in range(nrow):
        for j in range(1, ncol):
            cid = i*(ncol-1)+j-1
            if cid < data.shape[1]:
                data_show = data[:, cid:cid+1, slice_id, :, :]  # 1 1 h w
                data_show = F.interpolate(data_show, scale_factor=size_scale_factor, mode='bilinear', 
                    align_corners=False, recompute_scale_factor=True)[0, ...].detach().cpu().permute(1, 2, 0)  # h w 1
                axes[i, j].imshow(data_show)
            axes[i, j].invert_yaxis()
            axes[i, j].axis('off')

    plt.subplots_adjust(left=0., bottom=0., right=1., top=1., wspace=0.03, hspace=0.03)
    writer.add_figure(f'{suffix}', fig, global_step=step)

util_summary.py
oppo-us-research/NeuRBF