diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index f77c17f4d..462080287 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -16,8 +16,9 @@ Encoding functions """ +import itertools from abc import abstractmethod -from typing import Optional +from typing import Optional, Sequence import numpy as np import torch @@ -242,7 +243,6 @@ class HashEncoding(Encoding): implementation: Literal["tcnn", "torch"] = "tcnn", interpolation: Optional[Literal["Nearest", "Linear", "Smoothstep"]] = None, ) -> None: - super().__init__(in_dim=3) self.num_levels = num_levels self.features_per_level = features_per_level @@ -557,6 +557,87 @@ class TriplaneEncoding(Encoding): self.resolution = resolution +class KPlanesEncoding(Encoding): + """Learned K-Planes encoding + + A plane encoding supporting both 3D and 4D coordinates. With 3D coordinates this is similar to + :class:`TriplaneEncoding`. With 4D coordinates, the encoding at point ``[i,j,k,q]`` is + a n-dimensional vector computed as the elementwise product of 6 n-dimensional vectors at + ``planes[i,j]``, ``planes[i,k]``, ``planes[i,q]``, ``planes[j,k]``, ``planes[j,q]``, + ``planes[k,q]``. + + Unlike :class:`TriplaneEncoding` this class supports different resolution along each axis. + + This will return a tensor of shape (bs:..., num_components) + + Args: + resolution: Resolution of the grid. Can be a sequence of 3 or 4 integers. + num_components: The number of scalar planes to use (ie: output feature size) + init_a: The lower-bound of the uniform distribution used to initialize the spatial planes + init_b: The upper-bound of the uniform distribution used to initialize the spatial planes + reduce: Whether to use the element-wise product of the planes or the sum + """ + + def __init__( + self, + resolution: Sequence[int] = (128, 128, 128), + num_components: int = 64, + init_a: float = 0.1, + init_b: float = 0.5, + reduce: Literal["sum", "product"] = "product", + ) -> None: + super().__init__(in_dim=len(resolution)) + + self.resolution = resolution + self.num_components = num_components + self.reduce = reduce + if self.in_dim not in {3, 4}: + raise ValueError( + f"The dimension of coordinates must be either 3 (static scenes) " + f"or 4 (dynamic scenes). Found resolution with {self.in_dim} dimensions." + ) + has_time_planes = self.in_dim == 4 + + self.coo_combs = list(itertools.combinations(range(self.in_dim), 2)) + # Unlike the Triplane encoding, we use a parameter list instead of batching all planes + # together to support uneven resolutions (especially useful for time). + # Dynamic models (in_dim == 4) will have 6 planes: + # (y, x), (z, x), (t, x), (z, y), (t, y), (t, z) + # static models (in_dim == 3) will only have the 1st, 2nd and 4th planes. + self.plane_coefs = nn.ParameterList() + for coo_comb in self.coo_combs: + new_plane_coef = nn.Parameter( + torch.empty([self.num_components] + [self.resolution[cc] for cc in coo_comb[::-1]]) + ) + if has_time_planes and 3 in coo_comb: # Time planes initialized to 1 + nn.init.ones_(new_plane_coef) + else: + nn.init.uniform_(new_plane_coef, a=init_a, b=init_b) + self.plane_coefs.append(new_plane_coef) + + def get_out_dim(self) -> int: + return self.num_components + + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + """Sample features from this encoder. Expects ``in_tensor`` to be in range [-1, 1]""" + original_shape = in_tensor.shape + + output = 1.0 if self.reduce == "product" else 0.0 # identity for corresponding op + for ci, coo_comb in enumerate(self.coo_combs): + grid = self.plane_coefs[ci].unsqueeze(0) # [1, feature_dim, reso1, reso2] + coords = in_tensor[..., coo_comb].view(1, 1, -1, 2) # [1, 1, flattened_bs, 2] + interp = F.grid_sample( + grid, coords, align_corners=True, padding_mode="border" + ) # [1, output_dim, 1, flattened_bs] + interp = interp.view(self.num_components, -1).T # [flattened_bs, output_dim] + if self.reduce == "product": + output = output * interp + else: + output = output + interp + + return output.reshape(*original_shape[:-1], self.num_components) + + class SHEncoding(Encoding): """Spherical harmonic encoding diff --git a/tests/field_components/test_encodings.py b/tests/field_components/test_encodings.py index 271adddc9..dc59b79ff 100644 --- a/tests/field_components/test_encodings.py +++ b/tests/field_components/test_encodings.py @@ -171,6 +171,32 @@ def test_tensor_hash_encoder(): assert encoded.shape == (10, out_dim) +def test_kplane_encoder(): + """Test K-Planes encoder""" + + out_dim = 8 + + encoder = encodings.KPlanesEncoding( + resolution=(32, 32, 32, 16), + num_components=out_dim, + reduce="product", + ) + assert encoder.get_out_dim() == out_dim + + in_tensor = torch.randn((10, 4)) + encoded = encoder(in_tensor) + assert encoded.shape == (10, out_dim) + + encoder = encodings.KPlanesEncoding( + resolution=(32, 32, 32), + num_components=out_dim, + reduce="product", + ) + in_tensor = torch.randn((10, 3)) + encoded = encoder(in_tensor) + assert encoded.shape == (10, out_dim) + + if __name__ == "__main__": test_scaling_and_offset() test_nerf_encoder()