算法组件API

介绍算法组件相关的api，包括基础 Configs, BaseModels, Components 和 Ray Samplers 。

Configs

config中包含数据集配置、模型配置和训练配置，分别对应DatasetConfig、ModelConfig和TrainConfig。 定义模型配置类的时候，需要继承自 landmark.nerf_components.configs.ConfigClass 类， 并实现 check_args 函数，该函数用于检查参数是否有效。

class landmark.nerf_components.configs.ConfigClass(*args)[源代码]

Base Config

model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True, 'protected_namespaces': ()}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[dict[str, FieldInfo]] = {}

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

定义模型配置类的示例：

from landmark.nerf_components.configs import ConfigClass
class OctreeGSConfig(ConfigClass):
    """Octree-GS Config"""

    feat_dim: int = 32
    view_dim: int = 3
    n_offsets: int = 10
    fork: int = 2

    use_feat_bank: bool = False
    source_path: str = ""
    model_path: str = ""
    images: str = "images"
    resolution: int = 1
    white_background: bool = False
    random_background: bool = False
    resolution_scales: List[float] = [1.0]

    ...
class landmark.nerf_components.configs.DatasetConfig(*args, partition: str = 'all', datadir: str, subfolder: list = None, dataset_name: Literal['city', 'matrixcity', 'blender', 'Colmap'] = 'city', dataset_type: Literal['nerf', 'gaussian'] = 'nerf', use_preprocessed_data: bool = False, processed_data_type: str = 'file', preprocessed_dir: str = '/cpfs01/shared/landmarks-viewer/preprocessed_data/', downsample_train: int = 1, train_all_data: bool = False, img_uint8: bool = False, preload: bool = False, dynamic_load: bool = False, max_iter_each_block: int = -1, lb: List[float] | Tensor = [0, 0, 0], ub: List[float] | Tensor = [0, 0, 0], pose_rotation: List[float] | Tensor = [0, 0, 0], pose_translation: List[float] | Tensor = [0, 0, 0])[源代码]

Dataset Configuration class for setting up dataset parameters.

参数:

  • partition (str) – Specifies the dataset partition to use.

  • datadir (str) – Directory where the dataset is located.

  • subfolder (list, optional) – List of subfolders to include in the dataset. Defaults to None.

  • dataset_name (Literal["city", "matrixcity", "blender", "Colmap"], optional) – Name of the dataset. Defaults to “city”.

  • dataset_type (Literal["nerf", "gaussian"], optional) – Type of the dataset. Defaults to “nerf”.

  • use_preprocessed_data (bool, optional) – Flag to indicate if preprocessed data should be used. Defaults to False.

  • processed_data_type (str, optional) – Type of preprocessed data. Defaults to “file”.

  • preprocessed_dir (str, optional) – Directory for preprocessed data. Defaults to “/cpfs01/shared/landmarks-viewer/preprocessed_data/”.

  • downsample_train (int, optional) – Downsampling factor for training images. Defaults to 1.

  • train_all_data (bool, optional) – Flag to indicate if all data should be used for training. Defaults to False.

  • img_uint8 (bool, optional) – Flag to indicate if images are in uint8 format. Defaults to False.

  • preload (bool, optional) – Flag to indicate if data should be preloaded into memory. Defaults to False.

  • dynamic_load (bool, optional) – Flag to indicate if data should be dynamically loaded. Defaults to False.

  • max_iter_each_block (int, optional) – Maximum iterations per block for dynamic training. Defaults to -1.

  • lb (Union[List[float], torch.Tensor], optional) – Lower bound of the bounding box. Defaults to [0, 0, 0].

  • ub (Union[List[float], torch.Tensor], optional) – Upper bound of the bounding box. Defaults to [0, 0, 0].

  • pose_rotation (Union[List[float], torch.Tensor], optional) – Rotation part of the pose. Defaults to [0, 0, 0].

  • pose_translation (Union[List[float], torch.Tensor], optional) – Translation part of the pose. Defaults to [0, 0, 0].

check_args()[源代码]

Performs a sanity check on the dataset configuration.

class landmark.nerf_components.configs.RenderConfig(*args, model_name: str = None, expname: str = None, batch_size: int = 4096, basedir: str = './log', add_timestamp: bool = False, random_seed: int = 20211202, debug: bool = False, N_vis: int = 5, ckpt: str | None = None, ckpt_type: typing_extensions.Literal[full, sub] = 'full', kwargs: str | None = None, is_train: bool = False, distributed: bool = False, device: str | device = device(type='cuda'), DDP: bool = False, channel_parallel: bool = False, branch_parallel: bool = False, plane_division: list | None = None, channel_parallel_size: int | None = None, model_parallel_and_DDP: bool = False, test_iterations: List[int] = [], save_iterations: List[int] = [], quiet: bool = False, start_checkpoint: str | None = None, compute_extra_metrics: bool = False)[源代码]

Render Config

参数:

  • model_name (str) – The name of the model.

  • expname (str) – The name of the experiment.

  • batch_size (int) – The size of the batch for processing.

  • basedir (str) – The base directory for logging and saving outputs.

  • add_timestamp (bool) – Whether to add a timestamp to the output filenames.

  • random_seed (int) – The seed for random number generation to ensure reproducibility.

  • debug (bool) – Enables debug mode if set to True.

  • N_vis (int) – Number of visualizations to generate.

  • ckpt (Optional[str]) – Path to a checkpoint file.

  • ckpt_type (Literal["full", "sub"]) – Type of checkpoint, either ‘full’ or ‘sub’.

  • kwargs (Optional[str]) – Additional keyword arguments.

  • is_train (bool) – Flag indicating whether the configuration is for training.

  • distributed (bool) – Flag indicating whether distributed training is enabled.

  • device (Union[str, torch.device]) – The computing device to use (‘cuda’ or ‘cpu’).

  • DDP (bool) – Flag indicating whether Distributed Data Parallel (DDP) is used.

  • channel_parallel (bool) – Flag indicating whether channel parallelism is used.

  • branch_parallel (bool) – Flag indicating whether branch parallelism is used.

  • plane_division (Optional[list]) – Specifies how to divide the input data into planes.

  • channel_parallel_size (Optional[int]) – Size of the channel parallelism.

  • model_parallel_and_DDP (bool) – Flag indicating whether model parallelism and DDP are combined.

  • test_iterations (List[int]) – List of iterations at which to perform testing.

  • save_iterations (List[int]) – List of iterations at which to save outputs.

  • quiet (bool) – If True, reduces the verbosity of the output.

  • start_checkpoint (Optional[str]) – Path to the starting checkpoint for resuming training.

  • compute_extra_metrics (bool) – Flag indicating whether to compute additional metrics.

check_args()[源代码]

This method ensures that the configuration is valid for rendering

from_train_config(train_config: TrainConfig)[源代码]

Updates the configuration based on a given TrainConfig instance.

参数:

train_config (TrainConfig) – The training configuration to update this configuration with.

class landmark.nerf_components.configs.TrainConfig(*args, model_name: str, expname: str, start_iters: int = 0, n_iters: int, batch_size: int = 4096, progress_refresh_rate: int = 10, wandb: bool = False, tensorboard: bool = False, basedir: str = './log', add_timestamp: bool = False, random_seed: int = 20211202, debug: bool = False, N_vis: int = 5, vis_every: int = 10000, skip_save_imgs: bool = False, optim_dir: str | None = None, ckpt: str | None = None, ckpt_type: typing_extensions.Literal[full, sub] = 'full', kwargs: str | None = None, is_train: bool = True, device: str | device = device(type='cuda'), DDP: bool = False, channel_parallel: bool = False, branch_parallel: bool = False, plane_division: list | None = [1, 1], channel_parallel_size: int | None = None, model_parallel_and_DDP: bool = False, test_iterations: List[int] = [], save_iterations: List[int] = [], start_checkpoint: str | None = None, print_timestamp: bool = True)[源代码]

Train Configuration class for setting up training parameters.

参数:

  • model_name (str) – Name of the model to be trained.

  • expname (str) – Experiment name.

  • start_iters (int, optional) – Starting iteration number. Defaults to 0.

  • n_iters (int) – Total number of iterations.

  • batch_size (int, optional) – Batch size for training. Defaults to 4096.

  • progress_refresh_rate (int, optional) – Refresh rate for progress updates. Defaults to 10.

  • wandb (bool, optional) – Flag to indicate if Weights & Biases logging is enabled. Defaults to False.

  • tensorboard (bool, optional) – Flag to indicate if TensorBoard logging is enabled. Defaults to False.

  • basedir (str, optional) – Base directory for logging. Defaults to “./log”.

  • add_timestamp (bool, optional) – Flag to indicate if timestamp should be added to log filenames. Defaults to False.

  • random_seed (int, optional) – Random seed. Defaults to 20211202.

  • debug (bool, optional) – Flag to indicate if debug mode is enabled. Defaults to False.

  • N_vis (int, optional) – Number of images to visualize. Defaults to 5.

  • vis_every (int, optional) – Iteration interval for visualization. Defaults to 10000.

  • skip_save_imgs (bool, optional) – Flag to indicate if saving images is skipped. Defaults to False.

  • optim_dir (Optional[str], optional) – Directory for optimization checkpoints. Defaults to None.

  • ckpt (Optional[str], optional) – Path to a specific checkpoint. Defaults to None.

  • ckpt_type (Literal["full", "sub"], optional) – Type of checkpoint. Defaults to “full”.

  • kwargs (Optional[str], optional) – Additional arguments for checkpointing. Defaults to None.

  • is_train (bool, optional) – Flag to indicate if the configuration is for training. Defaults to True.

  • device (Union[str, torch.device], optional) – Device to use for training. Defaults to torch.device(“cuda”).

  • DDP (bool, optional) – Flag to indicate if Distributed Data Parallel is used. Defaults to False.

  • channel_parallel (bool, optional) – Flag to indicate if channel parallelism is used. Defaults to False.

  • branch_parallel (bool, optional) – Flag to indicate if branch parallelism is used. Defaults to False.

  • plane_division (Optional[list], optional) – Division of planes for parallelism. Defaults to [1, 1].

  • channel_parallel_size (Optional[int], optional) – Size for channel parallelism. Defaults to None.

  • model_parallel_and_DDP (bool, optional) – Flag to indicate if model parallelism and DDP are combined. Defaults to False.

  • test_iterations (List[int], optional) – Iterations for testing. Defaults to [].

  • save_iterations (List[int], optional) – Iterations for saving checkpoints. Defaults to [].

  • start_checkpoint (Optional[str], optional) – Path to the starting checkpoint. Defaults to None.

  • print_timestamp (bool, optional) – Flag to indicate if timestamps should be printed. Defaults to True.

check_args()[源代码]

Performs a sanity check on the training configuration.

Base Models

BaseModels中包含了NeRF的基础模型BaseNeRF，Gaussian的基础模型BaseGaussian，自定义模型需要继承基础模型，并实现其中方法。 NeRF模型可以参考benchmarks/nerf/gridnerf/gridnerf.py，Gaussian模型可以参考benchmarks/nerf/origin_gs/gs_model.py。

class landmark.nerf_components.model.base_model.Model(config=None, scene_manager=None, **kwargs)[源代码]

A basic model class for user-defined models, providing functionality for configuration management, scene management, and handling of additional keyword arguments.

This class is designed to be a flexible foundation for building complex models, allowing for easy extension and customization.

config

Configuration object containing model settings.

Type:

Config

scene_manager

Manager for scene information.

Type:

SceneManager

_kwargs_keys

Keys of additional keyword arguments passed during initialization.

Type:

list

kwargs()[源代码]

Returns a dictionary of additional keyword arguments.

state_dict(destination=None, prefix='', keep_vars=False)[源代码]

Returns the state dictionary of the model.

save_kwargs(file_path)[源代码]

Saves the keyword arguments to a file.

static get_kwargs_from_ckpt(file_path: str | PathLike, map_location: str | Dict[str, str]) → Mapping[str, Any][源代码]

Get kwargs from ckpt to initialize model class.

参数:

  • file_path – a string or os.PathLike object containing a file name

  • map_location – a string or a dict specifying how to remap storage locations

get_state_dict_from_ckpt(file_path: str | PathLike, map_location: str | Dict[str, str]) → Mapping[str, Any][源代码]

Get state_dict from the given file path.

参数:

  • file_path – a string or os.PathLike object containing a file name

  • map_location – a string or a dict specifying how to remap storage locations

kwargs()[源代码]

Returns a dictionary of the model’s keyword arguments.

返回:

A dictionary containing the keyword arguments.

返回类型:

dict

load_from_state_dict(state_dict: Mapping[str, Any], strict: bool = True)[源代码]

Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

参数:

  • state_dict (dict) – a dict containing parameters and persistent buffers.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s key

save_kwargs(file_path: str | PathLike)[源代码]
save_model(dir_path: str | PathLike, rank=None)[源代码]

Saves both the state dictionary and keyword arguments of the model to files in the specified directory.

This method is a convenience function that saves the model’s state dictionary and any additional keyword arguments to separate files within a given directory. It supports saving multiple copies of the model for different ranks in distributed training scenarios.

参数:

  • dir_path (Union[str, os.PathLike]) – The directory path where the model files will be saved.

  • rank (Optional[int]) – The rank of the process in distributed training. If provided, the method saves the files with a suffix indicating the rank.

抛出:

AssertionError – If there are no keyword arguments to save, indicating an empty keyword argument list.

save_state_dict(file_path: str | PathLike)[源代码]

Saves the state dictionary of the model to a file.

This method serializes the state dictionary of the model and saves it to the specified file path. It is useful for checkpointing models during training or for later restoration.

参数:

file_path (Union[str, os.PathLike]) – The path to the file where the state dictionary will be saved.

抛出:

AssertionError – If the state dictionary is empty, indicating there is nothing to save.

state_dict(destination=None, prefix='', keep_vars=False)[源代码]

Returns the state dictionary of the model, potentially modified by the destination and prefix.

参数:

  • destination – The destination for saving the state dictionary.

  • prefix (str) – A prefix to add to each key in the state dictionary.

  • keep_vars (bool) – Whether to keep variables that are not part of the model.

返回:

The modified state dictionary.

返回类型:

dict

class landmark.nerf_components.model.base_model.BaseNeRF(config=None, scene_manager=None, **kwargs)[源代码]

Base Class for volumetric rendered NeRF model

abstract init_field_components()[源代码]
abstract init_renderer()[源代码]
abstract init_sampler()[源代码]
render_all_rays(rays, chunk_size=None, N_samples=-1, idxs=None, app_code=None)[源代码]
class landmark.nerf_components.model.base_model.BaseGaussian(config=None, scene_manager=None, **kwargs)[源代码]

Base Class for Gaussian based model

get_state_dict_from_ckpt(file_path: str | PathLike, map_location: str | Dict[str, str]) → Mapping[str, Any][源代码]

Retrieves and merges the state dictionary from a checkpoint file, excluding the Gaussian encoding state.

This method is designed to load the state dictionary from a checkpoint file, specifically excluding the state of the Gaussian encoding module. It then attempts to merge this state dictionary with an existing state dictionary saved in the same directory as the checkpoint file. If duplicate keys are found during the merge, it raises a ValueError.

参数:

  • file_path (Union[str, os.PathLike]) – The path to the checkpoint file from which to load the state

  • dictionary.

  • map_location (Union[str, Dict[str, str]]) – The location where the model should be mapped if it is loaded on a different device than it was saved on.

返回:

The merged state dictionary, excluding the Gaussian encoding state.

返回类型:

Mapping[str, Any]

抛出:

ValueError – If duplicate keys are found between the loaded state dictionary and any existing state dictionary in the same directory.

abstract init_field_components()[源代码]
abstract init_renderer()[源代码]
load_from_state_dict(state_dict: Mapping[str, Any], strict: bool = True)[源代码]

Loads the model’s state dictionary from a given mapping, excluding the Gaussian encoding state.

This method loads the state dictionary into the model, specifically excluding the state of the Gaussian encoding module. It first extracts and removes any keys related to the Gaussian encoding from the provided state dictionary. Then, it calls the base class’s load_from_state_dict method to load the remaining state dictionary. Finally, it loads the Gaussian encoding state using its own load_from_state_dict method.

参数:

  • state_dict (Mapping[str, Any]) – The state dictionary to load into the model.

  • strict (bool, optional) – If True, the method will raise an error if keys in the state dictionary do not match the model’s keys. Defaults to True.

备注

This method modifies the input state_dict by removing keys related to the Gaussian encoding before passing it to the base class’s method. It also prints a warning if no keys related to the Gaussian encoding are found in the provided state dictionary.

save_model(dir_path: str | PathLike, rank=None)[源代码]

Saves the model’s state and point cloud data to the specified directory.

This method extends the base model’s save functionality by also saving the point cloud data associated with the model. It first saves the point cloud data using the save_ply method and then calls the base class’s save_model method to save the model’s state.

参数:

  • dir_path (Union[str, os.PathLike]) – The directory path where the model and point cloud data will be saved.

  • rank (Optional[int]) – The rank of the process in distributed training. If provided, the method saves the files with a suffix indicating the rank.

save_ply(dir_path: str | PathLike, rank=None)[源代码]

Saves the point cloud data to a PLY file in the specified directory.

This method finds the Gaussian encoding module within the model, uses it to generate the point cloud data, and saves this data to a PLY file. The file name is adjusted based on the rank in distributed training.

参数:

  • dir_path (Union[str, os.PathLike]) – The directory path where the point cloud data will be saved.

  • rank (Optional[int]) – The rank of the process in distributed training. If provided, the method saves the files with a suffix indicating the rank.

save_state_dict(file_path: str | PathLike)[源代码]

Saves the model’s state dictionary to a file, excluding the Gaussian encoding state.

This method modifies the base class’s save_state_dict method by excluding the state of the Gaussian encoding module from the saved state dictionary. This is useful for scenarios where the Gaussian encoding state does not need to be saved or restored, such as when the encoding is specific to the training process and not required for inference.

参数:

file_path (Union[str, os.PathLike]) – The path to the file where the state dictionary will be saved.

备注

This method directly modifies the state dictionary before saving it, removing any keys that start with the name of the Gaussian encoding module. This ensures that the saved state dictionary does not include the Gaussian encoding state.

Model Components

Components中提供了常用的算法组件，建议使用这里的组件进行模型定义。

class landmark.nerf_components.model_components.fields.alpha_mask.AlphaGridMask(device, aabb, alpha_volume)[源代码]

A class for the alpha grid mask.

参数:

  • device (str) – The device to use.

  • aabb (torch.Tensor) – The axis-aligned bounding box.

  • alpha_volume (torch.Tensor) – The alpha volume.

forward(xyz_sampled, filter_thresh: Tensor, mask: Tensor | None = None)[源代码]

Computes the forward pass of the AlphaGridMask.

参数:

  • xyz_sampled (torch.Tensor) – The sampled coordinates.

  • filter_thresh (torch.Tensor) – The filter threshold.

  • mask (torch.Tensor, optional) – The mask tensor. Defaults to None.

返回:

The computed mask.

返回类型:

torch.Tensor

generate_alpha_mask(xyz_sampled, filter_thresh=0, mask=None)[源代码]

Generates the alpha mask.

参数:

  • xyz_sampled (torch.Tensor) – The sampled coordinates.

  • filter_thresh (int, optional) – The filter threshold. Defaults to 0.

  • mask (torch.Tensor, optional) – The mask tensor. Defaults to None.

返回:

The generated alpha mask or a tuple of alpha mask and valid mask.

返回类型:

torch.Tensor or tuple

normalize_coord(xyz_sampled)[源代码]

Normalizes the sampled coordinates.

参数:

xyz_sampled (torch.Tensor) – The sampled coordinates.

返回:

The normalized coordinates.

返回类型:

torch.Tensor

sample_alpha(xyz_sampled)[源代码]

Samples the alpha values.

参数:

xyz_sampled (torch.Tensor) – The sampled coordinates.

返回:

The alpha values.

返回类型:

torch.Tensor

update_device(device)[源代码]

Updates the device used by the AlphaGridMask instance.

参数:

device (str) – The new device to use.

class landmark.nerf_components.model_components.fields.anchor_decoder.AnchorDecoder(use_feat_bank: bool, add_opacity_dist: bool, add_color_dist: bool, add_cov_dist: bool, view_dim: int, feat_dim: int, n_offsets: int, appearance_dim: int = 0, add_level: bool = False, dist2level: str | None = None)[源代码]

Decode the anchors to neural gaussians

参数:

  • use_feat_bank (bool) – Whether to use feature bank.

  • add_opacity_dist (bool) – Whether to add opacity distribution.

  • add_color_dist (bool) – Whether to add color distribution.

  • add_cov_dist (bool) – Whether to add covariance distribution.

  • view_dim (int) – Dimension of the view.

  • feat_dim (int) – Dimension of the features.

  • n_offsets (int) – Number of offsets.

  • appearance_dim (int, optional) – Dimension of the appearance code. Defaults to 0.

  • add_level (bool, optional) – Whether to add level. Defaults to False.

  • dist2level (str, optional) – Distance to level. Defaults to None.

use_feat_bank

Whether to use feature bank.

Type:

bool

add_opacity_dist

Whether to add opacity distribution.

Type:

bool

add_color_dist

Whether to add color distribution.

Type:

bool

add_cov_dist

Whether to add covariance distribution.

Type:

bool

add_level

Whether to add level.

Type:

bool

dist2level

Distance to level.

Type:

str

view_dim

Dimension of the view.

Type:

int

feat_dim

Dimension of the features.

Type:

int

n_offsets

Number of offsets.

Type:

int

appearance_dim

Dimension of the appearance code.

Type:

int

mlp_feature_bank

MLP for feature bank.

Type:

nn.Sequential

mlp_opacity

MLP for opacity.

Type:

nn.Sequential

mlp_cov

MLP for covariance.

Type:

nn.Sequential

mlp_color

MLP for color.

Type:

nn.Sequential

forward(camera_center, anchor, anchor_feat, offset, scaling, visible_mask, level, prog_ratio,
transition_mask, appearance_code)

Forward pass of the Anchor Decoder.

forward(camera_center: Tensor, anchor: Tensor, anchor_feat: Tensor, offset: Tensor, scaling: Tensor, visible_mask: Tensor | None, level: Tensor | None = None, prog_ratio: Tensor | None = None, transition_mask: Tensor | None = None, appearance_code: Tensor | None = None)[源代码]

Forward pass of the Anchor Decoder.

参数:

  • camera_center (torch.Tensor) – Camera center.

  • anchor (torch.Tensor) – Anchor.

  • anchor_feat (torch.Tensor) – Anchor features.

  • offset (torch.Tensor) – Offset.

  • scaling (torch.Tensor) – Scaling.

  • visible_mask (torch.Tensor, optional) – Visible mask. Defaults to None.

  • level (torch.Tensor, optional) – Level. Defaults to None.

  • prog_ratio (torch.Tensor, optional) – Progress ratio. Defaults to None.

  • transition_mask (torch.Tensor, optional) – Transition mask. Defaults to None.

  • appearance_code (torch.Tensor, optional) – Appearance code. Defaults to None.

返回:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:

  • xyz (torch.Tensor): XYZ coordinates.

  • color (torch.Tensor): Color.

  • opacity (torch.Tensor): Opacity.

  • scaling (torch.Tensor): Scaling.

  • rot (torch.Tensor): Rotation.

  • neural_opacity (torch.Tensor, optional): Neural opacity. Defaults to None.

  • mask (torch.Tensor, optional): Mask. Defaults to None.

class landmark.nerf_components.model_components.fields.embedding.AppearanceEmbedding(n_imgs: int, n_component: int, device: device)[源代码]

A class for the appearance embedding.

forward(idxs: Tensor | None = None, xyz_sampled: Tensor | None = None, app_code: int | None = None)[源代码]

Compute the appearance latent vector.

参数:

  • idxs (torch.Tensor, optional) – Indices of the sampled points.

  • xyz_sampled (torch.Tensor, optional) – Sampled points.

  • app_code (int, optional) – Given appearance code.

返回:

The appearance latent vector.

返回类型:

torch.Tensor

class landmark.nerf_components.model_components.MLPDecoder(inChanel, layer_width=128, num_layers: int | None = None, out_dim: int | None = None, out_activation: Module | None = None, mlp_bias: bool = True)[源代码]

A MLP for rendering feature

参数:

  • inChanel (int) – The number of input channels.

  • layer_width (int) – The number of output channels.

  • num_layers (int) – The number of layers in the MLP.

  • layer_width – The width of the MLP.

  • out_dim (int) – The output dimension of the MLP.

  • out_activation (torch.nn.Module) – The output activation function.

  • mlp_bias (bool) – The learnable bias of the module.

forward(features)[源代码]

Forward pass of the MLP Render Feature.

参数:

  • viewdirs (torch.Tensor) – The view direction tensor.

  • features (torch.Tensor) – The input feature tensor.

返回:

The output RGB tensor.

返回类型:

torch.Tensor

class landmark.nerf_components.model_components.NeRF(args, den_n_comp, app_n_comp, nfreqs=None)[源代码]

Class for NeRF branch

参数:

  • args (object) – The arguments object.

  • den_n_comp (int) – The number of components for the density field.

  • app_n_comp (int) – The number of components for the appearance field.

  • nfreqs (int, optional) – The number of frequencies for positional encoding. Defaults to None.

num_freq

The number of frequencies for positional encoding.

Type:

int

den_n_comp

The number of components for the density field.

Type:

int

app_n_comp

The number of components for the appearance field.

Type:

int

skips

The list of skip connections.

Type:

list

use_viewdirs

Flag indicating whether to use view directions.

Type:

bool

D

The value of Depth in the first stage.

Type:

int

D_a

The value of Depth in the second stage.

Type:

int

W

The value of Width.

Type:

int

encode_app

Flag indicating whether to encode appearance.

Type:

bool

init_module_v0()[源代码]

Initializes the module.

forward()[源代码]

Forward pass of the NeRF branch.

forward(pts, viewdir=None, den_feats=None, app_feats=None, app_latent=None)[源代码]

Forward pass of the NeRF branch.

参数:

  • pts (torch.Tensor) – The input points tensor.

  • viewdir (torch.Tensor, optional) – The view direction tensor. Defaults to None.

  • den_feats (torch.Tensor, optional) – The density field tensor. Defaults to None.

  • app_feats (torch.Tensor, optional) – The appearance field tensor. Defaults to None.

  • dists (torch.Tensor, optional) – The distances. Defaults to None.

返回:

The output tensor.

返回类型:

torch.Tensor

init_module_v0()[源代码]

Initializes the module.

class landmark.nerf_components.model_components.VolumeRenderer(white_bg: bool = False)[源代码]

Class for volume rendering.

参数:

white_bg (bool, optional) – Flag indicating whether to use a white background. Defaults to False.

forward(weights: Tensor, rgbs: Tensor, z_vals: Tensor, random_white_bg_ratio: float = 0.0)[源代码]

Compute the weighted sum of the RGB values.

参数:

  • weights (torch.Tensor) – Weights. Shape: [N_rays, N_samples]

  • rgbs (torch.Tensor) – RGB values. Shape: [N_rays, N_samples, 3]

  • z_vals (torch.Tensor) – Depth values. Shape: [N_rays, N_samples]

  • random_white_bg_ratio (float, optional) – Random white background ratio. Defaults to 0.0.

返回:

RGB map. Shape: [N_rays, 3] torch.Tensor: Depth map. Shape: [N_rays]

返回类型:

torch.Tensor

get_weights(sigma: Tensor, distances: Tensor) → Tensor[源代码]

Compute the weights for each sample.

参数:

  • sigma (torch.Tensor) – Sigma values. Shape: [N_rays, N_samples]

  • distances (torch.Tensor) – Distances. Shape: [N_rays, N_samples]

返回:

Weights. Shape: [N_rays, N_samples]

返回类型:

torch.Tensor

Ray Samplers

Ray Samplers中提供了常用的光线采样方法，一般用于NeRF的模型中，可以根据需要选择使用。实现自定义采样器建议继承BaseSampler，并实现其中的方法。

class landmark.nerf_components.ray_samplers.base_sampler.BaseSampler(*args, **kwargs)[源代码]

Sample points along the Rays.

abstract forward(*args, **kwargs) → Samples[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class landmark.nerf_components.ray_samplers.uniform_sampler.UniformSampler(num_samples: int, step_size: int, aabb: Tensor, near_far: list, *args, **kwargs)[源代码]

Sample N_samples number points on a ray. If is_train is True, use segmented random sampling. Otherwise use uniform sampling.

参数:

num_samples (int) – number of samples on each ray.

forward(rays: ~landmark.nerf_components.data.rays.Rays | ~torch.Tensor, num_samples: int = -1, random_sampling: bool = False, sample_within_hull: bool = False) -> (<class 'landmark.nerf_components.data.samples.Samples'>, <class 'torch.Tensor'>)[源代码]

Sample points along the rays. :param rays: rays to sample. :type rays: Rays :param num_samples: number of samples on each ray. If None, use self.num_samples. :type num_samples: int

返回:

sampled points. (torch.Tensor): mask of the points outside the bounding box.

返回类型:

(Samples)

class landmark.nerf_components.ray_samplers.pdf_sampler.PDFSampler(num_samples: int, *args, **kwargs)[源代码]

A sampler that generates samples based on a Probability Density Function (PDF).

This class extends the BaseSampler to implement sampling methods that are specifically designed to generate samples according to a given Probability Density Function. It is useful in scenarios where the distribution of the samples needs to follow a specific statistical distribution.

参数:

  • num_samples (int) – The number of samples to generate.

  • *args – Variable length argument list for the base class.

  • **kwargs – Arbitrary keyword arguments for the base class.

forward(rays: ~landmark.nerf_components.data.rays.Rays | ~torch.Tensor, z_vals, weights, num_samples: int | None = None, det=False) -> (<class 'landmark.nerf_components.data.samples.Samples'>, <class 'torch.Tensor'>)[源代码]

Sample points along the rays based on a probability density function.

参数:

  • rays (Union[Rays, torch.Tensor]) – Rays to sample from.

  • z_vals (torch.Tensor) – Values along the rays.

  • weights (torch.Tensor) – Weights for the PDF.

  • num_samples (Optional[int], optional) – Number of samples on each ray. If None, use self.num_samples.

  • det (bool, optional) – Deterministic sampling. Defaults to False.

返回:

A tuple containing the sampled points and a mask of the points outside the bounding box.

返回类型:

(Samples, torch.Tensor)

class landmark.nerf_components.ray_samplers.proposal_network_sampler.ProposalNetworkSampler(near_far, num_proposal_samples_per_ray: ~typing.Tuple[int, ...] = (64, ), num_nerf_samples_per_ray: int = 32, num_proposal_network_iterations: int = 2, config=None, aabb=None, single_jitter: bool = False, update_sched: ~typing.Callable = <function ProposalNetworkSampler.<lambda>>, initial_sampler=None, pdf_sampler=None)[源代码]

Sampler that uses a proposal network to generate samples.

参数:

  • num_proposal_samples_per_ray – Number of samples to generate per ray for each proposal step.

  • num_nerf_samples_per_ray – Number of samples to generate per ray for the NERF model.

  • num_proposal_network_iterations – Number of proposal network iterations to run.

  • single_jitter – Use a same random jitter for all samples along a ray.

  • update_sched – A function that takes the iteration number of steps between updates.

  • initial_sampler – Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set.

  • pdf_sampler – PDFSampler to use after the first iteration. Uses PDFSampler if not set.

forward(ray_bundle, render_stratified_sampling: bool, training: bool = False)[源代码]

Defines the computation performed at every call.

Should be overridden by all subclasses.

备注

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

set_anneal(anneal: float) → None[源代码]

Set the anneal value for the proposal network.

step_cb(step)[源代码]

Callback to register a training step has passed. This is used to keep track of the sampling schedule

class landmark.nerf_components.ray_samplers.volumetric_sampler.VolumetricSampler(occupancy_grid, density_fn, near_far, step_size, alpha_thre=0.01, cone_angle=0.0)[源代码]

Sampler inspired by the one proposed in the Instant-NGP paper. Generates samples along a ray by sampling the occupancy field. Optionally removes occluded samples if the density_fn is provided.

Args: occupancy_grid: Occupancy grid to sample from. density_fn: Function that evaluates density at a given point. scene_aabb: Axis-aligned bounding box of the scene, should be set to None if the scene is unbounded.

forward(ray_bundle: Rays | Tensor, training: bool, render_stratified_sampling: bool)[源代码]

Generate ray samples in a bounding box.

参数:

  • ray_bundle – Rays to generate samples for

  • render_step_size – Minimum step size to use for rendering

  • near_plane – Near plane for raymarching

  • far_plane – Far plane for raymarching

  • alpha_thre – Opacity threshold skipping samples.

  • cone_angle – Cone angle for raymarching, set to 0 for uniform marching.

返回:

a tuple of (ray_samples, packed_info, ray_indices) The ray_samples are packed, only storing the valid samples. The ray_indices contains the indices of the rays that each sample belongs to.

get_sigma_fn(origins, directions, training, times=None) → Callable | None[源代码]

Returns a function that returns the density of a point.

参数:

  • origins – Origins of rays

  • directions – Directions of rays

  • times – Times at which rays are sampled

返回:

Function that returns the density of a point or None if a density function is not provided.