FlexiCubes

repository·main·Indexed 21 days ago

A high-quality isosurface representation for gradient-based mesh optimization. FlexiCubes enables the reconstruction of 3D shapes by optimizing a Signed Distance Field (SDF) using geometric, visual, or physical objectives. Core functions are integrated into the Kaolin library (v0.15.0+) via `kaolin.ops.conversions.FlexiCubes`. The repository provides tools for mesh extraction, deformation, and regularization losses to ensure shape quality and training stability.

Tokens: 2.7K
Snippets: 5
Records: 10
Agent score: 25%

What's inside FlexiCubes

Optimize regularization loss weights in FlexiCubes
main
When performing mesh optimization, you can use regularizers to control the shape quality. It is recommended to start with low weights to avoid hindering convergence, then incrementally increase them if artifacts appear.
Common regularizers used in the examples/optimize.py pipeline:
- Floater Removal Loss: Helps remove floaters in un-supervised areas (e.g., internal faces).
- L_dev loss: Use this to reduce artifacts in flat areas (developability).
- L1 regularizer on flexible weights: Primarily used to stabilize training in generative pipelines like GET3D.
https://github.com/nv-tlabs/flexicubes/blob/main/README.md
Compare FlexiCubes resolution with DMTet
main
If migrating from DMTet to FlexiCubes, note that FlexiCubes uses a tetrahedral grid which is denser than a voxel grid for the same resolution n (where n is the edge length).
- DMTet (Voxel): (n+1)³ vertices.
- FlexiCubes (Tetrahedral): (n/2+1)³ grid vertices, but results in a denser output mesh.
Recommendation: To achieve a similar triangle count in the output mesh, use a 4:5 resolution ratio between the voxel grid and the tetrahedral grid. For example, a 64³ FlexiCubes grid produces approximately the same number of triangles as an 80³ DMTet grid.
https://github.com/nv-tlabs/flexicubes/blob/main/README.md
Run Gradient-Based Mesh Optimization
main
FlexiCubes can be used to reconstruct meshes by optimizing a Signed Distance Field (SDF) towards a reference mesh using geometric and depth losses.
1. Download Data: Run python download_data.py inside the examples directory to get the necessary datasets.
2. Run Optimization: Use optimize.py to start the process.
Command Line Arguments:
- --ref_mesh: Path to the reference .obj mesh.
- --out_dir: Directory where results will be saved.
- --develop_reg: (Optional) Set to True to add a developability regularizer to encourage fabricability from panels.
- --iter: (Optional) Number of optimization iterations.
```
# Basic optimization
python optimize.py --ref_mesh data/inputmodels/block.obj --out_dir out/block

# Optimization with developability regularization
python optimize.py --ref_mesh data/inputmodels/david.obj --out_dir out/david_dev --develop_reg True --iter=1250
```
https://github.com/nv-tlabs/flexicubes/blob/main/README.md

Set up the FlexiCubes Conda environment

main

To run the provided optimization examples, create a Conda environment with the following dependencies. This setup includes PyTorch, nvdiffrast, and Kaolin.

conda create -n flexicubes python=3.9
conda activate flexicubes
conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.3 -c pytorch
pip install imageio trimesh tqdm matplotlib torch_scatter ninja
pip install git+https://github.com/NVlabs/nvdiffrast/
pip install kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-1.12.0_cu113.html

Install FlexiCubes via Kaolin
main
The core functions of FlexiCubes are integrated into the Kaolin library starting from version v0.15.0. For production use, it is recommended to use the Kaolin implementation rather than the original flexicube.py file provided in this repository.
Refer to the Kaolin documentation for specific installation instructions and the API documentation for kaolin.ops.conversions.FlexiCubes.
https://github.com/nv-tlabs/flexicubes/blob/main/README.md
Extract mesh from a known Signed Distance Field (SDF)
main
While FlexiCubes is optimized for gradient-based reconstruction, you can also use it to extract a mesh directly from a known SDF without an optimization loop.
For a detailed walkthrough, refer to the tutorial in examples/extraction.ipynb.
https://github.com/nv-tlabs/flexicubes/blob/main/README.md
Apply regularization losses in FlexiCubes optimization
main
To ensure high-quality meshes and stable optimization, several regularization losses are recommended:
1. SDF Regularization: Uses loss.sdf_reg_loss(sdf, grid_edges) to smooth the scalar field. It is often useful to decay the weight of this loss over time.
2. Deformation Regularization: Uses the L_dev term returned by the FlexiCubes operator.
3. Weight Regularization: Penalizing the absolute value of the $\alpha$ and $\beta$ weights (weight[:,:20].abs().mean()) helps keep the optimization stable by encouraging the weights to stay near zero unless needed.
4. Internal Element Removal: A specific SDF regularization (as described in the nvdiffrec paper) can be used to remove internal floating elements that are not visible to the user.
https://github.com/nv-tlabs/flexicubes/blob/main/examples/optimization.ipynb
Perform gradient-based mesh optimization with FlexiCubes
main
FlexiCubes allows for gradient-based mesh optimization by representing a surface as the isosurface of a scalar field (SDF). This enables direct evaluation of objectives on the extracted surface while allowing for topological flexibility.
To optimize a mesh, you typically parameterize and optimize the following components:
1. SDF (sdf): The scalar field defining the surface.
2. Weights (weight): Per-cube learnable weights (e.g., $\beta$, $\alpha$, $\gamma$) that control the local geometry.
3. Deformations (deform): A displacement field applied to the voxel grid vertices.
In a typical optimization loop, you sample camera poses, render the current FlexiCubes mesh, compute reconstruction losses (like mask and depth loss) and regularization losses, and then backpropagate the total loss to update the parameters.
```
# Typical optimization loop structure
for it in tqdm.tqdm(range(iter)):
    optimizer.zero_grad()
    
    # 1. Sample cameras
    cameras = render.get_random_camera_batch(batch, iter_res=train_res, device=device)

    # 2. Extract mesh (use training=True during optimization)
    grid_verts = x_nx3 + (2-1e-8) / (voxel_grid_res * 2) * torch.tanh(deform)
    vertices, faces, L_dev = fc(
        grid_verts, sdf, cube_fx8, voxel_grid_res, 
        beta=weight[:,:12], alpha=weight[:,12:20], gamma_f=weight[:,20], 
        training=True
    )
    flexicubes_mesh = kal.rep.SurfaceMesh(vertices=vertices, faces=faces)

    # 3. Render and compute losses
    buffers = render.render_mesh(flexicubes_mesh, cameras, train_res)
    mask_loss = (buffers['mask'] - target['mask']).abs().mean()
    # ... compute other losses ...

    total_loss.backward()
    optimizer.step()
```
https://github.com/nv-tlabs/flexicubes/blob/main/examples/optimization.ipynb

Initialize FlexiCubes voxel grid and parameters

main

To use FlexiCubes, you first initialize the FlexiCubes operator and construct a voxel grid. You then define the learnable parameters for the SDF, weights, and deformations.

Note that the weight parameter is a tensor of shape (num_cubes, 21), where specific slices are used for $\beta$, $\alpha$, and $\gamma$ parameters during the forward pass.

import kaolin as kal

# Initialize operator
fc = kal.ops.conversions.FlexiCubes(device)

# Construct voxel grid
x_nx3, cube_fx8 = fc.construct_voxel_grid(voxel_grid_res)

# Initialize learnable parameters
sdf = torch.nn.Parameter(torch.rand_like(x_nx3[:,0]) - 0.1, requires_grad=True)
weight = torch.nn.Parameter(torch.zeros((cube_fx8.shape[0], 21), device='cuda'), requires_grad=True)
deform = torch.nn.Parameter(torch.zeros_like(x_nx3), requires_grad=True)

# Extract grid edges for regularization
all_edges = cube_fx8[:, fc.cube_edges].reshape(-1, 2)
grid_edges = torch.unique(all_edges, dim=0)

Use the FlexiCubes operator for mesh extraction
main
The FlexiCubes operator (accessed via fc(...)) extracts vertices and faces from a voxel grid and scalar field.
Key Arguments:
- grid_verts: The (N, 3) positions of the voxel grid vertices (can be deformed).
- sdf: The (N,) scalar field values.
- cube_fx8: The cube connectivity/topology data.
- voxel_grid_res: The resolution of the grid.
- beta, alpha, gamma_f: Slices of the learnable weights used to control local geometry.
- training: A boolean flag.
  When True: Each quadrilateral face is divided into four triangles (optimized for training stability).
  When False: Each quadrilateral face is divided into two triangles (standard mesh representation).
Returns:
- vertices: The extracted mesh vertices.
- faces: The extracted mesh faces.
- L_dev: A loss term related to deformation/geometry used for regularization.
```
# During training (4 triangles per quad)
vertices, faces, L_dev = fc(
    grid_verts, sdf, cube_fx8, voxel_grid_res, 
    beta=weight[:,:12], alpha=weight[:,12:20], gamma_f=weight[:,20], 
    training=True
)

# For final extraction (2 triangles per quad)
vertices, faces, L_dev = fc(
    grid_verts, sdf, cube_fx8, voxel_grid_res, 
    beta=weight[:,:12], alpha=weight[:,12:20], gamma_f=weight[:,20], 
    training=False
)
```
https://github.com/nv-tlabs/flexicubes/blob/main/examples/optimization.ipynb