mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
457788d689
* renaming dataloader to data. adding some dataloader classes * restructuring configs * fixed instant ngp tcnn direction rescaling * working through test cases * black and lint and tests
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""
|
|
Test the fields
|
|
"""
|
|
import torch
|
|
|
|
from pyrad.fields.modules.field_heads import FieldHeadNames
|
|
from pyrad.fields.instant_ngp_field import TCNNInstantNGPField
|
|
from pyrad.cameras.rays import Frustums, PointSamples
|
|
|
|
|
|
def test_tcnn_instant_ngp_field():
|
|
"""Test the tiny-cuda-nn field"""
|
|
# pylint: disable=import-outside-toplevel
|
|
# pylint: disable=unused-import
|
|
try:
|
|
import tinycudann as tcnn
|
|
except ImportError as e:
|
|
# tinycudann module doesn't exist
|
|
print(e)
|
|
return
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
aabb_scale = 1.0
|
|
aabb = torch.tensor(
|
|
[[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
|
|
).to(device)
|
|
field = TCNNInstantNGPField(aabb)
|
|
num_rays = 1024
|
|
num_samples = 256
|
|
positions = torch.rand((num_rays, num_samples, 3), dtype=torch.float32, device=device)
|
|
directions = torch.rand_like(positions)
|
|
frustums = Frustums(
|
|
origins=positions,
|
|
directions=directions,
|
|
frustum_starts=torch.zeros((*directions.shape[:-1], 1), device=device),
|
|
frustum_ends=torch.zeros((*directions.shape[:-1], 1), device=device),
|
|
pixel_area=torch.ones((*directions.shape[:-1], 1), device=device),
|
|
)
|
|
point_samples = PointSamples(frustums=frustums)
|
|
field_outputs = field.forward(point_samples)
|
|
|
|
positions_dtype = positions.dtype
|
|
field_outputs_density_dtype = field_outputs[FieldHeadNames.DENSITY].dtype
|
|
assert (
|
|
positions_dtype == field_outputs_density_dtype
|
|
), f"The input and output dtypes do not match: {positions_dtype} vs. {field_outputs_density_dtype}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_tcnn_instant_ngp_field()
|