From e186f7770920df074eb9dbe4ea6e3b687dc69dd8 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 31 Aug 2022 00:51:41 +0100 Subject: [PATCH] Add camera distortion (#312) --- docs/model_components/visualize_cameras.ipynb | 57367 +++++++++++++++- nerfactory/cameras/cameras.py | 40 +- nerfactory/cameras/utils.py | 112 + 3 files changed, 57095 insertions(+), 424 deletions(-) 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