diff --git a/.vscode/launch.json b/.vscode/launch.json index c7d3f4781..9f626f025 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -20,7 +20,7 @@ "justMyCode": true }, { - "name": "Python: run_train (gpu)", + "name": "Python: run_train", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -30,17 +30,7 @@ ] }, { - "name": "Python: run_train (cpu)", - "type": "python", - "request": "launch", - "program": "scripts/run_train.py", - "console": "integratedTerminal", - "args": [ - "machine.num_gpus=0" - ] - }, - { - "name": "Python: run_train ngp (gpu)", + "name": "Python: run_train ngp", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -50,7 +40,7 @@ ] }, { - "name": "Python: run_train ngp nerfactory (gpu)", + "name": "Python: run_train ngp nerfactory", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -60,16 +50,16 @@ "--experiment_name=nerfactory_experiment", "--logging.writer.2.max-log-size=0", // viewer - "--viewer.enable", - "--viewer.zmq-url=tcp://127.0.0.1:6001", - "--viewer.websocket-port=8014", - "--pipeline.model.randomize-background", - "--pipeline.model.cone-angle=0.00390625", + // "--viewer.enable", + // "--viewer.zmq-url=tcp://127.0.0.1:6001", + // "--viewer.websocket-port=8014", + // "--pipeline.model.randomize-background", + // "--pipeline.model.cone-angle=0.00390625", "nerfactory-data" ] }, { - "name": "Python: run_train ngp-friends (gpu)", + "name": "Python: run_train ngp-friends", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -88,7 +78,7 @@ ] }, { - "name": "Python: run_train mipnerf 360 (gpu)", + "name": "Python: run_train mipnerf 360", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -98,7 +88,7 @@ ] }, { - "name": "Python: run_train tensorf (gpu)", + "name": "Python: run_train tensorf", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -108,7 +98,7 @@ ] }, { - "name": "Python: run_train nerfw (gpu)", + "name": "Python: run_train nerfw", "type": "python", "request": "launch", "program": "scripts/run_train.py", @@ -118,7 +108,7 @@ ] }, { - "name": "Python: run_train semantic_nerf (gpu)", + "name": "Python: run_train semantic_nerf", "type": "python", "request": "launch", "program": "scripts/run_train.py", diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 05bf21e4d..f55abab2c 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -345,6 +345,10 @@ class VanillaDataManagerConfig(InstantiateConfig): """number of images to sample during training iteration""" eval_dataparser: dcargs.conf.Fixed[Optional[InstantiateConfig]] = None """optionally specify different dataparser to use during eval; if None, uses train_dataparser""" + eval_num_rays_per_batch: int = 1024 + """number of rays per batch to use per eval iteration""" + eval_num_images_to_sample_from: int = -1 + """number of images to sample during eval iteration""" eval_image_indices: Optional[Tuple[int, ...]] = (0,) """specifies the image indices to use during eval; if None, uses all""" eval_num_rays_per_chunk: int = 4096 @@ -407,8 +411,6 @@ class CompoundModelConfig(ModelConfig): field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) num_samples: int = 1024 # instead of course/fine samples - cone_angle: float = 0.0 - near_plane: float = 0.5 @dataclass diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py index fffa040d2..95bd500d5 100644 --- a/nerfactory/datamanagers/base.py +++ b/nerfactory/datamanagers/base.py @@ -235,7 +235,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """Sets up the data loaders for training""" assert self.train_input_dataset is not None if self.world_size > 1: - self.sampler = DistributedSampler( + sampler = DistributedSampler( self.train_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 ) self.train_image_dataloader = CacheImageDataloader( @@ -244,7 +244,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method device=self.device, num_workers=self.world_size * 4, pin_memory=True, - sampler=self.sampler, + sampler=sampler, ) # TODO(ethan): pass this in else: self.train_image_dataloader = CacheImageDataloader( @@ -261,6 +261,30 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def setup_eval(self): """Sets up the data loader for evaluation""" assert self.eval_input_dataset is not None + if self.world_size > 1: + sampler = DistributedSampler( + self.eval_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 + ) + self.eval_image_dataloader = CacheImageDataloader( + self.eval_input_dataset, + num_images_to_sample_from=self.config.eval_num_images_to_sample_from, + device=self.device, + num_workers=self.world_size * 4, + pin_memory=True, + sampler=sampler, + ) # TODO(ethan): pass this in + else: + self.eval_image_dataloader = CacheImageDataloader( + self.eval_input_dataset, + num_images_to_sample_from=self.config.eval_num_images_to_sample_from, + device=self.device, + num_workers=self.world_size * 4, + pin_memory=True, + ) + self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) + self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch) + self.eval_ray_generator = RayGenerator(self.eval_input_dataset.dataset_inputs.cameras.to(self.device)) + # for loading full images self.eval_dataloader = FixedIndicesEvalDataloader( input_dataset=self.eval_input_dataset, num_rays_per_chunk=self.config.eval_num_rays_per_chunk, @@ -270,8 +294,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) def next_train(self) -> Tuple[RayBundle, Dict]: - """Returns the next batch of data from the train dataloader. - The RayBundle can be shaped in whatever way.""" + """Returns the next batch of data from the train dataloader.""" self.train_count += 1 image_batch = next(self.iter_train_image_dataloader) batch = self.train_pixel_sampler.sample(image_batch) @@ -280,9 +303,10 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method return ray_bundle, batch def next_eval(self) -> Tuple[RayBundle, Dict]: - """Returns the next batch of data from the eval dataloader. - The RayBundle should be shaped like an image.""" + """Returns the next batch of data from the eval dataloader.""" self.eval_count += 1 - assert self.eval_dataloader is not None, "Must setup eval dataloader before calling next_eval" - camera_ray_bundle, batch = next(self.eval_dataloader) - return camera_ray_bundle, batch + image_batch = next(self.iter_eval_image_dataloader) + batch = self.eval_pixel_sampler.sample(image_batch) + ray_indices = batch["indices"] + ray_bundle = self.eval_ray_generator(ray_indices) + return ray_bundle, batch diff --git a/nerfactory/fields/instant_ngp_field.py b/nerfactory/fields/instant_ngp_field.py index 779563249..2c57a5c80 100644 --- a/nerfactory/fields/instant_ngp_field.py +++ b/nerfactory/fields/instant_ngp_field.py @@ -151,6 +151,26 @@ class TCNNInstantNGPField(Field): rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) return {FieldHeadNames.RGB: rgb} + def get_outputs_from_positions_and_direction(self, positions: TensorType["bs", 3], directions: TensorType["bs", 3]): + """Given positions and directions, return the outputs of the field. + + Args: + positions: positions of the points + directions: directions at the point locations + """ + # assume flat positions and directions + # TODO: normalize positions + positions = SceneBounds.get_normalized_positions(positions, self.aabb) + directions = get_normalized_directions(directions) + h = self.mlp_base(positions) + density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) + density = trunc_exp(density_before_activation.to(positions)) + # TODO: normalize directions + d = self.direction_encoding(directions) + h = torch.cat([d, base_mlp_out], dim=-1) + rgb = self.mlp_head(h).to(directions) + return {FieldHeadNames.RGB: rgb, FieldHeadNames.DENSITY: density} + class TorchInstantNGPField(NeRFField): """ diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index e659721f8..3f31b2af5 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -18,24 +18,25 @@ Implementation of Instant NGP. from __future__ import annotations +import math from dataclasses import dataclass, field from typing import Dict, List, Literal, Optional, Type +import nerfacc # pylint: disable=import-error import torch from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from torchtyping import TensorType -import nerfactory.cuda as nerfactory_cuda from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg from nerfactory.fields.instant_ngp_field import field_implementation_to_class from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.models.base import Model -from nerfactory.models.modules.ray_sampler import NGPSpacedSampler from nerfactory.optimizers.loss import MSELoss -from nerfactory.utils import colors, visualization, writer +from nerfactory.utils import visualization, writer from nerfactory.utils.callbacks import ( TrainingCallback, TrainingCallbackAttributes, @@ -57,7 +58,8 @@ class InstantNGPModelConfig(cfg.ModelConfig): """Instant NGP doesn't use a collider.""" field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'""" - enable_density_field: bool = True + # TODO(ethan): remove the density field specified here + enable_density_field: bool = False """Whether to create a density field to filter samples.""" num_samples: int = 1024 """Number of samples in field evaluation. Defaults to 1024,""" @@ -83,27 +85,39 @@ class NGPModel(Model): self.field = None super().__init__(config=config, **kwargs) - def get_training_callbacks( - self, training_callback_attributes: TrainingCallbackAttributes - ) -> List[TrainingCallback]: - assert self.density_field is not None - return [ - TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - update_every_num_iters=self.density_field.update_every_num_iters, - func=self.density_field.update_density_grid, - kwargs={"density_eval_func": self.field.density_fn}, # type: ignore - ) - ] - def populate_modules(self): """Set the fields and modules.""" super().populate_modules() + # torch or tiny-cuda-nn version self.field = field_implementation_to_class[self.config.field_implementation](self.scene_bounds.aabb) - # samplers - self.sampler = NGPSpacedSampler(num_samples=self.config.num_samples, density_field=self.density_field) + # to match Ruilong's code naming + self.scene_aabb = Parameter(self.scene_bounds.aabb.flatten(), requires_grad=False) + + # setup some rendering settings + render_n_samples = 1024 + self.render_step_size = ( + (self.scene_aabb[3:] - self.scene_aabb[:3]).max() * math.sqrt(3) / render_n_samples + ).item() + + # setup occupancy field with eval function + def occ_eval_fn(x: torch.Tensor) -> torch.Tensor: + """Evaluate occupancy given positions. + + Args: + x: positions with shape (N, 3). + Returns: + occupancy values with shape (N, 1). + """ + density_after_activation = self.field.density_fn(x) + # those two are similar when density is small. + # occupancy = 1.0 - torch.exp(-density_after_activation * render_step_size) + occupancy = density_after_activation * self.render_step_size + return occupancy + + # occupancy grid + self.occ_field = nerfacc.OccupancyField(occ_eval_fn=occ_eval_fn, aabb=self.scene_aabb, resolution=128) # losses self.rgb_loss = MSELoss() @@ -113,6 +127,17 @@ class NGPModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + return [ + TrainingCallback( + where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], + update_every_num_iters=1, + func=self.occ_field.every_n_step, # will take in step + ) + ] + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field is None: @@ -123,43 +148,34 @@ class NGPModel(Model): def get_outputs(self, ray_bundle: RayBundle): assert self.field is not None - # TODO(ruilongli) - # - train test difference - # - visualize "depth_density_grid" - num_rays = len(ray_bundle) - device = ray_bundle.origins.device + def query_fn(positions: TensorType["bs", 3], directions: TensorType["bs", 3], only_density=False): + if only_density: + return self.field.density_fn(positions) + field_outputs = self.field.get_outputs_from_positions_and_direction(positions, directions) + rgbs = field_outputs[FieldHeadNames.RGB] + sigmas = field_outputs[FieldHeadNames.DENSITY] + return rgbs, sigmas - ray_samples, packed_info, t_min, t_max = self.sampler( - ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle, near_plane=self.config.near_plane - ) - - field_outputs = self.field.forward(ray_samples) - rgbs = field_outputs[FieldHeadNames.RGB] - sigmas = field_outputs[FieldHeadNames.DENSITY] - - # accumulate all the rays start from zero opacity - opacities = torch.zeros((num_rays, 1), device=device) ( - accumulated_weight, - accumulated_depth, accumulated_color, - alive_ray_mask, - ) = nerfactory_cuda.VolumeRenderer.apply( - packed_info, - ray_samples.frustums.starts, - ray_samples.frustums.ends, - sigmas.contiguous(), - rgbs.contiguous(), - opacities, + accumulated_depth, + accumulated_weight, + steps_counter, # pylint: disable=unused-variable + compact_steps_counter, # pylint: disable=unused-variable + ) = nerfacc.volumetric_rendering( + query_fn=query_fn, + rays_o=ray_bundle.origins, + rays_d=ray_bundle.directions, + scene_aabb=self.scene_aabb, + scene_occ_binary=self.occ_field.occ_grid_binary, + scene_resolution=self.occ_field.resolution, + render_bkgd=torch.ones(3, device=self.device), + render_step_size=self.render_step_size, + near_plane=self.config.near_plane, + stratified=self.training, # only use stratified sampling during training ) - accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None]) - if self.config.randomize_background: - background_colors = torch.rand((num_rays, 3), device=device) # (num_rays, 3) - else: - background_colors = colors.WHITE.to(accumulated_color) # (3,) - - accumulated_color = accumulated_color + background_colors * (1.0 - accumulated_weight) + alive_ray_mask = accumulated_weight.squeeze(-1) > 0 outputs = { "rgb": accumulated_color, diff --git a/pyproject.toml b/pyproject.toml index 469019caf..946669cc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "nerfactory" -version = "0.0.0" +version = "0.0.1" description = "All-in-one repository for state-of-the-art NeRFs" readme = "README.md" license = { text="Apache 2.0"} @@ -33,7 +33,7 @@ dependencies = [ "mediapy==1.1.0", "msgpack==1.0.4", "msgpack_numpy==0.4.8", - "nerfacc==0.0.3", + "nerfacc==0.0.5", "numpy-quaternion", "opencv-python==4.6.0.66", "plotly==5.7.0",