diff --git a/docs/nerfology/methods/in2n.md b/docs/nerfology/methods/in2n.md new file mode 100644 index 000000000..3f42cecbc --- /dev/null +++ b/docs/nerfology/methods/in2n.md @@ -0,0 +1,96 @@ +# Instruct-NeRF2NeRF + +

Editing 3D Scenes with Instructions

+ +```{button-link} https://instruct-nerf2nerf.github.io/ +:color: primary +:outline: +Paper Website +``` + +```{button-link} https://github.com/ayaanzhaque/instruct-nerf2nerf +:color: primary +:outline: +Code +``` + + + +**Instruct-NeRF2NeRF enables instruction-based editing of NeRFs via a 2D diffusion model** + +## Installation + +First install nerfstudio dependencies. Then run: + +```bash +pip install git+https://github.com/ayaanzhaque/instruct-nerf2nerf +``` + +## Running Instruct-NeRF2NeRF + +Details for running Instruct-NeRF2NeRF (built with Nerfstudio!) can be found [here](https://github.com/ayaanzhaque/instruct-nerf2nerf). Once installed, run: + +```bash +ns-train in2n --help +``` + +Three varients of Instruct-NeRF2NeRF are provided: + +| Method | Description | Memory | Quality | +| ----------- | ---------------------------- | ------ | ------- | +| `in2n` | Full model, used in paper | ~15GB | Best | +| `in2n-small` | Half precision model | ~12GB | Good | +| `in2n-tiny` | Half prevision with no LPIPS | ~10GB | Ok | + +## Method + +### Overview + +Instruct-NeRF2NeRF is a method for editing NeRF scenes with text-instructions. Given a NeRF of a scene and the collection of images used to reconstruct it, the method uses an image-conditioned diffusion model ([InstructPix2Pix](https://www.timothybrooks.com/instruct-pix2pix)) to iteratively edit the input images while optimizing the underlying scene, resulting in an optimized 3D scene that respects the edit instruction. The paper demonstrates that their method is able to edit large-scale, real-world scenes, and is able to accomplish more realistic, targeted edits than prior work. + +## Pipeline + + + +This section will walk through each component of the Instruct-NeRF2NeRF method. + +### How it Works + +Instruct-NeRF2NeRF gradually updates a reconstructed NeRF scene by iteratively updating the dataset images while training the NeRF: + +1. An image is rendered from the scene at a training viewpoint. +2. It is edited by InstructPix2Pix given a global text instruction. +3. The training dataset image is replaced with the edited image. +4. The NeRF continues training as usual. + +### Editing Images with InstructPix2Pix + +InstructPix2Pix is an image-editing diffusion model which can be prompted using text instructions. More details on InstructPix2Pix can be found [here](https://www.timothybrooks.com/instruct-pix2pix). + +Traditionally, at inference time, InstructPix2Pix takes as input random noise and is conditioned on an image (the image to edit) and a text instruction. The strength of the edit can be controlled using the image and text classifier-free guidance scales. + +To update a dataset image a given viewpoint, Instruct-NeRF2NeRF first takes the original, unedited training image as image conditioning and uses the global text instruction as text conditioning. The main input to the diffusion model is a noised version of the current render from the given viewpoint. The noise is sampled from a normal distribution and scaled based on a randomly chosen timestep. Then InstructPix2Pix slowly denoises the rendered image by predicting the noised version of the image at previous timesteps until the image is fully denoised. This will produce an edited version of the input image. + +This process mixes the information of the diffusion model, which attempts to edit the image, the current 3D structure of the NeRF, and view-consistent information from the unedited, ground-truth images. By combining this set of information, the edit is respected while maintaining 3D consistency. + +The code snippet for how an image is edited in the pipeline can be found [here](https://github.com/ayaanzhaque/instruct-nerf2nerf/blob/main/in2n/ip2p.py). + +### Iterative Dataset Update + +When NeRF training starts, the dataset consists of the original, unedited images used to train the original scene. These images are saved separately to use as conditioning for InstructPix2Pix. At each optimization step, some number of NeRF optimization steps are performed, and then some number of images (often just one) are updated. The images are randomly ordered prior to training and then at each step, the images are chosen in order to edit. Once an image has been edited, it is replaced in the dataset. Importantly, at each NeRF step, rays are sampled across the entire dataset, meaning there is a mixed source of supervision between edited images and unedited images. This allows for a gradual optimization that balances maintaining the 3D structure and consistency of the NeRF as well as performing the edit. + +At early iterations of this process, the edited images may be inconistent with one another, as InstructPix2Pix often doesn't perform consistent edits across viewpoints. However, over time, since images are edited using the current render of the NeRF, the edits begin to converge towards a globally consistent depiction of the underlying scene. Here is an example of how the underlying dataset evolves and becomes more consistent. + + + +The traditional method for supervising NeRFs using diffusion models is to use a Score Distillation Sampling (SDS) loss, as proposed in [DreamFusion](https://dreamfusion3d.github.io/). The Iterative Dataset Update method can be viewed as a variant of SDS, as instead of updating a discrete set of images at each step, the loss is a mix of rays from various viewpoints which are edited to varying degrees. The results show that this leads to higher quality performance and more stable optimization. + +## Results + +For results, view the [project page](https://instruct-nerf2nerf.github.io/)! diff --git a/docs/nerfology/methods/index.md b/docs/nerfology/methods/index.md index 9a821b623..59b91ad93 100644 --- a/docs/nerfology/methods/index.md +++ b/docs/nerfology/methods/index.md @@ -1,8 +1,8 @@ # Methods -We provide a set of pre implemented nerfstudio methods. +We provide a set of pre implemented nerfstudio methods. -**The goal of nerfstudio is to modularize the various NeRF techniques as much as possible.** +**The goal of nerfstudio is to modularize the various NeRF techniques as much as possible.** As a result, many of the techniques from these pre-implemented methods can be mixed 🎨. @@ -20,9 +20,9 @@ To list the available methods run: ns-train --help ``` -## Guides +## Methods -In addition to their implementations, we have provided guides that walk through each of these method. +The following methods are supported in nerfstudio: ```{toctree} :maxdepth: 1 @@ -30,6 +30,7 @@ In addition to their implementations, we have provided guides that walk through Mip-NeRF Nerfacto Instant-NGP + Instruct-NeRF2NeRF Semantic NeRF-W CleaNeRF ``` diff --git a/nerfstudio/cameras/camera_paths.py b/nerfstudio/cameras/camera_paths.py index a313cad4d..9754a2bbf 100644 --- a/nerfstudio/cameras/camera_paths.py +++ b/nerfstudio/cameras/camera_paths.py @@ -125,10 +125,6 @@ def get_path_from_json(camera_path: Dict[str, Any]) -> Cameras: image_height = camera_path["render_height"] image_width = camera_path["render_width"] - if "camera_type" in camera_path: - camera_type = camera_path["camera_type"] - else: - camera_type = "perspective" if "camera_type" not in camera_path: camera_type = CameraType.PERSPECTIVE diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 79490ae67..8486f1097 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -164,8 +164,8 @@ def quaternion_matrix(quaternion: ArrayLike) -> np.ndarray: def get_interpolated_poses(pose_a: ArrayLike, pose_b: ArrayLike, steps: int = 10) -> List[float]: """Return interpolation of poses with specified number of steps. Args: - poseA: first pose - poseB: second pose + pose_a: first pose + pose_b: second pose steps: number of steps the interpolated pose path should contain """ @@ -190,8 +190,8 @@ def get_interpolated_k(k_a, k_b, steps: int = 10) -> TensorType[3, 4]: Returns interpolated path between two camera poses with specified number of steps. Args: - KA: camera matrix 1 - KB: camera matrix 2 + k_a: camera matrix 1 + k_b: camera matrix 2 steps: number of steps the interpolated pose path should contain """ Ks = [] diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index 0184b4c80..61746187c 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -57,7 +57,7 @@ class Frustums(TensorDataclass): return pos def get_start_positions(self) -> TensorType[..., 3]: - """Calulates "start" position of frustum. + """Calculates "start" position of frustum. Returns: xyz positions. @@ -193,7 +193,7 @@ class RayBundle(TensorDataclass): """Times at which rays are sampled""" def set_camera_indices(self, camera_index: int) -> None: - """Sets all of the the camera indices to a specific camera index. + """Sets all the camera indices to a specific camera index. Args: camera_index: Camera index. diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py index 9b692c65b..2b01f77d2 100644 --- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py @@ -51,7 +51,7 @@ class SDFStudioDataParserConfig(DataParserConfig): scene_scale: float = 2.0 """ Sets the bounding cube to have edge length of this size. - The longest dimension of the Friends axis-aligned bbox will be scaled to this value. + The longest dimension of the axis-aligned bbox will be scaled to this value. """ skip_every_for_val_split: int = 1 """sub sampling validation images""" @@ -141,8 +141,8 @@ class SDFStudio(DataParser): # TODO supports downsample # cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) - - assert meta["has_mono_prior"] == self.config.include_mono_prior, f"no mono prior in {self.config.data}" + if self.config.include_mono_prior: + assert meta["has_mono_prior"], f"no mono prior in {self.config.data}" dataparser_outputs = DataparserOutputs( image_filenames=image_filenames, diff --git a/nerfstudio/data/utils/nerfstudio_collate.py b/nerfstudio/data/utils/nerfstudio_collate.py index ffcd989b4..91774343a 100644 --- a/nerfstudio/data/utils/nerfstudio_collate.py +++ b/nerfstudio/data/utils/nerfstudio_collate.py @@ -23,7 +23,6 @@ from typing import Callable, Dict, Union import torch import torch.utils.data -from torch._six import string_classes from nerfstudio.cameras.cameras import Cameras @@ -119,7 +118,7 @@ def nerfstudio_collate( return torch.tensor(batch, dtype=torch.float64) elif isinstance(elem, int): return torch.tensor(batch) - elif isinstance(elem, string_classes): + elif isinstance(elem, (str, bytes)): return batch elif isinstance(elem, collections.abc.Mapping): try: diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 0f9015d45..763580bb4 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -242,6 +242,14 @@ class Trainer: writer.put_scalar(name="Train Loss", scalar=loss, step=step) writer.put_dict(name="Train Loss Dict", scalar_dict=loss_dict, step=step) writer.put_dict(name="Train Metrics Dict", scalar_dict=metrics_dict, step=step) + # The actual memory allocated by Pytorch. This is likely less than the amount + # shown in nvidia-smi since some unused memory can be held by the caching + # allocator and some context needs to be created on GPU. See Memory management + # (https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management) + # for more details about GPU memory management. + writer.put_scalar( + name="GPU Memory (MB)", scalar=torch.cuda.max_memory_allocated() / (1024**2), step=step + ) # Do not perform evaluation if there are no validation images if self.pipeline.datamanager.eval_dataset: diff --git a/nerfstudio/exporter/marching_cubes.py b/nerfstudio/exporter/marching_cubes.py new file mode 100644 index 000000000..03cd7da86 --- /dev/null +++ b/nerfstudio/exporter/marching_cubes.py @@ -0,0 +1,253 @@ +# Copyright 2022 The Nerfstudio Team. 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. + +""" +This module implements the Marching Cubes algorithm for extracting +isosurfaces +""" + +from typing import Callable, List, Tuple, Union + +import numpy as np +import torch +import trimesh +from skimage import measure +from torchtyping import TensorType + +avg_pool_3d = torch.nn.AvgPool3d(2, stride=2) +upsample = torch.nn.Upsample(scale_factor=2, mode="nearest") + + +def create_point_pyramid( + points: TensorType[3, "height", "width", "depth"] +) -> List[TensorType[3, "height", "width", "depth"]]: + """ + Create a point pyramid for multi-resolution evaluation. + + Args: + points: A torch tensor containing 3D points. + + Returns: + A list of torch tensors representing points at different resolutions. + """ + points_pyramid = [points] + for _ in range(3): + points = avg_pool_3d(points[None])[0] + points_pyramid.append(points) + points_pyramid = points_pyramid[::-1] + return points_pyramid + + +def evaluate_sdf(sdf: Callable, points: TensorType["batch", 3]) -> TensorType["batch"]: + """ + Evaluate a signed distance function (SDF) for a batch of points. + + Args: + sdf: A callable function that takes a tensor of size (N, 3) containing + 3D points and returns a tensor of size (N,) with the SDF values. + points: A torch tensor containing 3D points. + + Returns: + A torch tensor with the SDF values evaluated at the given points. + """ + z = [] + for _, pnts in enumerate(torch.split(points, 100000, dim=0)): + z.append(sdf(pnts)) + z = torch.cat(z, axis=0) + return z + + +def evaluate_multiresolution_sdf( + evaluate: Callable, + points_pyramid: List[TensorType[3, "height", "width", "depth"]], + coarse_mask: Union[TensorType[1, 1, "height", "width", "depth"], None], + x_max: float, + x_min: float, + crop_n: int, +) -> TensorType["batch"]: + + """ + Evaluate SDF values using a multi-resolution approach with a given point pyramid. + + Args: + evaluate: A callable function to evaluate SDF values at given points. + points_pyramid: A list of torch tensors representing points at different resolutions. + coarse_mask: A binary mask tensor indicating regions of the bounding box where the SDF + is expected to have a zero-crossing. + x_max: The maximum x-coordinate of the bounding box. + x_min: The minimum x-coordinate of the bounding box. + crop_n: The resolution of the grid used to sample the SDF. + + Returns: + A torch tensor with the SDF values evaluated at the given points. + """ + mask = None + threshold = 2 * (x_max - x_min) / crop_n * 8 + for pid, pts in enumerate(points_pyramid): + coarse_n = pts.shape[-1] + pts = pts.reshape(3, -1).permute(1, 0).contiguous() + + if mask is None: + # Only evaluate SDF + if coarse_mask is not None: + pts_sdf = torch.ones_like(pts[:, 1]) + valid_mask = torch.nn.functional.grid_sample(coarse_mask, pts[None, None, None])[0, 0, 0, 0] > 0 + if valid_mask.any(): + pts_sdf[valid_mask] = evaluate(pts[valid_mask].contiguous()) + else: + pts_sdf = evaluate(pts) + else: + mask = mask.reshape(-1) + pts_to_eval = pts[mask] + + if pts_to_eval.shape[0] > 0: + pts_sdf_eval = evaluate(pts_to_eval.contiguous()) + pts_sdf[mask] = pts_sdf_eval + + if pid < 3: + # Update mask + mask = torch.abs(pts_sdf) < threshold + mask = mask.reshape(coarse_n, coarse_n, coarse_n)[None, None] + mask = upsample(mask.float()).bool() + + pts_sdf = pts_sdf.reshape(coarse_n, coarse_n, coarse_n)[None, None] + pts_sdf = upsample(pts_sdf) + pts_sdf = pts_sdf.reshape(-1) + + threshold /= 2.0 + return pts_sdf + + +@torch.no_grad() +def generate_mesh_with_multires_marching_cubes( + geometry_callable_field: Callable, + resolution: int = 512, + 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), + isosurface_threshold: float = 0.0, + coarse_mask: Union[None, TensorType["height", "width", "depth"]] = None, +) -> trimesh.Trimesh: + """ + Computes the isosurface of a signed distance function (SDF) defined by the + callable `sdf` in a given bounding box with a specified resolution. The SDF + is sampled at a set of points within a regular grid, and the marching cubes + algorithm is used to generate a mesh that approximates the isosurface at a + specified isovalue `level`. + + Args: + sdf: A callable function that takes as input a tensor of size + (N, 3) containing 3D points, and returns a tensor of size (N,) containing + the signed distance function evaluated at those points. + output_path: The output directory where the resulting mesh will be saved. + resolution: The resolution of the grid used to sample the SDF. + bounding_box_min: The minimum coordinates of the bounding box in which the SDF + will be evaluated. + bounding_box_max: The maximum coordinates of the bounding box in which the SDF + will be evaluated. + isosurface_threshold: The isovalue at which to approximate the isosurface. + coarse_mask: A binary mask tensor of size ("height", "width", "depth") that indicates the regions + of the bounding box where the SDF is expected to have a zero-crossing. If + provided, the algorithm first evaluates the SDF at the coarse voxels where + the mask is True, and then refines the evaluation within these voxels using + a multi-scale approach. If None, evaluates the SDF at all points in the + bounding box. + Returns: + A torch tensor with the SDF values evaluated at the given points. + """ + # Check if resolution is divisible by 512 + assert ( + resolution % 512 == 0 + ), f"""resolution must be divisible by 512, got {resolution}. + This is important because the algorithm uses a multi-resolution approach + to evaluate the SDF where the mimimum resolution is 512.""" + # Prepare coarse mask if provided + if coarse_mask is not None: + coarse_mask = coarse_mask.permute(2, 1, 0)[None, None].cuda().float() + + # Initialize variables + crop_n = 512 + N = resolution // crop_n + grid_min = bounding_box_min + grid_max = bounding_box_max + xs = np.linspace(grid_min[0], grid_max[0], N + 1) + ys = np.linspace(grid_min[1], grid_max[1], N + 1) + zs = np.linspace(grid_min[2], grid_max[2], N + 1) + + # Initialize meshes list + meshes = [] + + # Iterate over the grid + for i in range(N): + for j in range(N): + for k in range(N): + # Calculate grid cell boundaries + x_min, x_max = xs[i], xs[i + 1] + y_min, y_max = ys[j], ys[j + 1] + z_min, z_max = zs[k], zs[k + 1] + + # Create point grid + x = np.linspace(x_min, x_max, crop_n) + y = np.linspace(y_min, y_max, crop_n) + z = np.linspace(z_min, z_max, crop_n) + xx, yy, zz = np.meshgrid(x, y, z, indexing="ij") + points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).cuda() + + # Function to evaluate SDF for a batch of points + def evaluate(points: torch.Tensor) -> torch.Tensor: + return evaluate_sdf(geometry_callable_field, points) + + # Construct point pyramids + points = points.reshape(crop_n, crop_n, crop_n, 3).permute(3, 0, 1, 2) + if coarse_mask is not None: + points_tmp = points.permute(1, 2, 3, 0)[None].cuda() + current_mask = torch.nn.functional.grid_sample(coarse_mask, points_tmp) + current_mask = (current_mask > 0.0).cpu().numpy()[0, 0] + else: + current_mask = None + + # Create point pyramid for multi-resolution evaluation + points_pyramid = create_point_pyramid(points) + + # Evaluate SDF using multi-resolution approach and mask + pts_sdf = evaluate_multiresolution_sdf(evaluate, points_pyramid, coarse_mask, x_max, x_min, crop_n) + + z = pts_sdf.detach().cpu().numpy() + + # Skip if no surface found + if current_mask is not None: + valid_z = z.reshape(crop_n, crop_n, crop_n)[current_mask] + if valid_z.shape[0] <= 0 or ( + np.min(valid_z) > isosurface_threshold or np.max(valid_z) < isosurface_threshold + ): + continue + + if not (np.min(z) > isosurface_threshold or np.max(z) < isosurface_threshold): + z = z.astype(np.float32) + verts, faces, normals, _ = measure.marching_cubes( + volume=z.reshape(crop_n, crop_n, crop_n), + level=isosurface_threshold, + spacing=( + (x_max - x_min) / (crop_n - 1), + (y_max - y_min) / (crop_n - 1), + (z_max - z_min) / (crop_n - 1), + ), + mask=current_mask, + ) + verts = verts + np.array([x_min, y_min, z_min]) + + meshcrop = trimesh.Trimesh(verts, faces, normals) + meshes.append(meshcrop) + + combined_mesh: trimesh.Trimesh = trimesh.util.concatenate(meshes) + return combined_mesh diff --git a/nerfstudio/exporter/texture_utils.py b/nerfstudio/exporter/texture_utils.py index 4001d1c60..d7e08c669 100644 --- a/nerfstudio/exporter/texture_utils.py +++ b/nerfstudio/exporter/texture_utils.py @@ -401,6 +401,7 @@ def export_textured_mesh( camera_indices = torch.zeros_like(origins[..., 0:1]) nears = torch.zeros_like(origins[..., 0:1]) fars = torch.ones_like(origins[..., 0:1]) * raylen + directions_norm = torch.ones_like(origins[..., 0:1]) # for surface model camera_ray_bundle = RayBundle( origins=origins, directions=directions, @@ -408,6 +409,7 @@ def export_textured_mesh( camera_indices=camera_indices, nears=nears, fars=fars, + metadata={"directions_norm": directions_norm}, ) CONSOLE.print("Creating texture image by rendering with NeRF...") diff --git a/pyproject.toml b/pyproject.toml index 969830d88..9ea93de89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,8 @@ dependencies = [ "u-msgpack-python>=2.4.1", "nuscenes-devkit>=1.1.1", "wandb>=0.13.3", - "xatlas" + "xatlas", + "trimesh>=3.20.2" ] [project.urls] diff --git a/scripts/exporter.py b/scripts/exporter.py index cd18aa354..5a58a99b3 100644 --- a/scripts/exporter.py +++ b/scripts/exporter.py @@ -27,6 +27,9 @@ from nerfstudio.exporter.exporter_utils import ( generate_point_cloud, get_mesh_from_filename, ) +from nerfstudio.exporter.marching_cubes import ( + generate_mesh_with_multires_marching_cubes, +) from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline from nerfstudio.utils.eval_utils import eval_setup @@ -313,14 +316,68 @@ class ExportPoissonMesh(Exporter): @dataclass class ExportMarchingCubesMesh(Exporter): - """ - NOT YET IMPLEMENTED - Export a mesh using marching cubes. - """ + """Export a mesh using marching cubes.""" + + isosurface_threshold: float = 0.0 + """The isosurface threshold for extraction. For SDF based methods the surface is the zero level set.""" + resolution: int = 1024 + """Marching cube resolution.""" + simplify_mesh: bool = False + """Whether to simplify the mesh.""" + bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0) + """Minimum of the bounding box.""" + bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0) + """Maximum of the bounding box.""" + px_per_uv_triangle: int = 4 + """Number of pixels per UV triangle.""" + unwrap_method: Literal["xatlas", "custom"] = "xatlas" + """The method to use for unwrapping the mesh.""" + num_pixels_per_side: int = 2048 + """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.""" def main(self) -> None: - """Export mesh""" - raise NotImplementedError("Marching cubes not implemented yet.") + """Main function.""" + if not self.output_dir.exists(): + self.output_dir.mkdir(parents=True) + + _, pipeline, _ = eval_setup(self.load_config) + + # TODO: Make this work with Density Field + assert hasattr(pipeline.model.config, "sdf_field"), "Model must have an SDF field." + + CONSOLE.print("Extracting mesh with marching cubes... which may take a while") + + assert ( + self.resolution % 512 == 0 + ), f"""resolution must be divisible by 512, got {self.resolution}. + This is important because the algorithm uses a multi-resolution approach + to evaluate the SDF where the mimimum resolution is 512.""" + + # Extract mesh using marching cubes for sdf at a multi-scale resolution. + multi_res_mesh = generate_mesh_with_multires_marching_cubes( + geometry_callable_field=lambda x: pipeline.model.field.forward_geonetwork(x)[:, 0].contiguous(), + resolution=self.resolution, + bounding_box_min=self.bounding_box_min, + bounding_box_max=self.bounding_box_max, + isosurface_threshold=self.isosurface_threshold, + coarse_mask=None, + ) + filename = self.output_dir / "sdf_marching_cubes_mesh.ply" + multi_res_mesh.export(filename) + + # load the mesh from the marching cubes export + mesh = get_mesh_from_filename(str(filename), target_num_faces=self.target_num_faces) + CONSOLE.print("Texturing mesh with NeRF...") + texture_utils.export_textured_mesh( + mesh, + pipeline, + self.output_dir, + px_per_uv_triangle=self.px_per_uv_triangle if self.unwrap_method == "custom" else None, + unwrap_method=self.unwrap_method, + num_pixels_per_side=self.num_pixels_per_side, + ) @dataclass