Skip to content

Add PSF-guided hashgrid auto-configuration for Nerfstudio models #3744

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
txdai wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add PSF-guided hashgrid auto-configuration for Nerfstudio models #3744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions nerfstudio/configs/method_configs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
method_configs: Dict[str, Union[TrainerConfig, ExternalMethodDummyTrainerConfig]] = {}
descriptions = {
"nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.",
"nerfacto-psf": "Nerfacto with PSF-guided hashgrid selection (no b/grid sweep needed).",
"nerfacto-huge": "Larger version of Nerfacto with higher quality.",
"depth-nerfacto": "Nerfacto with depth supervision.",
"instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.",
Expand Down Expand Up @@ -120,6 +121,45 @@
vis="viewer",
)

method_configs["nerfacto-psf"] = TrainerConfig(
method_name="nerfacto-psf",
steps_per_eval_batch=500,
steps_per_save=2000,
max_num_iterations=30000,
mixed_precision=True,
pipeline=VanillaPipelineConfig(
datamanager=ParallelDataManagerConfig(
dataparser=NerfstudioDataParserConfig(),
train_num_rays_per_batch=4096,
eval_num_rays_per_batch=4096,
),
model=NerfactoModelConfig(
eval_num_rays_per_chunk=1 << 15,
average_init_density=0.01,
psf_guided_hashgrid=True,
psf_target_fwhm_pixels=0.5,
psf_broadening=3.0,
camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
),
),
optimizers={
"proposal_networks": {
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
},
"fields": {
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
},
"camera_opt": {
"optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000),
},
},
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
vis="viewer",
)

method_configs["nerfacto-big"] = TrainerConfig(
method_name="nerfacto",
steps_per_eval_batch=500,
Expand Down
46 changes: 45 additions & 1 deletion nerfstudio/models/instant_ngp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@
from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer
from nerfstudio.models.base_model import Model, ModelConfig
from nerfstudio.utils import colormaps
from nerfstudio.utils.hashgrid_psf import compute_psf_guided_hashgrid_params
from nerfstudio.utils.rich_utils import CONSOLE


@dataclass
Expand All @@ -53,10 +55,20 @@ class InstantNGPModelConfig(ModelConfig):
"""Resolution of the grid used for the field."""
grid_levels: int = 4
"""Levels of the grid used for the field."""
num_levels: int = 16
"""Number of levels of the hashgrid encoding."""
base_res: int = 16
"""Resolution of the base level of the hashgrid encoding."""
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"""
psf_guided_hashgrid: bool = False
"""If True, solve hash growth factor from PSF target FWHM and dataset scale instead of sweeping."""
psf_target_fwhm_pixels: float = 0.5
"""Target empirical PSF FWHM measured in image pixels."""
psf_broadening: float = 3.0
"""Empirical broadening factor used to map idealized to observed PSF width."""
alpha_thre: float = 0.01
"""Threshold for opacity skipping."""
cone_angle: float = 0.004
Expand Down Expand Up @@ -97,6 +109,36 @@ def populate_modules(self):
"""Set the fields and modules."""
super().populate_modules()

active_max_res = int(self.config.max_res)
if self.config.psf_guided_hashgrid:
train_cameras = self.kwargs.get("train_cameras")
if train_cameras is not None and not hasattr(train_cameras, "camera_to_worlds"):
train_cameras = None
psf_params = compute_psf_guided_hashgrid_params(
train_cameras=train_cameras,
scene_aabb=self.scene_box.aabb,
base_resolution=self.config.base_res,
num_levels=self.config.num_levels,
target_fwhm_pixels=self.config.psf_target_fwhm_pixels,
dim=3,
broadening=self.config.psf_broadening,
)
if psf_params is None:
CONSOLE.print(
"[bold yellow]PSF-guided hashgrid could not compute parameters; "
f"falling back to max_res={active_max_res}[/bold yellow]"
)
else:
active_max_res = int(psf_params["max_res"])
self.config.max_res = active_max_res
growth_factor = float(psf_params["growth_factor"])
source = "scene-scale fallback" if psf_params["used_scene_scale_fallback"] > 0.5 else "camera scale"
CONSOLE.print(
"[green]PSF-guided hashgrid[/green] "
f"b={growth_factor:.4f} max_res={active_max_res} "
f"target_fwhm_world={psf_params['target_fwhm_world']:.4e} ({source})"
)

if self.config.disable_scene_contraction:
scene_contraction = None
else:
Expand All @@ -106,8 +148,10 @@ def populate_modules(self):
aabb=self.scene_box.aabb,
appearance_embedding_dim=0 if self.config.use_appearance_embedding else 32,
num_images=self.num_train_data,
num_levels=self.config.num_levels,
base_res=self.config.base_res,
log2_hashmap_size=self.config.log2_hashmap_size,
max_res=self.config.max_res,
max_res=active_max_res,
spatial_distortion=scene_contraction,
)

Expand Down
58 changes: 54 additions & 4 deletions nerfstudio/models/nerfacto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@
from nerfstudio.model_components.shaders import NormalsShader
from nerfstudio.models.base_model import Model, ModelConfig
from nerfstudio.utils import colormaps
from nerfstudio.utils.hashgrid_psf import compute_psf_guided_hashgrid_params, growth_factor_to_max_res
from nerfstudio.utils.rich_utils import CONSOLE


@dataclass
Expand Down Expand Up @@ -75,6 +77,14 @@ class NerfactoModelConfig(ModelConfig):
"""Size of the hashmap for the base mlp"""
features_per_level: int = 2
"""How many hashgrid features per level"""
psf_guided_hashgrid: bool = False
"""If True, solve hash growth factor from PSF target FWHM and dataset scale instead of sweeping."""
psf_target_fwhm_pixels: float = 0.5
"""Target empirical PSF FWHM measured in image pixels."""
psf_broadening: float = 3.0
"""Empirical broadening factor used to map idealized to observed PSF width."""
psf_adjust_proposal_nets: bool = True
"""If True, apply the solved growth factor to proposal hash pyramids as well."""
num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96)
"""Number of samples per ray for each proposal network."""
num_nerf_samples_per_ray: int = 48
Expand Down Expand Up @@ -145,6 +155,46 @@ def populate_modules(self):
"""Set the fields and modules."""
super().populate_modules()

active_max_res = int(self.config.max_res)
proposal_net_args_list = [dict(args) for args in self.config.proposal_net_args_list]
if self.config.psf_guided_hashgrid:
train_cameras = self.kwargs.get("train_cameras")
if train_cameras is not None and not hasattr(train_cameras, "camera_to_worlds"):
train_cameras = None
psf_params = compute_psf_guided_hashgrid_params(
train_cameras=train_cameras,
scene_aabb=self.scene_box.aabb,
base_resolution=self.config.base_res,
num_levels=self.config.num_levels,
target_fwhm_pixels=self.config.psf_target_fwhm_pixels,
dim=3,
broadening=self.config.psf_broadening,
)
if psf_params is None:
CONSOLE.print(
"[bold yellow]PSF-guided hashgrid could not compute parameters; "
f"falling back to max_res={active_max_res}[/bold yellow]"
)
else:
growth_factor = float(psf_params["growth_factor"])
active_max_res = int(psf_params["max_res"])
self.config.max_res = active_max_res
if self.config.psf_adjust_proposal_nets:
for proposal_args in proposal_net_args_list:
proposal_levels = int(proposal_args.get("num_levels", self.config.num_levels))
proposal_base_res = int(proposal_args.get("base_res", self.config.base_res))
proposal_args["max_res"] = growth_factor_to_max_res(
base_resolution=proposal_base_res,
num_levels=proposal_levels,
growth_factor=growth_factor,
)
source = "scene-scale fallback" if psf_params["used_scene_scale_fallback"] > 0.5 else "camera scale"
CONSOLE.print(
"[green]PSF-guided hashgrid[/green] "
f"b={growth_factor:.4f} max_res={active_max_res} "
f"target_fwhm_world={psf_params['target_fwhm_world']:.4e} ({source})"
)

if self.config.disable_scene_contraction:
scene_contraction = None
else:
Expand All @@ -157,7 +207,7 @@ def populate_modules(self):
self.scene_box.aabb,
hidden_dim=self.config.hidden_dim,
num_levels=self.config.num_levels,
max_res=self.config.max_res,
max_res=active_max_res,
base_res=self.config.base_res,
features_per_level=self.config.features_per_level,
log2_hashmap_size=self.config.log2_hashmap_size,
Expand All @@ -180,8 +230,8 @@ def populate_modules(self):
# Build the proposal network(s)
self.proposal_networks = torch.nn.ModuleList()
if self.config.use_same_proposal_network:
assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
prop_net_args = self.config.proposal_net_args_list[0]
assert len(proposal_net_args_list) == 1, "Only one proposal network is allowed."
prop_net_args = proposal_net_args_list[0]
network = HashMLPDensityField(
self.scene_box.aabb,
spatial_distortion=scene_contraction,
Expand All @@ -193,7 +243,7 @@ def populate_modules(self):
self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])
else:
for i in range(num_prop_nets):
prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]
prop_net_args = proposal_net_args_list[min(i, len(proposal_net_args_list) - 1)]
network = HashMLPDensityField(
self.scene_box.aabb,
spatial_distortion=scene_contraction,
Expand Down
1 change: 1 addition & 0 deletions nerfstudio/pipelines/base_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,7 @@ def __init__(
scene_box=self.datamanager.train_dataset.scene_box,
num_train_data=len(self.datamanager.train_dataset),
metadata=self.datamanager.train_dataset.metadata,
train_cameras=self.datamanager.train_dataset.cameras,
device=device,
grad_scaler=grad_scaler,
seed_points=seed_pts,
Expand Down
Loading
Loading