4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
nerfstudio/tests/nerf/test_sampler.py
Ethan Weber fb440fedf3
Occupancy grid (#79)
* occupancy grid in progress

* visualizing occupancy grid. currently restructing field and renderer head outputs

* black and pylint. broke nerfw and semantic nerf but it's okay

* get_project_root
2022-05-17 19:09:25 -07:00

58 lines
1.5 KiB
Python

"""
Test samplers
"""
import torch
from mattport.nerf.sampler import PDFSampler, UniformSampler
from mattport.structures.rays import RayBundle
from mattport.nerf.colliders import NearFarCollider
def test_uniform_sampler():
"""Test uniform sampler"""
near_plane = 2
far_plane = 4
num_samples = 15
sampler = UniformSampler(num_samples=num_samples)
origins = torch.zeros((10, 3))
directions = torch.ones_like(origins)
ray_bundle = RayBundle(origins=origins, directions=directions)
collider = NearFarCollider(near_plane, far_plane)
ray_bundle = collider(ray_bundle)
ray_samples = sampler(ray_bundle)
assert ray_samples.positions.shape[-2] == num_samples
# TODO Tancik: Add more precise tests
def test_pdf_sampler():
"""Test pdf sampler"""
near_plane = 2
far_plane = 4
num_samples = 15
origins = torch.zeros((10, 3))
directions = torch.ones_like(origins)
ray_bundle = RayBundle(origins=origins, directions=directions)
collider = NearFarCollider(near_plane, far_plane)
ray_bundle = collider(ray_bundle)
uniform_sampler = UniformSampler(num_samples=num_samples)
coarse_ray_samples = uniform_sampler(ray_bundle)
weights = torch.ones((10, num_samples))
# Just check that it doesn't crash. Need to add some actual tests.
pdf_sampler = PDFSampler(num_samples)
pdf_sampler(coarse_ray_samples, weights, num_samples)
# TODO Tancik: Add more precise tests
if __name__ == "__main__":
test_uniform_sampler()
test_pdf_sampler()