tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Add PyTorch jit-ing to the undistort computation for speed (#776)

Browse Source
This commit is contained in:
Justin Kerr 2022-10-13 18:33:36 -07:00 committed by GitHub
parent a48a81fc48
commit 7ebe1f88bb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 48 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

101
nerfstudio/cameras/camera_utils.py
View File

@ -271,60 +271,14 @@ def get_distortion_params(
return torch.Tensor([k1, k2, k3, k4, p1, p2])
def radial_and_tangential_undistort(
coords: TensorType["num_coords":..., 2],
distortion_params: TensorType["num_coords":..., 6],
eps: float = 1e-9,
max_iterations=10,
) -> TensorType["num_coords":..., 2]:
"""Computes undistorted coords given opencv distortion parameters.
Addapted from MultiNeRF
https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509
Args:
coords: The distorted coordinates.
distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].
eps: The epsilon for the convergence.
max_iterations: The maximum number of iterations to perform.
Returns:
The undistorted coordinates.
"""
# Initialize from the distorted point.
x = torch.clone(coords[..., 0])
y = torch.clone(coords[..., 1])
for _ in range(max_iterations):
fx, fy, fx_x, fx_y, fy_x, fy_y = _compute_residual_and_jacobian(
x=x, y=y, xd=coords[..., 0], yd=coords[..., 1], distortion_params=distortion_params
)
denominator = fy_x * fx_y - fx_x * fy_y
x_numerator = fx * fy_y - fy * fx_y
y_numerator = fy * fx_x - fx * fy_x
step_x = torch.where(torch.abs(denominator) > eps, x_numerator / denominator, torch.zeros_like(denominator))
step_y = torch.where(torch.abs(denominator) > eps, y_numerator / denominator, torch.zeros_like(denominator))
x = x + step_x
y = y + step_y
return torch.stack([x, y], dim=-1)
@torch.jit.script
def _compute_residual_and_jacobian(
x: TensorType["num_coords":...],
y: TensorType["num_coords":...],
xd: TensorType["num_coords":...],
yd: TensorType["num_coords":...],
distortion_params: TensorType["num_coords":..., 6],
) -> Tuple[
TensorType["num_coords":...],
TensorType["num_coords":...],
TensorType["num_coords":...],
TensorType["num_coords":...],
TensorType["num_coords":...],
TensorType["num_coords":...],
]:
x: torch.Tensor,
y: torch.Tensor,
xd: torch.Tensor,
yd: torch.Tensor,
distortion_params: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,]:
"""Auxiliary function of radial_and_tangential_undistort() that computes residuals and jacobians.
Adapted from MultiNeRF:
https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474
@ -383,6 +337,47 @@ def _compute_residual_and_jacobian(
return fx, fy, fx_x, fx_y, fy_x, fy_y
@torch.jit.script
def radial_and_tangential_undistort(
coords: torch.Tensor,
distortion_params: torch.Tensor,
eps: float = 1e-9,
max_iterations: int = 10,
) -> torch.Tensor:
"""Computes undistorted coords given opencv distortion parameters.
Addapted from MultiNeRF
https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509
Args:
coords: The distorted coordinates.
distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].
eps: The epsilon for the convergence.
max_iterations: The maximum number of iterations to perform.
Returns:
The undistorted coordinates.
"""
# Initialize from the distorted point.
x = coords[..., 0]
y = coords[..., 1]
for _ in range(max_iterations):
fx, fy, fx_x, fx_y, fy_x, fy_y = _compute_residual_and_jacobian(
x=x, y=y, xd=coords[..., 0], yd=coords[..., 1], distortion_params=distortion_params
)
denominator = fy_x * fx_y - fx_x * fy_y
x_numerator = fx * fy_y - fy * fx_y
y_numerator = fy * fx_x - fx * fy_x
step_x = torch.where(torch.abs(denominator) > eps, x_numerator / denominator, torch.zeros_like(denominator))
step_y = torch.where(torch.abs(denominator) > eps, y_numerator / denominator, torch.zeros_like(denominator))
x = x + step_x
y = y + step_y
return torch.stack([x, y], dim=-1)
def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:
"""Compute the rotation matrix that rotates vector a to vector b.