diff --git a/projects-lerf/lerf/data/utils/pyramid_interpolator.py b/projects-lerf/lerf/data/utils/pyramid_interpolator.py index 9cc6f76fd..7a95d9188 100644 --- a/projects-lerf/lerf/data/utils/pyramid_interpolator.py +++ b/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) ) ) diff --git a/projects-lerf/lerf/lerf.py b/projects-lerf/lerf/lerf.py index 836e2a85b..6932f360b 100644 --- a/projects-lerf/lerf/lerf.py +++ b/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 diff --git a/projects-lerf/lerf/lerf_config.py b/projects-lerf/lerf/lerf_config.py index d8df60219..73d203259 100644 --- a/projects-lerf/lerf/lerf_config.py +++ b/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" -) \ No newline at end of file + ), + description="TODO", +) diff --git a/projects-lerf/lerf/lerf_field.py b/projects-lerf/lerf/lerf_field.py index 91193de36..68178bc19 100644 --- a/projects-lerf/lerf/lerf_field.py +++ b/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, },