Awesome

Differentiable Iso-Surface Extraction Package (DISO)

[News] DISO now supports batch training and checkpointing! It is no longer necessary to allocate multiple iso-surface extractors.

This repository consists of a variety of differentiable iso-surface extraction algorithms implemented in cuda.

Currently, two algorithms are incorporated:

• Differentiable Marching Cubes [1] (DiffMC)
• Differentiable Dual Marching Cubes [2] (DiffDMC)

The differentiable iso-surface algorithms have multiple applications in gradient-based optimization, such as shape, texture, materials reconstruction from images.

Installation

Requirements: torch (must be compatible with CUDA version), trimesh

pip install diso

Quick Start

You can effortlessly try the following command, which converts a sphere SDF into triangle mesh using different algorithms. The generated results are saved in out/.

python test/example.py

Note:

• DiffMC and DiffDMC generate watertight manifold meshes when grid deformation is disabled. When enabling grid deformation, self-intersection may occur, but the face connectivity remains manifold.
• DiffDMC can produce a more uniform triangle distribution and smoother surfaces than DiffMC and supports generating quad meshes (in the example, the quad is automatically divided into two triangles by trimesh).

Usage

All the functions share the same interfaces. Firstly you should build an iso-surface extractor as follows:

from diso import DiffMC
from diso import DiffDMC

diffmc = DiffMC(dtype=torch.float32).cuda() # or dtype=torch.float64
diffdmc = DiffDMC(dtype=torch.float32).cuda() # or dtype=torch.float64

Then use its forward function to generate a single mesh:

verts, faces = diffmc(sdf, deform, normalize=True)  # or deform=None
verts, faces = diffdmc(sdf, deform, return_quads=False, normalize=True)  # or deform=None

Input

• sdf: queries SDF values on the grid vertices (see the test/example.py for how to create the grid). The gradient will be back-propagated to the source that generates the SDF values. ([N, N, N, 3])
• deform (optional): (learnable) deformation values on the grid vertices, the range must be [-0.5, 0.5], default=None. ([N, N, N, 3])
• normalize: whether to normalize the output vertices, default=True. If set to True, the vertices are normalized to [0, 1]. When False, the vertices remain unnormalized as [0, dim-1],
• return_quads: whether return quad meshes; only applicable for DiffDMC, default=False. If set to True, the function returns quad meshes ([F, 4]).

Output

• verts: mesh vertices within the range of [0, 1] or [0, dim-1]. ([V, 3])
• faces: mesh face indices (starting from 0). ([F, 3]).

The gradient will be automatically computed when backward function is called.

Batch Training

You can loop over the batch size to conduct an iso-surface extraction on each object, then compute the loss for the backward pass.

extractor = DiffMC(dtype=torch.float32).cuda()  # or DiffDMC
for bs in range(batchsize):
    verts, faces = extractor(sdf, deform)
    # compute and accumulate losses
loss.backward()

Note: It is suggested to create the extractors outside the epoch loop so that they can be reused by different iterations and avoid allocating memory every time.

Speed Comparison

We compare our library with DMTet [3] and FlexiCubes [4] on two examples: a simple round cube and a random initialized signed distance function.

The algorithms have been rigorously tested on an NVIDIA RTX 4090 GPU. Each algorithm underwent 100 repeated runs, and the table presents the time and CUDA memory consumption for a single run.

Citation

If you find this repository useful, please cite the following paper:

@article{wei2023neumanifold,
  title={NeuManifold: Neural Watertight Manifold Reconstruction with Efficient and High-Quality Rendering Support},
  author={Wei, Xinyue and Xiang, Fanbo and Bi, Sai and Chen, Anpei and Sunkavalli, Kalyan and Xu, Zexiang and Su, Hao},
  journal={arXiv preprint arXiv:2305.17134},
  year={2023}
}

Reference

[1] We L S. Marching cubes: A high resolution 3d surface construction algorithm[J]. Comput Graph, 1987, 21: 163-169.

[2] Nielson G M. Dual marching cubes[C]//IEEE visualization 2004. IEEE, 2004: 489-496.

[3] Shen T, Gao J, Yin K, et al. Deep marching tetrahedra: a hybrid representation for high-resolution 3d shape synthesis[J]. Advances in Neural Information Processing Systems, 2021, 34: 6087-6101.

[4] Shen T, Munkberg J, Hasselgren J, et al. Flexible isosurface extraction for gradient-based mesh optimization[J]. ACM Transactions on Graphics (TOG), 2023, 42(4): 1-16.