# 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"}
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 to pixel-space points with camera
intrinsics.
v_cam: N x V x 3
focal: N x 2 x 2
princpt: N x 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 = (focal[:, None] @ v_proj[..., None])[..., 0] + 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 to distorted pixel-space using the radial
and tangential model (4 parameters).
v_cam: N x V x 3
focal: N x 2 x 2
princpt: N x 2
D: N x 4
fov: N x 1
"""
# 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 = (focal[:, None] @ v_proj_dist[..., None])[..., 0] + 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 to distort pixel-space points using the
fisheye distortion model.
v_cam: N x V x 3
focal: N x 2 x 2
princpt: N x 2
D: N x 4
fov: N x 1
"""
# 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 = (focal[:, None] @ v_proj_dist[..., None])[..., 0] + 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 to distort pixel-space points using the
Fisheye62 distortion model. See the perception camera model implementation
where this was copied from here: https://fburl.com/code/oqpu8xdm
v_cam: N x V x 3
focal: N x 2 x 2
princpt: N x 2
D: N x 8
fov: N x 1
lut_vector_field: N x 2 x H_lut x W_lut
lut_spacing: N x 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 = (focal[:, None] @ v_proj_dist[..., None])[..., 0] + 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 the maximum field of view based on the assumption that the 5th order
polynomial for fish-eye effect is non-decreasing.
D: N x 4
"""
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(
[
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):
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):
fov = th.from_numpy(fov).to(D)
return fov
def estimate_fisheye_fov(D: Union[np.ndarray, th.Tensor]) -> th.Tensor:
"""Estimate the maximum field of view based on the assumption that the 9th order
polynomial is non-decreasing.
D: N x 4
"""
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(
[
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 the maximum field of view based on the assumption that the 13th
order fisheye62 polynomial (k0..k5) is non-decreasing.
`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.
D: N x >= 6 (only the first six entries are read)
"""
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(
[
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 3D world-space vertices to pixel-space, optionally applying a
distortion model with provided coefficients.
Returns v_pix, v_cam, both N x V x 3 since we preserve the camera-space
Z-coordinate for v_pix.
v: N x V x 3
camrot: N x 3 x 3
campos: N x 3
focal: N x 2 x 2
princpt: N x 2
distortion_coeff: N x 4
fov: N x 1
"""
if distortion_mode is not None:
assert distortion_coeff is not None, "Missing distortion coefficients."
v_cam = (camrot[:, None] @ (v - campos[:, None])[..., None])[..., 0]
# 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:
"""Computes the gradient of projected (pixel-space) vertex positions with
respect to the 3D world-space vertex positions given the gradient of the 3D
world-space vertex positions.
project_points_grad(dv, v) = d project_points(v) / dv * dv
Args:
v_grad: Gradient of 3D world-space vertices. Shape: N x V x 3
v: 3D world-space vertices. Shape: N x V x 3
camrot: Camera rotation. Shape: N x 3 x 3
camrot: Camera position. Shape: N x 3
focal: Focal length. Shape: N x 2 x 2
distortion_mode: Distortion currently not implemented and must be None.
distortion_coeff: Distortion currently not implemented and must be None.
Returns:
Gradient of 2D pixel-space vertices: N x V x 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 = (camrot[:, None] @ v_grad[..., None])[..., 0]
v_cam = (camrot[:, None] @ (v - campos[:, None])[..., None])[..., 0]
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 = (focal[:, None] @ v_proj_grad[..., None])[..., 0]
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