# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
from typing import List, Optional, Set, Tuple, Union
import numpy as np
import torch as th
DISTORTION_MODES: Set[Optional[str]] = {
None,
"pinhole",
"radial-tangential",
"fisheye",
}
_FISHEYE62_MODES: Set[str] = {"fisheye62", "fisheye62_lut"}
_ROCM: bool = th.version.hip is not None
_ROCM_FORM_DTYPES: tuple[th.dtype, ...] = (
th.float16,
th.bfloat16,
th.float32,
th.float64,
)
def _apply_matrix(m: th.Tensor, x: th.Tensor) -> th.Tensor:
"""Apply each batch item's matrix to each of its points: ``(m[:, None] @ x[..., None])[..., 0]``.
On ROCm, PyTorch runs that per-point product as a batched GEMM with one 3x3 (or 2x2) GEMM per
point, which made the projection most of the GPU time of a DRTK fitting loop on MI300X. There
``_apply_matrix_rocm`` replaces it for the inputs ``_rocm_form_applies`` accepts. CPU, CUDA and
TorchScript keep the matmul form, bit for bit.
"""
# TorchScript compiles only the matmul form, so scripted graphs are the same on every device.
if not th.jit.is_scripting():
if _ROCM and x.is_cuda and _rocm_form_applies(m, x):
return _apply_matrix_rocm(m, x)
return (m[:, None] @ x[..., None])[..., 0]
def _rocm_form_applies(m: th.Tensor, x: th.Tensor) -> bool:
"""Whether the broadcast form can stand in for the matmul: the matmul accepts ``m`` and ``x``,
and the broadcast form computes the same products in float32 or float64, without TF32 or
autocast lowering their precision (under autocast it returns float32, or float64 for float64
inputs). Every other input goes to the matmul, so it behaves exactly as before."""
if m.dim() < 2 or x.dim() < 1 or x.shape[-1] != m.shape[-1]:
return False
if m.dtype not in _ROCM_FORM_DTYPES or x.dtype not in _ROCM_FORM_DTYPES:
return False
# Products in float16 or bfloat16 would round before the sum (and float16 ones could overflow).
if th.promote_types(m.dtype, x.dtype) not in (th.float32, th.float64):
return False
if m.dtype == x.dtype:
return True
# Autocast casts both operands of the matmul to one dtype, but never float64 ones.
return th.is_autocast_enabled(x.device.type) and th.float64 not in (
m.dtype,
x.dtype,
)
def _apply_matrix_rocm(m: th.Tensor, x: th.Tensor) -> th.Tensor:
"""``_apply_matrix`` as a broadcast multiply and sum, with no GEMM: TF32 and autocast do not
lower its precision."""
return (x[..., None, :] * m[:, None]).sum(-1)
def _any_mode_is_fisheye62(
distortion_mode: Optional[Union[List[str], str]],
) -> bool:
# Check list/tuple first -- `list in set` raises TypeError because lists
# are unhashable.
if isinstance(distortion_mode, (list, tuple)):
return any(m in _FISHEYE62_MODES for m in distortion_mode)
return distortion_mode in _FISHEYE62_MODES
def project_pinhole(
v_cam: th.Tensor, focal: th.Tensor, princpt: th.Tensor
) -> th.Tensor:
"""Project camera-space points with an undistorted pinhole camera.
Args:
v_cam: Camera-space points with shape ``(N, V, 3)``.
focal: Camera intrinsic matrices with shape ``(N, 2, 2)``.
princpt: Principal points in pixels with shape ``(N, 2)``.
Returns:
Pixel coordinates with shape ``(N, V, 2)``.
"""
z = v_cam[:, :, 2:3]
z = th.where(z < 0, z.clamp(max=-1e-8), z.clamp(min=1e-8))
v_proj = v_cam[:, :, 0:2] / z
v_pix = _apply_matrix(focal, v_proj) + princpt[:, None]
return v_pix
def project_pinhole_distort_rt(
v_cam: th.Tensor,
focal: th.Tensor,
princpt: th.Tensor,
D: th.Tensor,
fov: Optional[th.Tensor] = None,
) -> th.Tensor:
"""Project camera-space points with an OpenCV radial-tangential model.
``D`` may contain OpenCV's 4-, 5-, or 8-parameter radial/tangential
coefficients. The normalized radius is clamped to ``fov`` before applying
distortion; when ``fov`` is omitted it is estimated from ``D``.
Args:
v_cam: Camera-space points with shape ``(N, V, 3)``.
focal: Camera intrinsic matrices with shape ``(N, 2, 2)``.
princpt: Principal points in pixels with shape ``(N, 2)``.
D: Distortion coefficients with shape ``(N, 4)``, ``(N, 5)``, or
``(N, 8)``.
fov: Optional maximum normalized radius with shape ``(N, 1)``.
Returns:
Distorted pixel coordinates with shape ``(N, V, 2)``.
"""
# See https://docs.opencv.org/2.4/doc/tutorials/calib3d/camera_calibration/camera_calibration.html
if fov is None:
with th.no_grad():
fov = estimate_rt_fov(D)
z = v_cam[:, :, 2:3]
z = th.where(z < 0, z.clamp(max=-1e-8), z.clamp(min=1e-8))
v_proj = v_cam[:, :, :2] / z
r2 = v_proj.pow(2).sum(-1)
# Clamp x, y and r to avoid wrapping behavior of the distortion model.
r2 = r2.clamp(max=fov.pow(2))
v_clamped = v_proj.clamp(min=-fov[..., None], max=fov[..., None])
assert D.shape[1] in [4, 5, 8]
# 4 param: R = (1 + k1 r^2 + k2 r^4)
R = 1 + D[:, 0:1] * r2 + D[:, 1:2] * r2.pow(2)
# 5 param: R = (1 + k1 r^2 + k2 r^4 + k3 r^6)
if D.shape[1] == 5:
R = R + D[:, 4:5] * r2.pow(3)
# 8 param: R = (1 + k1 r^2 + k2 r^4 + k3 r^6) / (1 + k4 r^2 + k5 r^4 + k6 r^6)
if D.shape[1] == 8:
R = R + D[:, 4:5] * r2.pow(3)
R = R / (1 + D[:, 5:6] * r2 + D[:, 6:7] * r2.pow(2) + D[:, 7:8] * r2.pow(3))
# [x' y'] * R
v_proj_dist = v_proj * R[..., None]
# [2 p1 x' y', 2 p2 x' y']
v_proj_dist += (
2
* v_clamped[..., 0:1]
* v_clamped[..., 1:2]
* th.stack((D[:, 2:3], D[:, 3:4]), dim=-1)
)
# [p2 r^2, p1 r^2]
v_proj_dist += r2[..., None] * th.stack((D[:, 3:4], D[:, 2:3]), dim=-1)
# [2 p2 x'^2, 2 p1 y'^2]
v_proj_dist += th.stack(
(
2 * D[:, 3:4] * v_clamped[..., 0].pow(2),
2 * D[:, 2:3] * v_clamped[..., 1].pow(2),
),
dim=-1,
)
v_pix_dist = _apply_matrix(focal, v_proj_dist) + princpt[:, None]
return v_pix_dist
def project_fisheye_distort(
v_cam: th.Tensor,
focal: th.Tensor,
princpt: th.Tensor,
D: th.Tensor,
fov: Optional[th.Tensor] = None,
) -> th.Tensor:
"""Project camera-space points with OpenCV's fisheye model.
Args:
v_cam: Camera-space points with shape ``(N, V, 3)``.
focal: Camera intrinsic matrices with shape ``(N, 2, 2)``.
princpt: Principal points in pixels with shape ``(N, 2)``.
D: Fisheye distortion coefficients with shape ``(N, 4)``.
fov: Optional maximum normalized radius with shape ``(N, 1)``. When
omitted, it is estimated from ``D``.
Returns:
Distorted pixel coordinates with shape ``(N, V, 2)``.
"""
# See https://github.com/opencv/opencv/blob/master/modules/calib3d/src/fisheye.cpp
if fov is None:
with th.no_grad():
fov = estimate_fisheye_fov(D)
z = v_cam[:, :, 2:3]
z = th.where(z < 0, z.clamp(max=-1e-8), z.clamp(min=1e-8))
v_proj = v_cam[:, :, :2] / z
r = v_proj.pow(2).sum(-1).sqrt()
r = r.clamp(max=fov, min=1e-8 * th.ones_like(fov))
theta = th.atan(r)
theta_d = theta * (
1
+ D[:, 0:1] * theta.pow(2)
+ D[:, 1:2] * theta.pow(4)
+ D[:, 2:3] * theta.pow(6)
+ D[:, 3:4] * theta.pow(8)
)
r = th.where(r < 0, r.clamp(max=-1e-8), r.clamp(min=1e-8))
v_proj_dist = v_proj * (theta_d / r)[..., None]
v_pix_dist = _apply_matrix(focal, v_proj_dist) + princpt[:, None]
return v_pix_dist
def project_fisheye_distort_62(
v_cam: th.Tensor,
focal: th.Tensor,
princpt: th.Tensor,
D: th.Tensor,
fov: Optional[th.Tensor] = None,
lut_vector_field: Optional[th.Tensor] = None,
lut_spacing: Optional[th.Tensor] = None,
) -> th.Tensor:
"""Project camera-space points with the Fisheye62 model.
Fisheye62 uses six radial coefficients followed by two tangential
coefficients. If ``lut_vector_field`` is provided, a pixel-space correction
is sampled from the lookup table and added to the projected coordinates.
Args:
v_cam: Camera-space points with shape ``(N, V, 3)``.
focal: Camera intrinsic matrices with shape ``(N, 2, 2)``.
princpt: Principal points in pixels with shape ``(N, 2)``.
D: Fisheye62 coefficients ``(k0, k1, k2, k3, k4, k5, p0, p1)`` with
shape ``(N, 8)``.
fov: Optional maximum normalized radius with shape ``(N, 1)``.
lut_vector_field: Optional pixel offset table with shape
``(N, 2, H_lut, W_lut)``.
lut_spacing: Pixel spacing of the lookup table with shape ``(N, 2)``.
Required when ``lut_vector_field`` is provided.
Returns:
Distorted pixel coordinates with shape ``(N, V, 2)``.
"""
assert D.shape[1] == 8, (
f"Fisheye62 model requires 8 distortion parameters: {D.shape}"
)
if fov is None:
with th.no_grad():
fov = estimate_fisheye_fov(D)
z = v_cam[:, :, 2:3]
z = th.where(z < 0, z.clamp(max=-1e-8), z.clamp(min=1e-8))
v_proj = v_cam[:, :, :2] / z
r = v_proj.pow(2).sum(-1).sqrt()
r = r.clamp(max=fov, min=1e-8 * th.ones_like(fov))
theta = th.atan(r)
theta2 = theta * theta
k0 = D[:, 0].unsqueeze(1)
k1 = D[:, 1].unsqueeze(1)
k2 = D[:, 2].unsqueeze(1)
k3 = D[:, 3].unsqueeze(1)
k4 = D[:, 4].unsqueeze(1)
k5 = D[:, 5].unsqueeze(1)
p0 = D[:, 6].unsqueeze(1)
p1 = D[:, 7].unsqueeze(1)
theta3 = theta2 * theta
theta5 = theta2 * theta3
theta7 = theta2 * theta5
theta9 = theta2 * theta7
theta11 = theta2 * theta9
theta13 = theta2 * theta11
thetad = (
theta
+ k0 * theta3
+ k1 * theta5
+ k2 * theta7
+ k3 * theta9
+ k4 * theta11
+ k5 * theta13
)
r = th.where(r < 0, r.clamp(max=-1e-8), r.clamp(min=1e-8))
v_proj_dist = v_proj * (thetad / r)[..., None]
if fov is not None:
v_proj_dist = v_proj_dist.clamp(min=-fov[..., None], max=fov[..., None])
x_r = v_proj_dist[:, :, 0]
y_r = v_proj_dist[:, :, 1]
rRadial2 = x_r * x_r + y_r * y_r
x_t = (2 * x_r * x_r + rRadial2) * p0 + (2 * x_r * y_r) * p1
y_t = (2 * x_r * y_r) * p0 + (2 * y_r * y_r + rRadial2) * p1
pTangential = th.cat([x_t[..., None], y_t[..., None]], dim=-1)
v_proj_dist = v_proj_dist + pTangential
v_pix_dist = _apply_matrix(focal, v_proj_dist) + princpt[:, None]
if lut_vector_field is not None:
assert lut_spacing is not None, (
"lookup table spacing must be provided along with vector field"
)
# Normalize pixel position to [-1, 1] range for torch.nn.functional.grid_sample
normalized_pixel_pos = v_pix_dist / lut_spacing[:, None, :]
lut_col, lut_row = lut_vector_field.shape[2:4]
normalized_pixel_pos[..., 0] = (
normalized_pixel_pos[..., 0] / (lut_col - 1) * 2.0 - 1.0
)
normalized_pixel_pos[..., 1] = (
normalized_pixel_pos[..., 1] / (lut_row - 1) * 2.0 - 1.0
)
# do grid sample
pixel_offset = th.nn.functional.grid_sample(
input=lut_vector_field,
grid=normalized_pixel_pos.unsqueeze(1),
align_corners=True,
).permute(0, 2, 3, 1)[:, 0, ...]
# mask for pixels out of the lookup table, whose offset need to be set to 0.0
out_bound_mask = (
(normalized_pixel_pos[..., 0] < -1.0)
| (normalized_pixel_pos[..., 0] > 1.0)
| (normalized_pixel_pos[..., 1] < -1.0)
| (normalized_pixel_pos[..., 1] > 1.0)
)
pixel_offset[out_bound_mask] = 0.0
v_pix_dist = v_pix_dist + pixel_offset
return v_pix_dist
def estimate_rt_fov(D: Union[np.ndarray, th.Tensor]) -> th.Tensor:
"""Estimate a radial-tangential distortion field-of-view bound.
The returned value is the smallest positive radius where the radial
polynomial can stop being monotonic. If no such root exists, the bound is
infinity. This helper uses NumPy root finding and is not differentiable with
respect to ``D``.
Args:
D: Radial-tangential coefficients with shape ``(N, 4)`` or larger. Only
the first two radial coefficients are read.
Returns:
Maximum normalized radius with shape ``(N, 1)``.
"""
if th.is_tensor(D):
coefs = D.cpu().numpy()
else:
coefs = D
ones = np.ones_like(coefs[:, 0])
zeros = np.zeros_like(coefs[:, 0])
coefs = np.stack(
# pyrefly: ignore [bad-argument-type]
[
5 * coefs[:, 1],
zeros,
3 * coefs[:, 0],
zeros,
ones,
],
axis=-1,
)
fov = []
for coef in coefs:
roots = np.roots(coef)
real_valued = roots.real[abs(roots.imag) < 1e-5]
positive_roots = real_valued[real_valued > 0]
if len(positive_roots) == 0:
fov.append(np.inf)
else:
fov.append(positive_roots.min())
fov = np.asarray(fov, dtype=np.float32)[..., None]
if th.is_tensor(D):
# pyrefly: ignore [no-matching-overload]
fov = th.from_numpy(fov).to(D)
return fov
def _solve_monotonic_fisheye_fov(
derivative_coefs: np.ndarray, D: Union[np.ndarray, th.Tensor]
) -> th.Tensor:
"""Shared fisheye FOV solver. For each row of `derivative_coefs`
(polynomial coefficients in descending order), returns
`tan(min(smallest positive real root, pi/2))`. Falls back to `tan(pi/2)`
when the derivative has no positive real root.
Return shape is `[N, 1]`, dtype float32, moved to `D`'s device if `D` is
a tensor.
"""
fov = []
for coef in derivative_coefs:
roots = np.roots(coef)
real_valued = roots.real[abs(roots.imag) < 1e-5]
positive_roots = real_valued[real_valued > 0]
if len(positive_roots) == 0:
fov.append(np.pi / 2)
else:
fov.append(min(positive_roots.min(), np.pi / 2))
fov = np.asarray(np.tan(fov), dtype=np.float32)[..., None]
if th.is_tensor(D):
# pyrefly: ignore [no-matching-overload]
fov = th.from_numpy(fov).to(D)
return fov
def estimate_fisheye_fov(D: Union[np.ndarray, th.Tensor]) -> th.Tensor:
"""Estimate an OpenCV fisheye field-of-view bound.
The returned value is ``tan(theta)`` at the first positive point where the
fisheye polynomial can stop being monotonic, capped at ``pi / 2``. This
helper uses NumPy root finding and is not differentiable with respect to
``D``.
Args:
D: Fisheye coefficients with shape ``(N, 4)`` or larger. Only the first
four radial coefficients are read.
Returns:
Maximum normalized radius with shape ``(N, 1)``.
"""
if th.is_tensor(D):
coefs = D.cpu().numpy()
else:
coefs = D
ones = np.ones_like(coefs[:, 0])
zeros = np.zeros_like(coefs[:, 0])
coefs = np.stack(
# pyrefly: ignore [bad-argument-type]
[
9 * coefs[:, 3],
zeros,
7 * coefs[:, 2],
zeros,
5 * coefs[:, 1],
zeros,
3 * coefs[:, 0],
zeros,
ones,
],
axis=-1,
)
return _solve_monotonic_fisheye_fov(coefs, D)
def estimate_fisheye62_fov(D: Union[np.ndarray, th.Tensor]) -> th.Tensor:
"""Estimate a Fisheye62 field-of-view bound.
`estimate_fisheye_fov` truncates at k0..k3 (9th order), which misses the
first root in `[0, pi/2]` for calibrations whose `k4*theta^11 + k5*theta^13`
tail pushes the polynomial non-monotonic before pi/2. Callers that use the
fisheye62 distortion model (8 coefficients: k0..k5, p0, p1) should prefer
this function.
This helper uses NumPy root finding and is not differentiable with respect
to ``D``.
Args:
D: Fisheye62 coefficients with shape ``(N, 6)`` or larger. Only the
first six radial coefficients are read.
Returns:
Maximum normalized radius with shape ``(N, 1)``.
"""
if th.is_tensor(D):
coefs = D.cpu().numpy()
else:
coefs = D
assert coefs.shape[-1] >= 6, (
f"fisheye62 FOV requires at least 6 coefficients, got shape {coefs.shape}"
)
ones = np.ones_like(coefs[:, 0])
zeros = np.zeros_like(coefs[:, 0])
coefs = np.stack(
# pyrefly: ignore [bad-argument-type]
[
13 * coefs[:, 5],
zeros,
11 * coefs[:, 4],
zeros,
9 * coefs[:, 3],
zeros,
7 * coefs[:, 2],
zeros,
5 * coefs[:, 1],
zeros,
3 * coefs[:, 0],
zeros,
ones,
],
axis=-1,
)
return _solve_monotonic_fisheye_fov(coefs, D)
[docs]
@th.jit.ignore
def project_points(
v: th.Tensor,
campos: th.Tensor,
camrot: th.Tensor,
focal: th.Tensor,
princpt: th.Tensor,
distortion_mode: Optional[Union[List[str], str]] = None,
distortion_coeff: Optional[th.Tensor] = None,
fov: Optional[th.Tensor] = None,
lut_vector_field: Optional[th.Tensor] = None,
lut_spacing: Optional[th.Tensor] = None,
) -> Tuple[th.Tensor, th.Tensor]:
"""Project world-space vertices to DRTK pixel coordinates.
The world-to-camera convention is ``v_cam = camrot @ (v - campos)``.
``v_pix`` stores pixel ``x, y`` followed by camera-space ``z`` so it can be
passed directly to :func:`drtk.rasterize`. Vertices with ``z <= 0`` are
behind the camera and are rejected by rasterization.
``distortion_mode`` may be ``None``/``"pinhole"``,
``"radial-tangential"``, ``"fisheye"``, ``"fisheye62"``, or
``"fisheye62_lut"``. A per-batch list of modes is supported for pinhole,
radial-tangential, and fisheye batches; Fisheye62 modes are currently
supported only when ``distortion_mode`` is a single string. For Fisheye62
with ``fov``, vertices outside the valid radius get ``z = -1`` so any
triangle touching them is culled by the rasterizer.
Args:
v: World-space vertices with shape ``(N, V, 3)``.
campos: Camera positions with shape ``(N, 3)``.
camrot: World-to-camera rotation matrices with shape ``(N, 3, 3)``.
focal: Camera intrinsic matrices with shape ``(N, 2, 2)``.
princpt: Principal points in pixels with shape ``(N, 2)``.
distortion_mode: Optional distortion mode string or per-batch list of
mode strings.
distortion_coeff: Distortion coefficients. Shape depends on
``distortion_mode``.
fov: Optional maximum normalized radius with shape ``(N, 1)``.
lut_vector_field: Optional Fisheye62 LUT offset field with shape
``(N, 2, H_lut, W_lut)``.
lut_spacing: Optional Fisheye62 LUT spacing with shape ``(N, 2)``.
Returns:
``(v_pix, v_cam)``. Both tensors have shape ``(N, V, 3)``. ``v_pix``
stores ``(x_pixels, y_pixels, z_camera)``.
"""
if distortion_mode is not None:
assert distortion_coeff is not None, "Missing distortion coefficients."
v_cam = _apply_matrix(camrot, v - campos[:, None])
# Fall back to single distortion mode if all the distortion modes are the same.
if isinstance(distortion_mode, (list, tuple)):
modes = list(set(distortion_mode))
if len(modes) == 0:
distortion_mode = None
elif len(modes) == 1:
distortion_mode = modes[0]
if distortion_mode is None or distortion_mode == "pinhole":
v_pix = project_pinhole(v_cam, focal, princpt)
elif isinstance(distortion_mode, str):
assert distortion_coeff is not None
# Single distortion model
if distortion_mode == "radial-tangential":
v_pix = project_pinhole_distort_rt(
v_cam, focal, princpt, distortion_coeff, fov
)
elif distortion_mode == "fisheye":
v_pix = project_fisheye_distort(
v_cam, focal, princpt, distortion_coeff, fov
)
elif distortion_mode == "fisheye62" or distortion_mode == "fisheye62_lut":
v_pix = project_fisheye_distort_62(
v_cam=v_cam,
focal=focal,
princpt=princpt,
D=distortion_coeff,
fov=fov,
lut_vector_field=lut_vector_field,
lut_spacing=lut_spacing,
)
else:
raise ValueError(
f"Invalid distortion mode: {distortion_mode}. Valid options: {DISTORTION_MODES}."
)
elif isinstance(distortion_mode, (list, tuple)):
assert distortion_coeff is not None
# A mix of multiple distortion modes
modes = set(distortion_mode)
if not modes <= DISTORTION_MODES:
raise ValueError(
f"Invalid distortion mode: {distortion_mode}. Valid options: {DISTORTION_MODES}."
)
v_pix = th.empty_like(v_cam[..., :2])
if None in modes or "pinhole" in modes:
idx = th.tensor(
[mode is None or mode == "pinhole" for mode in distortion_mode],
device=v_pix.device,
)
v_pix[idx] = project_pinhole(v_cam[idx], focal[idx], princpt[idx])
if "radial-tangential" in modes:
idx = th.tensor(
[mode == "radial-tangential" for mode in distortion_mode],
device=v_pix.device,
)
v_pix[idx] = project_pinhole_distort_rt(
v_cam[idx],
focal[idx],
princpt[idx],
distortion_coeff[idx],
fov[idx] if fov is not None else None,
)
if "fisheye" in modes:
idx = th.tensor(
[mode == "fisheye" for mode in distortion_mode], device=v_pix.device
)
v_pix[idx] = project_fisheye_distort(
v_cam[idx],
focal[idx],
princpt[idx],
distortion_coeff[idx],
fov[idx] if fov is not None else None,
)
else:
raise ValueError(
f"Invalid distortion mode: {distortion_mode}. Valid options: {DISTORTION_MODES}."
)
z_cam = v_cam[:, :, 2:3]
# fisheye62's per-axis clamp + tangential terms silently project rays
# whose `r = ||(x/z, y/z)|| > fov` into the image, giving garbage pixel
# coordinates. Drive `z = -1` on those vertices so DRTK's rasterizer
# near-plane check culls any triangle that touches them.
if fov is not None and _any_mode_is_fisheye62(distortion_mode):
z_safe = th.where(
z_cam.abs() < 1e-8,
th.where(z_cam < 0, z_cam.clamp(max=-1e-8), z_cam.clamp(min=1e-8)),
z_cam,
)
r_raw = (v_cam[:, :, :2] / z_safe).pow(2).sum(-1, keepdim=True).sqrt()
outside_fov = r_raw > fov.view(-1, 1, 1)
# Mixed-mode batches (defensive): only cull batch elements whose
# distortion mode is fisheye62 / fisheye62_lut. The other per-element
# branches above handle their own out-of-FOV clamping. Note the list
# path for fisheye62 is not wired through the dispatch above today --
# this guard keeps the cull safe if that's added later.
if isinstance(distortion_mode, (list, tuple)):
fisheye_mask = th.tensor(
[m in _FISHEYE62_MODES for m in distortion_mode],
device=z_cam.device,
).view(-1, 1, 1)
outside_fov = outside_fov & fisheye_mask
z_cam = th.where(outside_fov, th.full_like(z_cam, -1.0), z_cam)
v_pix = th.cat((v_pix[:, :, 0:2], z_cam), dim=-1)
return v_pix, v_cam
[docs]
@th.jit.ignore
def project_points_grad(
v_grad: th.Tensor,
v: th.Tensor,
campos: th.Tensor,
camrot: th.Tensor,
focal: th.Tensor,
distortion_mode: Optional[Union[List[str], str]] = None,
distortion_coeff: Optional[th.Tensor] = None,
) -> th.Tensor:
"""Apply the pinhole projection Jacobian to world-space vertex deltas.
This computes ``d project_points(v) / d v @ v_grad`` for the undistorted
pinhole projection path. Distortion gradients are not implemented here; pass
``distortion_mode=None``.
Args:
v_grad: World-space vertex perturbations with shape ``(N, V, 3)``.
v: World-space vertices with shape ``(N, V, 3)``.
campos: Camera positions with shape ``(N, 3)``.
camrot: World-to-camera rotation matrices with shape ``(N, 3, 3)``.
focal: Camera intrinsic matrices with shape ``(N, 2, 2)``.
distortion_mode: Distortion is not implemented and must be ``None``.
distortion_coeff: Distortion is not implemented and must be ``None``.
Returns:
Pixel-space perturbations with shape ``(N, V, 2)``.
"""
if distortion_mode is not None:
assert distortion_coeff is not None, "Missing distortion coefficients."
# d v_cam = d (Rv + T) = Rdv
v_cam_grad = _apply_matrix(camrot, v_grad)
v_cam = _apply_matrix(camrot, v - campos[:, None])
if distortion_mode is None:
z = v_cam[:, :, 2:3]
z_grad = v_cam_grad[:, :, 2:3]
z = th.where(z < 0, z.clamp(max=-1e-8), z.clamp(min=1e-8))
# Using quotient rule:
# d (v_cam / z) = (d v_cam * z - v_cam * dz) / z^2
v_proj_grad = (v_cam_grad[:, :, 0:2] * z - v_cam[:, :, 0:2] * z_grad) / z**2.0
# d v_pix = d (Kv + cp) = Kdv
v_pix_grad = _apply_matrix(focal, v_proj_grad)
elif distortion_mode == "radial-tangential":
raise NotImplementedError
elif distortion_mode == "fisheye":
raise NotImplementedError
else:
raise ValueError(
f"Invalid distortion mode: {distortion_mode}. Valid options: {DISTORTION_MODES}."
)
return v_pix_grad