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:
David Novotny 2021-02-02 05:42:59 -08:00 committed by Facebook GitHub Bot
parent f4f3d403f3
commit fba419b7f7
4 changed files with 104 additions and 0 deletions

View File

@ -0,0 +1 @@
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

View 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

View File

@ -0,0 +1 @@
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

View 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])
)