mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
73 lines
1.6 KiB
Python
73 lines
1.6 KiB
Python
"""
|
|
Field output tests
|
|
"""
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from nerfstudio.field_components.field_heads import (
|
|
DensityFieldHead,
|
|
FieldHead,
|
|
FieldHeadNames,
|
|
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, field_head_name=FieldHeadNames.RGB, activation=activation)
|
|
assert render_head.get_out_dim() == out_dim
|
|
|
|
x = torch.ones((9, in_dim))
|
|
render_head(x)
|
|
|
|
# Test in_dim not provided at construction
|
|
render_head = FieldHead(out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation)
|
|
with pytest.raises(SystemError):
|
|
render_head(x)
|
|
render_head.set_in_dim(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()
|