diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py index 82771a10e..e6d2d8c35 100644 --- a/nerfactory/cameras/rays.py +++ b/nerfactory/cameras/rays.py @@ -17,7 +17,7 @@ Some ray datastructures. """ import random from dataclasses import dataclass -from typing import Optional +from typing import Dict, Optional import torch from torchtyping import TensorType @@ -100,6 +100,7 @@ class RaySamples(TensorDataclass): camera_indices: Optional[TensorType["bs":..., 1]] = None valid_mask: Optional[TensorType["bs":..., 1]] = None deltas: Optional[TensorType["bs":..., 1]] = None + metadata: Optional[Dict[str, TensorType["bs":..., "latent_dims"]]] = None def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]: """Return weights based on predicted densities @@ -160,6 +161,7 @@ class RayBundle(TensorDataclass): fars: Rays Distance along ray to stop sampling valid_mask: Rays that are valid num_rays_per_chunk: Number of rays per chunk + metadata: Additional metadata or data needed for interpolation, will mimic shape of rays """ # TODO(ethan): make sure the sizes with ... are correct @@ -171,6 +173,7 @@ class RayBundle(TensorDataclass): fars: Optional[TensorType[..., 1]] = None valid_mask: Optional[TensorType[..., 1, bool]] = None num_rays_per_chunk: int = 128 + metadata: Optional[Dict[str, TensorType["num_rays", "latent_dims"]]] = None def set_camera_indices(self, camera_index: int) -> None: """Sets all of the the camera indices to a specific camera index. @@ -215,6 +218,8 @@ class RayBundle(TensorDataclass): ) -> RaySamples: """Produces samples for each ray by projection points along the ray direction. + Currently samples uniformly. + Args: bin_starts (TensorType[..., "num_samples", 1]): Distance from origin to start of bin. @@ -235,12 +240,14 @@ class RayBundle(TensorDataclass): else: camera_indices = None + shaped_raybundle_fields = self[..., None] + frustums = Frustums( - origins=self.origins[..., None, :], # [..., 1, 3] - directions=self.directions[..., None, :], # [..., 1, 3] + origins=shaped_raybundle_fields.origins, # [..., 1, 3] + directions=shaped_raybundle_fields.directions, # [..., 1, 3] starts=bin_starts, # [..., N_samples, 1] ends=bin_ends, # [..., N_samples, 1] - pixel_area=self.pixel_area[..., None, :], # [..., 1, 1] + pixel_area=shaped_raybundle_fields.pixel_area, # [..., 1, 1] ).to(device) ray_samples = RaySamples( @@ -248,6 +255,7 @@ class RayBundle(TensorDataclass): camera_indices=camera_indices, # [..., 1, 1] valid_mask=valid_mask, # [..., N_samples, 1] deltas=deltas, # [..., N_samples, 1] + metadata=shaped_raybundle_fields.metadata, ) return ray_samples diff --git a/nerfactory/fields/base.py b/nerfactory/fields/base.py index 802a34855..c4bae0ad2 100644 --- a/nerfactory/fields/base.py +++ b/nerfactory/fields/base.py @@ -81,6 +81,7 @@ class Field(nn.Module): if valid_mask is not None: # Hacky handling of empty masks. Tests on a single ray but doesn't use results if not valid_mask.any(): + # TODO: Can we just throw an exception here? If not we need to update this to include some fake metadata ray_samples = RaySamples(frustums=Frustums.get_mock_frustum().to(valid_mask.device)) else: ray_samples = ray_samples.apply_masks() diff --git a/nerfactory/utils/tensor_dataclass.py b/nerfactory/utils/tensor_dataclass.py index 47158a25e..863a2792d 100644 --- a/nerfactory/utils/tensor_dataclass.py +++ b/nerfactory/utils/tensor_dataclass.py @@ -15,7 +15,8 @@ """Tensor dataclass""" import dataclasses -from typing import Callable, NoReturn, Optional, Tuple, TypeVar, Union +from copy import deepcopy +from typing import Callable, Dict, List, NoReturn, Optional, Tuple, TypeVar, Union import numpy as np import torch @@ -24,8 +25,10 @@ TensorDataclassT = TypeVar("TensorDataclassT", bound="TensorDataclass") class TensorDataclass: - """Data class of tensors with the same size batch. Allows indexing and standard tensor ops. + """@dataclass of tensors with the same size batch. Allows indexing and standard tensor ops. Fields that are not Tensors will not be batched unless they are also a TensorDataclass. + Any fields that are dictionaries will have their Tensors or TensorDataclasses batched, and + dictionaries will have their tensors or TensorDataclasses considered in the initial broadcast. Example: @@ -54,32 +57,64 @@ class TensorDataclass: _shape: tuple def __post_init__(self) -> None: + """Finishes setting up the TensorDataclass + + This will 1) find the broadcasted shape and 2) broadcast all fields to this shape 3) + set _shape to be the broadcasted shape.""" if not dataclasses.is_dataclass(self): raise TypeError("TensorDataclass must be a dataclass") - field_names = [f.name for f in dataclasses.fields(self)] - batch_shapes = [] - for f in field_names: - v = self.__getattribute__(f) - if v is not None: - if isinstance(v, torch.Tensor): - batch_shapes.append(v.shape[:-1]) - elif isinstance(v, TensorDataclass): - batch_shapes.append(v.shape) + batch_shapes = self._get_dict_batch_shapes( + {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)} + ) if len(batch_shapes) == 0: raise ValueError("TensorDataclass must have at least one tensor") batch_shape = torch.broadcast_shapes(*batch_shapes) - for f in field_names: - v = self.__getattribute__(f) - if v is not None: - if isinstance(v, torch.Tensor): - self.__setattr__(f, v.broadcast_to((*batch_shape, v.shape[-1]))) - elif isinstance(v, TensorDataclass): - self.__setattr__(f, v.broadcast_to(batch_shape)) + broadcasted_fields = self._broadcast_dict_fields( + {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)}, batch_shape + ) + for f, v in broadcasted_fields.items(): + self.__setattr__(f, v) self.__setattr__("_shape", batch_shape) + def _get_dict_batch_shapes(self, dict_: Dict) -> List: + """Returns batch shapes of all tensors in a dictionary + + Args: + dict_ (Dict): The dictionary to get the batch shapes of. + + Returns: + list: The batch shapes of all tensors in the dictionary.""" + batch_shapes = [] + for v in dict_.values(): + if isinstance(v, torch.Tensor): + batch_shapes.append(v.shape[:-1]) + elif isinstance(v, TensorDataclass): + batch_shapes.append(v.shape) + elif isinstance(v, Dict): + batch_shapes.extend(self._get_dict_batch_shapes(v)) + return batch_shapes + + def _broadcast_dict_fields(self, dict_: Dict, batch_shape) -> Dict: + """Broadcasts all tensors in a dictionary according to batch_shape + + Args: + dict_ (Dict): The dictionary to broadcast. + + Returns: + Dict: The broadcasted dictionary.""" + new_dict = {} + for k, v in dict_.items(): + if isinstance(v, torch.Tensor): + new_dict[k] = v.broadcast_to((*batch_shape, v.shape[-1])) + elif isinstance(v, TensorDataclass): + new_dict[k] = v.broadcast_to(batch_shape) + elif isinstance(v, Dict): + new_dict[k] = self._broadcast_dict_fields(v, batch_shape) + return new_dict + def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT: if isinstance(indices, torch.Tensor): return self._apply_fn_to_fields(lambda x: x[indices]) @@ -123,6 +158,8 @@ class TensorDataclass: def reshape(self: TensorDataclassT, shape: Tuple[int, ...]) -> TensorDataclassT: """Returns a new TensorDataclass with the same data but with a new shape. + This should deepcopy as well. + Args: shape (Tuple[int, ...]): The new shape of the tensor dataclass. @@ -146,6 +183,9 @@ class TensorDataclass: def broadcast_to(self: TensorDataclassT, shape: Union[torch.Size, Tuple[int, ...]]) -> TensorDataclassT: """Returns a new TensorDataclass broadcast to new shape. + Changes to the original tensor dataclass should effect the returned tensor dataclass, + meaning it is NOT a deepcopy, and they are still linked. + Args: shape (Tuple[int, ...]): The new shape of the tensor dataclass. @@ -170,22 +210,52 @@ class TensorDataclass: ) -> TensorDataclassT: """Applies a function to all fields of the tensor dataclass. + TODO: Someone needs to make a high level design choice for whether not not we want this + to apply the function to any fields in arbitray superclasses. This is an edge case until we + upgrade to python 3.10 and dataclasses can actually be subclassed with vanilla python and no + janking, but if people try to jank some subclasses that are grandchildren of TensorDataclass + (imagine if someone tries to subclass the RayBundle) this will matter even before upgrading + to 3.10 . Currently we aren't going to be able to work properly for grandchildren, but you + want to use self.__dict__ if you want to apply this to grandchildren instead of our dictionary + from dataclasses.fields(self) as we do below and in other places. + Args: fn (Callable): The function to apply to tensor fields. - dataclass_fn (Optional[Callable]): The function to apply to TensorDataclass fields. Else use fn. + dataclass_fn (Optional[Callable]): The function to apply to TensorDataclass fields. Returns: TensorDataclass: A new TensorDataclass with the same data but with a new shape. """ - field_names = [f.name for f in dataclasses.fields(self)] - new_fields = {} - for f in field_names: - v = self.__getattribute__(f) - if v is not None: - if isinstance(v, TensorDataclass) and dataclass_fn is not None: - new_fields[f] = dataclass_fn(v) - elif isinstance(v, (torch.Tensor, TensorDataclass)): - new_fields[f] = fn(v) + new_fields = self._apply_fn_to_dict( + {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)}, fn, dataclass_fn + ) return dataclasses.replace(self, **new_fields) + + def _apply_fn_to_dict(self, dict_: Dict, fn: Callable, dataclass_fn: Optional[Callable] = None) -> Dict: + """A helper function for _apply_fn_to_fields, applying a function to all fields of dict_ + + Args: + dict_ (Dict): The dictionary to apply the function to. + fn (Callable): The function to apply to tensor fields. + dataclass_fn (Optional[Callable]): The function to apply to TensorDataclass fields. + + Returns: + Dict: A new dictionary with the same data but with a new shape. Will deep copy""" + + field_names = dict_.keys() + new_dict = {} + for f in field_names: + v = dict_[f] + if v is not None: + if isinstance(v, TensorDataclass) and dataclass_fn is not None: + new_dict[f] = dataclass_fn(v) + elif isinstance(v, (torch.Tensor, TensorDataclass)): + new_dict[f] = fn(v) + elif isinstance(v, Dict): + new_dict[f] = self._apply_fn_to_dict(v, fn, dataclass_fn) + else: + new_dict[f] = deepcopy(v) + + return new_dict diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py index 39121e224..286c06db5 100644 --- a/tests/utils/test_tensor_dataclass.py +++ b/tests/utils/test_tensor_dataclass.py @@ -2,6 +2,7 @@ Test tensor dataclass """ from dataclasses import dataclass +from typing import Dict import pytest import torch @@ -23,6 +24,7 @@ class DummyTensorDataclass(TensorDataclass): a: torch.Tensor b: torch.Tensor c: DummyNestedClass = None + d: Dict = None def test_init(): @@ -59,16 +61,19 @@ def test_broadcasting(): tensor_dataclass = DummyTensorDataclass(a=a, b=b) -def test_tensor_ops(): +def test_tensor_ops(): # pylint: disable=(too-many-statements) """Test tensor operations""" a = torch.ones((4, 6, 3)) b = torch.ones((6, 2)) - tensor_dataclass = DummyTensorDataclass(a=a, b=b) + d = {"t1": torch.ones((4, 6, 3)), "t2": {"t3": torch.ones((6, 4))}} + tensor_dataclass = DummyTensorDataclass(a=a, b=b, d=d) assert tensor_dataclass.shape == (4, 6) assert tensor_dataclass.a.shape == (4, 6, 3) assert tensor_dataclass.b.shape == (4, 6, 2) + assert tensor_dataclass.d["t1"].shape == (4, 6, 3) + assert tensor_dataclass.d["t2"]["t3"].shape == (4, 6, 4) assert tensor_dataclass.size == 24 assert tensor_dataclass.ndim == 2 assert len(tensor_dataclass) == 4 @@ -77,32 +82,52 @@ def test_tensor_ops(): assert reshaped.shape == (2, 12) assert reshaped.a.shape == (2, 12, 3) assert reshaped.b.shape == (2, 12, 2) + assert reshaped.d["t1"].shape == (2, 12, 3) + assert reshaped.d["t2"]["t3"].shape == (2, 12, 4) flattened = tensor_dataclass.flatten() assert flattened.shape == (24,) assert flattened.a.shape == (24, 3) assert flattened.b.shape == (24, 2) + assert flattened.d["t1"].shape == (24, 3) + assert flattened.d["t2"]["t3"].shape == (24, 4) assert flattened[0:4].shape == (4,) # Test indexing operations assert tensor_dataclass[:, 1].shape == (4,) assert tensor_dataclass[:, 1].a.shape == (4, 3) + assert tensor_dataclass[:, 1].d["t1"].shape == (4, 3) + assert tensor_dataclass[:, 1].d["t2"]["t3"].shape == (4, 4) assert tensor_dataclass[:, 0:2].shape == (4, 2) assert tensor_dataclass[:, 0:2].a.shape == (4, 2, 3) + assert tensor_dataclass[:, 0:2].d["t1"].shape == (4, 2, 3) + assert tensor_dataclass[:, 0:2].d["t2"]["t3"].shape == (4, 2, 4) assert tensor_dataclass[..., 1].shape == (4,) assert tensor_dataclass[..., 1].a.shape == (4, 3) assert tensor_dataclass[0].shape == (6,) assert tensor_dataclass[0].a.shape == (6, 3) + assert tensor_dataclass[0].d["t1"].shape == (6, 3) + assert tensor_dataclass[0].d["t2"]["t3"].shape == (6, 4) assert tensor_dataclass[0, ...].shape == (6,) assert tensor_dataclass[0, ...].a.shape == (6, 3) - tensor_dataclass = DummyTensorDataclass(a=torch.ones((2, 3, 4, 5)), b=torch.ones((4, 5))) + tensor_dataclass = DummyTensorDataclass( + a=torch.ones((2, 3, 4, 5)), + b=torch.ones((4, 5)), + d={"t1": torch.ones((2, 3, 4, 5)), "t2": {"t3": torch.ones((4, 5))}}, + ) assert tensor_dataclass[0, ...].shape == (3, 4) assert tensor_dataclass[0, ...].a.shape == (3, 4, 5) + assert tensor_dataclass[0, ...].d["t1"].shape == (3, 4, 5) + assert tensor_dataclass[0, ...].d["t2"]["t3"].shape == (3, 4, 5) assert tensor_dataclass[0, ..., 0].shape == (3,) assert tensor_dataclass[0, ..., 0].a.shape == (3, 5) + assert tensor_dataclass[0, ..., 0].d["t1"].shape == (3, 5) + assert tensor_dataclass[0, ..., 0].d["t2"]["t3"].shape == (3, 5) assert tensor_dataclass[..., 0].shape == (2, 3) assert tensor_dataclass[..., 0].a.shape == (2, 3, 5) + assert tensor_dataclass[..., 0].d["t1"].shape == (2, 3, 5) + assert tensor_dataclass[..., 0].d["t2"]["t3"].shape == (2, 3, 5) mask = torch.rand(size=(2, 3)) > 0.5 assert tensor_dataclass[mask].ndim == 2