Render Omni-directional stereo video (360 3D VR) (#1986)

* Render Omni-directional stereo (360 3D VR)

Adding the omni-directional stereo camera model to render 3D 360 video and images for VR. This includes the camera models for the left and right eyes and additions to the render script to stack the left and right renders for the final render.
I have also updated the documentation on the custom_dataset page to have information on ODS but let me know if there is a better place to put it.

* Updated ODS camera model and render paths

Updates based on suggested code changes including adding a helper function for ODS and improving path creation.
I also added a part where if a user tries to render a video with an output path that doesn't have an extension it will automatically add .mp4 to the end. Otherwise this usually results in an error when there is no file extension when rendering a video.

* Updated formatting for render and cam path scripts

Updated python formatting for render and camera_paths script

* Adding type checking for ODS function and reformatting camera type code

Added the type checking and doc string for the ODS camera rays function and reformatted the cam_type ray if statements.

docs/quickstart/custom_dataset.md

@@ -346,3 +346,49 @@ The `crop-factor` argument is optional but often very helpful. This is because e
 The pixels representing the distorted hand and head are obviously not useful in training a nerf so we can remove it by cropping the bottom 20% of the image. This can be done by using the `--crop-factor 0 0.2 0 0` argument.
 
 If cropping only needs to be done from the bottom, you can use the `--crop-bottom [num]` argument which would be the same as doing `--crop-factor 0.0 [num] 0.0 0.0`
+
+## 🥽 Render Omni-directional Stereo (360 VR)
+
+Stereo equirectangular rendering for VR video is supported as omni-directional stereo video and image rendering. This outputs two equirectangular renders vertically stacked, one for each eye. Omni-directional stereo (ODS) is a method to render VR 3D 360 videos, and may introduce slight depth distortions for close objects. For additional information on how ODS works, refer to this [writeup](https://developers.google.com/vr/jump/rendering-ods-content.pdf).
+
+<center>
+<img img width="300" src="https://github-production-user-asset-6210df.s3.amazonaws.com/9502341/240530353-5c05d12b-6e42-42bb-8ff5-f54aa30ca262.jpg">
+</center>
+
+To render in ODS it is essential to adjust the NeRF to have an approximately true-to-life real world scale (adjustable in the camera path) to ensure that the scene depth and IPD (distance between the eyes) is appropriate for the render to be viewable in VR. You can adjust the scene scale with the [Blender Add-on](https://docs.nerf.studio/en/latest/extensions/blender_addon.html) by appropriately scaling a point cloud representation of the NeRF.
+Results may be unviewable if the scale is not set appropriately. The IPD is set at 64mm by default but only is accurate when the NeRF scene is true to scale.
+
+For good quality renders, it is recommended to render at 4096x2048 per eye, or 2048x1024 per eye. Render resolutions must be specified in the camera path in a 2:1 aspect ratio. The final stacked render output will automatically be constructed as 1:1.
+
+:::{admonition} Note
+:class: info
+If you are rendering an image sequence, it is recommended to render as png instead of jpeg, since the png will appear clearer. However, file sizes can be significantly larger with png.
+:::
+
+To render with the ODS camera:
+1. Use the Nerfstudio Blender Add-on to set the scale of the NeRF scene and create the camera path
+    - Export a point cloud representation of the NeRF
+   - Import the point cloud representation in Blender and enable the Blender Add-on
+    - Create a reference object such as a cube which may be 1x1x1 meter. You could also create a cylinder and scale it to an appropriate height of a viewer.
+    - Now scale the point cloud representation accordingly to match the reference object. This is to ensure that the NeRF scene is scaled as close to real life.
+    - To place the camera at the correct height from the ground in the scene, you create a cylinder representing the viewer vertically scaled to the viewer’s height, and place the camera at eye level.
+    - Animate the camera movement as needed
+    - Create the camera path JSON file with the Nerfstudio Blender Add-on
+2. Edit the JSON camera path file
+    - Open the camera path JSON file and specify the `camera_type` as `omnidirectional`
+    - Specify the `render_height` and `render_width` to the resolution of a single eye. The width:height aspect ratio must be 2:1. Recommended resolutions are 4096x2048 or 2048x1024.
+    - Save the camera path and render the NeRF
+  <center>
+  <img img width="300" src="https://github-production-user-asset-6210df.s3.amazonaws.com/9502341/240530527-22d14276-ac2c-46a5-a4b0-4785b7413241.png">
+  </center>
+
+:::{admonition} Note
+:class: info
+If the depth of the scene is unviewable and looks too close or expanded when viewing the render in VR, the scale of the NeRF may be set too small. If there is almost no discernible depth, the scale of the NeRF may be too large. Getting the right scale may take some experimentation, so it is recommended to either render at a much lower resolution or just one frame to ensure the depth and render is viewable in the VR headset.
+:::
+
+### Additional Notes
+- Rendering with ODS can take significantly longer than traditional renders due to higher resolutions and needing to render a left and right eye view for each frame. Render times may grow exponentially with larger resolutions.
+- When rendering ODS content, Nerfstudio will first render the left eye, then the right eye, and finally vertically stack the renders. During this process, Nerfstudio will create a temporary folder to store the left and right eye renders and delete this folder once the final renders are stacked.
+- If rendering content where the camera is stationary for many frames, it is recommended to only render once at that position and extend the time in a video editor since ODS renders can take a lot of time to render.
+- It is recommended to render a preliminary render at a much lower resolution or frame rate to test and ensure that the depth and camera position look accurate in VR.

nerfstudio/cameras/camera_paths.py

@@ -139,6 +139,8 @@ def get_path_from_json(camera_path: Dict[str, Any]) -> Cameras:
         camera_type = CameraType.FISHEYE
     elif camera_path["camera_type"] == "equirectangular":
         camera_type = CameraType.EQUIRECTANGULAR
+    elif camera_path["camera_type"].lower() == "omnidirectional":
+        camera_type = CameraType.OMNIDIRECTIONALSTEREO_L
     else:
         camera_type = CameraType.PERSPECTIVE
 
@@ -149,7 +151,11 @@ def get_path_from_json(camera_path: Dict[str, Any]) -> Cameras:
         # pose
         c2w = torch.tensor(camera["camera_to_world"]).view(4, 4)[:3]
         c2ws.append(c2w)
-        if camera_type == CameraType.EQUIRECTANGULAR:
+        if (
+            camera_type == CameraType.EQUIRECTANGULAR
+            or camera_type == CameraType.OMNIDIRECTIONALSTEREO_L
+            or camera_type == CameraType.OMNIDIRECTIONALSTEREO_R
+        ):
             fxs.append(image_width / 2)
             fys.append(image_height)
         else:

nerfstudio/cameras/cameras.py

@@ -19,7 +19,7 @@ import base64
 import math
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple, Union
 
 import cv2
 import torch
@@ -44,6 +44,8 @@ class CameraType(Enum):
     PERSPECTIVE = auto()
     FISHEYE = auto()
     EQUIRECTANGULAR = auto()
+    OMNIDIRECTIONALSTEREO_L = auto()
+    OMNIDIRECTIONALSTEREO_R = auto()
 
 
 CAMERA_MODEL_TO_TYPE = {
@@ -54,6 +56,8 @@ CAMERA_MODEL_TO_TYPE = {
     "OPENCV": CameraType.PERSPECTIVE,
     "OPENCV_FISHEYE": CameraType.FISHEYE,
     "EQUIRECTANGULAR": CameraType.EQUIRECTANGULAR,
+    "OMNIDIRECTIONALSTEREO_L": CameraType.OMNIDIRECTIONALSTEREO_L,
+    "OMNIDIRECTIONALSTEREO_R": CameraType.OMNIDIRECTIONALSTEREO_R,
 }
 
 
@@ -643,48 +647,120 @@ class Cameras(TensorDataclass):
         # directions_stack[2] is the direction for ray in camera coordinates offset by 1 in y
         cam_types = torch.unique(self.camera_type, sorted=False)
         directions_stack = torch.empty((3,) + num_rays_shape + (3,), device=self.device)
-        if CameraType.PERSPECTIVE.value in cam_types:
-            mask = (self.camera_type[true_indices] == CameraType.PERSPECTIVE.value).squeeze(-1)  # (num_rays)
+
+        c2w = self.camera_to_worlds[true_indices]
+        assert c2w.shape == num_rays_shape + (3, 4)
+
+        def _compute_rays_for_omnidirectional_stereo(
+            eye: Literal["left", "right"]
+        ) -> Tuple[Float[Tensor, "num_rays_shape 3"], Float[Tensor, "3 num_rays_shape 3"]]:
+            """Compute the rays for an omnidirectional stereo camera
+
+            Args:
+                eye: Which eye to compute rays for.
+
+            Returns:
+                A tuple containing the origins and the directions of the rays.
+            """
+            # Directions calculated similarly to equirectangular
+            ods_cam_type = (
+                CameraType.OMNIDIRECTIONALSTEREO_R.value if eye == "right" else CameraType.OMNIDIRECTIONALSTEREO_L.value
+            )
+            mask = (self.camera_type[true_indices] == ods_cam_type).squeeze(-1)
             mask = torch.stack([mask, mask, mask], dim=0)
-            directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0], mask).float()
-            directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1], mask).float()
-            directions_stack[..., 2][mask] = -1.0
-
-        if CameraType.FISHEYE.value in cam_types:
-            mask = (self.camera_type[true_indices] == CameraType.FISHEYE.value).squeeze(-1)  # (num_rays)
-            mask = torch.stack([mask, mask, mask], dim=0)
-
-            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[..., 0][mask] = torch.masked_select(coord_stack[..., 0] * sin_theta / theta, mask).float()
-            directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1] * sin_theta / theta, mask).float()
-            directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask).float()
-
-        if CameraType.EQUIRECTANGULAR.value in cam_types:
-            mask = (self.camera_type[true_indices] == CameraType.EQUIRECTANGULAR.value).squeeze(-1)  # (num_rays)
-            mask = torch.stack([mask, mask, mask], dim=0)
-
-            # For equirect, fx = fy = height = width/2
-            # Then coord[..., 0] goes from -1 to 1 and coord[..., 1] goes from -1/2 to 1/2
-            theta = -torch.pi * coord_stack[..., 0]  # minus sign for right-handed
+            theta = -torch.pi * coord_stack[..., 0]
             phi = torch.pi * (0.5 - coord_stack[..., 1])
-            # use spherical in local camera coordinates (+y up, x=0 and z<0 is theta=0)
+
             directions_stack[..., 0][mask] = torch.masked_select(-torch.sin(theta) * torch.sin(phi), mask).float()
             directions_stack[..., 1][mask] = torch.masked_select(torch.cos(phi), mask).float()
             directions_stack[..., 2][mask] = torch.masked_select(-torch.cos(theta) * torch.sin(phi), mask).float()
 
-        for value in cam_types:
-            if value not in [CameraType.PERSPECTIVE.value, CameraType.FISHEYE.value, CameraType.EQUIRECTANGULAR.value]:
-                raise ValueError(f"Camera type {value} not supported.")
+            vr_ipd = 0.064  # IPD in meters (note: scale of NeRF must be true to life and can be adjusted with the Blender add-on)
+            isRightEye = 1 if eye == "right" else -1
+
+            # find ODS camera position
+            c2w = self.camera_to_worlds[true_indices]
+            assert c2w.shape == num_rays_shape + (3, 4)
+            transposedC2W = c2w[0][0].t()
+            ods_cam_position = transposedC2W[3].repeat(c2w.shape[1], 1)
+
+            rotation = c2w[..., :3, :3]
+
+            ods_theta = -torch.pi * ((x - cx) / fx)[0]
+
+            # local axes of ODS camera
+            ods_x_axis = torch.tensor([1, 0, 0], device=c2w.device)
+            ods_z_axis = torch.tensor([0, 0, -1], device=c2w.device)
+
+            # circle of ODS ray origins
+            ods_origins_circle = (
+                ods_cam_position
+                + isRightEye * (vr_ipd / 2.0) * (ods_x_axis.repeat(c2w.shape[1], 1)) * (torch.cos(ods_theta))[:, None]
+                + isRightEye * (vr_ipd / 2.0) * (ods_z_axis.repeat(c2w.shape[1], 1)) * (torch.sin(ods_theta))[:, None]
+            )
+
+            # rotate origins to match the camera rotation
+            for i in range(ods_origins_circle.shape[0]):
+                ods_origins_circle[i] = rotation[0][0] @ ods_origins_circle[i] + ods_cam_position[0]
+            ods_origins_circle = ods_origins_circle.unsqueeze(0).repeat(c2w.shape[0], 1, 1)
+
+            # assign final camera origins
+            c2w[..., :3, 3] = ods_origins_circle
+
+            return ods_origins_circle, directions_stack
+
+        for cam in cam_types:
+            if CameraType.PERSPECTIVE.value in cam_types:
+                mask = (self.camera_type[true_indices] == CameraType.PERSPECTIVE.value).squeeze(-1)  # (num_rays)
+                mask = torch.stack([mask, mask, mask], dim=0)
+                directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0], mask).float()
+                directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1], mask).float()
+                directions_stack[..., 2][mask] = -1.0
+
+            elif CameraType.FISHEYE.value in cam_types:
+                mask = (self.camera_type[true_indices] == CameraType.FISHEYE.value).squeeze(-1)  # (num_rays)
+                mask = torch.stack([mask, mask, mask], dim=0)
+
+                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[..., 0][mask] = torch.masked_select(
+                    coord_stack[..., 0] * sin_theta / theta, mask
+                ).float()
+                directions_stack[..., 1][mask] = torch.masked_select(
+                    coord_stack[..., 1] * sin_theta / theta, mask
+                ).float()
+                directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask).float()
+
+            elif CameraType.EQUIRECTANGULAR.value in cam_types:
+                mask = (self.camera_type[true_indices] == CameraType.EQUIRECTANGULAR.value).squeeze(-1)  # (num_rays)
+                mask = torch.stack([mask, mask, mask], dim=0)
+
+                # For equirect, fx = fy = height = width/2
+                # Then coord[..., 0] goes from -1 to 1 and coord[..., 1] goes from -1/2 to 1/2
+                theta = -torch.pi * coord_stack[..., 0]  # minus sign for right-handed
+                phi = torch.pi * (0.5 - coord_stack[..., 1])
+                # use spherical in local camera coordinates (+y up, x=0 and z<0 is theta=0)
+                directions_stack[..., 0][mask] = torch.masked_select(-torch.sin(theta) * torch.sin(phi), mask).float()
+                directions_stack[..., 1][mask] = torch.masked_select(torch.cos(phi), mask).float()
+                directions_stack[..., 2][mask] = torch.masked_select(-torch.cos(theta) * torch.sin(phi), mask).float()
+
+            elif CameraType.OMNIDIRECTIONALSTEREO_L.value in cam_types:
+                ods_origins_circle, directions_stack = _compute_rays_for_omnidirectional_stereo("left")
+                # assign final camera origins
+                c2w[..., :3, 3] = ods_origins_circle
+
+            elif CameraType.OMNIDIRECTIONALSTEREO_R.value in cam_types:
+                ods_origins_circle, directions_stack = _compute_rays_for_omnidirectional_stereo("right")
+                # assign final camera origins
+                c2w[..., :3, 3] = ods_origins_circle
+            else:
+                raise ValueError(f"Camera type {cam} not supported.")
 
         assert directions_stack.shape == (3,) + num_rays_shape + (3,)
 
-        c2w = self.camera_to_worlds[true_indices]
-        assert c2w.shape == num_rays_shape + (3, 4)
-
         if camera_opt_to_camera is not None:
             c2w = pose_utils.multiply(c2w, camera_opt_to_camera)
         rotation = c2w[..., :3, :3]  # (..., 3, 3)

nerfstudio/scripts/render.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 import json
 import os
 import struct
+import shutil
 import sys
 from contextlib import ExitStack
 from dataclasses import dataclass, field
@@ -59,6 +60,7 @@ from nerfstudio.pipelines.base_pipeline import Pipeline
 from nerfstudio.utils import colormaps, install_checks
 from nerfstudio.utils.eval_utils import eval_setup
 from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn
+from nerfstudio.utils.scripts import run_command
 
 
 def _render_trajectory_video(
@@ -342,6 +344,22 @@ class RenderCameraPath(BaseRender):
         crop_data = get_crop_from_json(camera_path)
         camera_path = get_path_from_json(camera_path)
 
+        if camera_path.camera_type[0] == CameraType.OMNIDIRECTIONALSTEREO_L.value:
+            # temp folder for writing left and right view renders
+            temp_folder_path = self.output_path.parent / (self.output_path.stem + "_temp")
+
+            Path(temp_folder_path).mkdir(parents=True, exist_ok=True)
+            left_eye_path = temp_folder_path / "ods_render_Left.mp4"
+
+            self.output_path = left_eye_path
+
+            CONSOLE.print("[bold green]:goggles: Omni-directional Stereo VR :goggles:")
+            CONSOLE.print("Rendering left eye view")
+
+        # add mp4 suffix to video output if none is specified
+        if self.output_format == "video" and str(self.output_path.suffix) == "":
+            self.output_path = self.output_path.with_suffix(".mp4")
+
         _render_trajectory_video(
             pipeline,
             camera_path,
@@ -356,6 +374,51 @@ class RenderCameraPath(BaseRender):
             colormap_options=self.colormap_options,
         )
 
+        if camera_path.camera_type[0] == CameraType.OMNIDIRECTIONALSTEREO_L.value:
+            # declare paths for left and right renders
+
+            left_eye_path = self.output_path
+            right_eye_path = left_eye_path.parent / "ods_render_Right.mp4"
+
+            self.output_path = right_eye_path
+            camera_path.camera_type[0] = CameraType.OMNIDIRECTIONALSTEREO_R.value
+
+            CONSOLE.print("Rendering right eye view")
+            _render_trajectory_video(
+                pipeline,
+                camera_path,
+                output_filename=self.output_path,
+                rendered_output_names=self.rendered_output_names,
+                rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
+                crop_data=crop_data,
+                seconds=seconds,
+                output_format=self.output_format,
+                image_format=self.image_format,
+                jpeg_quality=self.jpeg_quality,
+                colormap_options=self.colormap_options,
+            )
+
+            # stack the left and right eye renders for final output
+            self.output_path = Path(str(left_eye_path.parent)[:-5] + ".mp4")
+            ffmpeg_ods_command = ""
+            if self.output_format == "video":
+                ffmpeg_ods_command = f'ffmpeg -y -i "{left_eye_path}" -i "{right_eye_path}" -filter_complex "[0:v]pad=iw:2*ih[int];[int][1:v]overlay=0:h" -c:v libx264 -crf 23 -preset veryfast "{self.output_path}"'
+                run_command(ffmpeg_ods_command, verbose=False)
+            if self.output_format == "images":
+                # create a folder for the stacked renders
+                self.output_path = Path(str(left_eye_path.parent)[:-5])
+                self.output_path.mkdir(parents=True, exist_ok=True)
+                if self.image_format == "png":
+                    ffmpeg_ods_command = f'ffmpeg -y -pattern_type glob -i "{str(left_eye_path.with_suffix("") / "*.png")}"  -pattern_type glob -i "{str(right_eye_path.with_suffix("") / "*.png")}" -filter_complex vstack -start_number 0 "{str(self.output_path)+"//%05d.png"}"'
+                elif self.image_format == "jpeg":
+                    ffmpeg_ods_command = f'ffmpeg -y -pattern_type glob -i "{str(left_eye_path.with_suffix("") / "*.jpg")}"  -pattern_type glob -i "{str(right_eye_path.with_suffix("") / "*.jpg")}" -filter_complex vstack -start_number 0 "{str(self.output_path)+"//%05d.jpg"}"'
+                run_command(ffmpeg_ods_command, verbose=False)
+
+            # remove the temp files directory
+            if str(left_eye_path.parent)[-5:] == "_temp":
+                shutil.rmtree(left_eye_path.parent, ignore_errors=True)
+            CONSOLE.print("[bold green]Final ODS Render Complete")
+
 
 @dataclass
 class RenderInterpolated(BaseRender):