mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-10 18:17:31 +08:00
76 lines
1.9 KiB
Python
76 lines
1.9 KiB
Python
"""
|
|
Test the ray classes.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from nerfstudio.cameras.rays import Frustums
|
|
|
|
|
|
def test_frustum_get_position():
|
|
"""Test position calculation"""
|
|
|
|
origin = torch.Tensor([0, 1, 2])[None, ...]
|
|
direction = torch.Tensor([0, 1, 0])[None, ...]
|
|
frustum_start = torch.Tensor([2])[None, ...]
|
|
frustum_end = torch.Tensor([3])[None, ...]
|
|
|
|
target_position = torch.Tensor([0, 3.5, 2])[None, ...]
|
|
|
|
frustum = Frustums(
|
|
origins=origin,
|
|
directions=direction,
|
|
starts=frustum_start,
|
|
ends=frustum_end,
|
|
pixel_area=torch.ones((1, 1)),
|
|
)
|
|
|
|
positions = frustum.get_positions()
|
|
assert positions == pytest.approx(target_position, abs=1e-6)
|
|
|
|
|
|
def test_frustum_get_gaussian_blob():
|
|
"""Test gaussian blob calculation"""
|
|
|
|
frustum = Frustums(
|
|
origins=torch.ones((5, 3)),
|
|
directions=torch.ones((5, 3)),
|
|
starts=torch.ones((5, 1)),
|
|
ends=torch.ones((5, 1)),
|
|
pixel_area=torch.ones((5, 1)),
|
|
)
|
|
|
|
gaussian_blob = frustum.get_gaussian_blob()
|
|
assert gaussian_blob.mean.shape == (5, 3)
|
|
assert gaussian_blob.cov.shape == (5, 3, 3)
|
|
|
|
|
|
def test_frustum_apply_masks():
|
|
"""Test masking frustum"""
|
|
frustum = Frustums(
|
|
origins=torch.ones((5, 3)),
|
|
directions=torch.ones((5, 3)),
|
|
starts=torch.ones((5, 1)),
|
|
ends=torch.ones((5, 1)),
|
|
pixel_area=torch.ones((5, 1)),
|
|
)
|
|
|
|
mask = torch.tensor([False, True, False, True, True], dtype=torch.bool)
|
|
frustum = frustum[mask]
|
|
|
|
assert frustum.origins.shape == (3, 3)
|
|
assert frustum.directions.shape == (3, 3)
|
|
assert frustum.starts.shape == (3, 1)
|
|
assert frustum.ends.shape == (3, 1)
|
|
assert frustum.pixel_area.shape == (3, 1)
|
|
|
|
|
|
def test_get_mock_frustum():
|
|
"""Test creation of mock frustum"""
|
|
Frustums.get_mock_frustum()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_frustum_get_gaussian_blob()
|