From ea2993d5b1b7f3c74d54b80b86a4b6f3288bcab7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ruilong=20Li=28=E6=9D=8E=E7=91=9E=E9=BE=99=29?= Date: Mon, 24 Apr 2023 17:11:30 -0700 Subject: [PATCH] Update NGP related models with nerfacc 0.5.2 (#1809) * profiling * set pose_optimizer to None if it's mode is "off" * radial_and_tangential_undistort in nerfacc * fix camera distortions * update to nerfacc 0.5.1. Upgrade instant-ngp model. * revert nerfacc to 0.5.1 * format * allows to fall back to the pytorch undistortion * format * fix linting * fix test * fix test * fix nerfplayer-ngp * enable_collider=False for NGP * nerfplayer verified. * zero cone_angle for ngp bounded * fix format * fix pylint --- nerfstudio/cameras/cameras.py | 2 +- nerfstudio/configs/method_configs.py | 26 +++++---- nerfstudio/fields/instant_ngp_field.py | 52 ++++++++--------- nerfstudio/fields/nerfacto_field.py | 2 +- nerfstudio/fields/nerfplayer_ngp_field.py | 34 ++++++----- nerfstudio/model_components/ray_samplers.py | 43 +++++++------- nerfstudio/model_components/renderers.py | 20 +++++-- nerfstudio/models/instant_ngp.py | 63 +++++++++++---------- nerfstudio/models/nerfplayer_ngp.py | 62 +++++++++++++------- nerfstudio/utils/math.py | 5 +- pyproject.toml | 2 +- tests/utils/test_aabb_intersection.py | 6 +- 12 files changed, 182 insertions(+), 135 deletions(-) diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index bda664642..0dec377c9 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -633,7 +633,7 @@ class Cameras(TensorDataclass): if distortion_params is not None: mask = (self.camera_type[true_indices] != CameraType.EQUIRECTANGULAR.value).squeeze(-1) # (num_rays) coord_mask = torch.stack([mask, mask, mask], dim=0) - if mask.any(): + if mask.any() and (distortion_params != 0).any(): coord_stack[coord_mask, :] = camera_utils.radial_and_tangential_undistort( coord_stack[coord_mask, :].reshape(3, -1, 2), distortion_params[mask, :], diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 8842f4e46..429607115 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -21,7 +21,6 @@ from __future__ import annotations from typing import Dict import tyro -from nerfacc import ContractionType from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig @@ -235,16 +234,20 @@ method_configs["instant-ngp"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=DynamicBatchPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192), + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + ), model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), ), optimizers={ "fields": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), } }, - viewer=ViewerConfig(num_rays_per_chunk=64000), + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), vis="viewer", ) @@ -259,9 +262,11 @@ method_configs["instant-ngp-bounded"] = TrainerConfig( datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192), model=InstantNGPModelConfig( eval_num_rays_per_chunk=8192, - contraction_type=ContractionType.AABB, + grid_levels=1, + alpha_thre=0.0, + cone_angle=0.0, render_step_size=0.001, - max_num_samples_per_ray=48, + disable_scene_contraction=True, near_plane=0.01, background_color="black", ), @@ -269,10 +274,10 @@ method_configs["instant-ngp-bounded"] = TrainerConfig( optimizers={ "fields": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), } }, - viewer=ViewerConfig(num_rays_per_chunk=64000), + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), vis="viewer", ) @@ -470,9 +475,10 @@ method_configs["nerfplayer-ngp"] = TrainerConfig( datamanager=DepthDataManagerConfig(dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=8192), model=NerfplayerNGPModelConfig( eval_num_rays_per_chunk=8192, - contraction_type=ContractionType.AABB, + grid_levels=1, + alpha_thre=0.0, render_step_size=0.001, - max_num_samples_per_ray=48, + disable_scene_contraction=True, near_plane=0.01, ), ), diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py index d1b3e2fa7..e205c73f0 100644 --- a/nerfstudio/fields/instant_ngp_field.py +++ b/nerfstudio/fields/instant_ngp_field.py @@ -21,7 +21,6 @@ from typing import Dict, Optional, Tuple import numpy as np import torch -from nerfacc import ContractionType, contract from torch.nn.parameter import Parameter from torchtyping import TensorType @@ -30,6 +29,10 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.field_components.activations import trunc_exp from nerfstudio.field_components.embedding import Embedding from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.spatial_distortions import ( + SceneContraction, + SpatialDistortion, +) from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn try: @@ -69,16 +72,16 @@ class TCNNInstantNGPField(Field): use_appearance_embedding: Optional[bool] = False, num_images: Optional[int] = None, appearance_embedding_dim: int = 32, - contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE, num_levels: int = 16, log2_hashmap_size: int = 19, max_res: int = 2048, + spatial_distortion: Optional[SpatialDistortion] = SceneContraction(), ) -> None: super().__init__() self.aabb = Parameter(aabb, requires_grad=False) self.geo_feat_dim = geo_feat_dim - self.contraction_type = contraction_type + self.spatial_distortion = spatial_distortion self.use_appearance_embedding = use_appearance_embedding if use_appearance_embedding: @@ -88,7 +91,8 @@ class TCNNInstantNGPField(Field): # TODO: set this properly based on the aabb base_res: int = 16 - per_level_scale = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) + features_per_level: int = 2 + growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) self.direction_encoding = tcnn.Encoding( n_input_dims=3, @@ -104,10 +108,10 @@ class TCNNInstantNGPField(Field): encoding_config={ "otype": "HashGrid", "n_levels": num_levels, - "n_features_per_level": 2, + "n_features_per_level": features_per_level, "log2_hashmap_size": log2_hashmap_size, "base_resolution": base_res, - "per_level_scale": per_level_scale, + "per_level_scale": growth_factor, }, network_config={ "otype": "FullyFusedMLP", @@ -134,9 +138,16 @@ class TCNNInstantNGPField(Field): ) def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: - positions = ray_samples.frustums.get_positions() + if self.spatial_distortion is not None: + positions = ray_samples.frustums.get_positions() + positions = self.spatial_distortion(positions) + positions = (positions + 2.0) / 4.0 + else: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + # Make sure the tcnn gets inputs between 0 and 1. + selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1) + positions = positions * selector[..., None] positions_flat = positions.view(-1, 3) - positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type) h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) @@ -145,20 +156,18 @@ class TCNNInstantNGPField(Field): # softplus, because it enables high post-activation (float32) density outputs # from smaller internal (float16) parameters. density = trunc_exp(density_before_activation.to(positions)) + density = density * selector[..., None] return density, base_mlp_out def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: + assert density_embedding is not None directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) - d = self.direction_encoding(directions_flat) - if density_embedding is None: - positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - h = torch.cat([d, positions.view(-1, 3)], dim=-1) - else: - h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + + h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) if self.use_appearance_embedding: if ray_samples.camera_indices is None: @@ -183,20 +192,5 @@ class TCNNInstantNGPField(Field): step_size: the step size to use for the opacity evaluation. """ density = self.density_fn(positions) - ## TODO: We should scale step size based on the distortion. Currently it uses too much memory. - # aabb_min, aabb_max = self.aabb[0], self.aabb[1] - # if self.contraction_type is not ContractionType.AABB: - # x = (positions - aabb_min) / (aabb_max - aabb_min) - # x = x * 2 - 1 # aabb is at [-1, 1] - # mag = x.norm(dim=-1, keepdim=True) - # mask = mag.squeeze(-1) > 1 - - # dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (1 / mag**3 - (2 * mag - 1) / mag**4) - # dev[~mask] = 1.0 - # dev = torch.clamp(dev, min=1e-6) - # step_size = step_size / dev.norm(dim=-1, keepdim=True) - # else: - # step_size = step_size * (aabb_max - aabb_min) - opacity = density * step_size return opacity diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index de0fd7c7d..acc868648 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -103,7 +103,7 @@ class TCNNNerfactoField(Field): pass_semantic_gradients: bool = False, use_pred_normals: bool = False, use_average_appearance_embedding: bool = False, - spatial_distortion: SpatialDistortion = None, + spatial_distortion: Optional[SpatialDistortion] = None, ) -> None: super().__init__() diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py index 2952a4369..bbf7b72e6 100644 --- a/nerfstudio/fields/nerfplayer_ngp_field.py +++ b/nerfstudio/fields/nerfplayer_ngp_field.py @@ -20,7 +20,6 @@ NeRFPlayer (https://arxiv.org/abs/2210.15947) field implementations with Instant from typing import Dict, Optional, Tuple import torch -from nerfacc import ContractionType, contract from torch.nn.parameter import Parameter from torchtyping import TensorType @@ -29,6 +28,10 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.field_components.activations import trunc_exp from nerfstudio.field_components.embedding import Embedding from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.spatial_distortions import ( + SceneContraction, + SpatialDistortion, +) from nerfstudio.field_components.temporal_grid import TemporalGridEncoder from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn @@ -82,13 +85,13 @@ class NerfplayerNGPField(Field): disable_viewing_dependent: bool = False, num_images: Optional[int] = None, appearance_embedding_dim: int = 32, - contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE, + spatial_distortion: Optional[SpatialDistortion] = SceneContraction(), ) -> None: super().__init__() self.aabb = Parameter(aabb, requires_grad=False) self.geo_feat_dim = geo_feat_dim - self.contraction_type = contraction_type + self.spatial_distortion = spatial_distortion self.use_appearance_embedding = use_appearance_embedding if use_appearance_embedding: @@ -144,9 +147,16 @@ class NerfplayerNGPField(Field): ) def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: - positions = ray_samples.frustums.get_positions() + if self.spatial_distortion is not None: + positions = ray_samples.frustums.get_positions() + positions = self.spatial_distortion(positions) + positions = (positions + 2.0) / 4.0 + else: + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + # Make sure the tcnn gets inputs between 0 and 1. + selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1) + positions = positions * selector[..., None] positions_flat = positions.view(-1, 3) - positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type) assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer" times_flat = ray_samples.times.view(-1, 1) @@ -158,28 +168,22 @@ class NerfplayerNGPField(Field): # softplus, because it enables high post-activation (float32) density outputs # from smaller internal (float16) parameters. density = trunc_exp(density_before_activation.to(positions)) + density = density * selector[..., None] return density, base_mlp_out def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: + assert density_embedding is not None directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) if self.direction_encoding is not None: d = self.direction_encoding(directions_flat) - if density_embedding is None: - positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - h = torch.cat([d, positions.view(-1, 3)], dim=-1) - else: - h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) else: # viewing direction is disabled - if density_embedding is None: - positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) - h = positions.view(-1, 3) - else: - h = density_embedding.view(-1, self.geo_feat_dim) + h = density_embedding.view(-1, self.geo_feat_dim) if self.use_appearance_embedding: if ray_samples.camera_indices is None: diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index f5a96255b..7b7219f15 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -19,9 +19,8 @@ Collection of sampling strategies from abc import abstractmethod from typing import Callable, List, Optional, Tuple, Union -import nerfacc import torch -from nerfacc import OccupancyGrid +from nerfacc import OccGridEstimator from torch import nn from torchtyping import TensorType @@ -382,24 +381,22 @@ class VolumetricSampler(Sampler): def __init__( self, - occupancy_grid: Optional[OccupancyGrid] = None, + occupancy_grid: OccGridEstimator, density_fn: Optional[Callable[[TensorType[..., 3]], TensorType[..., 1]]] = None, - scene_aabb: Optional[TensorType[2, 3]] = None, ) -> None: super().__init__() - self.scene_aabb = scene_aabb + assert occupancy_grid is not None self.density_fn = density_fn self.occupancy_grid = occupancy_grid - if self.scene_aabb is not None: - self.scene_aabb = self.scene_aabb.to("cuda").flatten() - def get_sigma_fn(self, origins, directions) -> Optional[Callable]: + def get_sigma_fn(self, origins, directions, times=None) -> Optional[Callable]: """Returns a function that returns the density of a point. Args: origins: Origins of rays directions: Directions of rays + times: Times at which rays are sampled Returns: Function that returns the density of a point or None if a density function is not provided. """ @@ -412,8 +409,10 @@ class VolumetricSampler(Sampler): def sigma_fn(t_starts, t_ends, ray_indices): t_origins = origins[ray_indices] t_dirs = directions[ray_indices] - positions = t_origins + t_dirs * (t_starts + t_ends) / 2.0 - return density_fn(positions) + positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0 + if times is None: + return density_fn(positions).squeeze(-1) + return density_fn(positions, times[ray_indices]).squeeze(-1) return sigma_fn @@ -429,6 +428,7 @@ class VolumetricSampler(Sampler): render_step_size: float, near_plane: float = 0.0, far_plane: Optional[float] = None, + alpha_thre: float = 0.01, cone_angle: float = 0.0, ) -> Tuple[RaySamples, TensorType["total_samples",]]: """Generate ray samples in a bounding box. @@ -438,6 +438,7 @@ class VolumetricSampler(Sampler): 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. Returns: @@ -448,6 +449,7 @@ class VolumetricSampler(Sampler): rays_o = ray_bundle.origins.contiguous() rays_d = ray_bundle.directions.contiguous() + times = ray_bundle.times if ray_bundle.nears is not None and ray_bundle.fars is not None: t_min = ray_bundle.nears.contiguous().reshape(-1) @@ -457,34 +459,33 @@ class VolumetricSampler(Sampler): t_min = None t_max = None + if far_plane is None: + far_plane = 1e10 + if ray_bundle.camera_indices is not None: camera_indices = ray_bundle.camera_indices.contiguous() else: camera_indices = None - - ray_indices, starts, ends = nerfacc.ray_marching( + ray_indices, starts, ends = self.occupancy_grid.sampling( rays_o=rays_o, rays_d=rays_d, t_min=t_min, t_max=t_max, - scene_aabb=self.scene_aabb, - grid=self.occupancy_grid, - # this is a workaround - using density causes crash and damage quality. should be fixed - sigma_fn=None, # self.get_sigma_fn(rays_o, rays_d), + sigma_fn=self.get_sigma_fn(rays_o, rays_d, times), render_step_size=render_step_size, near_plane=near_plane, far_plane=far_plane, stratified=self.training, cone_angle=cone_angle, - alpha_thre=1e-2, + alpha_thre=alpha_thre, ) num_samples = starts.shape[0] if num_samples == 0: # create a single fake sample and update packed_info accordingly # this says the last ray in packed_info has 1 sample, which starts and ends at 1 ray_indices = torch.zeros((1,), dtype=torch.long, device=rays_o.device) - starts = torch.ones((1, 1), dtype=starts.dtype, device=rays_o.device) - ends = torch.ones((1, 1), dtype=ends.dtype, device=rays_o.device) + starts = torch.ones((1,), dtype=starts.dtype, device=rays_o.device) + ends = torch.ones((1,), dtype=ends.dtype, device=rays_o.device) origins = rays_o[ray_indices] dirs = rays_d[ray_indices] @@ -496,8 +497,8 @@ class VolumetricSampler(Sampler): frustums=Frustums( origins=origins, directions=dirs, - starts=starts, - ends=ends, + starts=starts[..., None], + ends=ends[..., None], pixel_area=zeros, ), camera_indices=camera_indices, diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 43f81404c..5bc4be9b1 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -91,8 +91,12 @@ class RGBRenderer(nn.Module): # Necessary for packed samples from volumetric ray sampler if background_color == "last_sample": raise NotImplementedError("Background color 'last_sample' not implemented for packed samples.") - comp_rgb = nerfacc.accumulate_along_rays(weights, ray_indices, rgb, num_rays) - accumulated_weight = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays) + comp_rgb = nerfacc.accumulate_along_rays( + weights[..., 0], values=rgb, ray_indices=ray_indices, n_rays=num_rays + ) + accumulated_weight = nerfacc.accumulate_along_rays( + weights[..., 0], values=None, ray_indices=ray_indices, n_rays=num_rays + ) else: comp_rgb = torch.sum(weights * rgb, dim=-2) accumulated_weight = torch.sum(weights, dim=-2) @@ -217,7 +221,9 @@ class AccumulationRenderer(nn.Module): if ray_indices is not None and num_rays is not None: # Necessary for packed samples from volumetric ray sampler - accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays) + accumulation = nerfacc.accumulate_along_rays( + weights[..., 0], values=None, ray_indices=ray_indices, n_rays=num_rays + ) else: accumulation = torch.sum(weights, dim=-2) return accumulation @@ -274,8 +280,12 @@ class DepthRenderer(nn.Module): if ray_indices is not None and num_rays is not None: # Necessary for packed samples from volumetric ray sampler - depth = nerfacc.accumulate_along_rays(weights, ray_indices, steps, num_rays) - accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays) + depth = nerfacc.accumulate_along_rays( + weights[..., 0], values=steps, ray_indices=ray_indices, n_rays=num_rays + ) + accumulation = nerfacc.accumulate_along_rays( + weights[..., 0], values=None, ray_indices=ray_indices, n_rays=num_rays + ) depth = depth / (accumulation + eps) else: depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps) diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 2bdf85b6a..d8119db5c 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -23,7 +23,6 @@ from typing import Dict, List, Optional, Tuple, Type import nerfacc import torch -from nerfacc import ContractionType from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure @@ -37,7 +36,8 @@ from nerfstudio.engine.callbacks import ( TrainingCallbackLocation, ) from nerfstudio.field_components.field_heads import FieldHeadNames -from nerfstudio.fields.instant_ngp_field import TCNNInstantNGPField +from nerfstudio.field_components.spatial_distortions import SceneContraction +from nerfstudio.fields.nerfacto_field import TCNNNerfactoField from nerfstudio.model_components.losses import MSELoss from nerfstudio.model_components.ray_samplers import VolumetricSampler from nerfstudio.model_components.renderers import ( @@ -61,19 +61,19 @@ class InstantNGPModelConfig(ModelConfig): """Whether to create a scene collider to filter rays.""" collider_params: Optional[Dict[str, float]] = None """Instant NGP doesn't use a collider.""" - max_num_samples_per_ray: int = 24 - """Number of samples in field evaluation.""" grid_resolution: int = 128 """Resolution of the grid used for the field.""" + grid_levels: int = 4 + """Levels of the grid used for the field.""" max_res: int = 2048 """Maximum resolution of the hashmap for the base mlp.""" log2_hashmap_size: int = 19 """Size of the hashmap for the base mlp""" - contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE - """Contraction type used for spatial deformation of the field.""" + alpha_thre: float = 0.01 + """Threshold for opacity skipping.""" cone_angle: float = 0.004 """Should be set to 0.0 for blender scenes but 1./256 for real scenes.""" - render_step_size: float = 0.01 + render_step_size: Optional[float] = None """Minimum step size for rendering.""" near_plane: float = 0.05 """How far along ray to start sampling.""" @@ -83,6 +83,8 @@ class InstantNGPModelConfig(ModelConfig): """Whether to use an appearance embedding.""" background_color: Literal["random", "black", "white"] = "random" """The color that is given to untrained areas.""" + disable_scene_contraction: bool = False + """Whether to disable scene contraction or not.""" class NGPModel(Model): @@ -93,7 +95,7 @@ class NGPModel(Model): """ config: InstantNGPModelConfig - field: TCNNInstantNGPField + field: TCNNNerfactoField def __init__(self, config: InstantNGPModelConfig, **kwargs) -> None: super().__init__(config=config, **kwargs) @@ -102,28 +104,33 @@ class NGPModel(Model): """Set the fields and modules.""" super().populate_modules() - self.field = TCNNInstantNGPField( + if self.config.disable_scene_contraction: + scene_contraction = None + else: + scene_contraction = SceneContraction(order=float("inf")) + + self.field = TCNNNerfactoField( aabb=self.scene_box.aabb, - contraction_type=self.config.contraction_type, - use_appearance_embedding=self.config.use_appearance_embedding, num_images=self.num_train_data, log2_hashmap_size=self.config.log2_hashmap_size, max_res=self.config.max_res, + spatial_distortion=scene_contraction, ) self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False) - # Occupancy Grid - self.occupancy_grid = nerfacc.OccupancyGrid( + if self.config.render_step_size is None: + # auto step size: ~1000 samples in the base level grid + self.config.render_step_size = ((self.scene_aabb[3:] - self.scene_aabb[:3]) ** 2).sum().sqrt().item() / 1000 + # Occupancy Grid. + self.occupancy_grid = nerfacc.OccGridEstimator( roi_aabb=self.scene_aabb, resolution=self.config.grid_resolution, - contraction_type=self.config.contraction_type, + levels=self.config.grid_levels, ) # Sampler - vol_sampler_aabb = self.scene_box.aabb if self.config.contraction_type == ContractionType.AABB else None self.sampler = VolumetricSampler( - scene_aabb=vol_sampler_aabb, occupancy_grid=self.occupancy_grid, density_fn=self.field.density_fn, ) @@ -145,11 +152,9 @@ class NGPModel(Model): self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: def update_occupancy_grid(step: int): - # TODO: needs to get access to the sampler, on how the step size is determinated at each x. See - # https://github.com/KAIR-BAIR/nerfacc/blob/127223b11401125a9fce5ce269bb0546ee4de6e8/examples/train_ngp_nerf.py#L190-L213 - self.occupancy_grid.every_n_step( + self.occupancy_grid.update_every_n_steps( step=step, - occ_eval_fn=lambda x: self.field.get_opacity(x, self.config.render_step_size), + occ_eval_fn=lambda x: self.field.density_fn(x) * self.config.render_step_size, ) return [ @@ -177,6 +182,7 @@ class NGPModel(Model): near_plane=self.config.near_plane, far_plane=self.config.far_plane, render_step_size=self.config.render_step_size, + alpha_thre=self.config.alpha_thre, cone_angle=self.config.cone_angle, ) @@ -185,11 +191,12 @@ class NGPModel(Model): # accumulation packed_info = nerfacc.pack_info(ray_indices, num_rays) weights = nerfacc.render_weight_from_density( + t_starts=ray_samples.frustums.starts[..., 0], + t_ends=ray_samples.frustums.ends[..., 0], + sigmas=field_outputs[FieldHeadNames.DENSITY][..., 0], packed_info=packed_info, - sigmas=field_outputs[FieldHeadNames.DENSITY], - t_starts=ray_samples.frustums.starts, - t_ends=ray_samples.frustums.ends, - ) + )[0] + weights = weights[..., None] rgb = self.renderer_rgb( rgb=field_outputs[FieldHeadNames.RGB], @@ -201,13 +208,11 @@ class NGPModel(Model): weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays ) accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays) - alive_ray_mask = accumulation.squeeze(-1) > 0 outputs = { "rgb": rgb, "accumulation": accumulation, "depth": depth, - "alive_ray_mask": alive_ray_mask, # the rays we kept from sampler "num_samples_per_ray": packed_info[:, 1], } return outputs @@ -221,8 +226,7 @@ class NGPModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) - mask = outputs["alive_ray_mask"] - rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) + rgb_loss = self.rgb_loss(image, outputs["rgb"]) loss_dict = {"rgb_loss": rgb_loss} return loss_dict @@ -237,12 +241,10 @@ class NGPModel(Model): outputs["depth"], accumulation=outputs["accumulation"], ) - alive_ray_mask = colormaps.apply_colormap(outputs["alive_ray_mask"]) combined_rgb = torch.cat([image, rgb], dim=1) combined_acc = torch.cat([acc], dim=1) combined_depth = torch.cat([depth], dim=1) - combined_alive_ray_mask = torch.cat([alive_ray_mask], dim=1) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] @@ -260,7 +262,6 @@ class NGPModel(Model): "img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth, - "alive_ray_mask": combined_alive_ray_mask, } return metrics_dict, images_dict diff --git a/nerfstudio/models/nerfplayer_ngp.py b/nerfstudio/models/nerfplayer_ngp.py index 09f68d67f..4f884674f 100644 --- a/nerfstudio/models/nerfplayer_ngp.py +++ b/nerfstudio/models/nerfplayer_ngp.py @@ -19,11 +19,10 @@ Implementation of NeRFPlayer (https://arxiv.org/abs/2210.15947) with InstantNGP from __future__ import annotations from dataclasses import dataclass, field -from typing import Type +from typing import List, Type import nerfacc import torch -from nerfacc import ContractionType from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure @@ -31,7 +30,13 @@ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from typing_extensions import Literal from nerfstudio.cameras.rays import RayBundle +from nerfstudio.engine.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.fields.nerfplayer_ngp_field import NerfplayerNGPField from nerfstudio.model_components.losses import MSELoss from nerfstudio.model_components.ray_samplers import VolumetricSampler @@ -80,6 +85,8 @@ class NerfplayerNGPModelConfig(InstantNGPModelConfig): """The training background color that is given to untrained areas.""" disable_viewing_dependent: bool = True """Disable viewing dependent effects.""" + disable_scene_contraction: bool = False + """Whether to disable scene contraction or not.""" class NerfplayerNGPModel(NGPModel): @@ -96,9 +103,13 @@ class NerfplayerNGPModel(NGPModel): """Set the fields and modules.""" Model.populate_modules(self) + if self.config.disable_scene_contraction: + scene_contraction = None + else: + scene_contraction = SceneContraction(order=float("inf")) + self.field = NerfplayerNGPField( aabb=self.scene_box.aabb, - contraction_type=self.config.contraction_type, use_appearance_embedding=self.config.use_appearance_embedding, num_images=self.num_train_data, temporal_dim=self.config.temporal_dim, @@ -107,21 +118,20 @@ class NerfplayerNGPModel(NGPModel): log2_hashmap_size=self.config.log2_hashmap_size, base_resolution=self.config.base_resolution, disable_viewing_dependent=self.config.disable_viewing_dependent, + spatial_distortion=scene_contraction, ) self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False) # Occupancy Grid - self.occupancy_grid = nerfacc.OccupancyGrid( + self.occupancy_grid = nerfacc.OccGridEstimator( roi_aabb=self.scene_aabb, resolution=self.config.grid_resolution, - contraction_type=self.config.contraction_type, + levels=self.config.grid_levels, ) # Sampler - vol_sampler_aabb = self.scene_box.aabb if self.config.contraction_type == ContractionType.AABB else None self.sampler = VolumetricSampler( - scene_aabb=vol_sampler_aabb, occupancy_grid=self.occupancy_grid, density_fn=self.field.density_fn, ) # need to update the density_fn later during forward (for input time) @@ -140,11 +150,26 @@ class NerfplayerNGPModel(NGPModel): self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) self.temporal_distortion = True # for viewer + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + def update_occupancy_grid(step: int): + self.occupancy_grid.update_every_n_steps( + step=step, + occ_eval_fn=lambda x: self.field.get_opacity(x, self.config.render_step_size), + ) + + return [ + TrainingCallback( + where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], + update_every_num_iters=1, + func=update_occupancy_grid, + ), + ] + def get_outputs(self, ray_bundle: RayBundle): num_rays = len(ray_bundle) - # update the density_fn of the sampler so that the density is time aware - self.sampler.density_fn = lambda x: self.field.density_fn(x, ray_bundle.times) with torch.no_grad(): ray_samples, ray_indices = self.sampler( ray_bundle=ray_bundle, @@ -152,6 +177,7 @@ class NerfplayerNGPModel(NGPModel): far_plane=self.config.far_plane, render_step_size=self.config.render_step_size, cone_angle=self.config.cone_angle, + alpha_thre=self.config.alpha_thre, ) field_outputs = self.field(ray_samples) @@ -160,10 +186,11 @@ class NerfplayerNGPModel(NGPModel): packed_info = nerfacc.pack_info(ray_indices, num_rays) weights = nerfacc.render_weight_from_density( packed_info=packed_info, - sigmas=field_outputs[FieldHeadNames.DENSITY], - t_starts=ray_samples.frustums.starts, - t_ends=ray_samples.frustums.ends, - ) + sigmas=field_outputs[FieldHeadNames.DENSITY][..., 0], + t_starts=ray_samples.frustums.starts[..., 0], + t_ends=ray_samples.frustums.ends[..., 0], + )[0] + weights = weights[..., None] # update bgcolor in the renderer; usually random color for training and fixed color for inference if self.training: @@ -180,13 +207,11 @@ class NerfplayerNGPModel(NGPModel): weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays ) accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays) - alive_ray_mask = accumulation.squeeze(-1) > 0 outputs = { "rgb": rgb, "accumulation": accumulation, "depth": depth, - "alive_ray_mask": alive_ray_mask, # the rays we kept from sampler "num_samples_per_ray": packed_info[:, 1], } # adding them to outputs for calculating losses @@ -199,10 +224,9 @@ class NerfplayerNGPModel(NGPModel): def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) - mask = outputs["alive_ray_mask"] - rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) + rgb_loss = self.rgb_loss(image, outputs["rgb"]) loss_dict = {"rgb_loss": rgb_loss} - if "depth_image" in batch.keys() and self.config.depth_weight > 0: + if "depth_image" in batch.keys() and self.config.depth_weight > 0 and self.training: mask = batch["depth_image"] != 0 # First we calculate the depth value, just like most of the papers. loss_dict["depth_loss"] = (outputs["depth"][mask] - batch["depth_image"][mask]).abs().mean() @@ -228,6 +252,6 @@ class NerfplayerNGPModel(NGPModel): density_min_mask = (gt_depth_packed - steps > margin) & (gt_depth_packed != 0) loss_dict["depth_loss"] += (outputs["sigmas"][density_min_mask[..., 0]].pow(2)).mean() * 1e-2 loss_dict["depth_loss"] *= self.config.depth_weight - if self.config.temporal_tv_weight > 0: + if self.config.temporal_tv_weight > 0 and self.training: loss_dict["temporal_tv_loss"] = self.config.temporal_tv_weight * self.field.mlp_base.get_temporal_tv_loss() return loss_dict diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index f54b6509a..89a4d3250 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -262,7 +262,10 @@ def intersect_aabb( try: import nerfacc # pylint: disable=import-outside-toplevel - t_min, t_max = nerfacc.ray_aabb_intersect(origins, directions, aabb) + t_min, t_max, _ = nerfacc.ray_aabb_intersect( + origins, directions, aabb[None, :], near_plane=0, far_plane=1e10, miss_value=1e10 + ) + t_min, t_max = t_min.squeeze(-1), t_max.squeeze(-1) except: # pylint: disable=bare-except t_min, t_max = _intersect_aabb(origins, directions, aabb, max_bound=1e10, invalid_value=1e10) _USE_NERFACC = False diff --git a/pyproject.toml b/pyproject.toml index ffcbcf19e..12b6288ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ "mediapy>=1.1.0", "msgpack>=1.0.4", "msgpack_numpy>=0.4.8", - "nerfacc==0.3.5", + "nerfacc==0.5.2", "open3d>=0.16.0", "opencv-python==4.6.0.66", "Pillow>=9.3.0", diff --git a/tests/utils/test_aabb_intersection.py b/tests/utils/test_aabb_intersection.py index 1d6ba34b9..ab8e2ce0e 100644 --- a/tests/utils/test_aabb_intersection.py +++ b/tests/utils/test_aabb_intersection.py @@ -168,7 +168,11 @@ def test_equall_nerfacc(): # time2 = time.time() # time3 = time.time() - t_min_nerfacc, t_max_nerfacc = nerfacc.ray_aabb_intersect(origins, directions, aabb) + t_min_nerfacc, t_max_nerfacc, _ = nerfacc.ray_aabb_intersect( + origins, directions, aabb[None, :], near_plane=0, far_plane=1e10, miss_value=1e10 + ) + t_max_nerfacc = t_max_nerfacc.squeeze(-1) + t_min_nerfacc = t_min_nerfacc.squeeze(-1) # time4 = time.time() # print("pytorch ", time2-time1)