tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
Code

Update NGP related models with nerfacc 0.5.2 (#1809)

Browse Source 
* 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
This commit is contained in:
Ruilong Li(李瑞龙) 2023-04-24 17:11:30 -07:00 committed by GitHub
parent c6abd3abcb
commit ea2993d5b1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
12 changed files with 182 additions and 135 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfstudio/cameras/cameras.py
View File

@ -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, :],

26
nerfstudio/configs/method_configs.py
View File

@ -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,
),
),

52
nerfstudio/fields/instant_ngp_field.py
View File

@ -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

2
nerfstudio/fields/nerfacto_field.py
View File

@ -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__()

34
nerfstudio/fields/nerfplayer_ngp_field.py
View File

@ -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:

43
nerfstudio/model_components/ray_samplers.py
View File

@ -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,

20
nerfstudio/model_components/renderers.py
View File

@ -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)

63
nerfstudio/models/instant_ngp.py
View File

@ -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

62
nerfstudio/models/nerfplayer_ngp.py
View File

@ -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

5
nerfstudio/utils/math.py
View File

@ -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

2
pyproject.toml
View File

@ -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",

6
tests/utils/test_aabb_intersection.py
View File

@ -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)