Flesh out viewer beta

.vscode/settings.json

@@ -24,15 +24,15 @@
   "typescript.suggestionActions.enabled": false,
   "javascript.suggestionActions.enabled": false,
   "[python]": {
-    "editor.defaultFormatter": "ms-python.python",
+    "editor.defaultFormatter": "ms-python.black-formatter",
     "editor.codeActionsOnSave": {
       "source.organizeImports": true,
       "source.fixAll": true
     }
   },
   "editor.formatOnSave": true,
   "python.envFile": "${workspaceFolder}/.env",
-  "python.formatting.provider": "black",
+  "python.formatting.provider": "none",
   "python.formatting.blackArgs": ["--line-length=120"],
   "python.linting.pylintEnabled": false,
   "python.linting.flake8Enabled": false,

nerfstudio/cameras/cameras.py

@@ -32,7 +32,7 @@
 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
+from nerfstudio.data.scene_box import SceneBox, OrientedBox
 from nerfstudio.utils.tensor_dataclass import TensorDataclass
 
 TORCH_DEVICE = Union[torch.device, str]
@@ -323,6 +323,7 @@ def generate_rays(
         keep_shape: Optional[bool] = None,
         disable_distortion: bool = False,
         aabb_box: Optional[SceneBox] = None,
+        obb_box: Optional[OrientedBox] = None,
     ) -> RayBundle:
         """Generates rays for the given camera indices.
 
@@ -465,20 +466,24 @@ def generate_rays(
         if keep_shape is False:
             raybundle = raybundle.flatten()
 
-        if aabb_box:
+        if aabb_box is not None or obb_box is not None:
             with torch.no_grad():
-                tensor_aabb = Parameter(aabb_box.aabb.flatten(), requires_grad=False)
-
                 rays_o = raybundle.origins.contiguous()
                 rays_d = raybundle.directions.contiguous()
 
-                tensor_aabb = tensor_aabb.to(rays_o.device)
                 shape = rays_o.shape
 
                 rays_o = rays_o.reshape((-1, 3))
                 rays_d = rays_d.reshape((-1, 3))
 
-                t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb)
+                if aabb_box is not None:
+                    tensor_aabb = Parameter(aabb_box.aabb.flatten(), requires_grad=False)
+                    tensor_aabb = tensor_aabb.to(rays_o.device)
+                    t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb)
+                elif obb_box is not None:
+                    t_min, t_max = nerfstudio.utils.math.intersect_obb(rays_o, rays_d, obb_box)
+                else:
+                    assert False
 
                 t_min = t_min.reshape([shape[0], shape[1], 1])
                 t_max = t_max.reshape([shape[0], shape[1], 1])
@@ -747,8 +752,6 @@ def _compute_rays_for_vr180(
 
             rotation = c2w[..., :3, :3]
 
-            -torch.pi * ((x - cx) / fx)[0]
-
             # interocular axis of the VR180 camera
             vr180_x_axis = torch.tensor([1, 0, 0], device=c2w.device)
 

nerfstudio/data/scene_box.py

@@ -17,7 +17,8 @@
 """
 
 from dataclasses import dataclass
-from typing import Union
+from typing import Union, Tuple
+import viser.transforms as vtf
 
 import torch
 from jaxtyping import Float
@@ -76,3 +77,38 @@ def from_camera_poses(poses: Float[Tensor, "*batch 3 4"], scale_factor: float) -
         xyzs = poses[..., :3, -1]
         aabb = torch.stack([torch.min(xyzs, dim=0)[0], torch.max(xyzs, dim=0)[0]])
         return SceneBox(aabb=aabb * scale_factor)
+
+
+@dataclass
+class OrientedBox:
+    R: Float[Tensor, "3 3"]
+    """R: rotation matrix."""
+    T: Float[Tensor, "3"]
+    """T: translation vector."""
+    S: Float[Tensor, "3"]
+    """S: scale vector."""
+
+    def within(self, pts: Float[Tensor, "n 3"]):
+        """Returns a boolean mask indicating whether each point is within the box."""
+        R, T, S = self.R, self.T, self.S.to(pts)
+        H = torch.eye(4, device=pts.device, dtype=pts.dtype)
+        H[:3, :3] = R
+        H[:3, 3] = T
+        H_world2bbox = torch.inverse(H)
+        pts = torch.cat((pts, torch.ones_like(pts[..., :1])), dim=-1)
+        pts = torch.matmul(H_world2bbox, pts.T).T[..., :3]
+
+        comp_l = torch.tensor(-S / 2)
+        comp_m = torch.tensor(S / 2)
+        mask = torch.all(torch.concat([pts > comp_l, pts < comp_m], dim=-1), dim=-1)
+        return mask
+
+    @staticmethod
+    def from_params(
+        pos: Tuple[float, float, float], rpy: Tuple[float, float, float], scale: Tuple[float, float, float]
+    ):
+        """Construct a box from position, rotation, and scale parameters."""
+        R = torch.tensor(vtf.SO3.from_rpy_radians(rpy[0], rpy[1], rpy[2]).as_matrix())
+        T = torch.tensor(pos)
+        S = torch.tensor(scale)
+        return OrientedBox(R=R, T=T, S=S)

nerfstudio/exporter/exporter_utils.py

@@ -33,6 +33,7 @@
 
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.data.scene_box import OrientedBox
 from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline
 from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn
 
@@ -84,8 +85,9 @@ def generate_point_cloud(
     depth_output_name: str = "depth",
     normal_output_name: Optional[str] = None,
     use_bounding_box: bool = True,
-    bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0),
-    bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0),
+    bounding_box_min: Optional[Tuple[float, float, float]] = None,
+    bounding_box_max: Optional[Tuple[float, float, float]] = None,
+    crop_obb: Optional[OrientedBox] = None,
     std_ratio: float = 10.0,
 ) -> o3d.geometry.PointCloud:
     """Generate a point cloud from a nerf.
@@ -119,7 +121,8 @@ def generate_point_cloud(
     rgbs = []
     normals = []
     view_directions = []
-
+    if use_bounding_box and (crop_obb is not None and bounding_box_max is not None):
+        CONSOLE.print("Provided aabb and crop_obb at the same time, using only the obb", style="bold yellow")
     with progress as progress_bar:
         task = progress_bar.add_task("Generating Point Cloud", total=num_points)
         while not progress_bar.finished:
@@ -163,12 +166,15 @@ def generate_point_cloud(
                 normal = normal[mask]
 
             if use_bounding_box:
-                comp_l = torch.tensor(bounding_box_min, device=point.device)
-                comp_m = torch.tensor(bounding_box_max, device=point.device)
-                assert torch.all(
-                    comp_l < comp_m
-                ), f"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}"
-                mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
+                if crop_obb is None:
+                    comp_l = torch.tensor(bounding_box_min, device=point.device)
+                    comp_m = torch.tensor(bounding_box_max, device=point.device)
+                    assert torch.all(
+                        comp_l < comp_m
+                    ), f"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}"
+                    mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
+                else:
+                    mask = crop_obb.within(point)
                 point = point[mask]
                 rgb = rgb[mask]
                 view_direction = view_direction[mask]

nerfstudio/model_components/renderers.py

@@ -219,6 +219,7 @@ def forward(
         Returns:
             Outputs of rgb values.
         """
+
         if background_color is None:
             background_color = self.background_color
 

nerfstudio/scripts/exporter.py

@@ -34,6 +34,7 @@
 
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
+from nerfstudio.data.scene_box import OrientedBox
 from nerfstudio.exporter import texture_utils, tsdf_utils
 from nerfstudio.exporter.exporter_utils import (
     collect_camera_poses,
@@ -109,10 +110,17 @@ class ExportPointCloud(Exporter):
     """Name of the RGB output."""
     use_bounding_box: bool = True
     """Only query points within the bounding box"""
-    bounding_box_min: Tuple[float, float, float] = (-1, -1, -1)
+    bounding_box_min: Optional[Tuple[float, float, float]] = (-1, -1, -1)
     """Minimum of the bounding box, used if use_bounding_box is True."""
-    bounding_box_max: Tuple[float, float, float] = (1, 1, 1)
+    bounding_box_max: Optional[Tuple[float, float, float]] = (1, 1, 1)
     """Maximum of the bounding box, used if use_bounding_box is True."""
+
+    obb_center: Optional[Tuple[float, float, float]] = None
+    """Center of the oriented bounding box."""
+    obb_rotation: Optional[Tuple[float, float, float]] = None
+    """Rotation of the oriented bounding box. Expressed as RPY Euler angles in radians"""
+    obb_scale: Optional[Tuple[float, float, float]] = None
+    """Scale of the oriented bounding box along each axis."""
     num_rays_per_batch: int = 32768
     """Number of rays to evaluate per batch. Decrease if you run out of memory."""
     std_ratio: float = 10.0
@@ -135,7 +143,9 @@ def main(self) -> None:
 
         # Whether the normals should be estimated based on the point cloud.
         estimate_normals = self.normal_method == "open3d"
-
+        crop_obb = None
+        if self.obb_center is not None and self.obb_rotation is not None and self.obb_scale is not None:
+            crop_obb = OrientedBox.from_params(self.obb_center, self.obb_rotation, self.obb_scale)
         pcd = generate_point_cloud(
             pipeline=pipeline,
             num_points=self.num_points,
@@ -148,6 +158,7 @@ def main(self) -> None:
             use_bounding_box=self.use_bounding_box,
             bounding_box_min=self.bounding_box_min,
             bounding_box_max=self.bounding_box_max,
+            crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
         torch.cuda.empty_cache()
@@ -264,6 +275,12 @@ class ExportPoissonMesh(Exporter):
     """Minimum of the bounding box, used if use_bounding_box is True."""
     bounding_box_max: Tuple[float, float, float] = (1, 1, 1)
     """Minimum of the bounding box, used if use_bounding_box is True."""
+    obb_center: Optional[Tuple[float, float, float]] = None
+    """Center of the oriented bounding box."""
+    obb_rotation: Optional[Tuple[float, float, float]] = None
+    """Rotation of the oriented bounding box. Expressed as RPY Euler angles in radians"""
+    obb_scale: Optional[Tuple[float, float, float]] = None
+    """Scale of the oriented bounding box along each axis."""
     num_rays_per_batch: int = 32768
     """Number of rays to evaluate per batch. Decrease if you run out of memory."""
     texture_method: Literal["point_cloud", "nerf"] = "nerf"
@@ -296,6 +313,10 @@ def main(self) -> None:
 
         # Whether the normals should be estimated based on the point cloud.
         estimate_normals = self.normal_method == "open3d"
+        if self.obb_center is not None and self.obb_rotation is not None and self.obb_scale is not None:
+            crop_obb = OrientedBox.from_params(self.obb_center, self.obb_rotation, self.obb_scale)
+        else:
+            crop_obb = None
 
         pcd = generate_point_cloud(
             pipeline=pipeline,
@@ -309,6 +330,7 @@ def main(self) -> None:
             use_bounding_box=self.use_bounding_box,
             bounding_box_min=self.bounding_box_min,
             bounding_box_max=self.bounding_box_max,
+            crop_obb=crop_obb,
             std_ratio=self.std_ratio,
         )
         torch.cuda.empty_cache()

nerfstudio/utils/math.py

@@ -21,7 +21,7 @@
 from jaxtyping import Bool, Float
 from torch import Tensor
 
-from nerfstudio.utils.misc import torch_compile
+from nerfstudio.data.scene_box import OrientedBox
 
 
 def components_from_spherical_harmonics(
@@ -199,7 +199,7 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
     return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
 
 
-@torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
+# @torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
 def intersect_aabb(
     origins: torch.Tensor,
     directions: torch.Tensor,
@@ -240,6 +240,42 @@ def intersect_aabb(
     return t_min, t_max
 
 
+def intersect_obb(
+    origins: torch.Tensor,
+    directions: torch.Tensor,
+    obb: OrientedBox,
+    max_bound: float = 1e10,
+    invalid_value: float = 1e10,
+):
+    """
+    Ray intersection with an oriented bounding box (OBB)
+
+    Args:
+        origins: [N,3] tensor of 3d positions
+        directions: [N,3] tensor of normalized directions
+        R: [3,3] rotation matrix
+        T: [3] translation vector
+        S: [3] extents of the bounding box
+        max_bound: Maximum value of t_max
+        invalid_value: Value to return in case of no intersection
+    """
+    # Transform ray to OBB space
+    R, T, S = obb.R, obb.T, obb.S.to(origins.device)
+    H = torch.eye(4, device=origins.device, dtype=origins.dtype)
+    H[:3, :3] = R
+    H[:3, 3] = T
+    H_world2bbox = torch.inverse(H)
+    origins = torch.cat((origins, torch.ones_like(origins[..., :1])), dim=-1)
+    origins = torch.matmul(H_world2bbox, origins.T).T[..., :3]
+    directions = torch.matmul(H_world2bbox[:3, :3], directions.T).T
+
+    # Compute intersection with axis-aligned bounding box with min as -S and max as +S
+    aabb = torch.concat((-S / 2, S / 2))
+    t_min, t_max = intersect_aabb(origins, directions, aabb, max_bound=max_bound, invalid_value=invalid_value)
+
+    return t_min, t_max
+
+
 def safe_normalize(
     vectors: Float[Tensor, "*batch_dim N"],
     eps: float = 1e-10,