From f48fa6a9238f367d53261746ce6cc2e3e9445266 Mon Sep 17 00:00:00 2001
From: Georgy Ponimatkin <38923619+ponimatkin@users.noreply.github.com>
Date: Thu, 2 Mar 2023 23:49:30 +0100
Subject: [PATCH] Instant ngp args (#1549)

* Expand instant-ngp args.

* Fixed misplaced string.
---
 nerfstudio/fields/instant_ngp_field.py | 8 ++++++--
 nerfstudio/models/instant_ngp.py       | 6 ++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py
index 102912ba7..b0f4e97e4 100644
--- a/nerfstudio/fields/instant_ngp_field.py
+++ b/nerfstudio/fields/instant_ngp_field.py
@@ -19,6 +19,7 @@ Instant-NGP field implementations using tiny-cuda-nn, torch, ....
 
 from typing import Dict, Optional, Tuple
 
+import numpy as np
 import torch
 from nerfacc import ContractionType, contract
 from torch.nn.parameter import Parameter
@@ -63,6 +64,7 @@ class TCNNInstantNGPField(Field):
         contraction_type: type of contraction
         num_levels: number of levels of the hashmap for the base mlp
         log2_hashmap_size: size of the hashmap for the base mlp
+        max_res: maximum resolution of the hashmap for the base mlp
     """
 
     def __init__(
@@ -79,6 +81,7 @@ class TCNNInstantNGPField(Field):
         contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE,
         num_levels: int = 16,
         log2_hashmap_size: int = 19,
+        max_res: int = 2048,
     ) -> None:
         super().__init__()
 
@@ -93,7 +96,8 @@ class TCNNInstantNGPField(Field):
             self.appearance_embedding = Embedding(num_images, appearance_embedding_dim)
 
         # TODO: set this properly based on the aabb
-        per_level_scale = 1.4472692012786865
+        base_res: int = 16
+        per_level_scale = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
         self.direction_encoding = tcnn.Encoding(
             n_input_dims=3,
@@ -111,7 +115,7 @@ class TCNNInstantNGPField(Field):
                 "n_levels": num_levels,
                 "n_features_per_level": 2,
                 "log2_hashmap_size": log2_hashmap_size,
-                "base_resolution": 16,
+                "base_resolution": base_res,
                 "per_level_scale": per_level_scale,
             },
             network_config={
diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py
index ace4b77c6..2bdf85b6a 100644
--- a/nerfstudio/models/instant_ngp.py
+++ b/nerfstudio/models/instant_ngp.py
@@ -65,6 +65,10 @@ class InstantNGPModelConfig(ModelConfig):
     """Number of samples in field evaluation."""
     grid_resolution: int = 128
     """Resolution 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."""
     cone_angle: float = 0.004
@@ -103,6 +107,8 @@ class NGPModel(Model):
             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,
         )
 
         self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False)