4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
nerfstudio/tests/model_components/test_renderers.py
Matthew Tancik ca69adaaec
Add median depth calculation (#877)
* Add median depth calculation

* lint

* Fix test
2022-11-01 15:05:16 -07:00

90 lines
2.4 KiB
Python

"""
Test renderers
"""
import pytest
import torch
from nerfstudio.cameras.rays import Frustums, RaySamples
from nerfstudio.model_components import renderers
def test_rgb_renderer():
"""Test RGB volumetric rendering"""
num_samples = 10
rgb_samples = torch.ones((3, num_samples, 3))
weights = torch.ones((3, num_samples, 1))
weights /= torch.sum(weights, dim=-2, keepdim=True)
rgb_renderer = renderers.RGBRenderer()
rgb = rgb_renderer(rgb=rgb_samples, weights=weights)
assert torch.max(rgb) > 0.9
rgb = rgb_renderer(rgb=rgb_samples * 0, weights=weights)
assert torch.max(rgb) == pytest.approx(0, abs=1e-6)
def test_sh_renderer():
"""Test SH volumetric rendering"""
levels = 4
num_samples = 10
sh = torch.ones((3, num_samples, 3 * levels**2))
weights = torch.ones((3, num_samples, 1))
weights /= torch.sum(weights, dim=-2, keepdim=True)
directions = torch.zeros((3, num_samples, 3))
directions[..., 0] = 1
sh_renderer = renderers.SHRenderer()
rgb = sh_renderer(sh=sh, directions=directions, weights=weights)
assert torch.max(rgb) > 0.9
def test_acc_renderer():
"""Test accumulation rendering"""
num_samples = 10
weights = torch.ones((3, num_samples, 1))
weights /= torch.sum(weights, dim=-2, keepdim=True)
acc_renderer = renderers.AccumulationRenderer()
accumulation = acc_renderer(weights=weights)
assert torch.max(accumulation) > 0.9
def test_depth_renderer():
"""Test depth rendering"""
num_samples = 10
weights = torch.ones((num_samples, 1))
weights /= torch.sum(weights, dim=-2, keepdim=True)
frustums = Frustums.get_mock_frustum()
frustums.starts = torch.linspace(0, 100, num_samples)[..., None]
frustums.ends = torch.linspace(1, 101, num_samples)[..., None]
ray_samples = RaySamples(
frustums=frustums,
camera_indices=torch.ones((num_samples, 1)),
deltas=torch.ones((num_samples, 1)),
)
depth_renderer = renderers.DepthRenderer(method="median")
depth = depth_renderer(weights=weights, ray_samples=ray_samples)
assert torch.min(depth) > 0
depth_renderer = renderers.DepthRenderer(method="expected")
depth = depth_renderer(weights=weights, ray_samples=ray_samples)
assert torch.min(depth) > 0
if __name__ == "__main__":
test_rgb_renderer()
test_sh_renderer()
test_acc_renderer()
test_depth_renderer()