put field params in config

projects-lerf/lerf/data/utils/pyramid_interpolator.py

@@ -1,9 +1,9 @@
 import numpy as np
 import torch
+from lerf.data.utils.feature_dataloader import FeatureDataloader
+from lerf.encoders.image_encoder import BaseImageEncoder
 from tqdm import tqdm
 
-from lerf.encoders.image_encoder import BaseImageEncoder
-from lerf.data.utils.feature_dataloader import FeatureDataloader
 
 class PyramidInterpolator(FeatureDataloader):
     def __init__(
@@ -18,17 +18,22 @@ class PyramidInterpolator(FeatureDataloader):
         assert "stride_ratio" in cfg
         assert "image_shape" in cfg
 
-        self.kernel_size = int(cfg['image_shape'][0] * cfg["tile_ratio"])
+        self.kernel_size = int(cfg["image_shape"][0] * cfg["tile_ratio"])
         self.stride = int(self.kernel_size * cfg["stride_ratio"])
         self.padding = self.kernel_size // 2
         self.center_x = (
-            (self.kernel_size - 1) / 2 - self.padding + self.stride
+            (self.kernel_size - 1) / 2
+            - self.padding
+            + self.stride
             * np.arange(
                 np.floor((cfg["image_shape"][0] + 2 * self.padding - (self.kernel_size - 1) - 1) / self.stride + 1)
             )
         )
         self.center_y = (
-            (self.kernel_size - 1) / 2 - self.padding + self.stride * np.arange(
+            (self.kernel_size - 1) / 2
+            - self.padding
+            + self.stride
+            * np.arange(
                 np.floor((cfg["image_shape"][1] + 2 * self.padding - (self.kernel_size - 1) - 1) / self.stride + 1)
             )
         )

projects-lerf/lerf/lerf.py

@@ -5,6 +5,10 @@ from typing import Dict, List, Tuple, Type
 import numpy as np
 import open_clip
 import torch
+from lerf.encoders.image_encoder import BaseImageEncoder
+from lerf.lerf_field import LERFField
+from lerf.lerf_fieldheadnames import LERFFieldHeadNames
+from lerf.lerf_renderers import CLIPRenderer, MeanRenderer
 from torch.nn import Parameter
 
 from nerfstudio.cameras.rays import RayBundle, RaySamples
@@ -15,11 +19,6 @@ from nerfstudio.model_components.ray_samplers import PDFSampler
 from nerfstudio.model_components.renderers import DepthRenderer
 from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
 
-from lerf.encoders.image_encoder import BaseImageEncoder
-from lerf.lerf_field import LERFField
-from lerf.lerf_renderers import CLIPRenderer, MeanRenderer
-from lerf.lerf_fieldheadnames import LERFFieldHeadNames
-
 
 @dataclass
 class LERFModelConfig(NerfactoModelConfig):
@@ -30,6 +29,9 @@ class LERFModelConfig(NerfactoModelConfig):
     """maximum scale used to compute relevancy with"""
     specify_scale: bool = False
     num_lerf_samples: int = 12
+    hashgrid_layers: Tuple[int] = (12, 12)
+    hashgrid_resolutions: Tuple[Tuple[int]] = ((16, 128), (128, 512))
+    hashgrid_sizes: Tuple[int] = (19, 19)
 
 
 class LERFModel(NerfactoModel):
@@ -42,7 +44,12 @@ class LERFModel(NerfactoModel):
         self.renderer_mean = MeanRenderer()
 
         self.image_encoder: BaseImageEncoder = self.kwargs["image_encoder"]
-        self.lerf_field = LERFField(clip_n_dims=self.image_encoder.embedding_dim)
+        self.lerf_field = LERFField(
+            self.config.hashgrid_layers,
+            self.config.hashgrid_sizes,
+            self.config.hashgrid_resolutions,
+            clip_n_dims=self.image_encoder.embedding_dim,
+        )
 
     def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, preset_scales=None):
         # TODO smoothen this out

projects-lerf/lerf/lerf_config.py

@@ -2,6 +2,10 @@
 LERF configuration file.
 """
 
+from lerf.data.lerf_datamanager import LERFDataManagerConfig
+from lerf.lerf import LERFModelConfig
+from lerf.lerf_pipeline import LERFPipelineConfig
+
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
@@ -10,18 +14,14 @@ from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.plugins.types import MethodSpecification
 
-from lerf.data.lerf_datamanager import LERFDataManagerConfig
-from lerf.lerf import LERFModelConfig
-from lerf.lerf_pipeline import LERFPipelineConfig
-
 """
 Swap out the network config to use OpenCLIP or CLIP here.
 """
-from lerf.encoders.openclip_encoder import OpenCLIPNetworkConfig
 from lerf.encoders.clip_encoder import CLIPNetworkConfig
+from lerf.encoders.openclip_encoder import OpenCLIPNetworkConfig
 
 lerf_method = MethodSpecification(
-  config=TrainerConfig(
+    config=TrainerConfig(
         method_name="lerf",
         steps_per_eval_batch=500,
         steps_per_save=2000,
@@ -36,12 +36,15 @@ lerf_method = MethodSpecification(
                     mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
                 ),
             ),
-            model=LERFModelConfig(eval_num_rays_per_chunk=1 << 16),
+            model=LERFModelConfig(
+                eval_num_rays_per_chunk=1 << 16,
+                hashgrid_sizes=(19, 19),
+                hashgrid_layers=(12, 12),
+                hashgrid_resolutions=((16, 128), (128, 512)),
+            ),
             network=OpenCLIPNetworkConfig(
-                clip_model_type = "ViT-B-16",
-                clip_model_pretrained = "laion2b_s34b_b88k",
-                clip_n_dims= 512
-            )
+                clip_model_type="ViT-B-16", clip_model_pretrained="laion2b_s34b_b88k", clip_n_dims=512
+            ),
         ),
         optimizers={
             "proposal_networks": {
@@ -59,6 +62,6 @@ lerf_method = MethodSpecification(
         },
         viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
         vis="viewer",
-  ),
-  description="TODO"
-)
+    ),
+    description="TODO",
+)

projects-lerf/lerf/lerf_field.py

@@ -2,6 +2,7 @@ from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
+from lerf.lerf_fieldheadnames import LERFFieldHeadNames
 from torch import nn
 from torch.nn.parameter import Parameter
 from torchtyping import TensorType
@@ -16,8 +17,6 @@ from nerfstudio.field_components.spatial_distortions import (
 )
 from nerfstudio.fields.base_field import Field
 
-from lerf.lerf_fieldheadnames import LERFFieldHeadNames
-
 try:
     import tinycudann as tcnn
 except ImportError:
@@ -25,15 +24,22 @@ except ImportError:
 
 
 class LERFField(Field):
-    def __init__(self, clip_n_dims: int, spatial_distortion: SpatialDistortion = SceneContraction()):
+    def __init__(
+        self,
+        grid_layers,
+        grid_sizes,
+        grid_resolutions,
+        clip_n_dims: int,
+        spatial_distortion: SpatialDistortion = SceneContraction(),
+    ):
         super().__init__()
-
+        assert len(grid_layers) == len(grid_sizes) and len(grid_resolutions) == len(grid_layers)
         self.spatial_distortion = spatial_distortion
         self.clip_encs = torch.nn.ModuleList(
             [
-                # LERFField._get_encoding(16, 512, 16, indim=3)
-                LERFField._get_encoding(16, 128, 12, indim=3),
-                LERFField._get_encoding(128, 512, 12, indim=3),
+                LERFField._get_encoding(
+                    grid_resolutions[i][0], grid_resolutions[i][1], grid_layers[i], indim=3, hash_size=grid_sizes[i]
+                ) for i in range(len(grid_layers))
             ]
         )
         tot_out_dims = sum([e.n_output_dims for e in self.clip_encs])
@@ -63,7 +69,7 @@ class LERFField(Field):
         )
 
     @staticmethod
-    def _get_encoding(start_res, end_res, levels, indim=3):
+    def _get_encoding(start_res, end_res, levels, indim=3, hash_size=19):
         growth = np.exp((np.log(end_res) - np.log(start_res)) / (levels - 1))
         enc = tcnn.Encoding(
             n_input_dims=indim,
@@ -71,7 +77,7 @@ class LERFField(Field):
                 "otype": "HashGrid",
                 "n_levels": levels,
                 "n_features_per_level": 8,
-                "log2_hashmap_size": 19,
+                "log2_hashmap_size": hash_size,
                 "base_resolution": start_res,
                 "per_level_scale": growth,
             },