Ray bundle latents add (#292)

* Making TensorDataclass reshape nested dict attributes.

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

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()

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

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