Models

Base

Base Model implementation which takes in RayBundles or Cameras

class nerfstudio.models.base_model.Model(config: ModelConfig, scene_box: SceneBox, num_train_data: int, **kwargs)

Bases: Module

Model class Where everything (Fields, Optimizers, Samplers, Visualization, etc) is linked together. This should be subclassed for custom NeRF model.

Parameters:

• config – configuration for instantiating model

• scene_box – dataset scene box

property device

Returns the device that the model is on.

forward(ray_bundle: Union[RayBundle, Cameras]) → Dict[str, Union[Tensor, List]]

Run forward starting with a ray bundle. This outputs different things depending on the configuration of the model and whether or not the batch is provided (whether or not we are training basically)

Parameters:

ray_bundle – containing all the information needed to render that ray latents included

abstract get_image_metrics_and_images(outputs: Dict[str, Tensor], batch: Dict[str, Tensor]) → Tuple[Dict[str, float], Dict[str, Tensor]]

Writes the test image outputs. TODO: This shouldn’t return a loss

Parameters:

• image_idx – Index of the image.

• step – Current step.

• batch – Batch of data.

• outputs – Outputs of the model.

Returns:

A dictionary of metrics.

abstract get_loss_dict(outputs, batch, metrics_dict=None) → Dict[str, Tensor]

Computes and returns the losses dict.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

• metrics_dict – dictionary of metrics, some of which we can use for loss

get_metrics_dict(outputs, batch) → Dict[str, Tensor]

Compute and returns metrics.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

abstract get_outputs(ray_bundle: Union[RayBundle, Cameras]) → Dict[str, Union[Tensor, List]]

Takes in a Ray Bundle and returns a dictionary of outputs.

Parameters:

• ray_bundle – Input bundle of rays. This raybundle should have all the

• outputs. (needed information to compute the) –

Returns:

Outputs of model. (ie. rendered colors)

get_outputs_for_camera(camera: Cameras, obb_box: Optional[OrientedBox] = None) → Dict[str, Tensor]

Takes in a camera, generates the raybundle, and computes the output of the model. Assumes a ray-based model.

Parameters:

camera – generates raybundle

get_outputs_for_camera_ray_bundle(camera_ray_bundle: RayBundle) → Dict[str, Tensor]

Takes in camera parameters and computes the output of the model.

Parameters:

camera_ray_bundle – ray bundle to calculate outputs over

abstract get_param_groups() → Dict[str, List[Parameter]]

Obtain the parameter groups for the optimizers

Returns:

Mapping of different parameter groups

get_rgba_image(outputs: Dict[str, Tensor], output_name: str = 'rgb') → Tensor

Returns the RGBA image from the outputs of the model.

Parameters:

outputs – Outputs of the model.

Returns:

RGBA image.

get_training_callbacks(training_callback_attributes: TrainingCallbackAttributes) → List[TrainingCallback]

Returns a list of callbacks that run functions at the specified training iterations.

load_model(loaded_state: Dict[str, Any]) → None

Load the checkpoint from the given path

Parameters:

loaded_state – dictionary of pre-trained model states

populate_modules()

Set the necessary modules to get the network working.

update_to_step(step: int) → None

Called when loading a model from a checkpoint. Sets any model parameters that change over training to the correct value, based on the training step of the checkpoint.

Parameters:

step – training step of the loaded checkpoint

class nerfstudio.models.base_model.ModelConfig(_target: ~typing.Type = <factory>, enable_collider: bool = True, collider_params: ~typing.Optional[~typing.Dict[str, float]] = <factory>, loss_coefficients: ~typing.Dict[str, float] = <factory>, eval_num_rays_per_chunk: int = 4096, prompt: ~typing.Optional[str] = None)

Bases: InstantiateConfig

Configuration for model instantiation

collider_params: Optional[Dict[str, float]]

parameters to instantiate scene collider with

enable_collider: bool = True

Whether to create a scene collider to filter rays.

eval_num_rays_per_chunk: int = 4096

specifies number of rays per chunk during eval

loss_coefficients: Dict[str, float]

parameters to instantiate density field with

prompt: Optional[str] = None

A prompt to be used in text to NeRF models

Instant NGP

Implementation of Instant NGP.

class nerfstudio.models.instant_ngp.InstantNGPModelConfig(_target: ~typing.Type = <factory>, enable_collider: bool = False, collider_params: ~typing.Optional[~typing.Dict[str, float]] = None, loss_coefficients: ~typing.Dict[str, float] = <factory>, eval_num_rays_per_chunk: int = 4096, prompt: ~typing.Optional[str] = None, grid_resolution: ~typing.Union[int, ~typing.List[int]] = 128, grid_levels: int = 4, max_res: int = 2048, log2_hashmap_size: int = 19, alpha_thre: float = 0.01, cone_angle: float = 0.004, render_step_size: ~typing.Optional[float] = None, near_plane: float = 0.05, far_plane: float = 1000.0, use_gradient_scaling: bool = False, use_appearance_embedding: bool = False, background_color: ~typing.Literal['random', 'black', 'white'] = 'random', disable_scene_contraction: bool = False)

Bases: ModelConfig

Instant NGP Model Config

alpha_thre: float = 0.01

Threshold for opacity skipping.

background_color: Literal['random', 'black', 'white'] = 'random'

The color that is given to masked areas. These areas are used to force the density in those regions to be zero.

collider_params: Optional[Dict[str, float]] = None

Instant NGP doesn’t use a collider.

cone_angle: float = 0.004

Should be set to 0.0 for blender scenes but 1./256 for real scenes.

disable_scene_contraction: bool = False

Whether to disable scene contraction or not.

enable_collider: bool = False

Whether to create a scene collider to filter rays.

far_plane: float = 1000.0

How far along ray to stop sampling.

grid_levels: int = 4

Levels of the grid used for the field.

grid_resolution: Union[int, List[int]] = 128

Resolution of the grid used for the field.

log2_hashmap_size: int = 19

Size of the hashmap for the base mlp

max_res: int = 2048

Maximum resolution of the hashmap for the base mlp.

near_plane: float = 0.05

How far along ray to start sampling.

render_step_size: Optional[float] = None

Minimum step size for rendering.

use_appearance_embedding: bool = False

Whether to use an appearance embedding.

use_gradient_scaling: bool = False

Use gradient scaler where the gradients are lower for points closer to the camera.

class nerfstudio.models.instant_ngp.NGPModel(config: InstantNGPModelConfig, **kwargs)

Bases: Model

Instant NGP model

Parameters:

config – instant NGP configuration to instantiate model

get_image_metrics_and_images(outputs: Dict[str, Tensor], batch: Dict[str, Tensor]) → Tuple[Dict[str, float], Dict[str, Tensor]]

Writes the test image outputs. TODO: This shouldn’t return a loss

Parameters:

• image_idx – Index of the image.

• step – Current step.

• batch – Batch of data.

• outputs – Outputs of the model.

Returns:

A dictionary of metrics.

get_loss_dict(outputs, batch, metrics_dict=None)

Computes and returns the losses dict.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

• metrics_dict – dictionary of metrics, some of which we can use for loss

get_metrics_dict(outputs, batch)

Compute and returns metrics.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

get_outputs(ray_bundle: RayBundle)

Takes in a Ray Bundle and returns a dictionary of outputs.

Parameters:

• ray_bundle – Input bundle of rays. This raybundle should have all the

• outputs. (needed information to compute the) –

Returns:

Outputs of model. (ie. rendered colors)

get_param_groups() → Dict[str, List[Parameter]]

Obtain the parameter groups for the optimizers

Returns:

Mapping of different parameter groups

get_training_callbacks(training_callback_attributes: TrainingCallbackAttributes) → List[TrainingCallback]

Returns a list of callbacks that run functions at the specified training iterations.

populate_modules()

Set the fields and modules.

Semantic NeRF-W

Semantic NeRF-W implementation which should be fast enough to view in the viewer.

class nerfstudio.models.semantic_nerfw.SemanticNerfWModel(config: SemanticNerfWModelConfig, metadata: Dict, **kwargs)

Bases: Model

Nerfacto model

Parameters:

config – Nerfacto configuration to instantiate model

get_image_metrics_and_images(outputs: Dict[str, Tensor], batch: Dict[str, Tensor]) → Tuple[Dict[str, float], Dict[str, Tensor]]

Writes the test image outputs. TODO: This shouldn’t return a loss

Parameters:

• image_idx – Index of the image.

• step – Current step.

• batch – Batch of data.

• outputs – Outputs of the model.

Returns:

A dictionary of metrics.

get_loss_dict(outputs, batch, metrics_dict=None)

Computes and returns the losses dict.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

• metrics_dict – dictionary of metrics, some of which we can use for loss

get_metrics_dict(outputs, batch)

Compute and returns metrics.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

get_outputs(ray_bundle: RayBundle)

Takes in a Ray Bundle and returns a dictionary of outputs.

Parameters:

• ray_bundle – Input bundle of rays. This raybundle should have all the

• outputs. (needed information to compute the) –

Returns:

Outputs of model. (ie. rendered colors)

get_param_groups() → Dict[str, List[Parameter]]

Obtain the parameter groups for the optimizers

Returns:

Mapping of different parameter groups

get_training_callbacks(training_callback_attributes: TrainingCallbackAttributes) → List[TrainingCallback]

Returns a list of callbacks that run functions at the specified training iterations.

populate_modules()

Set the fields and modules.

class nerfstudio.models.semantic_nerfw.SemanticNerfWModelConfig(_target: Type = <factory>, enable_collider: bool = True, collider_params: Optional[Dict[str, float]] = <factory>, loss_coefficients: Dict[str, float] = <factory>, eval_num_rays_per_chunk: int = 4096, prompt: Optional[str] = None, near_plane: float = 0.05, far_plane: float = 1000.0, background_color: Literal['random', 'last_sample', 'black', 'white'] = 'last_sample', hidden_dim: int = 64, hidden_dim_color: int = 64, hidden_dim_transient: int = 64, num_levels: int = 16, base_res: int = 16, max_res: int = 2048, log2_hashmap_size: int = 19, features_per_level: int = 2, num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96), num_nerf_samples_per_ray: int = 48, proposal_update_every: int = 5, proposal_warmup: int = 5000, num_proposal_iterations: int = 2, use_same_proposal_network: bool = False, proposal_net_args_list: List[Dict] = <factory>, proposal_initial_sampler: Literal['piecewise', 'uniform'] = 'piecewise', interlevel_loss_mult: float = 1.0, distortion_loss_mult: float = 0.002, orientation_loss_mult: float = 0.0001, pred_normal_loss_mult: float = 0.001, use_proposal_weight_anneal: bool = True, use_appearance_embedding: bool = True, use_average_appearance_embedding: bool = True, proposal_weights_anneal_slope: float = 10.0, proposal_weights_anneal_max_num_iters: int = 1000, use_single_jitter: bool = True, predict_normals: bool = False, disable_scene_contraction: bool = False, use_gradient_scaling: bool = False, implementation: Literal['tcnn', 'torch'] = 'tcnn', appearance_embed_dim: int = 32, average_init_density: float = 1.0, camera_optimizer: CameraOptimizerConfig = <factory>, use_transient_embedding: bool = False, semantic_loss_weight: float = 1.0, pass_semantic_gradients: bool = False)

Bases: NerfactoModelConfig

Nerfacto Model Config

use_transient_embedding: bool = False

Whether to use transient embedding.

NeRF

Implementation of vanilla nerf.

class nerfstudio.models.vanilla_nerf.NeRFModel(config: VanillaModelConfig, **kwargs)

Bases: Model

Vanilla NeRF model

Parameters:

config – Basic NeRF configuration to instantiate model

get_image_metrics_and_images(outputs: Dict[str, Tensor], batch: Dict[str, Tensor]) → Tuple[Dict[str, float], Dict[str, Tensor]]

Writes the test image outputs. TODO: This shouldn’t return a loss

Parameters:

• image_idx – Index of the image.

• step – Current step.

• batch – Batch of data.

• outputs – Outputs of the model.

Returns:

A dictionary of metrics.

get_loss_dict(outputs, batch, metrics_dict=None) → Dict[str, Tensor]

Computes and returns the losses dict.

Parameters:

• outputs – the output to compute loss dict to

• batch – ground truth batch corresponding to outputs

• metrics_dict – dictionary of metrics, some of which we can use for loss

get_outputs(ray_bundle: RayBundle)

Takes in a Ray Bundle and returns a dictionary of outputs.

Parameters:

• ray_bundle – Input bundle of rays. This raybundle should have all the

• outputs. (needed information to compute the) –

Returns:

Outputs of model. (ie. rendered colors)

get_param_groups() → Dict[str, List[Parameter]]

Obtain the parameter groups for the optimizers

Returns:

Mapping of different parameter groups

populate_modules()

Set the fields and modules

class nerfstudio.models.vanilla_nerf.VanillaModelConfig(_target: Type = <factory>, enable_collider: bool = True, collider_params: Optional[Dict[str, float]] = <factory>, loss_coefficients: Dict[str, float] = <factory>, eval_num_rays_per_chunk: int = 4096, prompt: Optional[str] = None, num_coarse_samples: int = 64, num_importance_samples: int = 128, enable_temporal_distortion: bool = False, temporal_distortion_params: Dict[str, Any] = <factory>, use_gradient_scaling: bool = False, background_color: Literal['random', 'last_sample', 'black', 'white'] = 'white')

Bases: ModelConfig

Vanilla Model Config

background_color: Literal['random', 'last_sample', 'black', 'white'] = 'white'

Whether to randomize the background color.

enable_temporal_distortion: bool = False

Specifies whether or not to include ray warping based on time.

num_coarse_samples: int = 64

Number of samples in coarse field evaluation

num_importance_samples: int = 128

Number of samples in fine field evaluation

temporal_distortion_params: Dict[str, Any]

Parameters to instantiate temporal distortion with

use_gradient_scaling: bool = False

Use gradient scaler where the gradients are lower for points closer to the camera.