test_grid.py 1.12 KB
Newer Older
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
1
2
3
import pytest
import torch

4
from nerfacc import OccupancyGrid
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
5
6
7
8
9
10
11

device = "cuda:0"


@pytest.mark.skipif(not torch.cuda.is_available, reason="No CUDA device")
def occ_eval_fn(x: torch.Tensor) -> torch.Tensor:
    """Pesudo occupancy function: (N, 3) -> (N, 1)."""
12
    return ((x - 0.5).norm(dim=-1, keepdim=True) < 0.5).float()
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
13
14
15
16


@pytest.mark.skipif(not torch.cuda.is_available, reason="No CUDA device")
def test_occ_grid():
17
18
    roi_aabb = [0, 0, 0, 1, 1, 1]
    occ_grid = OccupancyGrid(roi_aabb=roi_aabb, resolution=128).to(device)
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
19
20
    occ_grid.every_n_step(0, occ_eval_fn, occ_thre=0.1)
    assert occ_grid.roi_aabb.shape == (6,)
21
    assert occ_grid.binary.shape == (1, 128, 128, 128)
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
22
23


24
25
26
27
28
29
30
31
32
33
@pytest.mark.skipif(not torch.cuda.is_available, reason="No CUDA device")
def test_query_grid():
    roi_aabb = [0, 0, 0, 1, 1, 1]
    occ_grid = OccupancyGrid(roi_aabb=roi_aabb, resolution=128).to(device)
    occ_grid.every_n_step(0, occ_eval_fn, occ_thre=0.1)
    samples = torch.rand((100, 3), device=device)
    occs = occ_grid.query_occ(samples)
    assert occs.shape == (100,)


Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
34
35
if __name__ == "__main__":
    test_occ_grid()
36
    test_query_grid()