From 7ebe1f88bb2ba127d8684ea4f61646bc1f12ddb7 Mon Sep 17 00:00:00 2001
From: Justin Kerr <justin.g.kerr@gmail.com>
Date: Thu, 13 Oct 2022 18:33:36 -0700
Subject: [PATCH] Add PyTorch jit-ing to the undistort computation for speed
 (#776)

---
 nerfstudio/cameras/camera_utils.py | 101 ++++++++++++++---------------
 1 file changed, 48 insertions(+), 53 deletions(-)

diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py
index 973dbec7b..6fd85b358 100644
--- a/nerfstudio/cameras/camera_utils.py
+++ b/nerfstudio/cameras/camera_utils.py
@@ -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.