mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
NeRF Raymarcher
Summary: An initial NeRF diff which sets up the folder structure and implements the raymarching algorithm of NeRF. Reviewed By: nikhilaravi Differential Revision: D25623990 fbshipit-source-id: ac6b05a9b866358bd4bbf44858f06859d8a6ebd1
This commit is contained in:
parent
f4f3d403f3
commit
fba419b7f7
1
projects/nerf/nerf/__init__.py
Normal file
1
projects/nerf/nerf/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
68
projects/nerf/nerf/raymarcher.py
Normal file
68
projects/nerf/nerf/raymarcher.py
Normal file
@ -0,0 +1,68 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
import torch
|
||||
from pytorch3d.renderer import EmissionAbsorptionRaymarcher
|
||||
from pytorch3d.renderer.implicit.raymarching import (
|
||||
_check_density_bounds,
|
||||
_check_raymarcher_inputs,
|
||||
_shifted_cumprod,
|
||||
)
|
||||
|
||||
|
||||
class EmissionAbsorptionNeRFRaymarcher(EmissionAbsorptionRaymarcher):
|
||||
"""
|
||||
This is essentially the `pytorch3d.renderer.EmissionAbsorptionRaymarcher`
|
||||
which additionally returns the rendering weights. It also skips returning
|
||||
the computation of the alpha-mask which is, in case of NeRF, equal to 1
|
||||
everywhere.
|
||||
|
||||
The weights are later used in the NeRF pipeline to carry out the importance
|
||||
ray-sampling for the fine rendering pass.
|
||||
|
||||
For more details about the EmissionAbsorptionRaymarcher please refer to
|
||||
the documentation of `pytorch3d.renderer.EmissionAbsorptionRaymarcher`.
|
||||
"""
|
||||
|
||||
def forward(
|
||||
self,
|
||||
rays_densities: torch.Tensor,
|
||||
rays_features: torch.Tensor,
|
||||
eps: float = 1e-10,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
rays_densities: Per-ray density values represented with a tensor
|
||||
of shape `(..., n_points_per_ray, 1)` whose values range in [0, 1].
|
||||
rays_features: Per-ray feature values represented with a tensor
|
||||
of shape `(..., n_points_per_ray, feature_dim)`.
|
||||
eps: A lower bound added to `rays_densities` before computing
|
||||
the absorbtion function (cumprod of `1-rays_densities` along
|
||||
each ray). This prevents the cumprod to yield exact 0
|
||||
which would inhibit any gradient-based learning.
|
||||
|
||||
Returns:
|
||||
features: A tensor of shape `(..., feature_dim)` containing
|
||||
the rendered features for each ray.
|
||||
weights: A tensor of shape `(..., n_points_per_ray)` containing
|
||||
the ray-specific emission-absorbtion distribution.
|
||||
Each ray distribution `(..., :)` is a valid probability
|
||||
distribution, i.e. it contains non-negative values that integrate
|
||||
to 1, such that `weights.sum(dim=-1)==1).all()` yields `True`.
|
||||
"""
|
||||
_check_raymarcher_inputs(
|
||||
rays_densities,
|
||||
rays_features,
|
||||
None,
|
||||
z_can_be_none=True,
|
||||
features_can_be_none=False,
|
||||
density_1d=True,
|
||||
)
|
||||
_check_density_bounds(rays_densities)
|
||||
rays_densities = rays_densities[..., 0]
|
||||
absorption = _shifted_cumprod(
|
||||
(1.0 + eps) - rays_densities, shift=self.surface_thickness
|
||||
)
|
||||
weights = rays_densities * absorption
|
||||
features = (weights[..., None] * rays_features).sum(dim=-2)
|
||||
|
||||
return features, weights
|
1
projects/nerf/tests/__init__.py
Normal file
1
projects/nerf/tests/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
34
projects/nerf/tests/test_raymarcher.py
Normal file
34
projects/nerf/tests/test_raymarcher.py
Normal file
@ -0,0 +1,34 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from nerf.raymarcher import EmissionAbsorptionNeRFRaymarcher
|
||||
from pytorch3d.renderer import EmissionAbsorptionRaymarcher
|
||||
|
||||
|
||||
class TestRaymarcher(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
torch.manual_seed(42)
|
||||
|
||||
def test_raymarcher(self):
|
||||
"""
|
||||
Checks that the nerf raymarcher outputs are identical to the
|
||||
EmissionAbsorptionRaymarcher.
|
||||
"""
|
||||
|
||||
feat_dim = 3
|
||||
rays_densities = torch.rand(100, 10, 1)
|
||||
rays_features = torch.randn(100, 10, feat_dim)
|
||||
|
||||
out, out_nerf = [
|
||||
raymarcher(rays_densities, rays_features)
|
||||
for raymarcher in (
|
||||
EmissionAbsorptionRaymarcher(),
|
||||
EmissionAbsorptionNeRFRaymarcher(),
|
||||
)
|
||||
]
|
||||
|
||||
self.assertTrue(
|
||||
torch.allclose(out[..., :feat_dim], out_nerf[0][..., :feat_dim])
|
||||
)
|
Loading…
x
Reference in New Issue
Block a user