DeepSDFStruct.geom_reconstruction#

Shape Reconstruction#

Fit a spatially varying field of DeepSDF latent codes to a target mesh.

The decoder of a trained DeepSDF model turns a single latent code into one shape. To represent a whole part rather than a single microtile, this module tiles that decoder over a domain and lets the latent code vary in space: the codes are the control points of a B-spline, and fitting the shape means optimizing those control points against SDF samples taken from a target mesh. Because neighbouring tiles read nearby points of the same spline, the result is a continuous structure – a field of local shapes.

LocalShapesReconstructor wraps that whole procedure:

import torch, trimesh
from DeepSDFStruct.geom_reconstruction import LocalShapesReconstructor

recon = LocalShapesReconstructor(output_dir="output")
torch.manual_seed(42)

mesh_orig = trimesh.load_mesh("part.stl")
assert mesh_orig.is_watertight, "target mesh must be watertight"

struct, scaling, gt_sdf, params = recon.fit_mesh(
    mesh=mesh_orig, tiling=[32, 32, 32]
)
recon.export(struct, scaling)

struct is a LatticeSDFStruct defined in parameter space; scaling maps it back onto the original mesh’s coordinates, and is what you hand to create_3D_mesh() as deformation_function.

Functions#

build_parameter_spline

Build the B-spline whose control points carry the latent codes.

sample_gt_sdf

Draw uniform and near-surface SDF samples from a target mesh.

Classes#

LocalShapesReconstructor

Fits a latent-code field to a mesh.

StructBuild

Named tuple returned by LocalShapesReconstructor.build_struct().

FitResult

Named tuple returned by LocalShapesReconstructor.fit_mesh().

Functions

build_parameter_spline(spline_degrees, ...)

Build the B-spline that carries one latent code per control point.

sample_gt_sdf(gt_sdf, mesh, bounds, *, ...)

Draw uniform and near-surface SDF samples from a target mesh.

Classes

FitResult(struct, scaling, gt_sdf, params)

Outcome of LocalShapesReconstructor.fit_mesh().

LocalShapesReconstructor([model, ...])

Fits a field of DeepSDF latent codes to a target mesh.

StructBuild(struct, scaling, mesh_norm, ...)

What LocalShapesReconstructor.build_struct() produces.

class DeepSDFStruct.geom_reconstruction.FitResult(struct: DeepSDFStruct.lattice_structure.LatticeSDFStruct, scaling: DeepSDFStruct.torch_spline.TorchScaling, gt_sdf: DeepSDFStruct.SDF.SDFBase, params: list[torch.Tensor])#

Bases: NamedTuple

Outcome of LocalShapesReconstructor.fit_mesh().

Deliberately just the four headline outputs, so the common case unpacks directly:

struct, scaling, gt_sdf, params = recon.fit_mesh(...)

Fit diagnostics are not carried here. For the loss curve, pass loss_plot_path / loss_csv_path to write them out, or drive the steps yourself (build_struct(), sample_gt_sdf(), fit_samples()) – fit_samples returns loss_history, final_loss and num_steps.

Parameters:
struct#

The fitted structure, in parameter space, with its parametrization already set to the optimized control points. struct.bounds holds the evaluation bounds.

Type:

LatticeSDFStruct

scaling#

Parameter space -> original mesh coordinates. Hand this to create_3D_mesh() as deformation_function.

Type:

TorchScaling

gt_sdf#

Ground-truth SDF of the normalized target mesh, reusable for error metrics.

Type:

SDFBase

params#

Optimized parameters, i.e. [control_points] of shape (n_control_points, latent_dim).

Type:

list of torch.Tensor

gt_sdf: DeepSDFStruct.SDF.SDFBase#

Alias for field number 2

params: list[torch.Tensor]#

Alias for field number 3

scaling: DeepSDFStruct.torch_spline.TorchScaling#

Alias for field number 1

struct: DeepSDFStruct.lattice_structure.LatticeSDFStruct#

Alias for field number 0

class DeepSDFStruct.geom_reconstruction.LocalShapesReconstructor(model=PretrainedModels.Primitives, checkpoint='latest', device=None, output_dir='tests/tmp_outputs')#

Bases: object

Fits a field of DeepSDF latent codes to a target mesh.

Parameters:
  • model (str, PretrainedModels or DeepSDFModel, optional) – The decoder providing the microtile. Accepts a PretrainedModels member, a path to a checkpoint directory, or an already loaded DeepSDFModel. Defaults to the bundled primitives decoder.

  • checkpoint (str, default "latest") – Checkpoint name inside the model directory. Ignored when model is an already loaded model.

  • device (str or torch.device, optional) – Compute device. Defaults to CUDA when available.

  • output_dir (path-like or None, default "tests/tmp_outputs") –

    Where every file this reconstructor writes goes: the loss curve and CSV of fit_mesh(), and whatever export() produces. Relative paths are resolved against the working directory. The directory is created on first write, not here. Pass None to keep the reconstructor from writing anything unless an explicit path is given.

    The default points at the repository’s gitignored scratch directory, so a run started from a checkout does not leave artifacts in the repository root. It is relative to the working directory, so callers running from outside a checkout should pass an explicit path.

Examples

>>> recon = LocalShapesReconstructor()
>>> struct, scaling, gt_sdf, params = recon.fit_mesh(
...     mesh=mesh, tiling=[8, 8, 8]
... )
>>> recon.export(struct, scaling)

Notes

fit_mesh() is the one-call path. When you need to interleave work – exporting intermediate fields, or reusing one structure across several fits – use build_struct(), sample_gt_sdf() and fit_samples() separately; fit_mesh is just their composition.

build_struct(mesh, tiling, *, spline_degree=(1, 1, 1), shrink_factor=1.0)#

Normalize mesh and build the latent-code spline and structure.

Parameters:
  • mesh (trimesh.Trimesh) – Target mesh. Copied before normalization, so the caller’s mesh is left untouched.

  • tiling (list of 3 int) – Microtiles per dimension.

  • spline_degree (sequence of 3 int, default (1, 1, 1)) – Degree of the latent-code spline per dimension.

  • shrink_factor (float, default 1.0) – Passed to normalize_mesh_to_unit_cube(). Values below 1 leave a margin between the mesh and the domain border.

Returns:

Named tuple with the structure and everything derived alongside it.

Return type:

StructBuild

ensure_output_dir()#

Create output_dir if needed and return it.

Raises:

ValueError – If the reconstructor was built with output_dir=None, i.e. was explicitly told not to write files.

Return type:

Path

export(struct, scaling=None, *, mesh_resolution=32, bounds=None, output_dir=None, **kwargs)#

Write struct to disk as an SDF grid plus surface meshes.

Thin wrapper around export_reconstructed_artifacts() that fills in what the reconstructor already knows: the output directory (created if missing), the device, and – unless overridden – the structure’s own bounds.

Parameters:
  • struct (LatticeSDFStruct) – Structure to export, normally the one fit_mesh() returned.

  • scaling (TorchScaling, optional) – Parameter-to-physical-space map. Without it only the parameter-space mesh is written.

  • mesh_resolution (int, default 32) – FlexiCubes grid resolution per dimension.

  • bounds (torch.Tensor, optional) – Evaluation bounds. Defaults to struct.bounds.

  • output_dir (path-like, optional) – Destination, overriding output_dir for this call.

  • **kwargs – Forwarded to export_reconstructed_artifacts(), e.g. sdf_grid_N or the file-name arguments.

Returns:

Path of the physical-space mesh, or of the parameter-space mesh when no scaling was given.

Return type:

pathlib.Path

fit_mesh(mesh, tiling, *, num_iterations=10, lr=0.005, batch_size=4096, n_uniform=100000, n_surface=500000, spline_degree=(1, 1, 1), shrink_factor=1.0, samples_surface_stds=(0.005, 0.0001), box_constrained=False, code_reg_lambda=0.0, code_bound=1.0, grad_clip=1.0, eikonal_lambda=0.0, loss_plot_path=None, loss_csv_path=None, step_callback=None)#

Fit a latent-code field to mesh, end to end.

Normalizes the mesh into parameter space, builds the structure, samples the ground-truth SDF, runs the fit, and writes the optimized codes back into the structure.

Parameters:
  • mesh (trimesh.Trimesh) – Target mesh. Should be watertight, otherwise the sign of the ground-truth SDF is not well defined.

  • tiling (list of 3 int) – Microtiles per dimension.

  • num_iterations (int, float, int) – Fitting hyperparameters, see fit_samples().

  • lr (int, float, int) – Fitting hyperparameters, see fit_samples().

  • batch_size (int, float, int) – Fitting hyperparameters, see fit_samples().

  • n_uniform (int) – Sample counts, see sample_gt_sdf().

  • n_surface (int) – Sample counts, see sample_gt_sdf().

  • spline_degree – Structure construction, see build_struct().

  • shrink_factor (float) – Structure construction, see build_struct().

  • samples_surface_stds – Sampling behaviour, see sample_gt_sdf().

  • box_constrained (bool) – Sampling behaviour, see sample_gt_sdf().

  • code_reg_lambda (float) – Regularization, see fit_samples().

  • code_bound (float | None) – Regularization, see fit_samples().

  • grad_clip (float | None) – Regularization, see fit_samples().

  • eikonal_lambda (float) – Regularization, see fit_samples().

  • loss_plot_path (path-like, optional) – Where the loss curve and its raw values go. Both default to reconstruction_loss.png / .csv inside output_dir; with output_dir=None they are not written.

  • loss_csv_path (path-like, optional) – Where the loss curve and its raw values go. Both default to reconstruction_loss.png / .csv inside output_dir; with output_dir=None they are not written.

  • step_callback (Optional[Callable[[int, int, int], None]]) – Diagnostics, see fit_samples().

Returns:

Unpacks as struct, scaling, gt_sdf, params. Fit diagnostics are not included – see FitResult for how to get them.

Return type:

FitResult

static fit_samples(struct, samples, *, num_iterations=10, lr=0.005, batch_size=4096, code_reg_lambda=0.0, code_bound=1.0, grad_clip=1.0, eikonal_lambda=0.0, loss_plot_path=None, loss_csv_path=None, step_callback=None)#

Optimize struct’s latent codes against samples.

Static: everything needed is already inside struct, so this can be called as LocalShapesReconstructor.fit_samples(struct, samples, ...) to refit a lattice without loading a model again.

Parameters:
  • struct (LatticeSDFStruct) – Structure to fit; its parametrization parameters are optimized in place.

  • samples (SampledSDF) – Target samples, in the same (parameter) space as struct.

  • num_iterations (int, default 10) – Number of epochs over samples. One epoch is len(samples) // batch_size optimizer steps, so with the default sample counts of fit_mesh() this is already a few thousand steps.

  • lr (float, default 5e-3) – Adam learning rate.

  • batch_size (int, default 4096) – Samples per optimizer step.

  • code_reg_lambda (float, default 0.0) – Weight of the L2 penalty on evaluated latent codes.

  • code_bound (float or None, default 1.0) – If set, control points are clamped to [-code_bound, code_bound] after every step, keeping them inside the range the decoder was trained on. None disables the clamp.

  • grad_clip (float or None, default 1.0) – If set, gradient-norm clipping threshold. None disables it.

  • eikonal_lambda (float, default 0.0) – Weight of the near-surface Eikonal penalty (|grad SDF| - 1)^2.

  • loss_plot_path (path-like, optional) – Where to write the loss curve and its raw values.

  • loss_csv_path (path-like, optional) – Where to write the loss curve and its raw values.

  • step_callback (callable, optional) – Called as step_callback(epoch, batch_idx, n_batches) after every optimizer step, for progress exports.

Returns:

Keys params, loss_history, final_loss, num_steps.

Return type:

dict

class DeepSDFStruct.geom_reconstruction.StructBuild(struct: DeepSDFStruct.lattice_structure.LatticeSDFStruct, scaling: DeepSDFStruct.torch_spline.TorchScaling, mesh_norm: trimesh.base.Trimesh, bounds: torch.Tensor, param_spline: DeepSDFStruct.parametrization.SplineParametrization, param_spline_sp: splinepy.bspline.BSpline, scale: float, shift: numpy.ndarray)#

Bases: NamedTuple

What LocalShapesReconstructor.build_struct() produces.

Parameters:
struct#

Lattice over the normalized mesh, latent codes initialized but not yet fitted.

Type:

LatticeSDFStruct

scaling#

Parameter space -> original mesh coordinates.

Type:

TorchScaling

mesh_norm#

The target mesh normalized into parameter space.

Type:

trimesh.Trimesh

bounds#

(2, 3) bounds of struct, taken from mesh_norm.

Type:

torch.Tensor

param_spline#

The latent-code spline; its control points are the free variables.

Type:

SplineParametrization

param_spline_sp#

The underlying splinepy spline, for knot-grid exports.

Type:

splinepy.BSpline

scale#

Inverse of the normalization scale, as consumed by TorchScaling.

Type:

float

shift#

Center of the original mesh’s bounding box.

Type:

numpy.ndarray

bounds: torch.Tensor#

Alias for field number 3

mesh_norm: trimesh.base.Trimesh#

Alias for field number 2

param_spline: DeepSDFStruct.parametrization.SplineParametrization#

Alias for field number 4

param_spline_sp: splinepy.bspline.BSpline#

Alias for field number 5

scale: float#

Alias for field number 6

scaling: DeepSDFStruct.torch_spline.TorchScaling#

Alias for field number 1

shift: numpy.ndarray#

Alias for field number 7

struct: DeepSDFStruct.lattice_structure.LatticeSDFStruct#

Alias for field number 0

DeepSDFStruct.geom_reconstruction.build_parameter_spline(spline_degrees, tiling, latent_dim, bounds=None)#

Build the B-spline that carries one latent code per control point.

Starts from a single clamped span per dimension and inserts n_box - 1 uniformly spaced interior knots, so the spline has exactly one knot span per microtile. Control points are zero-initialized; callers are expected to set them (see LocalShapesReconstructor.build_struct(), which starts from the mean trained latent vector).

Parameters:
  • spline_degrees (list of 3 int) – Polynomial degree per spatial dimension. [1, 1, 1] gives trilinear interpolation between neighbouring latent codes.

  • tiling (list of 3 int) – Number of knot spans (microtiles) per dimension.

  • latent_dim (int) – Width of each control point, i.e. the model’s latent code length.

  • bounds ((2, 3) array-like, optional) – [[xmin, ymin, zmin], [xmax, ymax, zmax]] spanned by the spline. Defaults to the unit cube [0, 1]^3.

Returns:

Spline with prod(tiling + 1) control points of width latent_dim for degree 1.

Return type:

splinepy.BSpline

DeepSDFStruct.geom_reconstruction.sample_gt_sdf(gt_sdf, mesh, bounds, *, n_uniform, n_surface, device='cpu', stds=(0.025, 0.0001), box_constrained=False)#

Draw uniform and near-surface SDF samples from a target mesh.

Combines a uniform fill of bounds (which teaches the fit where the material is not) with a dense band hugging the surface (which sharpens the zero level set).

Parameters:
  • gt_sdf (SDFBase) – Ground-truth SDF, normally SDFfromMesh(mesh, scale=False). Passed in rather than built here so the caller can reuse the same object for error metrics afterwards.

  • mesh (trimesh.Trimesh) – Target mesh, already in the coordinate system of bounds.

  • bounds (torch.Tensor) – (2, 3) box for the uniform samples.

  • n_uniform (int) – Number of uniform samples inside bounds.

  • n_surface (int) – Number of surface points; each is perturbed once per entry of stds, so the surface band holds n_surface * len(stds) samples (or exactly n_surface when box_constrained is set).

  • device (str or torch.device) – Device the samples are created on.

  • stds (sequence of float) – Standard deviations of the Gaussian noise added to surface points – one coarse, one fine by convention.

  • box_constrained (bool, default False) – If True, reject surface samples outside bounds and re-sample until n_surface accepted points are collected. Needed when the lattice domain is smaller than the mesh’s own bounding box, as in a shape optimization whose design domain is a sub-box.

Returns:

Concatenation of the uniform and surface samples.

Return type:

SampledSDF