mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-14 11:54:01 +08:00
Ethan Weber
fb440fedf3
* 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
64 lines
1.3 KiB
Python
64 lines
1.3 KiB
Python
"""
|
|
Field output tests
|
|
"""
|
|
import torch
|
|
from torch import nn
|
|
|
|
from mattport.nerf.field_modules.field_heads import (
|
|
DensityFieldHead,
|
|
FieldHead,
|
|
RGBFieldHead,
|
|
SHFieldHead,
|
|
)
|
|
|
|
|
|
def test_field_output():
|
|
"""Test render output"""
|
|
in_dim = 6
|
|
out_dim = 4
|
|
activation = nn.ReLU()
|
|
render_head = FieldHead(in_dim=in_dim, out_dim=out_dim, activation=activation)
|
|
assert render_head.get_out_dim() == out_dim
|
|
|
|
x = torch.ones((9, in_dim))
|
|
render_head(x)
|
|
|
|
|
|
def test_density_output():
|
|
"""Test rgb output"""
|
|
in_dim = 6
|
|
density_head = DensityFieldHead(in_dim)
|
|
assert density_head.get_out_dim() == 1
|
|
|
|
x = torch.ones((9, in_dim))
|
|
density_head(x)
|
|
|
|
|
|
def test_rgb_output():
|
|
"""Test rgb output"""
|
|
in_dim = 6
|
|
rgb_head = RGBFieldHead(in_dim)
|
|
assert rgb_head.get_out_dim() == 3
|
|
|
|
x = torch.ones((9, in_dim))
|
|
rgb_head(x)
|
|
|
|
|
|
def test_sh_output():
|
|
"""Test sh output"""
|
|
in_dim = 6
|
|
levels = 4
|
|
channels = 3
|
|
rgb_head = SHFieldHead(in_dim, levels=levels, channels=channels)
|
|
assert rgb_head.get_out_dim() == channels * levels**2
|
|
|
|
x = torch.ones((9, in_dim))
|
|
rgb_head(x)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_field_output()
|
|
test_density_output()
|
|
test_rgb_output()
|
|
test_sh_output()
|