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 the B-spline that carries one latent code per control point. |
|
Draw uniform and near-surface SDF samples from a target mesh. |
Classes
|
Outcome of |
|
Fits a field of DeepSDF latent codes to a target mesh. |
|
What |
- 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:
NamedTupleOutcome 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_pathto write them out, or drive the steps yourself (build_struct(),sample_gt_sdf(),fit_samples()) –fit_samplesreturnsloss_history,final_lossandnum_steps.- Parameters:
scaling (DeepSDFStruct.torch_spline.TorchScaling)
gt_sdf (DeepSDFStruct.SDF.SDFBase)
params (list[torch.Tensor])
- struct#
The fitted structure, in parameter space, with its parametrization already set to the optimized control points.
struct.boundsholds the evaluation bounds.- Type:
- scaling#
Parameter space -> original mesh coordinates. Hand this to
create_3D_mesh()asdeformation_function.- Type:
- 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:
objectFits a field of DeepSDF latent codes to a target mesh.
- Parameters:
model (str, PretrainedModels or DeepSDFModel, optional) – The decoder providing the microtile. Accepts a
PretrainedModelsmember, a path to a checkpoint directory, or an already loadedDeepSDFModel. Defaults to the bundled primitives decoder.checkpoint (str, default "latest") – Checkpoint name inside the model directory. Ignored when
modelis 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 whateverexport()produces. Relative paths are resolved against the working directory. The directory is created on first write, not here. PassNoneto 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 – usebuild_struct(),sample_gt_sdf()andfit_samples()separately;fit_meshis 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:
- ensure_output_dir()#
Create
output_dirif 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_dirfor this call.**kwargs – Forwarded to
export_reconstructed_artifacts(), e.g.sdf_grid_Nor the file-name arguments.
- Returns:
Path of the physical-space mesh, or of the parameter-space mesh when no
scalingwas 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, seebuild_struct().samples_surface_stds – Sampling behaviour, see
sample_gt_sdf().box_constrained (
bool) – Sampling behaviour, seesample_gt_sdf().code_reg_lambda (
float) – Regularization, seefit_samples().code_bound (
float|None) – Regularization, seefit_samples().grad_clip (
float|None) – Regularization, seefit_samples().eikonal_lambda (
float) – Regularization, seefit_samples().loss_plot_path (path-like, optional) – Where the loss curve and its raw values go. Both default to
reconstruction_loss.png/.csvinsideoutput_dir; withoutput_dir=Nonethey are not written.loss_csv_path (path-like, optional) – Where the loss curve and its raw values go. Both default to
reconstruction_loss.png/.csvinsideoutput_dir; withoutput_dir=Nonethey are not written.step_callback (
Optional[Callable[[int,int,int],None]]) – Diagnostics, seefit_samples().
- Returns:
Unpacks as
struct, scaling, gt_sdf, params. Fit diagnostics are not included – seeFitResultfor how to get them.- Return type:
- 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 asLocalShapesReconstructor.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 islen(samples) // batch_sizeoptimizer steps, so with the default sample counts offit_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.Nonedisables the clamp.grad_clip (float or None, default 1.0) – If set, gradient-norm clipping threshold.
Nonedisables 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:
NamedTupleWhat
LocalShapesReconstructor.build_struct()produces.- Parameters:
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)
- struct#
Lattice over the normalized mesh, latent codes initialized but not yet fitted.
- Type:
- scaling#
Parameter space -> original mesh coordinates.
- Type:
- mesh_norm#
The target mesh normalized into parameter space.
- Type:
trimesh.Trimesh
- bounds#
(2, 3) bounds of
struct, taken frommesh_norm.- Type:
torch.Tensor
- param_spline#
The latent-code spline; its control points are the free variables.
- Type:
- 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 - 1uniformly spaced interior knots, so the spline has exactly one knot span per microtile. Control points are zero-initialized; callers are expected to set them (seeLocalShapesReconstructor.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 widthlatent_dimfor 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 holdsn_surface * len(stds)samples (or exactlyn_surfacewhenbox_constrainedis 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
boundsand re-sample untiln_surfaceaccepted 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: