diff --git a/docs/model_components/visualize_cameras.ipynb b/docs/model_components/visualize_cameras.ipynb
index bcb15f758..4e7aaff58 100644
--- a/docs/model_components/visualize_cameras.ipynb
+++ b/docs/model_components/visualize_cameras.ipynb
@@ -52,13 +52,99 @@
},
{
"cell_type": "code",
- "execution_count": 10,
+ "execution_count": 2,
"metadata": {
"tags": [
"hide-input"
]
},
"outputs": [
+ {
+ "data": {
+ "text/html": [
+ " \n",
+ " "
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
{
"data": {
"application/vnd.plotly.v1+json": {
@@ -13732,9 +13818,9 @@
}
},
"text/html": [
- ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# OUTPUT_ONLY\n",
+ "height = 15\n",
+ "width = 15\n",
+ "\n",
+ "from plotly.subplots import make_subplots\n",
+ "\n",
+ "distortion_names = (\"k1\", \"k2\", \"p1\", \"k3\", \"k4\", \"p2\")\n",
+ "distortion_min = (0.01, 0.001, 0.01, 0.001, -0.05, -0.05)\n",
+ "distortion_max = (0.05, 0.05, 0.05, 0.05, 0.05, 0.05)\n",
+ "fig = make_subplots(rows=2, cols=3, vertical_spacing=0.1, subplot_titles=distortion_names)\n",
+ "\n",
+ "num_steps = 19\n",
+ "all_steps = []\n",
+ "for i in range(len(distortion_names)):\n",
+ " for step in torch.linspace(distortion_min[i], distortion_max[i], num_steps):\n",
+ " distortion_params = torch.zeros(6)\n",
+ " distortion_params[i] = step\n",
+ "\n",
+ " from nerfactory.cameras import utils\n",
+ "\n",
+ " coords = torch.meshgrid(torch.linspace(-1, 1, height), torch.linspace(-1, 1, width), indexing=\"ij\")\n",
+ " coords = torch.stack(coords, dim=-1)\n",
+ "\n",
+ " coords = utils.radial_and_tangential_undistort(coords, distortion_params)\n",
+ "\n",
+ " fig.add_trace(\n",
+ " go.Scatter(\n",
+ " x=coords[..., 0].flatten(),\n",
+ " y=coords[..., 1].flatten(),\n",
+ " mode=\"markers\",\n",
+ " marker=dict(\n",
+ " size=10,\n",
+ " ),\n",
+ " visible=False,\n",
+ " ),\n",
+ " row=i % 2 + 1,\n",
+ " col=i // 2 + 1,\n",
+ " )\n",
+ "\n",
+ " fig.data[num_steps * i + num_steps // 2].visible = True\n",
+ "\n",
+ "# Create and add slider\n",
+ "steps = []\n",
+ "for i in range(num_steps):\n",
+ " step = dict(method=\"update\", args=[{\"visible\": [False] * len(fig.data)}], label=\"\")\n",
+ " for d in range(len(distortion_names)):\n",
+ " step[\"args\"][0][\"visible\"][num_steps * d + i] = True\n",
+ " steps.append(step)\n",
+ "\n",
+ "sliders = [\n",
+ " dict(\n",
+ " active=num_steps // 2,\n",
+ " pad={\"t\": 10},\n",
+ " steps=steps,\n",
+ " len=0.5,\n",
+ " x=0.5,\n",
+ " xanchor=\"center\",\n",
+ " )\n",
+ "]\n",
+ "\n",
+ "webdocs_layout = go.Layout(\n",
+ " height=600,\n",
+ " sliders=sliders,\n",
+ " margin=dict(r=0, b=0, l=0, t=20),\n",
+ " hovermode=False,\n",
+ " showlegend=False,\n",
+ " paper_bgcolor=\"rgba(0,0,0,0)\",\n",
+ " plot_bgcolor=\"rgba(0,0,0,0)\",\n",
+ " font=dict(color=\"#a3a3a3\", size=18),\n",
+ ")\n",
+ "fig.for_each_yaxis(lambda x: x.update(range=[-1.3, 1.3], constrain=\"domain\", showgrid=False, visible=False))\n",
+ "fig.for_each_xaxis(lambda x: x.update(scaleanchor=\"y\", scaleratio=1, showgrid=False, visible=False))\n",
+ "fig.update_layout(webdocs_layout)\n",
+ "\n",
+ "fig.show()"
+ ]
}
],
"metadata": {
diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py
index e76c71a33..33c8ac654 100644
--- a/nerfactory/cameras/cameras.py
+++ b/nerfactory/cameras/cameras.py
@@ -163,19 +163,11 @@ class Cameras:
fx, fy = self.fx[camera_indices], self.fy[camera_indices]
cx, cy = self.cx, self.cy
- directions = torch.stack([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1)
- directions_x_offset = torch.stack([(x - cx + 1) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1)
- directions_y_offset = torch.stack([(x - cx) / fx, -(y - cy + 1) / fy, -torch.ones_like(x)], -1)
+ coord = torch.stack([(x - cx) / fx, -(y - cy) / fy], -1)
+ coord_x_offset = torch.stack([(x - cx + 1) / fx, -(y - cy) / fy], -1)
+ coord_y_offset = torch.stack([(x - cx) / fx, -(y - cy + 1) / fy], -1)
- directions_stack = torch.stack([directions, directions_x_offset, directions_y_offset], dim=0)
-
- c2w = self.camera_to_worlds[camera_indices]
- if camera_to_world_delta is not None:
- c2w = c2w + camera_to_world_delta
- rotation = c2w[..., :3, :3] # (..., 3, 3)
- directions_stack = torch.sum(
- directions_stack[..., None, :] * rotation, dim=-1
- ) # (..., 1, 3) * (..., 3, 3) -> (..., 3)
+ coord_stack = torch.stack([coord, coord_x_offset, coord_y_offset], dim=0)
distortion_params = None
if self.distortion_params is not None:
@@ -187,19 +179,33 @@ class Cameras:
if distortion_params is not None:
raise NotImplementedError("Camera distortion not implemented.")
- if self.camera_type == CameraType.FISHEYE:
- theta = torch.sqrt(torch.sum(directions_stack[..., :2] ** 2, dim=-1))
+ if self.camera_type == CameraType.PERSPECTIVE:
+ directions_stack = torch.stack(
+ [coord_stack[..., 0], coord_stack[..., 1], -torch.ones_like(coord_stack[..., 1])], dim=-1
+ )
+ elif self.camera_type == CameraType.FISHEYE:
+ theta = torch.sqrt(torch.sum(coord_stack**2, dim=-1))
theta = torch.clip(theta, 0.0, math.pi)
sin_theta = torch.sin(theta)
directions_stack = torch.stack(
[
- directions_stack[..., 0] * sin_theta / theta,
- directions_stack[..., 1] * sin_theta / theta,
- torch.cos(theta),
+ coord_stack[..., 0] * sin_theta / theta,
+ coord_stack[..., 1] * sin_theta / theta,
+ -torch.cos(theta),
],
dim=-1,
)
+ else:
+ raise ValueError(f"Camera type {self.camera_type} not supported.")
+
+ c2w = self.camera_to_worlds[camera_indices]
+ if camera_to_world_delta is not None:
+ c2w = c2w + camera_to_world_delta
+ rotation = c2w[..., :3, :3] # (..., 3, 3)
+ directions_stack = torch.sum(
+ directions_stack[..., None, :] * rotation, dim=-1
+ ) # (..., 1, 3) * (..., 3, 3) -> (..., 3)
directions_stack = normalize(directions_stack, dim=-1)
diff --git a/nerfactory/cameras/utils.py b/nerfactory/cameras/utils.py
index 56f726f2f..0007680e3 100644
--- a/nerfactory/cameras/utils.py
+++ b/nerfactory/cameras/utils.py
@@ -324,3 +324,115 @@ def viewmatrix(lookat, up, pos):
vec1 = normalize(torch.cross(vec2, vec0))
m = torch.stack([vec0, vec1, vec2, pos], 1)
return m
+
+
+# Addapted from MultiNeRF
+# https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509
+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.
+
+ 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)
+
+
+# Adapted from MultiNeRF:
+# https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474
+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":...],
+]:
+ """Auxiliary function of radial_and_tangential_undistort() that computes residuals and jacobians.
+
+ Args:
+ x: The updated x coordinates.
+ y: The updated y coordinates.
+ xd: The distorted x coordinates.
+ yd: The distorted y coordinates.
+ distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].
+
+ Returns:
+ The residuals (fx, fy) and jacobians (fx_x, fx_y, fy_x, fy_y).
+ """
+
+ k1 = distortion_params[..., 0]
+ k2 = distortion_params[..., 1]
+ k3 = distortion_params[..., 2]
+ k4 = distortion_params[..., 3]
+ p1 = distortion_params[..., 4]
+ p2 = distortion_params[..., 5]
+
+ # let r(x, y) = x^2 + y^2;
+ # d(x, y) = 1 + k1 * r(x, y) + k2 * r(x, y) ^2 + k3 * r(x, y)^3 +
+ # k4 * r(x, y)^4;
+ r = x * x + y * y
+ d = 1.0 + r * (k1 + r * (k2 + r * (k3 + r * k4)))
+
+ # The perfect projection is:
+ # xd = x * d(x, y) + 2 * p1 * x * y + p2 * (r(x, y) + 2 * x^2);
+ # yd = y * d(x, y) + 2 * p2 * x * y + p1 * (r(x, y) + 2 * y^2);
+ #
+ # Let's define
+ #
+ # fx(x, y) = x * d(x, y) + 2 * p1 * x * y + p2 * (r(x, y) + 2 * x^2) - xd;
+ # fy(x, y) = y * d(x, y) + 2 * p2 * x * y + p1 * (r(x, y) + 2 * y^2) - yd;
+ #
+ # We are looking for a solution that satisfies
+ # fx(x, y) = fy(x, y) = 0;
+ fx = d * x + 2 * p1 * x * y + p2 * (r + 2 * x * x) - xd
+ fy = d * y + 2 * p2 * x * y + p1 * (r + 2 * y * y) - yd
+
+ # Compute derivative of d over [x, y]
+ d_r = k1 + r * (2.0 * k2 + r * (3.0 * k3 + r * 4.0 * k4))
+ d_x = 2.0 * x * d_r
+ d_y = 2.0 * y * d_r
+
+ # Compute derivative of fx over x and y.
+ fx_x = d + d_x * x + 2.0 * p1 * y + 6.0 * p2 * x
+ fx_y = d_y * x + 2.0 * p1 * x + 2.0 * p2 * y
+
+ # Compute derivative of fy over x and y.
+ fy_x = d_x * y + 2.0 * p2 * y + 2.0 * p1 * x
+ fy_y = d + d_y * y + 2.0 * p2 * x + 6.0 * p1 * y
+
+ return fx, fy, fx_x, fx_y, fy_x, fy_y