Update RaySamples to TensorDataclass

pyrad/cameras/rays.py

@@ -86,13 +86,13 @@ class Frustums(TensorDataclass):
 
 
 @dataclass
-class RaySamples:
+class RaySamples(TensorDataclass):
     """Samples along a ray
 
     Args:
         frustums (Frustums): Frustums along ray.
         camera_indices (TensorType[..., 1]): Camera index.
-        valid_mask (TensorType[...]): Rays that are valid.
+        valid_mask (TensorType[..., 1]): Rays that are valid.
         bin_starts (TensorType[..., 1]): frustum starts along ray.
         bin_ends (TensorType[..., 1]): frustum ends along ray.
         deltas )TensorType[..., 1]): "width" of each sample.
@@ -105,23 +105,23 @@ class RaySamples:
     bin_ends: TensorType[..., 1] = None
     deltas: TensorType[..., 1] = None
 
-    def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]:
+    def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]:
         """Return weights based on predicted densities
 
         Args:
             densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray
 
         Returns:
-            TensorType[..., "num_samples"]: Weights for each sample
+            TensorType[..., "num_samples", 1]: Weights for each sample
         """
 
-        delta_density = self.deltas * densities[..., 0]
+        delta_density = self.deltas * densities
         alphas = 1 - torch.exp(-delta_density)
 
         # mip-nerf version of transmittance calculation:
-        transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
+        transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)
         transmittance = torch.cat(
-            [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1
+            [torch.zeros((*transmittance.shape[:1], 1, 1)).to(densities.device), transmittance], axis=-2
         )
         transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
 
@@ -146,20 +146,7 @@ class RaySamples:
             RaySamples: New set of masked samples.
         """
         if is_not_none(self.valid_mask):
-            frustums = self.frustums[self.valid_mask]
-            camera_indices = (
-                self.camera_indices[self.valid_mask] if is_not_none(self.camera_indices) else self.camera_indices
-            )
-            bin_starts = self.bin_starts[self.valid_mask] if is_not_none(self.bin_starts) else self.bin_starts
-            bin_ends = self.bin_ends[self.valid_mask] if is_not_none(self.bin_ends) else self.bin_ends
-            deltas = self.deltas[self.valid_mask] if is_not_none(self.deltas) else self.deltas
-            return RaySamples(
-                frustums=frustums,
-                camera_indices=camera_indices,
-                bin_starts=bin_starts,
-                bin_ends=bin_ends,
-                deltas=deltas,
-            )
+            return self[self.valid_mask[..., 0]]
         return self
 
 
@@ -257,25 +244,25 @@ class RayBundle(TensorDataclass):
         deltas = dists * torch.norm(self.directions[:, None, :], dim=-1)
 
         if is_not_none(self.camera_indices):
-            camera_indices = self.camera_indices.repeat(1, num_samples)
+            camera_indices = self.camera_indices.repeat(1, num_samples)[..., None]
         else:
             camera_indices = None
 
         frustums = Frustums(
-            origins=self.origins[:, None, :],
-            directions=self.directions[:, None, :],
-            frustum_starts=bin_starts[:, :, None],
-            frustum_ends=bin_ends[:, :, None],
-            pixel_area=self.pixel_area[:, None, :],
+            origins=self.origins[:, None, :],  # [N_rays, 1, 3]
+            directions=self.directions[:, None, :],  # [N_rays, 1, 3]
+            frustum_starts=bin_starts[:, :, None],  # [N_rays, N_samples, 1]
+            frustum_ends=bin_ends[:, :, None],  # [N_rays, N_samples, 1]
+            pixel_area=self.pixel_area[:, None, :],  # [N_rays, 2, 1]
         ).to(device)
 
         ray_samples = RaySamples(
             frustums=frustums,
-            camera_indices=camera_indices,
-            valid_mask=valid_mask,
-            bin_starts=bin_starts,
-            bin_ends=bin_ends,
-            deltas=deltas,
+            camera_indices=camera_indices,  # [N_rays, N_samples, 1]
+            valid_mask=valid_mask[..., None],  # [N_rays, N_samples, 1]
+            bin_starts=bin_starts[..., None],  # [N_rays, N_samples, 1]
+            bin_ends=bin_ends[..., None],  # [N_rays, N_samples, 1]
+            deltas=deltas[..., None],  # [N_rays, N_samples, 1]
         )
 
         return ray_samples

pyrad/fields/base.py

@@ -90,15 +90,17 @@ class Field(nn.Module):
 
             field_outputs = {}
             for k, value in field_outputs_masked.items():
-                zeros = torch.zeros(*valid_mask.shape, value.shape[-1], dtype=torch.float32, device=valid_mask.device)
+                zeros = torch.zeros(
+                    *valid_mask.shape[:-1], value.shape[-1], dtype=torch.float32, device=valid_mask.device
+                )
                 if valid_mask.any():
-                    zeros[valid_mask] = value
+                    zeros[valid_mask[..., 0]] = value
                 else:
                     zeros[0, :] = value
                 field_outputs[k] = zeros
-            density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=valid_mask.device)
+            density = torch.zeros(valid_mask.shape, dtype=torch.float32, device=valid_mask.device)
             if valid_mask.any():
-                density[valid_mask] = density_masked
+                density[valid_mask[..., 0]] = density_masked
         else:
             density, density_embedding = self.get_density(ray_samples)
             field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding)

pyrad/graphs/modules/ray_sampler.py

@@ -264,7 +264,7 @@ class PDFSampler(Sampler):
         self,
         ray_bundle: RayBundle = None,
         ray_samples: RaySamples = None,
-        weights: TensorType[..., "num_samples"] = None,
+        weights: TensorType[..., "num_samples", 1] = None,
         num_samples: Optional[int] = None,
         eps: float = 1e-5,
     ) -> RaySamples:
@@ -273,7 +273,7 @@ class PDFSampler(Sampler):
         Args:
             ray_bundle (RayBundle): Rays to generate samples for
             ray_samples (RaySamples): Existing ray samples
-            weights: (TensorType[..., "num_samples"]): Weights for each bin
+            weights: (TensorType[..., "num_samples", 1]): Weights for each bin
             num_samples (Optional[int]): Number of samples per ray
             eps: float: Small value to prevent numerical issues. Defaults to 1e-5
 
@@ -283,7 +283,7 @@ class PDFSampler(Sampler):
         num_samples = num_samples or self.num_samples
         num_bins = num_samples + 1
 
-        weights = weights + self.histogram_padding
+        weights = weights[..., 0] + self.histogram_padding
 
         # Add small offset to rays with zero weight to prevent NaNs
         weights_sum = torch.sum(weights, dim=-1, keepdim=True)
@@ -337,9 +337,9 @@ class PDFSampler(Sampler):
         # Force bins to not have a gap between them. Kinda hacky, should reconsider.
         existing_bins = torch.cat(
             [
-                ray_samples.bin_starts[..., :1],
-                (ray_samples.bin_starts[..., 1:] + ray_samples.bin_ends[..., :-1]) / 2.0,
-                ray_samples.bin_ends[..., -1:],
+                ray_samples.bin_starts[..., :1, 0],
+                (ray_samples.bin_starts[..., 1:, 0] + ray_samples.bin_ends[..., :-1, 0]) / 2.0,
+                ray_samples.bin_ends[..., -1:, 0],
             ],
             axis=-1,
         )

pyrad/renderers/renderers.py

@@ -52,38 +52,38 @@ class RGBRenderer(nn.Module):
     def combine_rgb(
         cls,
         rgb: TensorType[..., "num_samples", 3],
-        weights: TensorType[..., "num_samples"],
+        weights: TensorType[..., "num_samples", 1],
         background_color: Optional[TensorType[3]] = None,
     ) -> TensorType[..., 3]:
         """Composite samples along ray and render color image
 
         Args:
-            rgb (TensorType[..., "num_samples", -1]): RGB for each sample
-            weights (TensorType[..., "num_samples"]): Weights for each sample
+            rgb (TensorType[..., "num_samples", 3]): RGB for each sample
+            weights (TensorType[..., "num_samples", 1]): Weights for each sample
             background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
 
         Returns:
             TensorType[..., 3]: Outputs rgb values.
         """
-        rgb = torch.sum(weights[..., None] * rgb, dim=-2)
+        rgb = torch.sum(weights * rgb, dim=-2)
 
         if background_color is None:
             background_color = torch.rand_like(rgb).to(rgb.device)
 
-        rgb = rgb + background_color.to(weights.device) * (1.0 - torch.sum(weights, dim=-1)[..., None])
+        rgb = rgb + background_color.to(weights.device) * (1.0 - torch.sum(weights, dim=-2))
 
         return rgb
 
     def forward(
         self,
         rgb: TensorType[..., "num_samples", 3],
-        weights: TensorType[..., "num_samples"],
+        weights: TensorType[..., "num_samples", 1],
     ) -> TensorType[..., 3]:
         """Composite samples along ray and render color image
 
         Args:
-            rgb (TensorType[..., "num_samples", -1]): RGB for each sample
-            weights (TensorType[..., "num_samples"]): Weights for each sample
+            rgb (TensorType[..., "num_samples", 3]): RGB for each sample
+            weights (TensorType[..., "num_samples", 1]): Weights for each sample
 
         Returns:
             TensorType[..., 3]: Outputs of rgb values.
@@ -114,14 +114,14 @@ class SHRenderer(nn.Module):
         self,
         sh: TensorType[..., "num_samples", "coeffs"],
         directions: TensorType[..., "num_samples", 3],
-        weights: TensorType[..., "num_samples"],
+        weights: TensorType[..., "num_samples", 1],
     ) -> TensorType[..., 3]:
         """Composite samples along ray and render color image
 
         Args:
             sh (TensorType[..., "num_samples", "coeffs"]): Spherical hamonics coefficients for each sample
             directions: (TensorType[..., "num_samples", 3]): Sample direction
-            weights (TensorType[..., "num_samples"]): Weights for each sample
+            weights (TensorType[..., "num_samples", 1]): Weights for each sample
 
         Returns:
             TensorType[..., 3]: Outputs of rgb values.
@@ -149,7 +149,7 @@ class AccumulationRenderer(nn.Module):
     @classmethod
     def forward(
         cls,
-        weights: TensorType[..., "num_samples"],
+        weights: TensorType[..., "num_samples", 1],
     ) -> TensorType:
         """Composite samples along ray and calculate accumulation.
 
@@ -160,7 +160,7 @@ class AccumulationRenderer(nn.Module):
             TensorType: Outputs of accumulated values.
         """
 
-        accumulation = torch.sum(weights, dim=-1)[..., None]
+        accumulation = torch.sum(weights, dim=-2)
         return accumulation
 
 
@@ -177,11 +177,11 @@ class DepthRenderer(nn.Module):
             raise ValueError(f"{method} is an invalid depth calculation method")
         self.method = method
 
-    def forward(self, weights: TensorType[..., "num_samples"], ray_samples: RaySamples) -> TensorType[..., 1]:
+    def forward(self, weights: TensorType[..., "num_samples", 1], ray_samples: RaySamples) -> TensorType[..., 1]:
         """Composite samples along ray and calculate disparities.
 
         Args:
-            weights (TensorType[..., "num_samples"]): Weights for each sample.
+            weights (TensorType[..., "num_samples", 1]): Weights for each sample.
             ray_samples (RaySamples): Set of ray samples.
 
         Returns:
@@ -191,11 +191,11 @@ class DepthRenderer(nn.Module):
         if self.method == "expected":
             eps = 1e-10
             steps = (ray_samples.bin_starts + ray_samples.bin_ends) / 2
-            depth = torch.sum(weights * steps, dim=-1) / (torch.sum(weights, -1) + eps)
+            depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps)
 
-            depth = torch.clip(depth, steps[..., 0], steps[..., -1])
+            depth = torch.clip(depth, steps[..., 0, :], steps[..., -1, :])
 
-            return depth[..., None]
+            return depth
 
         raise NotImplementedError(f"Method {self.method} not implemented")
 
@@ -204,17 +204,19 @@ class UncertaintyRenderer(nn.Module):
     """Calculate uncertainty along the ray."""
 
     @classmethod
-    def forward(cls, betas: TensorType[..., "num_samples", 1], weights: TensorType[..., "num_samples"]) -> TensorType:
+    def forward(
+        cls, betas: TensorType[..., "num_samples", 1], weights: TensorType[..., "num_samples", 1]
+    ) -> TensorType:
         """_summary_
 
         Args:
-            betas (TensorType[..., "num_samples", 1]): _description_
-            weights (TensorType[..., "num_samples"]): _description_
+            betas (TensorType[..., num_samples, 1]): _description_
+            weights (TensorType[..., num_samples, 1]): _description_
 
         Returns:
             TensorType: _description_
         """
-        uncertainty = torch.sum(weights[..., None] * betas, dim=-2)
+        uncertainty = torch.sum(weights * betas, dim=-2)
         return uncertainty
 
 
@@ -223,8 +225,8 @@ class SemanticRenderer(nn.Module):
 
     @classmethod
     def forward(
-        cls, semantics: TensorType[..., "num_samples", "num_classes"], weights: TensorType[..., "num_samples"]
+        cls, semantics: TensorType[..., "num_samples", "num_classes"], weights: TensorType[..., "num_samples", 1]
     ) -> TensorType[..., "num_classes"]:
         """_summary_"""
-        sem = torch.sum(weights[..., None] * semantics, dim=-2)
+        sem = torch.sum(weights * semantics, dim=-2)
         return sem

tests/graphs/modules/test_ray_sampler.py

@@ -102,7 +102,7 @@ def test_pdf_sampler():
     uniform_sampler = UniformSampler(num_samples=num_samples)
     coarse_ray_samples = uniform_sampler(ray_bundle)
 
-    weights = torch.ones((10, num_samples))
+    weights = torch.ones((10, num_samples, 1))
 
     # Just check that it doesn't crash. Need to add some actual tests.
     pdf_sampler = PDFSampler(num_samples)

tests/renderers/test_renderers.py

@@ -13,8 +13,8 @@ def test_rgb_renderer():
     num_samples = 10
 
     rgb_samples = torch.ones((3, num_samples, 3))
-    weights = torch.ones((3, num_samples))
-    weights /= torch.sum(weights, axis=-1, keepdim=True)
+    weights = torch.ones((3, num_samples, 1))
+    weights /= torch.sum(weights, axis=-2, keepdim=True)
 
     rgb_renderer = renderers.RGBRenderer()
 
@@ -32,8 +32,8 @@ def test_sh_renderer():
     num_samples = 10
 
     sh = torch.ones((3, num_samples, 3 * levels**2))
-    weights = torch.ones((3, num_samples))
-    weights /= torch.sum(weights, axis=-1, keepdim=True)
+    weights = torch.ones((3, num_samples, 1))
+    weights /= torch.sum(weights, axis=-2, keepdim=True)
     directions = torch.zeros((3, num_samples, 3))
     directions[..., 0] = 1
 
@@ -47,8 +47,8 @@ def test_acc_renderer():
     """Test accumulation rendering"""
 
     num_samples = 10
-    weights = torch.ones((3, num_samples))
-    weights /= torch.sum(weights, axis=-1, keepdim=True)
+    weights = torch.ones((3, num_samples, 1))
+    weights /= torch.sum(weights, axis=-2, keepdim=True)
 
     acc_renderer = renderers.AccumulationRenderer()
 
@@ -60,16 +60,16 @@ def test_depth_renderer():
     """Test depth rendering"""
 
     num_samples = 10
-    weights = torch.ones((3, num_samples))
-    weights /= torch.sum(weights, axis=-1, keepdim=True)
+    weights = torch.ones((3, num_samples, 1))
+    weights /= torch.sum(weights, axis=-2, keepdim=True)
 
     ray_samples = RaySamples(
         frustums=Frustums.get_mock_frustum(),
-        camera_indices=torch.ones((num_samples)),
-        valid_mask=torch.ones((num_samples)),
-        bin_starts=torch.linspace(0, 100, num_samples),
-        bin_ends=torch.linspace(1, 101, num_samples),
-        deltas=torch.ones((num_samples)),
+        camera_indices=torch.ones((num_samples, 1)),
+        valid_mask=torch.ones((num_samples, 1)),
+        bin_starts=torch.linspace(0, 100, num_samples)[..., None],
+        bin_ends=torch.linspace(1, 101, num_samples)[..., None],
+        deltas=torch.ones((num_samples, 1)),
     )
 
     depth_renderer = renderers.DepthRenderer()