diff --git a/README.md b/README.md index ab1874768..1ea4ca8fa 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,7 @@ Have feedback? We'd love for you to fill out our [Nerfstudio Feedback Form](http We hope nerfstudio enables you to build faster :hammer: learn together :books: and contribute to our NeRF community :sparkling_heart:. ## Sponsors + Sponsors of this work includes [Luma AI](https://lumalabs.ai/) and the [BAIR commons](https://bcommons.berkeley.edu/home).

@@ -97,7 +98,6 @@ Sponsors of this work includes [Luma AI](https://lumalabs.ai/) and the [BAIR com

- # Quickstart The quickstart will help you get started with the default vanilla NeRF trained on the classic Blender Lego scene. @@ -233,18 +233,19 @@ ns-export pointcloud --help Using an existing dataset is great, but likely you want to use your own data! We support various methods for using your own data. Before it can be used in nerfstudio, the camera location and orientations must be determined and then converted into our format using `ns-process-data`. We rely on external tools for this, instructions and information can be found in the documentation. -| Data | Capture Device | Requirements | `ns-process-data` Speed | -| ---------------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------------------------- | ----------------------- | -| 📷 [Images](https://docs.nerf.studio/quickstart/custom_dataset.html#images-or-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | -| 📹 [Video](https://docs.nerf.studio/quickstart/custom_dataset.html#images-or-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | -| 🌎 [360 Data](https://docs.nerf.studio/quickstart/custom_dataset.html#data-equirectangular) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | -| 📱 [Polycam](https://docs.nerf.studio/quickstart/custom_dataset.html#polycam-capture) | IOS with LiDAR | [Polycam App](https://poly.cam/) | 🐇 | -| 📱 [KIRI Engine](https://docs.nerf.studio/quickstart/custom_dataset.html#kiri-engine-capture) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | 🐇 | -| 📱 [Record3D](https://docs.nerf.studio/quickstart/custom_dataset.html#record3d-capture) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | -| 🖥 [Metashape](https://docs.nerf.studio/quickstart/custom_dataset.html#metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | -| 🖥 [RealityCapture](https://docs.nerf.studio/quickstart/custom_dataset.html#realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | -| 🖥 [ODM](https://docs.nerf.studio/quickstart/custom_dataset.html#ODM) | Any | [ODM](https://github.com/OpenDroneMap/ODM) | 🐇 | -| 🛠 [Custom](https://docs.nerf.studio/quickstart/data_conventions.html) | Any | Camera Poses | 🐇 | +| Data | Capture Device | Requirements | `ns-process-data` Speed | +| --------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------------------------- | ----------------------- | +| 📷 [Images](https://docs.nerf.studio/quickstart/custom_dataset.html#images-or-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | +| 📹 [Video](https://docs.nerf.studio/quickstart/custom_dataset.html#images-or-video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | +| 🌎 [360 Data](https://docs.nerf.studio/quickstart/custom_dataset.html#data-equirectangular) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | +| 📱 [Polycam](https://docs.nerf.studio/quickstart/custom_dataset.html#polycam-capture) | IOS with LiDAR | [Polycam App](https://poly.cam/) | 🐇 | +| 📱 [KIRI Engine](https://docs.nerf.studio/quickstart/custom_dataset.html#kiri-engine-capture) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | 🐇 | +| 📱 [Record3D](https://docs.nerf.studio/quickstart/custom_dataset.html#record3d-capture) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | +| 🖥 [Metashape](https://docs.nerf.studio/quickstart/custom_dataset.html#metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | +| 🖥 [RealityCapture](https://docs.nerf.studio/quickstart/custom_dataset.html#realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | +| 🖥 [ODM](https://docs.nerf.studio/quickstart/custom_dataset.html#ODM) | Any | [ODM](https://github.com/OpenDroneMap/ODM) | 🐇 | +| 👓 [Aria](https://docs.nerf.studio/quickstart/custom_dataset.html#Aria) | Aria glasses | [Project Aria](https://projectaria.com/) | 🐇 | +| 🛠 [Custom](https://docs.nerf.studio/quickstart/data_conventions.html) | Any | Camera Poses | 🐇 | ## 5. Advanced Options @@ -276,25 +277,25 @@ And that's it for getting started with the basics of nerfstudio. If you're interested in learning more on how to create your own pipelines, develop with the viewer, run benchmarks, and more, please check out some of the quicklinks below or visit our [documentation](https://docs.nerf.studio/) directly. -| Section | Description | -| -------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | +| Section | Description | +| ---------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | | [Documentation](https://docs.nerf.studio/) | Full API documentation and tutorials | -| [Viewer](https://viewer.nerf.studio/) | Home page for our web viewer | -| 🎒 **Educational** | +| [Viewer](https://viewer.nerf.studio/) | Home page for our web viewer | +| 🎒 **Educational** | | [Model Descriptions](https://docs.nerf.studio/nerfology/methods/index.html) | Description of all the models supported by nerfstudio and explanations of component parts. | | [Component Descriptions](https://docs.nerf.studio/nerfology/model_components/index.html) | Interactive notebooks that explain notable/commonly used modules in various models. | -| 🏃 **Tutorials** | +| 🏃 **Tutorials** | | [Getting Started](https://docs.nerf.studio/quickstart/installation.html) | A more in-depth guide on how to get started with nerfstudio from installation to contributing. | | [Using the Viewer](https://docs.nerf.studio/quickstart/viewer_quickstart.html) | A quick demo video on how to navigate the viewer. | -| [Using Record3D](https://www.youtube.com/watch?v=XwKq7qDQCQk) | Demo video on how to run nerfstudio without using COLMAP. | -| 💻 **For Developers** | +| [Using Record3D](https://www.youtube.com/watch?v=XwKq7qDQCQk) | Demo video on how to run nerfstudio without using COLMAP. | +| 💻 **For Developers** | | [Creating pipelines](https://docs.nerf.studio/developer_guides/pipelines/index.html) | Learn how to easily build new neural rendering pipelines by using and/or implementing new modules. | | [Creating datasets](https://docs.nerf.studio/quickstart/custom_dataset.html) | Have a new dataset? Learn how to run it with nerfstudio. | | [Contributing](https://docs.nerf.studio/reference/contributing.html) | Walk-through for how you can start contributing now. | -| 💖 **Community** | -| [Discord](https://discord.gg/uMbNqcraFc) | Join our community to discuss more. We would love to hear from you! | -| [Twitter](https://twitter.com/nerfstudioteam) | Follow us on Twitter @nerfstudioteam to see cool updates and announcements | -| [Feedback Form](TODO) | We welcome any feedback! This is our chance to learn what you all are using Nerfstudio for. | +| 💖 **Community** | +| [Discord](https://discord.gg/uMbNqcraFc) | Join our community to discuss more. We would love to hear from you! | +| [Twitter](https://twitter.com/nerfstudioteam) | Follow us on Twitter @nerfstudioteam to see cool updates and announcements | +| [Feedback Form](TODO) | We welcome any feedback! This is our chance to learn what you all are using Nerfstudio for. | # Supported Features diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index 46aaa907d..fddbe1d56 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -10,7 +10,7 @@ ns-process-data {video,images,polycam,record3d} --data {DATA_PATH} --output-dir A full set of arguments can be found {doc}`here`. -We Currently support the following custom data types: +We currently support the following custom data types: | Data | Capture Device | Requirements | `ns-process-data` Speed | | ----------------------------- | -------------- | ----------------------------------------------- | ----------------------- | | 📷 [Images](images_and_video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 | @@ -22,6 +22,7 @@ We Currently support the following custom data types: | 🖥 [Metashape](metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | | 🖥 [RealityCapture](realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | | 🖥 [ODM](odm) | Any | [ODM](https://github.com/OpenDroneMap/ODM) | 🐇 | +| 👓 [Aria](aria) | Aria glasses | [Project Aria](https://projectaria.com/) | 🐇 | (images_and_video)= @@ -348,6 +349,31 @@ ns-process-data odm --data /path/to/dataset --output-dir {output directory} ns-train nerfacto --data {output directory} ``` +(aria)= + +## Aria + +1. Install projectaria_tools: + +```bash +conda activate nerfstudio +pip install projectaria-tools'[all]' +``` + +2. Download a VRS file from Project Aria glasses, and run Machine Perception Services to extract poses. + +3. Convert to nerfstudio format. + +```bash +ns-process-data aria --vrs-file /path/to/vrs/file --mps-data-dir /path/to/mps/data --output-dir {output directory} +``` + +4. Train! + +```bash +ns-train nerfacto --data {output directory} +``` + (360_data)= ## 360 Data (Equirectangular) @@ -459,4 +485,4 @@ If the depth of the scene is unviewable and looks too close or expanded when vie - The IPD can be modified in the `cameras.py` script as the variable `vr_ipd` (default is 64 mm). - Compositing with Blender Objects and VR180 or ODS Renders - Configure the Blender camera as panoramic and equirectangular. For the VR180 Blender camera, set the panoramic longitude min and max to -90 and 90. - - Change the Stereoscopy mode to "Parallel" set the Interocular Distance to 0.064 m. \ No newline at end of file + - Change the Stereoscopy mode to "Parallel" set the Interocular Distance to 0.064 m. diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index ffb557589..445656623 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -25,7 +25,6 @@ from jaxtyping import Float from numpy.typing import NDArray from torch import Tensor - _EPS = np.finfo(float).eps * 4.0 @@ -622,3 +621,224 @@ def auto_orient_and_center_poses( raise ValueError(f"Unknown value for method: {method}") return oriented_poses, transform + + +@torch.jit.script +def fisheye624_project(xyz, params): + """ + Batched implementation of the FisheyeRadTanThinPrism (aka Fisheye624) camera + model project() function. + Inputs: + xyz: BxNx3 tensor of 3D points to be projected + params: Bx16 tensor of Fisheye624 parameters formatted like this: + [f_u f_v c_u c_v {k_0 ... k_5} {p_0 p_1} {s_0 s_1 s_2 s_3}] + or Bx15 tensor of Fisheye624 parameters formatted like this: + [f c_u c_v {k_0 ... k_5} {p_0 p_1} {s_0 s_1 s_2 s_3}] + Outputs: + uv: BxNx2 tensor of 2D projections of xyz in image plane + Model for fisheye cameras with radial, tangential, and thin-prism distortion. + This model allows fu != fv. + Specifically, the model is: + uvDistorted = [x_r] + tangentialDistortion + thinPrismDistortion + [y_r] + proj = diag(fu,fv) * uvDistorted + [cu;cv]; + where: + a = x/z, b = y/z, r = (a^2+b^2)^(1/2) + th = atan(r) + cosPhi = a/r, sinPhi = b/r + [x_r] = (th+ k0 * th^3 + k1* th^5 + ...) [cosPhi] + [y_r] [sinPhi] + the number of terms in the series is determined by the template parameter numK. + tangentialDistortion = [(2 x_r^2 + rd^2)*p_0 + 2*x_r*y_r*p_1] + [(2 y_r^2 + rd^2)*p_1 + 2*x_r*y_r*p_0] + where rd^2 = x_r^2 + y_r^2 + thinPrismDistortion = [s0 * rd^2 + s1 rd^4] + [s2 * rd^2 + s3 rd^4] + Author: Daniel DeTone (ddetone@meta.com) + """ + + assert xyz.ndim == 3 + assert params.ndim == 2 + assert params.shape[-1] == 16 or params.shape[-1] == 15, "This model allows fx != fy" + eps = 1e-9 + B, N = xyz.shape[0], xyz.shape[1] + + # Radial correction. + z = xyz[:, :, 2].reshape(B, N, 1) + z = torch.where(torch.abs(z) < eps, eps * torch.sign(z), z) + ab = xyz[:, :, :2] / z + r = torch.norm(ab, dim=-1, p=2, keepdim=True) + th = torch.atan(r) + th_divr = torch.where(r < eps, torch.ones_like(ab), ab / r) + th_k = th.reshape(B, N, 1).clone() + for i in range(6): + th_k = th_k + params[:, -12 + i].reshape(B, 1, 1) * torch.pow(th, 3 + i * 2) + xr_yr = th_k * th_divr + uv_dist = xr_yr + + # Tangential correction. + p0 = params[:, -6].reshape(B, 1) + p1 = params[:, -5].reshape(B, 1) + xr = xr_yr[:, :, 0].reshape(B, N) + yr = xr_yr[:, :, 1].reshape(B, N) + xr_yr_sq = torch.square(xr_yr) + xr_sq = xr_yr_sq[:, :, 0].reshape(B, N) + yr_sq = xr_yr_sq[:, :, 1].reshape(B, N) + rd_sq = xr_sq + yr_sq + uv_dist_tu = uv_dist[:, :, 0] + ((2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1) + uv_dist_tv = uv_dist[:, :, 1] + ((2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0) + uv_dist = torch.stack([uv_dist_tu, uv_dist_tv], dim=-1) # Avoids in-place complaint. + + # Thin Prism correction. + s0 = params[:, -4].reshape(B, 1) + s1 = params[:, -3].reshape(B, 1) + s2 = params[:, -2].reshape(B, 1) + s3 = params[:, -1].reshape(B, 1) + rd_4 = torch.square(rd_sq) + uv_dist[:, :, 0] = uv_dist[:, :, 0] + (s0 * rd_sq + s1 * rd_4) + uv_dist[:, :, 1] = uv_dist[:, :, 1] + (s2 * rd_sq + s3 * rd_4) + + # Finally, apply standard terms: focal length and camera centers. + if params.shape[-1] == 15: + fx_fy = params[:, 0].reshape(B, 1, 1) + cx_cy = params[:, 1:3].reshape(B, 1, 2) + else: + fx_fy = params[:, 0:2].reshape(B, 1, 2) + cx_cy = params[:, 2:4].reshape(B, 1, 2) + result = uv_dist * fx_fy + cx_cy + + return result + + +# Core implementation of fisheye 624 unprojection. More details are documented here: +# https://facebookresearch.github.io/projectaria_tools/docs/tech_insights/camera_intrinsic_models#the-fisheye62-model +@torch.jit.script +def fisheye624_unproject_helper(uv, params, max_iters: int = 5): + """ + Batched implementation of the FisheyeRadTanThinPrism (aka Fisheye624) camera + model. There is no analytical solution for the inverse of the project() + function so this solves an optimization problem using Newton's method to get + the inverse. + Inputs: + uv: BxNx3 tensor of 2D pixels to be projected + params: Bx16 tensor of Fisheye624 parameters formatted like this: + [f_u f_v c_u c_v {k_0 ... k_5} {p_0 p_1} {s_0 s_1 s_2 s_3}] + or Bx15 tensor of Fisheye624 parameters formatted like this: + [f c_u c_v {k_0 ... k_5} {p_0 p_1} {s_0 s_1 s_2 s_3}] + Outputs: + xyz: BxNx3 tensor of 3D rays of uv points with z = 1. + Model for fisheye cameras with radial, tangential, and thin-prism distortion. + This model assumes fu=fv. This unproject function holds that: + X = unproject(project(X)) [for X=(x,y,z) in R^3, z>0] + and + x = project(unproject(s*x)) [for s!=0 and x=(u,v) in R^2] + Author: Daniel DeTone (ddetone@meta.com) + """ + + assert uv.ndim == 3, "Expected batched input shaped BxNx3" + assert params.ndim == 2 + assert params.shape[-1] == 16 or params.shape[-1] == 15, "This model allows fx != fy" + eps = 1e-6 + B, N = uv.shape[0], uv.shape[1] + + if params.shape[-1] == 15: + fx_fy = params[:, 0].reshape(B, 1, 1) + cx_cy = params[:, 1:3].reshape(B, 1, 2) + else: + fx_fy = params[:, 0:2].reshape(B, 1, 2) + cx_cy = params[:, 2:4].reshape(B, 1, 2) + + uv_dist = (uv - cx_cy) / fx_fy + + # Compute xr_yr using Newton's method. + xr_yr = uv_dist.clone() # Initial guess. + for _ in range(max_iters): + uv_dist_est = xr_yr.clone() + # Tangential terms. + p0 = params[:, -6].reshape(B, 1) + p1 = params[:, -5].reshape(B, 1) + xr = xr_yr[:, :, 0].reshape(B, N) + yr = xr_yr[:, :, 1].reshape(B, N) + xr_yr_sq = torch.square(xr_yr) + xr_sq = xr_yr_sq[:, :, 0].reshape(B, N) + yr_sq = xr_yr_sq[:, :, 1].reshape(B, N) + rd_sq = xr_sq + yr_sq + uv_dist_est[:, :, 0] = uv_dist_est[:, :, 0] + ((2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1) + uv_dist_est[:, :, 1] = uv_dist_est[:, :, 1] + ((2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0) + # Thin Prism terms. + s0 = params[:, -4].reshape(B, 1) + s1 = params[:, -3].reshape(B, 1) + s2 = params[:, -2].reshape(B, 1) + s3 = params[:, -1].reshape(B, 1) + rd_4 = torch.square(rd_sq) + uv_dist_est[:, :, 0] = uv_dist_est[:, :, 0] + (s0 * rd_sq + s1 * rd_4) + uv_dist_est[:, :, 1] = uv_dist_est[:, :, 1] + (s2 * rd_sq + s3 * rd_4) + # Compute the derivative of uv_dist w.r.t. xr_yr. + duv_dist_dxr_yr = uv.new_ones(B, N, 2, 2) + duv_dist_dxr_yr[:, :, 0, 0] = 1.0 + 6.0 * xr_yr[:, :, 0] * p0 + 2.0 * xr_yr[:, :, 1] * p1 + offdiag = 2.0 * (xr_yr[:, :, 0] * p1 + xr_yr[:, :, 1] * p0) + duv_dist_dxr_yr[:, :, 0, 1] = offdiag + duv_dist_dxr_yr[:, :, 1, 0] = offdiag + duv_dist_dxr_yr[:, :, 1, 1] = 1.0 + 6.0 * xr_yr[:, :, 1] * p1 + 2.0 * xr_yr[:, :, 0] * p0 + xr_yr_sq_norm = xr_yr_sq[:, :, 0] + xr_yr_sq[:, :, 1] + temp1 = 2.0 * (s0 + 2.0 * s1 * xr_yr_sq_norm) + duv_dist_dxr_yr[:, :, 0, 0] = duv_dist_dxr_yr[:, :, 0, 0] + (xr_yr[:, :, 0] * temp1) + duv_dist_dxr_yr[:, :, 0, 1] = duv_dist_dxr_yr[:, :, 0, 1] + (xr_yr[:, :, 1] * temp1) + temp2 = 2.0 * (s2 + 2.0 * s3 * xr_yr_sq_norm) + duv_dist_dxr_yr[:, :, 1, 0] = duv_dist_dxr_yr[:, :, 1, 0] + (xr_yr[:, :, 0] * temp2) + duv_dist_dxr_yr[:, :, 1, 1] = duv_dist_dxr_yr[:, :, 1, 1] + (xr_yr[:, :, 1] * temp2) + # Compute 2x2 inverse manually here since torch.inverse() is very slow. + # Because this is slow: inv = duv_dist_dxr_yr.inverse() + # About a 10x reduction in speed with above line. + mat = duv_dist_dxr_yr.reshape(-1, 2, 2) + a = mat[:, 0, 0].reshape(-1, 1, 1) + b = mat[:, 0, 1].reshape(-1, 1, 1) + c = mat[:, 1, 0].reshape(-1, 1, 1) + d = mat[:, 1, 1].reshape(-1, 1, 1) + det = 1.0 / ((a * d) - (b * c)) + top = torch.cat([d, -b], dim=2) + bot = torch.cat([-c, a], dim=2) + inv = det * torch.cat([top, bot], dim=1) + inv = inv.reshape(B, N, 2, 2) + # Manually compute 2x2 @ 2x1 matrix multiply. + # Because this is slow: step = (inv @ (uv_dist - uv_dist_est)[..., None])[..., 0] + diff = uv_dist - uv_dist_est + a = inv[:, :, 0, 0] + b = inv[:, :, 0, 1] + c = inv[:, :, 1, 0] + d = inv[:, :, 1, 1] + e = diff[:, :, 0] + f = diff[:, :, 1] + step = torch.stack([a * e + b * f, c * e + d * f], dim=-1) + # Newton step. + xr_yr = xr_yr + step + + # Compute theta using Newton's method. + xr_yr_norm = xr_yr.norm(p=2, dim=2).reshape(B, N, 1) + th = xr_yr_norm.clone() + for _ in range(max_iters): + th_radial = uv.new_ones(B, N, 1) + dthd_th = uv.new_ones(B, N, 1) + for k in range(6): + r_k = params[:, -12 + k].reshape(B, 1, 1) + th_radial = th_radial + (r_k * torch.pow(th, 2 + k * 2)) + dthd_th = dthd_th + ((3.0 + 2.0 * k) * r_k * torch.pow(th, 2 + k * 2)) + th_radial = th_radial * th + step = (xr_yr_norm - th_radial) / dthd_th + # handle dthd_th close to 0. + step = torch.where(dthd_th.abs() > eps, step, torch.sign(step) * eps * 10.0) + th = th + step + # Compute the ray direction using theta and xr_yr. + close_to_zero = torch.logical_and(th.abs() < eps, xr_yr_norm.abs() < eps) + ray_dir = torch.where(close_to_zero, xr_yr, torch.tan(th) / xr_yr_norm * xr_yr) + ray = torch.cat([ray_dir, uv.new_ones(B, N, 1)], dim=2) + return ray + + +# unproject 2D point to 3D with fisheye624 model +def fisheye624_unproject(coords: torch.Tensor, distortion_params: torch.Tensor) -> torch.Tensor: + dirs = fisheye624_unproject_helper(coords.unsqueeze(0), distortion_params[0].unsqueeze(0)) + # correct for camera space differences: + dirs[..., 1] = -dirs[..., 1] + dirs[..., 2] = -dirs[..., 2] + return dirs diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 467c8fcc0..46ec42bdc 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -32,7 +32,7 @@ import nerfstudio.utils.math import nerfstudio.utils.poses as pose_utils from nerfstudio.cameras import camera_utils from nerfstudio.cameras.rays import RayBundle -from nerfstudio.data.scene_box import SceneBox, OrientedBox +from nerfstudio.data.scene_box import OrientedBox, SceneBox from nerfstudio.utils.tensor_dataclass import TensorDataclass TORCH_DEVICE = Union[torch.device, str] @@ -48,6 +48,7 @@ class CameraType(Enum): OMNIDIRECTIONALSTEREO_R = auto() VR180_L = auto() VR180_R = auto() + FISHEYE624 = auto() CAMERA_MODEL_TO_TYPE = { @@ -62,6 +63,7 @@ CAMERA_MODEL_TO_TYPE = { "OMNIDIRECTIONALSTEREO_R": CameraType.OMNIDIRECTIONALSTEREO_R, "VR180_L": CameraType.VR180_L, "VR180_R": CameraType.VR180_R, + "FISHEYE624": CameraType.FISHEYE624, } @@ -79,7 +81,7 @@ class Cameras(TensorDataclass): cy: Principal point y width: Image width height: Image height - distortion_params: OpenCV 6 radial distortion coefficients + distortion_params: distortion coefficients (OpenCV 6 radial or 6-2-4 radial, tangential, thin-prism for Fisheye624) camera_type: Type of camera model. This will be an int corresponding to the CameraType enum. times: Timestamps for each camera metadata: Additional metadata or data needed for interpolation, will mimic shape of the cameras @@ -629,8 +631,8 @@ class Cameras(TensorDataclass): assert coord_stack.shape == (3,) + num_rays_shape + (2,) # Undistorts our images according to our distortion parameters + distortion_params = None if not disable_distortion: - distortion_params = None if self.distortion_params is not None: distortion_params = self.distortion_params[true_indices] if distortion_params_delta is not None: @@ -832,6 +834,34 @@ class Cameras(TensorDataclass): # assign final camera origins c2w[..., :3, 3] = vr180_origins + elif CameraType.FISHEYE624.value in cam_types: + mask = (self.camera_type[true_indices] == CameraType.FISHEYE624.value).squeeze(-1) # (num_rays) + coord_mask = torch.stack([mask, mask, mask], dim=0) + + # fisheye624 requires pixel coordinates to unproject, so we need to recomput the offsets in pixel coords. + pcoord = torch.stack([x, y], -1) # (num_rays, 2) + pcoord_x_offset = torch.stack([x + 1, y], -1) # (num_rays, 2) + pcoord_y_offset = torch.stack([x, y + 1], -1) # (num_rays, 2) + + # Stack image coordinates and image coordinates offset by 1, check shapes too + pcoord_stack = torch.stack([pcoord, pcoord_x_offset, pcoord_y_offset], dim=0) # (3, num_rays, 2) + + assert distortion_params is not None + masked_coords = pcoord_stack[coord_mask, :] + # The fisheye unprojection does not rely on planar/pinhold unprojection, thus the method needs + # to access the focal length and principle points directly. + camera_params = torch.cat( + [ + fx[mask].unsqueeze(1), + fy[mask].unsqueeze(1), + cx[mask].unsqueeze(1), + cy[mask].unsqueeze(1), + distortion_params[mask, :], + ], + dim=1, + ) + directions_stack[coord_mask] = camera_utils.fisheye624_unproject(masked_coords, camera_params) + else: raise ValueError(f"Camera type {cam} not supported.") diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 7971f46e5..47144c00e 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -475,8 +475,14 @@ class VanillaDataManager(DataManager, Generic[TDataset]): is_equirectangular = (dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value).all() if is_equirectangular.any(): CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") + + fisheye_crop_radius = ( + None if dataset.cameras.metadata is None else dataset.cameras.metadata["fisheye_crop_radius"] + ) return self.config.pixel_sampler.setup( - is_equirectangular=is_equirectangular, num_rays_per_batch=num_rays_per_batch + is_equirectangular=is_equirectangular, + num_rays_per_batch=num_rays_per_batch, + fisheye_crop_radius=fisheye_crop_radius, ) def setup_train(self): diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index eaeda0bce..9b4c2951a 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -104,10 +104,11 @@ class Nerfstudio(DataParser): height_fixed = "h" in meta width_fixed = "w" in meta distort_fixed = False - for distort_key in ["k1", "k2", "k3", "p1", "p2"]: + for distort_key in ["k1", "k2", "k3", "p1", "p2", "distortion_params"]: if distort_key in meta: distort_fixed = True break + fisheye_crop_radius = meta.get("fisheye_crop_radius", None) fx = [] fy = [] cx = [] @@ -149,7 +150,9 @@ class Nerfstudio(DataParser): width.append(int(frame["w"])) if not distort_fixed: distort.append( - camera_utils.get_distortion_params( + torch.tensor(frame["distortion_params"], dtype=torch.float32) + if "distortion_params" in frame + else camera_utils.get_distortion_params( k1=float(frame["k1"]) if "k1" in frame else 0.0, k2=float(frame["k2"]) if "k2" in frame else 0.0, k3=float(frame["k3"]) if "k3" in frame else 0.0, @@ -274,17 +277,22 @@ class Nerfstudio(DataParser): height = int(meta["h"]) if height_fixed else torch.tensor(height, dtype=torch.int32)[idx_tensor] width = int(meta["w"]) if width_fixed else torch.tensor(width, dtype=torch.int32)[idx_tensor] if distort_fixed: - distortion_params = camera_utils.get_distortion_params( - k1=float(meta["k1"]) if "k1" in meta else 0.0, - k2=float(meta["k2"]) if "k2" in meta else 0.0, - k3=float(meta["k3"]) if "k3" in meta else 0.0, - k4=float(meta["k4"]) if "k4" in meta else 0.0, - p1=float(meta["p1"]) if "p1" in meta else 0.0, - p2=float(meta["p2"]) if "p2" in meta else 0.0, + distortion_params = ( + torch.tensor(meta["distortion_params"], dtype=torch.float32) + if "distortion_params" in meta + else camera_utils.get_distortion_params( + k1=float(meta["k1"]) if "k1" in meta else 0.0, + k2=float(meta["k2"]) if "k2" in meta else 0.0, + k3=float(meta["k3"]) if "k3" in meta else 0.0, + k4=float(meta["k4"]) if "k4" in meta else 0.0, + p1=float(meta["p1"]) if "p1" in meta else 0.0, + p2=float(meta["p2"]) if "p2" in meta else 0.0, + ) ) else: distortion_params = torch.stack(distort, dim=0)[idx_tensor] + metadata = {"fisheye_crop_radius": fisheye_crop_radius} if fisheye_crop_radius is not None else None cameras = Cameras( fx=fx, fy=fy, @@ -295,6 +303,7 @@ class Nerfstudio(DataParser): width=width, camera_to_worlds=poses[:, :3, :4], camera_type=camera_type, + metadata=metadata, ) assert self.downscale_factor is not None diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 07bab5d82..c7ce9e073 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -47,6 +47,8 @@ class PixelSamplerConfig(InstantiateConfig): """Whether or not to include a reference to the full image in returned batch.""" is_equirectangular: bool = False """List of whether or not camera i is equirectangular.""" + fisheye_crop_radius: Optional[float] = None + """Set to the radius (in pixels) for fisheye cameras.""" class PixelSampler: @@ -65,6 +67,7 @@ class PixelSampler: self.config.num_rays_per_batch = self.kwargs.get("num_rays_per_batch", self.config.num_rays_per_batch) self.config.keep_full_image = self.kwargs.get("keep_full_image", self.config.keep_full_image) self.config.is_equirectangular = self.kwargs.get("is_equirectangular", self.config.is_equirectangular) + self.config.fisheye_crop_radius = self.kwargs.get("fisheye_crop_radius", self.config.fisheye_crop_radius) self.set_num_rays_per_batch(self.config.num_rays_per_batch) def set_num_rays_per_batch(self, num_rays_per_batch: int): @@ -134,6 +137,36 @@ class PixelSampler: return indices + def sample_method_fisheye( + self, + batch_size: int, + num_images: int, + image_height: int, + image_width: int, + mask: Optional[Tensor] = None, + device: Union[torch.device, str] = "cpu", + ) -> Int[Tensor, "batch_size 3"]: + if isinstance(mask, torch.Tensor): + indices = self.sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + else: + rand_samples = torch.rand((batch_size, 3), device=device) + # convert random samples tto radius and theta + radii = self.config.fisheye_crop_radius * torch.sqrt(rand_samples[:, 1]) + theta = 2.0 * torch.pi * rand_samples[:, 2] + + # convert radius and theta to x and y between -radii and radii + x = radii * torch.cos(theta) + y = radii * torch.sin(theta) + + # Multiply by the batch size and height/width to get pixel indices. + indices = torch.floor( + torch.stack([rand_samples[:, 0], y, x], dim=1) + * torch.tensor([num_images, image_height // 2, image_width // 2], device=device) + + torch.tensor([0, image_height // 2, image_width // 2], device=device) + ).long() + + return indices + def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. @@ -154,6 +187,10 @@ class PixelSampler: indices = self.sample_method_equirectangular( num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device ) + elif self.config.fisheye_crop_radius is not None: + indices = self.sample_method_fisheye( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device + ) else: indices = self.sample_method( num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device @@ -163,6 +200,10 @@ class PixelSampler: indices = self.sample_method_equirectangular( num_rays_per_batch, num_images, image_height, image_width, device=device ) + elif self.config.fisheye_crop_radius is not None: + indices = self.sample_method_fisheye( + num_rays_per_batch, num_images, image_height, image_width, device=device + ) else: indices = self.sample_method(num_rays_per_batch, num_images, image_height, image_width, device=device) diff --git a/nerfstudio/scripts/datasets/process_project_aria.py b/nerfstudio/scripts/datasets/process_project_aria.py new file mode 100644 index 000000000..f2fdcc94c --- /dev/null +++ b/nerfstudio/scripts/datasets/process_project_aria.py @@ -0,0 +1,231 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import sys +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List + +import numpy as np +import tyro +from PIL import Image + +try: + from projectaria_tools.core import mps + from projectaria_tools.core.data_provider import VrsDataProvider, create_vrs_data_provider + from projectaria_tools.core.sophus import SE3 +except ImportError: + print("projectaria_tools import failed, please install with pip3 install projectaria-tools'[all]'") + sys.exit(1) + +ARIA_CAMERA_MODEL = "FISHEYE624" + +# The Aria coordinate system is different than the Blender/NerfStudio coordinate system. +# Blender / Nerfstudio: +Z = back, +Y = up, +X = right +# Surreal: +Z = forward, +Y = down, +X = right +T_ARIA_NERFSTUDIO = SE3.from_matrix( + np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) +) + + +@dataclass +class AriaCameraCalibration: + fx: float + fy: float + cx: float + cy: float + distortion_params: np.ndarray + width: int + height: int + t_device_camera: SE3 + + +@dataclass +class AriaImageFrame: + camera: AriaCameraCalibration + file_path: str + t_world_camera: SE3 + timestamp_ns: float + + +@dataclass +class TimedPoses: + timestamps_ns: np.ndarray + t_world_devices: List[SE3] + + +def get_camera_calibs(provider: VrsDataProvider) -> Dict[str, AriaCameraCalibration]: + """Retrieve the per-camera factory calibration from within the VRS.""" + + factory_calib = {} + name = "camera-rgb" + device_calib = provider.get_device_calibration() + assert device_calib is not None, "Could not find device calibration" + sensor_calib = device_calib.get_camera_calib(name) + assert sensor_calib is not None, f"Could not find sensor calibration for {name}" + + width = sensor_calib.get_image_size()[0].item() + height = sensor_calib.get_image_size()[1].item() + intrinsics = sensor_calib.projection_params() + + factory_calib[name] = AriaCameraCalibration( + fx=intrinsics[0], + fy=intrinsics[0], + cx=intrinsics[1], + cy=intrinsics[2], + distortion_params=intrinsics[3:15], + width=width, + height=height, + t_device_camera=sensor_calib.get_transform_device_camera(), + ) + + return factory_calib + + +def read_trajectory_csv_to_dict(file_iterable_csv: str) -> TimedPoses: + closed_loop_traj = mps.read_closed_loop_trajectory(file_iterable_csv) + + timestamps_secs, poses = zip( + *[(it.tracking_timestamp.total_seconds(), it.transform_world_device) for it in closed_loop_traj] + ) + + SEC_TO_NANOSEC = 1e9 + return TimedPoses( + timestamps_ns=(np.array(timestamps_secs) * SEC_TO_NANOSEC).astype(int), + t_world_devices=poses, + ) + + +def to_aria_image_frame( + provider: VrsDataProvider, + index: int, + name_to_camera: Dict[str, AriaCameraCalibration], + t_world_devices: TimedPoses, + output_dir: Path, +) -> AriaImageFrame: + name = "camera-rgb" + + camera_calibration = name_to_camera[name] + stream_id = provider.get_stream_id_from_label(name) + assert stream_id is not None, f"Could not find stream {name}" + + # Get the image corresponding to this index + image_data = provider.get_image_data_by_index(stream_id, index) + img = Image.fromarray(image_data[0].to_numpy_array()) + capture_time_ns = image_data[1].capture_timestamp_ns + + file_path = f"{output_dir}/{name}_{capture_time_ns}.jpg" + threading.Thread(target=lambda: img.save(file_path)).start() + + # Find the nearest neighbor pose with the closest timestamp to the capture time. + nearest_pose_idx = np.searchsorted(t_world_devices.timestamps_ns, capture_time_ns) + nearest_pose_idx = np.minimum(nearest_pose_idx, len(t_world_devices.timestamps_ns) - 1) + assert nearest_pose_idx != -1, f"Could not find pose for {capture_time_ns}" + t_world_device = t_world_devices.t_world_devices[nearest_pose_idx] + + # Compute the world to camera transform. + t_world_camera = t_world_device @ camera_calibration.t_device_camera @ T_ARIA_NERFSTUDIO + + return AriaImageFrame( + camera=camera_calibration, + file_path=file_path, + t_world_camera=t_world_camera, + timestamp_ns=capture_time_ns, + ) + + +def to_nerfstudio_frame(frame: AriaImageFrame) -> Dict: + return { + "fl_x": frame.camera.fx, + "fl_y": frame.camera.fy, + "cx": frame.camera.cx, + "cy": frame.camera.cy, + "distortion_params": frame.camera.distortion_params.tolist(), + "w": frame.camera.width, + "h": frame.camera.height, + "file_path": frame.file_path, + "transform_matrix": frame.t_world_camera.to_matrix().tolist(), + "timestamp": frame.timestamp_ns, + } + + +@dataclass +class ProcessProjectAria: + """Processes Project Aria data i.e. a VRS of the raw recording streams and the MPS attachments + that provide poses, calibration, and 3d points. More information on MPS data can be found at: + https://facebookresearch.github.io/projectaria_tools/docs/ARK/mps. + """ + + vrs_file: Path + """Path to the VRS file.""" + mps_data_dir: Path + """Path to Project Aria Machine Perception Services (MPS) attachments.""" + output_dir: Path + """Path to the output directory.""" + + def main(self) -> None: + """Generate a nerfstudio dataset from ProjectAria data (VRS) and MPS attachments.""" + # Create output directory if it doesn't exist. + self.output_dir = self.output_dir.absolute() + self.output_dir.mkdir(parents=True, exist_ok=True) + + provider = create_vrs_data_provider(str(self.vrs_file.absolute())) + assert provider is not None, "Cannot open file" + + name_to_camera = get_camera_calibs(provider) + + print("Getting poses from closed loop trajectory CSV...") + trajectory_csv = self.mps_data_dir / "closed_loop_trajectory.csv" + t_world_devices = read_trajectory_csv_to_dict(str(trajectory_csv.absolute())) + + name = "camera-rgb" + stream_id = provider.get_stream_id_from_label(name) + + # create an AriaImageFrame for each image in the VRS. + print("Creating Aria frames...") + aria_frames = [ + to_aria_image_frame(provider, index, name_to_camera, t_world_devices, self.output_dir) + for index in range(0, provider.get_num_data(stream_id)) + ] + + # create the NerfStudio frames from the AriaImageFrames. + print("Creating NerfStudio frames...") + CANONICAL_RGB_VALID_RADIUS = 707.5 + CANONICAL_RGB_WIDTH = 1408 + rgb_valid_radius = CANONICAL_RGB_VALID_RADIUS * (aria_frames[0].camera.width / CANONICAL_RGB_WIDTH) + nerfstudio_frames = { + "camera_model": ARIA_CAMERA_MODEL, + "frames": [to_nerfstudio_frame(frame) for frame in aria_frames], + "fisheye_crop_radius": rgb_valid_radius, + } + + # write the json out to disk as transforms.json + print("Writing transforms.json") + transform_file = self.output_dir / "transforms.json" + with open(transform_file, "w", encoding="UTF-8"): + transform_file.write_text(json.dumps(nerfstudio_frames)) + + +if __name__ == "__main__": + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(ProcessProjectAria).main() diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py index 83cdcca47..2c2cd7a38 100644 --- a/nerfstudio/scripts/process_data.py +++ b/nerfstudio/scripts/process_data.py @@ -475,6 +475,12 @@ class ProcessODM(BaseConverterToNerfstudioDataset): CONSOLE.rule() +@dataclass +class NotInstalled: + def main(self) -> None: + ... + + Commands = Union[ Annotated[ImagesToNerfstudioDataset, tyro.conf.subcommand(name="images")], Annotated[VideoToNerfstudioDataset, tyro.conf.subcommand(name="video")], @@ -485,6 +491,29 @@ Commands = Union[ Annotated[ProcessODM, tyro.conf.subcommand(name="odm")], ] +# Add aria subcommand if projectaria_tools is installed. +try: + import projectaria_tools +except ImportError: + projectaria_tools = None + +if projectaria_tools is not None: + from nerfstudio.scripts.datasets.process_project_aria import ProcessProjectAria + + # Note that Union[A, Union[B, C]] == Union[A, B, C]. + Commands = Union[Commands, Annotated[ProcessProjectAria, tyro.conf.subcommand(name="aria")]] +else: + Commands = Union[ + Commands, + Annotated[ + NotInstalled, + tyro.conf.subcommand( + name="aria", + description="**Not installed.** Processing Project Aria data requires `pip install projectaria_tools'[all]'`.", + ), + ], + ] + def entrypoint(): """Entrypoint for use with pyproject scripts.""" diff --git a/pyproject.toml b/pyproject.toml index b6d11c60c..3b6d69a84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,7 @@ dev = [ "opencv-stubs==0.0.7", "transformers==4.29.2", "pyright==1.1.331", + "projectaria_tools[all]>=1.2.0", ] # Documentation related packages