nerfstudio-project / nerfstudio

A collaboration friendly studio for NeRFs
https://docs.nerf.studio
Apache License 2.0
9.52k stars 1.3k forks source link

The marching cubes extraction requires an SDF field #2212

Open bhavikajalli opened 1 year ago

bhavikajalli commented 1 year ago

Describe the bug Hi, I am trying to extract a mesh from a trained perfecto model using marching cubes. I use the code: ns-export marching-cubes --load-config outputs/output/nerfacto/2023-07-12_193830/config.yml --output-dir outputs/output/nerfacto/2023-07-12_193830/

However I get the following error: AssertionError: Model must have an SDF field. To Reproduce Steps to reproduce the behavior:

  1. Train a nerfacto model ns-train nerfacto --data data/ --pipeline.model.predict-normals True
  2. Extract a marching cubes mesh ns-export marching-cubes --load-config outputs/output/nerfacto/2023-07-12_193830/config.yml --output-dir outputs/output/nerfacto/2023-07-12_193830/

Expected behavior A mesh created using marching cubes

Additional context I do see that the code(scripts/exporter.py) has

TODO: Make this work with Density Field

assert hasattr(pipeline.model.config, "sdf_field"), "Model must have an SDF field."

How can I train a nerfacto model with an SDF field or is there an intermediary step?

bhavikajalli commented 1 year ago

I think maybe this can be solved by extracting a dense grid but I am not sure how.

Zhuoyang-Pan commented 1 year ago

Instead of obtaining an SDF field, I believe there is a naive solution to tackle this problem. Here's what I've come up with:

    # TODO: Make this work with Density Field
-   assert hasattr(pipeline.model.config, "sdf_field"), "Model must have an SDF field."
+   if hasattr(pipeline.model.config, "sdf_field"):
+       geometry_callable_field_fn = lambda x: cast(SDFField, pipeline.model.field).forward_geonetwork(x)[:, 0].contiguous()
+   else:
+       geometry_callable_field_fn = lambda x: pipeline.model.field.density_fn(x)[:, 0].contiguous()

    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 minimum 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: cast(SDFField, pipeline.model.field)
-        .forward_geonetwork(x)[:, 0]
-        .contiguous(),
+       geometry_callable_field=geometry_callable_field_fn,
        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,
    )

where you could use density as a threshold for marching cubes. In my experiments, I've found 2.5 to be a promising threshold. You can implement it as:

ns-export marching-cubes --isosurface-threshold 2.5 --load-config xxx --output-dir xxx

However, the result of this naive method might have relatively poor quality sometimes. If you're in search of higher-quality outcomes, Poisson surface reconstruction tends to yield better results. You may also consider using sdfstudio to obtain a sdf field first.

Hope that this helps~