From bd78a68756036a4970c4ae6eb61399f539c6dfbd Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 6 Jul 2022 11:03:04 -0700 Subject: [PATCH] Update RaySamples to TensorDataclass --- pyrad/cameras/rays.py | 51 +++++++++--------------- pyrad/fields/base.py | 10 +++-- pyrad/graphs/modules/ray_sampler.py | 12 +++--- pyrad/renderers/renderers.py | 48 +++++++++++----------- tests/graphs/modules/test_ray_sampler.py | 2 +- tests/renderers/test_renderers.py | 26 ++++++------ 6 files changed, 70 insertions(+), 79 deletions(-) diff --git a/pyrad/cameras/rays.py b/pyrad/cameras/rays.py index bd58118c6..095ebbe89 100644 --- a/pyrad/cameras/rays.py +++ b/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 diff --git a/pyrad/fields/base.py b/pyrad/fields/base.py index 1ba11cb2c..e8c8a2c11 100644 --- a/pyrad/fields/base.py +++ b/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) diff --git a/pyrad/graphs/modules/ray_sampler.py b/pyrad/graphs/modules/ray_sampler.py index 5752060eb..08cc507c5 100644 --- a/pyrad/graphs/modules/ray_sampler.py +++ b/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, ) diff --git a/pyrad/renderers/renderers.py b/pyrad/renderers/renderers.py index c43cc5170..2b5e98edb 100644 --- a/pyrad/renderers/renderers.py +++ b/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 diff --git a/tests/graphs/modules/test_ray_sampler.py b/tests/graphs/modules/test_ray_sampler.py index f47944b52..641dc8f26 100644 --- a/tests/graphs/modules/test_ray_sampler.py +++ b/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) diff --git a/tests/renderers/test_renderers.py b/tests/renderers/test_renderers.py index 85d7121f2..58aea7f33 100644 --- a/tests/renderers/test_renderers.py +++ b/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()