mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
194 lines
6.6 KiB
Python
194 lines
6.6 KiB
Python
# Copyright 2022 The Plenoptix Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
""" Math Helper Functions """
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torchtyping import TensorType
|
|
|
|
|
|
def components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., "components"]:
|
|
"""
|
|
Returns value for each component of spherical harmonics.
|
|
|
|
Args:
|
|
levels: Number of spherical harmonic levels to compute.
|
|
directions: Spherical hamonic coefficients
|
|
"""
|
|
num_components = levels**2
|
|
components = torch.zeros((*directions.shape[:-1], num_components), device=directions.device)
|
|
|
|
assert 1 <= levels <= 5, f"SH levels must be in [1,4], got {levels}"
|
|
assert directions.shape[-1] == 3, f"Direction input should have three dimensions. Got {directions.shape[-1]}"
|
|
|
|
x = directions[..., 0]
|
|
y = directions[..., 1]
|
|
z = directions[..., 2]
|
|
|
|
xx = x**2
|
|
yy = y**2
|
|
zz = z**2
|
|
|
|
# l0
|
|
components[..., 0] = 0.28209479177387814
|
|
|
|
# l1
|
|
if levels > 1:
|
|
components[..., 1] = 0.4886025119029199 * y
|
|
components[..., 2] = 0.4886025119029199 * z
|
|
components[..., 3] = 0.4886025119029199 * x
|
|
|
|
# l2
|
|
if levels > 2:
|
|
components[..., 4] = 1.0925484305920792 * x * y
|
|
components[..., 5] = 1.0925484305920792 * y * z
|
|
components[..., 6] = 0.9461746957575601 * zz
|
|
components[..., 7] = 1.0925484305920792 * x * z
|
|
components[..., 8] = 0.5462742152960396 * (xx - yy)
|
|
|
|
# l3
|
|
if levels > 3:
|
|
components[..., 9] = 0.5900435899266435 * y * (3 * xx - yy)
|
|
components[..., 10] = 2.890611442640554 * x * y * z
|
|
components[..., 11] = 0.4570457994644658 * y * (5 * zz - 1)
|
|
components[..., 12] = 0.3731763325901154 * z * (5 * zz - 3)
|
|
components[..., 13] = 0.4570457994644658 * x * (5 * zz - 1)
|
|
components[..., 14] = 1.445305721320277 * z * (xx - yy)
|
|
components[..., 15] = 0.5900435899266435 * x * (xx - 3 * yy)
|
|
|
|
# l4
|
|
if levels > 4:
|
|
components[..., 16] = 2.5033429417967046 * x * y * (xx - yy)
|
|
components[..., 17] = 1.7701307697799304 * y * z * (3 * xx - yy)
|
|
components[..., 18] = 0.9461746957575601 * x * y * (7 * zz - 1)
|
|
components[..., 19] = 0.6690465435572892 * y * (7 * zz - 3 * z)
|
|
components[..., 20] = 0.10578554691520431 * (35 * zz * zz - 30 * zz + 3)
|
|
components[..., 21] = 0.6690465435572892 * x * z * (7 * zz - 3)
|
|
components[..., 22] = 0.47308734787878004 * (xx - yy) * (7 * zz - 1)
|
|
components[..., 23] = 1.7701307697799304 * x * z * (xx - 3 * yy)
|
|
components[..., 24] = 0.4425326924449826 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy))
|
|
|
|
return components
|
|
|
|
|
|
@dataclass
|
|
class Gaussians:
|
|
"""Stores Gaussians
|
|
|
|
Args:
|
|
mean: Mean of multivariate Gaussian
|
|
cov: Covariance of multivariate Gaussian.
|
|
"""
|
|
|
|
mean: TensorType[..., "dim"]
|
|
cov: TensorType[..., "dim", "dim"]
|
|
|
|
|
|
def compute_3d_gaussian(
|
|
directions: TensorType[..., 3],
|
|
means: TensorType[..., 3],
|
|
dir_variance: TensorType[..., 1],
|
|
radius_variance: TensorType[..., 1],
|
|
) -> Gaussians:
|
|
"""Compute guassian along ray.
|
|
|
|
Args:
|
|
directions: Axis of Gaussian.
|
|
means: Mean of Gaussian.
|
|
dir_variance: Variance along direction axis.
|
|
radius_variance: Variance tangent to direction axis.
|
|
|
|
Returns:
|
|
Gaussians: Oriented 3D gaussian.
|
|
"""
|
|
|
|
dir_outer_product = directions[..., :, None] * directions[..., None, :]
|
|
eye = torch.eye(directions.shape[-1], device=directions.device)
|
|
dir_mag_sq = torch.clamp(torch.sum(directions**2, dim=-1, keepdim=True), min=1e-10)
|
|
null_outer_product = eye - directions[..., :, None] * (directions / dir_mag_sq)[..., None, :]
|
|
dir_cov_diag = dir_variance[..., None] * dir_outer_product[..., :, :]
|
|
radius_cov_diag = radius_variance[..., None] * null_outer_product[..., :, :]
|
|
cov = dir_cov_diag + radius_cov_diag
|
|
return Gaussians(mean=means, cov=cov)
|
|
|
|
|
|
def cylinder_to_gaussian(
|
|
origins: TensorType[..., 3],
|
|
directions: TensorType[..., 3],
|
|
starts: TensorType[..., 1],
|
|
ends: TensorType[..., 1],
|
|
radius: TensorType[..., 1],
|
|
) -> Gaussians:
|
|
"""Approximates cylinders with a Gaussian distributions.
|
|
|
|
Args:
|
|
origins: Origins of cylinders.
|
|
directions: Direction (axis) of cylinders.
|
|
starts: Start of cylinders.
|
|
ends: End of cylinders.
|
|
radius: Radii of cylinders.
|
|
|
|
Returns:
|
|
Gaussians: Approximation of cylinders
|
|
"""
|
|
means = origins + directions * ((starts + ends) / 2.0)
|
|
dir_variance = (ends - starts) ** 2 / 12
|
|
radius_variance = radius**2 / 4.0
|
|
return compute_3d_gaussian(directions, means, dir_variance, radius_variance)
|
|
|
|
|
|
def conical_frustum_to_gaussian(
|
|
origins: TensorType[..., 3],
|
|
directions: TensorType[..., 3],
|
|
starts: TensorType[..., 1],
|
|
ends: TensorType[..., 1],
|
|
radius: TensorType[..., 1],
|
|
) -> Gaussians:
|
|
"""Approximates conical frustums with a Gaussian distributions.
|
|
|
|
Uses stable parameterization described in mip-NeRF publication.
|
|
|
|
Args:
|
|
origins: Origins of cones.
|
|
directions: Direction (axis) of frustums.
|
|
starts: Start of conical frustums.
|
|
ends: End of conical frustums.
|
|
radius: Radii of cone a distance of 1 from the origin.
|
|
|
|
Returns:
|
|
Gaussians: Approximation of conical frustums
|
|
"""
|
|
mu = (starts + ends) / 2.0
|
|
hw = (ends - starts) / 2.0
|
|
means = origins + directions * (mu + (2.0 * mu * hw**2.0) / (3.0 * mu**2.0 + hw**2.0))
|
|
dir_variance = (hw**2) / 3 - (4 / 15) * ((hw**4 * (12 * mu**2 - hw**2)) / (3 * mu**2 + hw**2) ** 2)
|
|
radius_variance = radius**2 * ((mu**2) / 4 + (5 / 12) * hw**2 - 4 / 15 * (hw**4) / (3 * mu**2 + hw**2))
|
|
return compute_3d_gaussian(directions, means, dir_variance, radius_variance)
|
|
|
|
|
|
def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
|
|
"""Computes the expected value of sin(y) where y ~ N(x_means, x_vars)
|
|
|
|
Args:
|
|
x_means: Mean values.
|
|
x_vars: Variance of values.
|
|
|
|
Returns:
|
|
torch.Tensor: The expected value of sin.
|
|
"""
|
|
|
|
return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
|