tests/test_measurements.py
import torch
from mcss.geometry import look_at, make_intrinsics, make_voxel_centers
from mcss.losses import MeasurementLoss
from mcss.measurements import FixedMeasurementRenderer, _sample_volume, _state_normal_grid
from mcss.types import Cameras, SceneState, StateAppearance
def _sphere_state(resolution: int = 12) -> SceneState:
bounds = torch.tensor([[[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]])
points = make_voxel_centers(bounds, resolution)
radius = torch.linalg.vector_norm(points, dim=-1)
density = (12.0 * (0.65 - radius)).unsqueeze(1).requires_grad_()
color = torch.zeros(1, 3, resolution, resolution, resolution)
color[:, 0] = 0.8
color[:, 1] = 0.2
color[:, 2] = 0.1
log_variance = torch.full_like(density, -3.0)
return SceneState(density, color, log_variance, bounds)
def _cameras() -> Cameras:
intrinsics = make_intrinsics((8, 8), 55.0).view(1, 1, 3, 3)
c2w = look_at(torch.tensor([0.0, 0.0, -3.0]), torch.zeros(3)).view(1, 1, 4, 4)
return Cameras(intrinsics, c2w, (8, 8))
def test_fixed_renderer_is_parameter_free_and_differentiable() -> None:
renderer = FixedMeasurementRenderer(n_samples=24, ray_chunk_size=31)
state = _sphere_state()
outputs = renderer(state, _cameras())
assert sum(parameter.numel() for parameter in renderer.parameters()) == 0
assert set(outputs) == {"depth", "normal", "point", "rgb", "uncertainty", "visibility"}
assert outputs["rgb"].shape == (1, 1, 3, 8, 8)
assert outputs["depth"].shape == (1, 1, 1, 8, 8)
assert all(torch.isfinite(value).all() for value in outputs.values())
outputs["depth"].mean().backward()
assert state.density_logits.grad is not None
assert torch.isfinite(state.density_logits.grad).all()
def test_measurement_selection_and_query_permutation_are_exact() -> None:
renderer = FixedMeasurementRenderer(n_samples=16, ray_chunk_size=64)
state = _sphere_state(10)
camera = _cameras()
second_pose = look_at(torch.tensor([2.5, 0.0, -1.5]), torch.zeros(3))
cameras = Cameras(
camera.intrinsics.expand(1, 2, 3, 3).clone(),
torch.stack((camera.c2w[0, 0], second_pose), dim=0).unsqueeze(0),
camera.image_size,
)
reference = renderer(state, cameras, {"depth", "visibility"})
permuted = renderer(state, cameras.select_views(torch.tensor([1, 0])), {"depth", "visibility"})
assert set(reference) == {"depth", "visibility"}
torch.testing.assert_close(reference["depth"][:, [1, 0]], permuted["depth"])
torch.testing.assert_close(reference["visibility"][:, [1, 0]], permuted["visibility"])
def test_cell_centered_volume_samples_reproduce_voxel_values() -> None:
bounds = torch.tensor([[[-3.0, -2.0, 1.0], [5.0, 4.0, 9.0]]])
resolution = (4, 3, 5)
values = torch.arange(4 * 3 * 5, dtype=torch.float32).reshape(1, 1, *resolution)
centers = make_voxel_centers(bounds, resolution).reshape(1, -1, 1, 3)
sampled = _sample_volume(values, centers, bounds)
torch.testing.assert_close(sampled.reshape(-1), values.reshape(-1), atol=1e-5, rtol=1e-5)
def test_density_gradient_normal_uses_world_space_cell_spacing() -> None:
bounds = torch.tensor([[[-4.0, -3.0, 1.0], [8.0, 4.0, 11.0]]])
resolution = (5, 7, 9)
centers = make_voxel_centers(bounds, resolution)
density = 20.0 + 0.1 * (centers[..., 0] + 2.0 * centers[..., 1] + 3.0 * centers[..., 2])
density_logits = torch.log(torch.expm1(density)).unsqueeze(1)
state = SceneState(
density_logits,
torch.zeros(1, 3, *resolution),
torch.zeros_like(density_logits),
bounds,
)
normal = _state_normal_grid(state)[0, :, 2, 3, 4]
expected = -torch.nn.functional.normalize(torch.tensor([1.0, 2.0, 3.0]), dim=0)
torch.testing.assert_close(normal, expected, atol=1e-5, rtol=1e-5)
def test_highres_typed_appearance_changes_only_rgb_and_receives_gradients() -> None:
renderer = FixedMeasurementRenderer(n_samples=24, ray_chunk_size=64)
native = _sphere_state(8)
confidence = torch.zeros(1, 1, 16, 16, 16)
provenance = torch.zeros(1, 2, 16, 16, 16)
red = torch.tensor([0.9, 0.1, 0.1]).view(1, 3, 1, 1, 1)
base_color = red.expand(1, 3, 16, 16, 16).clone()
residual = torch.zeros_like(base_color, requires_grad=True)
appearance = StateAppearance(
confidence=confidence,
unknown_probability=torch.ones_like(confidence),
completion_gate=torch.ones_like(confidence),
provenance=provenance,
base_color=base_color,
color_logit_residual=residual,
)
typed = SceneState(
native.density_logits,
native.color,
native.log_variance,
native.bounds,
appearance=appearance,
)
native_outputs = renderer(native, _cameras(), {"rgb", "depth"})
typed_outputs = renderer(typed, _cameras(), {"rgb", "depth"})
torch.testing.assert_close(typed_outputs["depth"], native_outputs["depth"])
assert not torch.allclose(typed_outputs["rgb"], native_outputs["rgb"])
assert typed_outputs["rgb"][:, :, 0].mean() > typed_outputs["rgb"][:, :, 1].mean()
typed_outputs["rgb"].mean().backward()
assert residual.grad is not None
assert torch.isfinite(residual.grad).all()
assert residual.grad.abs().sum() > 0
def test_nearly_flat_density_normal_backward_is_finite_under_bfloat16_cuda_amp() -> None:
if not torch.cuda.is_available():
return
device = torch.device("cuda")
resolution = 8
bounds = torch.tensor([[[-6.0, -4.0, 0.05], [6.0, 4.0, 12.05]]], device=device)
ramp = torch.linspace(-1e-3, 1e-3, resolution, device=device)
density_parameter = torch.nn.Parameter(
torch.full(
(1, 1, resolution, resolution, resolution),
-2.0,
device=device,
)
+ ramp.view(1, 1, 1, 1, resolution)
)
density = density_parameter.to(torch.bfloat16)
state = SceneState(
density,
torch.zeros(1, 3, resolution, resolution, resolution, device=device, dtype=torch.bfloat16),
torch.full_like(density, -3.0),
bounds,
)
camera = Cameras(
make_intrinsics((8, 8), 55.0, device=device).view(1, 1, 3, 3),
torch.eye(4, device=device).view(1, 1, 4, 4),
(8, 8),
)
renderer = FixedMeasurementRenderer(n_samples=8, ray_chunk_size=64).to(device)
target = torch.tensor([0.0, 0.0, 1.0], device=device).view(1, 1, 3, 1, 1)
target = target.expand(1, 1, 3, 8, 8)
objective = MeasurementLoss({"normal": 1.0})
with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
normal = renderer(state, camera, {"normal"})["normal"]
loss, _ = objective({"normal": normal}, {"normal": target})
loss.backward()
assert torch.isfinite(normal).all()
assert density_parameter.grad is not None
assert torch.isfinite(density_parameter.grad).all()