diff --git a/docs/extensions/imgs/desolation_unreal.png b/docs/extensions/imgs/desolation_unreal.png new file mode 100644 index 000000000..3214d14a1 Binary files /dev/null and b/docs/extensions/imgs/desolation_unreal.png differ diff --git a/docs/extensions/imgs/export_nvol.png b/docs/extensions/imgs/export_nvol.png new file mode 100644 index 000000000..d00be381c Binary files /dev/null and b/docs/extensions/imgs/export_nvol.png differ diff --git a/docs/extensions/imgs/nvol_ready.png b/docs/extensions/imgs/nvol_ready.png new file mode 100644 index 000000000..4d8d94715 Binary files /dev/null and b/docs/extensions/imgs/nvol_ready.png differ diff --git a/docs/extensions/unreal_engine.md b/docs/extensions/unreal_engine.md new file mode 100644 index 000000000..f533f2a56 --- /dev/null +++ b/docs/extensions/unreal_engine.md @@ -0,0 +1,40 @@ +# Exporting to Unreal Engine + + ```{image} imgs/desolation_unreal.png + :width: 800 + :align: center + :alt: NeRF in Unreal Engine + ``` + +## Overview + +NeRFStudio models can be used in Unreal Engine if they are converted to an NVOL file. NVOL is a new standard file format to store NeRFs in a fast and efficient way. NVOL files can be obtained from NeRFStudio checkpoints files (.ckpt) using the [Volinga Suite](https://volinga.ai/). + +:::{admonition} Note +:class: warning +Volinga Suite is not yet publicly available. If you are interested in converting your models to NVOL, sign up for the beta at [volinga.ai](https://volinga.ai). +::: + + +## Exporting your model to NVOL +Currently NVOL file only supports Volinga model (which is based on nerfacto). You can train your model using the following command: + +```bash +ns-train volinga --data /path/to/your/data --vis viewer +``` + +Once the training is done, you can find your checkpoint file in the `outputs/path-to-your-data/volinga` folder. Then, you can drag it to Volinga Suite to export it to NVOL. + + ```{image} imgs/export_nvol.png + :width: 400 + :align: center + :alt: Nvol export in Voliga Suite + ``` + +Once the NVOL is ready, you can download it and use it in Unreal Engine. + + ```{image} imgs/nvol_ready.png + :width: 800 + :align: center + :alt: NVOL ready to use + ``` \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index d421535e9..b03e8712b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -58,8 +58,8 @@ Contributing ```{toctree} :hidden: :caption: Extensions - extensions/blender_addon +extensions/unreal_engine ``` ```{toctree} diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 6e4fdab87..e888747b4 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -63,6 +63,7 @@ method_configs: Dict[str, TrainerConfig] = {} descriptions = { "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.", "depth-nerfacto": "Nerfacto with depth supervision.", + "volinga": "Real-time rendering model from Volinga. Directly exportable to NVOL format at https://volinga.ai/", "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.", "instant-ngp-bounded": "Implementation of Instant-NGP. Recommended for bounded real and synthetic scenes", "mipnerf": "High quality model for bounded scenes. (slow)", @@ -137,6 +138,47 @@ method_configs["depth-nerfacto"] = TrainerConfig( vis="viewer", ) +method_configs["volinga"] = TrainerConfig( + method_name="volinga", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=NerfactoModelConfig( + eval_num_rays_per_chunk=1 << 15, + hidden_dim=32, + hidden_dim_color=32, + hidden_dim_transient=32, + num_nerf_samples_per_ray=24, + proposal_net_args_list=[ + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": True}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": True}, + ], + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + method_configs["instant-ngp"] = TrainerConfig( method_name="instant-ngp", steps_per_eval_batch=500, diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index 9f4fb1571..443ca92c3 100644 --- a/nerfstudio/fields/density_fields.py +++ b/nerfstudio/fields/density_fields.py @@ -21,7 +21,6 @@ from typing import Optional import numpy as np import torch -from torch.nn.parameter import Parameter from torchtyping import TensorType from nerfstudio.cameras.rays import RaySamples @@ -62,11 +61,15 @@ class HashMLPDensityField(Field): features_per_level=2, ) -> None: super().__init__() - self.aabb = Parameter(aabb, requires_grad=False) + self.register_buffer("aabb", aabb) self.spatial_distortion = spatial_distortion self.use_linear = use_linear growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) + self.register_buffer("max_res", torch.tensor(max_res)) + self.register_buffer("num_levels", torch.tensor(num_levels)) + self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size)) + config = { "encoding": { "otype": "HashGrid", diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 0ed9069e1..7e4908e0b 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -116,9 +116,13 @@ class TCNNNerfactoField(Field): ) -> None: super().__init__() - self.aabb = Parameter(aabb, requires_grad=False) + self.register_buffer("aabb", aabb) self.geo_feat_dim = geo_feat_dim + self.register_buffer("max_res", torch.tensor(max_res)) + self.register_buffer("num_levels", torch.tensor(num_levels)) + self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size)) + self.spatial_distortion = spatial_distortion self.num_images = num_images self.appearance_embedding_dim = appearance_embedding_dim diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 024e3a7a6..53173640a 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -73,6 +73,12 @@ class NerfactoModelConfig(ModelConfig): """How far along the ray to stop sampling.""" background_color: Literal["random", "last_sample", "black", "white"] = "last_sample" """Whether to randomize the background color.""" + hidden_dim: int = 64 + """Dimension of hidden layers""" + hidden_dim_color: int = 64 + """Dimension of hidden layers for color network""" + hidden_dim_transient: int = 64 + """Dimension of hidden layers for transient network""" num_levels: int = 16 """Number of levels of the hashmap for the base mlp.""" max_res: int = 2048 @@ -93,8 +99,8 @@ class NerfactoModelConfig(ModelConfig): """Use the same proposal network. Otherwise use different ones.""" proposal_net_args_list: List[Dict] = field( default_factory=lambda: [ - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128}, - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": False}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": False}, ] ) """Arguments for the proposal density fields.""" @@ -145,9 +151,12 @@ class NerfactoModel(Model): # Fields self.field = TCNNNerfactoField( self.scene_box.aabb, + hidden_dim=self.config.hidden_dim, num_levels=self.config.num_levels, max_res=self.config.max_res, log2_hashmap_size=self.config.log2_hashmap_size, + hidden_dim_color=self.config.hidden_dim_color, + hidden_dim_transient=self.config.hidden_dim_transient, spatial_distortion=scene_contraction, num_images=self.num_train_data, use_pred_normals=self.config.predict_normals, diff --git a/tests/test_train.py b/tests/test_train.py index 5c52e50e8..89aa1c29a 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -21,6 +21,7 @@ BLACKLIST = [ "instant-ngp", "instant-ngp-bounded", "nerfacto", + "volinga", "phototourism", "depth-nerfacto", "nerfplayer-ngp",