From 3ba27e723a5c8c2a4c7eabd27f6aac57a10475e3 Mon Sep 17 00:00:00 2001 From: nepfaff Date: Wed, 2 Oct 2024 09:50:04 -0400 Subject: [PATCH] Add option for refining the TSDF using the mesh's AABB --- nerfstudio/exporter/tsdf_utils.py | 29 +++++++++++++++++++++++++++++ nerfstudio/scripts/exporter.py | 7 +++++++ 2 files changed, 36 insertions(+) diff --git a/nerfstudio/exporter/tsdf_utils.py b/nerfstudio/exporter/tsdf_utils.py index 9be7b3d9fe..f3f1e53571 100644 --- a/nerfstudio/exporter/tsdf_utils.py +++ b/nerfstudio/exporter/tsdf_utils.py @@ -285,6 +285,8 @@ def export_tsdf_mesh( 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), + refine_mesh_using_initial_aabb_estimate: bool = False, + refinement_epsilon: float = 1e-2, ) -> None: """Export a TSDF mesh from a pipeline. @@ -299,6 +301,9 @@ def export_tsdf_mesh( use_bounding_box: Whether to use a bounding box for the TSDF volume. bounding_box_min: Minimum coordinates of the bounding box. bounding_box_max: Maximum coordinates of the bounding box. + refine_mesh_using_initial_aabb_estimate: Whether to refine the TSDF using the initial AABB estimate. + refinement_epsilon: Epsilon for refining the TSDF. This is the distance in meters that the refined AABB/OBB will + be expanded by in each direction. """ device = pipeline.device @@ -356,5 +361,29 @@ def export_tsdf_mesh( CONSOLE.print("Computing Mesh") mesh = tsdf.get_mesh() + + if refine_mesh_using_initial_aabb_estimate: + CONSOLE.print("Refining the TSDF based on the Mesh AABB") + + # Compute the AABB of the mesh and use it to initialize a new TSDF + vertices_min = torch.min(mesh.vertices, dim=0).values - refinement_epsilon + vertices_max = torch.max(mesh.vertices, dim=0).values + refinement_epsilon + aabb = torch.stack([vertices_min, vertices_max]).cpu() + tsdf = TSDF.from_aabb(aabb, volume_dims=volume_dims) + # move TSDF to device + tsdf.to(device) + + CONSOLE.print("Integrating the updated TSDF") + for i in range(0, len(c2w), batch_size): + tsdf.integrate_tsdf( + c2w[i : i + batch_size], + K[i : i + batch_size], + depth_images[i : i + batch_size], + color_images=color_images[i : i + batch_size], + ) + + CONSOLE.print("Computing the updated Mesh") + mesh = tsdf.get_mesh() + CONSOLE.print("Saving TSDF Mesh") tsdf.export_mesh(mesh, filename=str(output_dir / "tsdf_mesh.ply")) diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index 650eaebff9..fccb70ffba 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -218,6 +218,11 @@ class ExportTSDFMesh(Exporter): """If using xatlas for unwrapping, the pixels per side of the texture image.""" target_num_faces: Optional[int] = 50000 """Target number of faces for the mesh to texture.""" + refine_mesh_using_initial_aabb_estimate: bool = False + """Refine the mesh using the initial AABB estimate.""" + refinement_epsilon: float = 1e-2 + """Refinement epsilon for the mesh. This is the distance in meters that the refined AABB/OBB will be expanded by + in each direction.""" def main(self) -> None: """Export mesh""" @@ -238,6 +243,8 @@ 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, + refine_mesh_using_initial_aabb_estimate=self.refine_mesh_using_initial_aabb_estimate, + refinement_epsilon=self.refinement_epsilon, ) # possibly