From 5a1c8a1c86bd6ad31f437c19bac9dda4f0d58c36 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 1 Jul 2022 09:32:43 -0700 Subject: [PATCH] Remove point samples --- pyrad/cameras/rays.py | 68 +++++++++++-------------------- pyrad/fields/base.py | 32 +++++++-------- pyrad/fields/instant_ngp_field.py | 24 +++++------ pyrad/fields/nerf_field.py | 12 +++--- pyrad/fields/nerfw_field.py | 14 +++---- pyrad/graphs/instant_ngp.py | 4 +- pyrad/graphs/mipnerf.py | 4 +- pyrad/graphs/mipnerf_360.py | 4 +- pyrad/graphs/nerfw.py | 4 +- pyrad/graphs/semantic_nerf.py | 14 +++---- pyrad/graphs/vanilla_nerf.py | 4 +- pyrad/renderers/renderers.py | 2 +- tests/fields/test_fields.py | 6 +-- 13 files changed, 85 insertions(+), 107 deletions(-) diff --git a/pyrad/cameras/rays.py b/pyrad/cameras/rays.py index 9053e19c4..19384e9ce 100644 --- a/pyrad/cameras/rays.py +++ b/pyrad/cameras/rays.py @@ -101,42 +101,6 @@ class Frustums: ) -@dataclass -class PointSamples: - """Samples in space. - - Args: - frustims (Frustums): Frustums along ray. - directions (TensorType[..., 3]): Unit direction vector. - camera_indices (TensorType[..., 1]): Camera index. - valid_mask (TensorType[...]): Rays that are valid. - """ - - frustums: Frustums - camera_indices: TensorType[..., 1] = None - valid_mask: TensorType[...] = None - - def apply_masks(self) -> "PointSamples": - """Use valid_mask to mask samples. - - Returns: - PointSamples: New set of masked samples. - """ - if is_not_none(self.valid_mask): - frustums = self.frustums.apply_masks(self.valid_mask) - camera_indices = ( - self.camera_indices[self.valid_mask] if is_not_none(self.camera_indices) else self.camera_indices - ) - return PointSamples( - frustums=frustums, - camera_indices=camera_indices, - ) - return PointSamples( - frustums=self.frustums, - camera_indices=self.camera_indices, - ) - - @dataclass class RaySamples: """Samples along a ray @@ -157,15 +121,6 @@ class RaySamples: bin_ends: TensorType[..., 1] = None deltas: TensorType[..., 1] = None - def to_point_samples(self) -> PointSamples: - """Convert to PointSamples instance and return.""" - # TODO: make this more interpretable - return PointSamples( - frustums=self.frustums, - camera_indices=self.camera_indices, - valid_mask=self.valid_mask, - ) - def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: """Return weights based on predicted densities @@ -200,6 +155,29 @@ class RaySamples: """Sets valid mask""" self.valid_mask = valid_mask + def apply_masks(self) -> "RaySamples": + """Use valid_mask to mask samples. + + Returns: + RaySamples: New set of masked samples. + """ + if is_not_none(self.valid_mask): + frustums = self.frustums.apply_masks(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 + @dataclass class RayBundle: diff --git a/pyrad/fields/base.py b/pyrad/fields/base.py index 7c892839a..01d00ce6b 100644 --- a/pyrad/fields/base.py +++ b/pyrad/fields/base.py @@ -24,7 +24,7 @@ import torch from torchtyping import TensorType from pyrad.fields.modules.field_heads import FieldHeadNames -from pyrad.cameras.rays import Frustums, PointSamples +from pyrad.cameras.rays import Frustums, RaySamples from pyrad.utils.misc import is_not_none @@ -34,7 +34,7 @@ class Field(nn.Module): def density_fn(self, positions): """Returns only the density. Used primarily with the occupancy grid.""" # Need to figure out a better way to descibe positions with a ray. - point_samples = PointSamples( + ray_samples = RaySamples( frustums=Frustums( origins=positions, directions=torch.ones_like(positions), @@ -43,15 +43,15 @@ class Field(nn.Module): pixel_area=torch.ones_like((positions[..., :1])), ) ) - density, _ = self.get_density(point_samples) + density, _ = self.get_density(ray_samples) return density @abstractmethod - def get_density(self, point_samples: PointSamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: + def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: """Computes and returns the densities. Args: - point_samples (PointSamples): Samples locations to compute density. + ray_samples (RaySamples): Samples locations to compute density. Returns: Tuple[TensorType[...,1], TensorType[...,"num_features"]]: A tensor of densities and a tensor of features. @@ -59,34 +59,34 @@ class Field(nn.Module): @abstractmethod def get_outputs( - self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: """Computes and returns the colors. Args: - point_samples (PointSamples): Samples locations to compute outputs. + ray_samples (RaySamples): Samples locations to compute outputs. density_embedding (TensorType, optional): Density embeddings to condition on. Returns: Dict[FieldHeadNames, TensorType]: Output field values. """ - def forward(self, point_samples: PointSamples): + def forward(self, ray_samples: RaySamples): """Evaluates the field at points along the ray. Args: - point_samples (PointSamples): Samples to evaluate field on. + ray_samples (RaySamples): Samples to evaluate field on. """ - valid_mask = point_samples.valid_mask + valid_mask = ray_samples.valid_mask if is_not_none(valid_mask): # Hacky handling of empty masks. Tests on a single ray but doesn't use results if not valid_mask.any(): - point_samples = PointSamples(frustums=Frustums.get_mock_frustum()) + ray_samples = RaySamples(frustums=Frustums.get_mock_frustum()) else: - point_samples = point_samples.apply_masks() - density_masked, density_embedding_masked = self.get_density(point_samples) - field_outputs_masked = self.get_outputs(point_samples, density_embedding=density_embedding_masked) + ray_samples = ray_samples.apply_masks() + density_masked, density_embedding_masked = self.get_density(ray_samples) + field_outputs_masked = self.get_outputs(ray_samples, density_embedding=density_embedding_masked) field_outputs = {} for k, value in field_outputs_masked.items(): @@ -100,8 +100,8 @@ class Field(nn.Module): if valid_mask.any(): density[valid_mask] = density_masked else: - density, density_embedding = self.get_density(point_samples) - field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding) + density, density_embedding = self.get_density(ray_samples) + field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding) field_outputs[FieldHeadNames.DENSITY] = density return field_outputs diff --git a/pyrad/fields/instant_ngp_field.py b/pyrad/fields/instant_ngp_field.py index 2098e7580..9870444e3 100644 --- a/pyrad/fields/instant_ngp_field.py +++ b/pyrad/fields/instant_ngp_field.py @@ -27,7 +27,7 @@ from pyrad.fields.modules.encoding import Encoding, HashEncoding, SHEncoding from pyrad.fields.modules.field_heads import FieldHeadNames from pyrad.fields.base import Field from pyrad.fields.nerf_field import NeRFField -from pyrad.cameras.rays import PointSamples +from pyrad.cameras.rays import RaySamples try: import tinycudann as tcnn @@ -107,26 +107,26 @@ class TCNNInstantNGPField(Field): }, ) - def get_density(self, point_samples: PointSamples): + def get_density(self, ray_samples: RaySamples): """Computes and returns the densities.""" - positions = get_normalized_positions(point_samples.frustums.get_positions(), self.aabb) + positions = get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions_flat = positions.view(-1, 3) dtype = positions_flat.dtype x = self.position_encoding(positions_flat) - h = self.mlp_base(x).view(*point_samples.frustums.get_positions().shape[:-1], -1).to(dtype) + h = self.mlp_base(x).view(*ray_samples.frustums.get_positions().shape[:-1], -1).to(dtype) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) density = F.softplus(density_before_activation) return density, base_mlp_out - def get_outputs(self, point_samples: PointSamples, density_embedding=None): + def get_outputs(self, ray_samples: RaySamples, density_embedding=None): # TODO: add valid_mask masking! # tcnn requires directions in the range [0,1] - directions = get_normalized_directions(point_samples.frustums.directions) + directions = get_normalized_directions(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) dtype = directions_flat.dtype d = self.direction_encoding(directions_flat) h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) - h = self.mlp_head(h).view(*point_samples.frustums.directions.shape[:-1], -1).to(dtype) + h = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(dtype) rgb = torch.sigmoid(h) return {FieldHeadNames.RGB: rgb} @@ -158,12 +158,12 @@ class TorchInstantNGPField(NeRFField): ) self.aabb = Parameter(aabb, requires_grad=False) - def get_density(self, point_samples: PointSamples): - normalized_point_samples = point_samples - normalized_point_samples.positions = get_normalized_positions( - normalized_point_samples.frustums.get_positions(), self.aabb + def get_density(self, ray_samples: RaySamples): + normalized_ray_samples = ray_samples + normalized_ray_samples.positions = get_normalized_positions( + normalized_ray_samples.frustums.get_positions(), self.aabb ) - return super().get_density(normalized_point_samples) + return super().get_density(normalized_ray_samples) field_implementation_to_class = {"tcnn": TCNNInstantNGPField, "torch": TorchInstantNGPField} diff --git a/pyrad/fields/nerf_field.py b/pyrad/fields/nerf_field.py index 6a865fe84..95dbbe703 100644 --- a/pyrad/fields/nerf_field.py +++ b/pyrad/fields/nerf_field.py @@ -25,7 +25,7 @@ from pyrad.fields.modules.encoding import Encoding, Identity from pyrad.fields.modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead from pyrad.fields.modules.mlp import MLP from pyrad.fields.base import Field -from pyrad.cameras.rays import PointSamples +from pyrad.cameras.rays import RaySamples from pyrad.fields.modules.spatial_distortions import SpatialDistortion @@ -81,14 +81,14 @@ class NeRFField(Field): for field_head in self.field_heads: field_head.set_in_dim(self.mlp_head.get_out_dim()) - def get_density(self, point_samples: PointSamples): + def get_density(self, ray_samples: RaySamples): if self.use_integrated_encoding: - gaussian_samples = point_samples.frustums.get_gaussian_blob() + gaussian_samples = ray_samples.frustums.get_gaussian_blob() if self.spatial_distortion is not None: gaussian_samples = self.spatial_distortion(gaussian_samples) encoded_xyz = self.position_encoding(gaussian_samples.mean, covs=gaussian_samples.cov) else: - positions = point_samples.frustums.get_positions() + positions = ray_samples.frustums.get_positions() if self.spatial_distortion is not None: positions = self.spatial_distortion(positions) encoded_xyz = self.position_encoding(positions) @@ -97,11 +97,11 @@ class NeRFField(Field): return density, base_mlp_out def get_outputs( - self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: outputs = {} for field_head in self.field_heads: - encoded_dir = self.direction_encoding(point_samples.frustums.directions) + encoded_dir = self.direction_encoding(ray_samples.frustums.directions) mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) outputs[field_head.field_head_name] = field_head(mlp_out) return outputs diff --git a/pyrad/fields/nerfw_field.py b/pyrad/fields/nerfw_field.py index 4e1ea4563..55728a96a 100644 --- a/pyrad/fields/nerfw_field.py +++ b/pyrad/fields/nerfw_field.py @@ -32,7 +32,7 @@ from pyrad.fields.modules.field_heads import ( ) from pyrad.fields.modules.mlp import MLP from pyrad.fields.base import Field -from pyrad.cameras.rays import PointSamples +from pyrad.cameras.rays import RaySamples class VanillaNerfWField(Field): @@ -106,22 +106,22 @@ class VanillaNerfWField(Field): self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.get_out_dim()) self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.get_out_dim()) - def get_density(self, point_samples: PointSamples): + def get_density(self, ray_samples: RaySamples): """Computes and returns the densities.""" - encoded_xyz = self.position_encoding(point_samples.frustums.positions()) + encoded_xyz = self.position_encoding(ray_samples.frustums.positions()) base_mlp_out = self.mlp_base(encoded_xyz) density = self.field_head_density(base_mlp_out) return density, base_mlp_out def get_outputs( - self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: outputs = {} - encoded_dir = self.direction_encoding(point_samples.directions) - embedded_appearance = self.embedding_appearance(point_samples.camera_indices) + encoded_dir = self.direction_encoding(ray_samples.frustums.directions) + embedded_appearance = self.embedding_appearance(ray_samples.camera_indices) mlp_head_out = self.mlp_head(torch.cat([density_embedding, encoded_dir, embedded_appearance], dim=-1)) outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_head_out) # static rgb - embedded_transient = self.embedding_transient(point_samples.camera_indices) + embedded_transient = self.embedding_transient(ray_samples.camera_indices) transient_mlp_out = self.mlp_transient(torch.cat([density_embedding, embedded_transient], dim=-1)) outputs[self.field_head_transient_uncertainty.field_head_name] = self.field_head_transient_uncertainty( transient_mlp_out diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index ac8283b21..5ff49aa01 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/pyrad/graphs/instant_ngp.py @@ -93,12 +93,12 @@ class NGPGraph(Graph): # uniform sampling ray_samples_uniform = self.sampler_occupancy_grid(ray_bundle) - field_outputs_uniform = self.field.forward(ray_samples_uniform.to_point_samples()) + field_outputs_uniform = self.field.forward(ray_samples_uniform) weights_uniform = ray_samples_uniform.get_weights(field_outputs_uniform[FieldHeadNames.DENSITY]) # pdf sampling ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_uniform) - field_outputs_pdf = self.field.forward(ray_samples_pdf.to_point_samples()) + field_outputs_pdf = self.field.forward(ray_samples_pdf) # Hacky treatment of bins as points to allow us to merge uniform and pdf. ts_uniform = (ray_samples_uniform.bin_starts + ray_samples_uniform.bin_ends) / 2.0 diff --git a/pyrad/graphs/mipnerf.py b/pyrad/graphs/mipnerf.py index 6dd68e5b8..b8b6958a8 100644 --- a/pyrad/graphs/mipnerf.py +++ b/pyrad/graphs/mipnerf.py @@ -98,7 +98,7 @@ class MipNerfGraph(Graph): ray_samples_uniform = self.sampler_uniform(ray_bundle) # First pass: - field_outputs_coarse = self.field.forward(ray_samples_uniform.to_point_samples()) + field_outputs_coarse = self.field.forward(ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -111,7 +111,7 @@ class MipNerfGraph(Graph): ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # Second pass: - field_outputs_fine = self.field.forward(ray_samples_pdf.to_point_samples()) + field_outputs_fine = self.field.forward(ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB], diff --git a/pyrad/graphs/mipnerf_360.py b/pyrad/graphs/mipnerf_360.py index e769d5998..35343057a 100644 --- a/pyrad/graphs/mipnerf_360.py +++ b/pyrad/graphs/mipnerf_360.py @@ -102,7 +102,7 @@ class MipNerf360Graph(Graph): ray_samples_uniform = self.sampler_uniform(ray_bundle) # First pass: - field_outputs_coarse = self.field.forward(ray_samples_uniform.to_point_samples()) + field_outputs_coarse = self.field.forward(ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -115,7 +115,7 @@ class MipNerf360Graph(Graph): ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # Second pass: - field_outputs_fine = self.field.forward(ray_samples_pdf.to_point_samples()) + field_outputs_fine = self.field.forward(ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB], diff --git a/pyrad/graphs/nerfw.py b/pyrad/graphs/nerfw.py index 1f4c234fc..7f6b01969 100644 --- a/pyrad/graphs/nerfw.py +++ b/pyrad/graphs/nerfw.py @@ -105,7 +105,7 @@ class NerfWGraph(Graph): ray_samples_uniform = self.sampler_uniform(ray_bundle) # coarse field - field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples()) + field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -117,7 +117,7 @@ class NerfWGraph(Graph): ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field - field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) + field_outputs_fine = self.field_fine.forward(ray_samples_pdf) # fine weights weights_fine = ray_samples_pdf.get_weights( diff --git a/pyrad/graphs/semantic_nerf.py b/pyrad/graphs/semantic_nerf.py index e3121e5ac..6f3176add 100644 --- a/pyrad/graphs/semantic_nerf.py +++ b/pyrad/graphs/semantic_nerf.py @@ -30,7 +30,7 @@ from pyrad.fields.base import Field from pyrad.fields.nerf_field import NeRFField from pyrad.graphs.vanilla_nerf import NeRFGraph from pyrad.renderers.renderers import SemanticRenderer -from pyrad.cameras.rays import PointSamples, RayBundle +from pyrad.cameras.rays import RaySamples, RayBundle from pyrad.utils import writer @@ -75,16 +75,16 @@ class SemanticNerfField(Field): in_dim=self.mlp_semantic.get_out_dim(), num_classes=self.num_semantic_classes ) - def get_density(self, point_samples: PointSamples): - encoded_xyz = self.position_encoding(point_samples.frustums.get_positions()) + def get_density(self, ray_samples: RaySamples): + encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions()) base_mlp_out = self.mlp_base(encoded_xyz) density = self.field_head_density(base_mlp_out) return density, base_mlp_out def get_outputs( - self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: - encoded_dir = self.direction_encoding(point_samples.directions) + encoded_dir = self.direction_encoding(ray_samples.frustums.directions) mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) outputs = {} # rgb @@ -128,7 +128,7 @@ class SemanticNerfGraph(NeRFGraph): ray_samples_uniform = self.sampler_uniform(ray_bundle) # coarse field - field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples()) + field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -141,7 +141,7 @@ class SemanticNerfGraph(NeRFGraph): ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field - field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) + field_outputs_fine = self.field_fine.forward(ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB], diff --git a/pyrad/graphs/vanilla_nerf.py b/pyrad/graphs/vanilla_nerf.py index e5beeff1f..db1782d80 100644 --- a/pyrad/graphs/vanilla_nerf.py +++ b/pyrad/graphs/vanilla_nerf.py @@ -107,7 +107,7 @@ class NeRFGraph(Graph): ray_samples_uniform = self.sampler_uniform(ray_bundle) # coarse field: - field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples()) + field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform) weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( rgb=field_outputs_coarse[FieldHeadNames.RGB], @@ -120,7 +120,7 @@ class NeRFGraph(Graph): ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field: - field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) + field_outputs_fine = self.field_fine.forward(ray_samples_pdf) weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) rgb_fine = self.renderer_rgb( rgb=field_outputs_fine[FieldHeadNames.RGB], diff --git a/pyrad/renderers/renderers.py b/pyrad/renderers/renderers.py index e4ece1677..c43cc5170 100644 --- a/pyrad/renderers/renderers.py +++ b/pyrad/renderers/renderers.py @@ -19,7 +19,7 @@ Example: .. code-block:: python - field_outputs = field(ray_sampler.to_point_samples()) + field_outputs = field(ray_sampler) weights = ray_sampler.get_weights(field_outputs[FieldHeadNames.DENSITY]) rgb_renderer = RGBRenderer() diff --git a/tests/fields/test_fields.py b/tests/fields/test_fields.py index ad76fa4e0..28d898290 100644 --- a/tests/fields/test_fields.py +++ b/tests/fields/test_fields.py @@ -5,7 +5,7 @@ import torch from pyrad.fields.modules.field_heads import FieldHeadNames from pyrad.fields.instant_ngp_field import TCNNInstantNGPField -from pyrad.cameras.rays import Frustums, PointSamples +from pyrad.cameras.rays import Frustums, RaySamples def test_tcnn_instant_ngp_field(): @@ -35,8 +35,8 @@ def test_tcnn_instant_ngp_field(): frustum_ends=torch.zeros((*directions.shape[:-1], 1), device=device), pixel_area=torch.ones((*directions.shape[:-1], 1), device=device), ) - point_samples = PointSamples(frustums=frustums) - field_outputs = field.forward(point_samples) + ray_samples = RaySamples(frustums=frustums) + field_outputs = field.forward(ray_samples) positions_dtype = positions.dtype field_outputs_density_dtype = field_outputs[FieldHeadNames.DENSITY].dtype