DeepSDFStruct 1.7.2__py3-none-any.whl

DeepSDFStruct/__init__.py

@@ -0,0 +1,66 @@
+"""
+DeepSDFStruct - Differentiable Framework for 3D Microstructured Materials
+===========================================================================
+
+DeepSDFStruct is a comprehensive library for creating, manipulating, and optimizing
+3D microstructured geometries using Signed Distance Functions (SDFs) and spline-based
+lattices. The framework integrates classical geometric representations with deep
+learning approaches, enabling differentiable design and optimization workflows.
+
+Key Components
+--------------
+
+SDF Representations
+    - ``DeepSDFStruct.SDF``: Abstract base class and core SDF utilities
+    - ``DeepSDFStruct.sdf_primitives``: Geometric primitives (spheres, cylinders, etc.)
+    - ``DeepSDFStruct.lattice_structure``: Periodic lattice microstructures
+
+Mesh Operations
+    - ``DeepSDFStruct.mesh``: Mesh generation, processing, and export
+    - ``DeepSDFStruct.flexicubes``: Advanced dual contouring (3D)
+    - ``DeepSDFStruct.flexisquares``: Dual contouring for 2D cross-sections
+
+Deep Learning
+    - ``DeepSDFStruct.deep_sdf``: Neural network models and training
+    - ``DeepSDFStruct.sampling``: Data generation and sampling strategies
+
+Optimization
+    - ``DeepSDFStruct.optimization``: MMA and gradient-based optimization
+    - ``DeepSDFStruct.parametrization``: Spatially-varying parameter functions
+
+Utilities
+    - ``DeepSDFStruct.torch_spline``: Differentiable B-spline operations
+    - ``DeepSDFStruct.plotting``: Visualization tools
+    - ``DeepSDFStruct.utils``: General utility functions
+
+Examples
+--------
+Create a simple sphere and generate a mesh::
+
+    from DeepSDFStruct.sdf_primitives import SphereSDF
+    from DeepSDFStruct.mesh import create_3D_mesh
+
+    sphere = SphereSDF(center=[0, 0, 0], radius=0.5)
+    mesh = create_3D_mesh(sphere, resolution=64)
+
+Create a lattice structure::
+
+    from DeepSDFStruct.lattice_structure import LatticeSDFStruct
+    from DeepSDFStruct.torch_spline import TorchSpline
+
+    # Define deformation spline and unit cell
+    # ... (see documentation for details)
+
+    lattice = LatticeSDFStruct(
+        tiling=[3, 3, 3],
+        microtile=unit_cell
+    )
+
+For comprehensive examples, see the example notebook in the repository.
+"""
+
+import DeepSDFStruct.utils
+
+DeepSDFStruct.utils.configure_logging()
+
+__author__ = "Michael Kofler"

DeepSDFStruct/deep_sdf/__init__.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+# Copyright 2004-present Facebook. All Rights Reserved.
+
+"""
+Deep Learning for Signed Distance Functions (DeepSDF)
+======================================================
+
+This submodule implements the DeepSDF approach for learning implicit neural
+representations of 3D geometry. It provides complete workflows for:
+
+- Training neural networks to represent geometry as learned SDFs
+- Generating datasets from explicit geometry
+- Reconstructing shapes from latent codes
+- Optimizing latent codes for shape fitting
+
+The implementation is based on "DeepSDF: Learning Continuous Signed Distance
+Functions for Shape Representation" (Park et al., CVPR 2019) with extensions
+for microstructured materials and lattice geometries.
+
+Key Components
+--------------
+
+models.py
+    DeepSDFModel class wrapping trained decoder networks and latent codes.
+
+training.py
+    Complete training pipeline including loss functions, learning rate
+    schedules, and training loops.
+
+data.py
+    Dataset classes for loading and batching SDF samples.
+
+reconstruction.py
+    Methods for fitting latent codes to target geometries.
+
+workspace.py
+    Utilities for managing experiments, checkpoints, and results.
+
+networks/
+    Neural network architectures (decoders, hierarchical models).
+
+Typical Workflow
+----------------
+
+1. Generate training data from explicit geometries::
+
+    from DeepSDFStruct.sampling import generate_dataset
+    generate_dataset(geometries, output_dir, n_samples=500000)
+
+2. Train a DeepSDF model::
+
+    from DeepSDFStruct.deep_sdf import training
+    training.main(specs)  # specs define architecture, training params
+
+3. Use the trained model::
+
+    from DeepSDFStruct.SDF import SDFfromDeepSDF
+    from DeepSDFStruct.deep_sdf.models import DeepSDFModel
+
+    model = DeepSDFModel(decoder, latent_vectors, device)
+    sdf = SDFfromDeepSDF(model, latent_code=latent_vectors[0])
+
+4. Reconstruct new shapes::
+
+    from DeepSDFStruct.deep_sdf import reconstruction
+    latent_code = reconstruction.reconstruct(model, target_sdf)
+
+For complete examples, see the example notebook and test files.
+"""

DeepSDFStruct/deep_sdf/create_screenshots_from_plyfiles.py

@@ -0,0 +1,65 @@
+"""
+PLY File Screenshot Generator
+=============================
+
+This utility script generates screenshots and animated GIFs from PLY mesh files.
+It uses Vedo for 3D rendering with a fixed camera configuration to ensure
+consistent viewpoints across multiple meshes.
+
+Usage
+-----
+From command line::
+
+    python create_screenshots_from_plyfiles.py /path/to/ply/files
+
+The script will:
+1. Find all .ply files in the specified folder
+2. Generate a .png screenshot for each mesh
+3. Create an animated GIF combining all screenshots
+
+This is useful for:
+- Visualizing reconstruction sequences
+- Creating animations of optimization progress
+- Generating consistent documentation images
+- Batch processing of mesh files
+"""
+
+import vedo
+import pathlib
+from natsort import natsorted
+import imageio.v3 as imageio
+
+cam = dict(
+    position=(3.51484, 3.27242, 4.06787),
+    focal_point=(-3.98290e-3, -2.36815e-3, 1.02887e-3),
+    viewup=(-0.308731, 0.853375, -0.420043),
+    roll=2.37119,
+    distance=6.29647,
+    clipping_range=(3.00412, 10.4561),
+)
+
+
+def main(ply_file_folder):
+    images = []
+    plt = vedo.Plotter(interactive=False)
+    for ply_filename in natsorted(pathlib.Path(ply_file_folder).rglob("*.ply")):
+        print(f"Processing {ply_filename}")
+        mesh = vedo.load(str(ply_filename))
+        plt.show(mesh, interactive=False, camera=vedo.camera_from_dict(cam))
+        vedo.screenshot(ply_filename.with_suffix(".png").as_posix())
+        image = imageio.imread(str(ply_filename.with_suffix(".png")))
+        images.append(image)
+        plt.clear()
+
+    plt.close()
+    imageio.imwrite(ply_filename.parent / "reconstruction.gif", images, duration=300)
+
+
+if __name__ == "__main__":
+    # argument parser
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Create screenshots from ply files")
+    parser.add_argument("ply_file_folder", type=str, help="Folder containing ply files")
+    args = parser.parse_args()
+    main(args.ply_file_folder)

DeepSDFStruct/deep_sdf/data.py

@@ -0,0 +1,258 @@
+"""
+DeepSDF Data Loading and Preprocessing
+======================================
+
+This module provides PyTorch Dataset classes and utilities for loading
+and preprocessing SDF training data. It handles:
+
+- Loading sampled SDF values from .npz files
+- Splitting samples into positive and negative regions
+- Batching and shuffling for training
+- Data normalization and clamping
+- Multi-class and multi-dataset support
+
+The module is designed to work with the dataset format produced by the
+sampling module, where each geometry is stored as a .npz file containing
+positive and negative SDF samples.
+
+Classes
+-------
+SDFSamples
+    PyTorch Dataset for loading SDF samples with lazy loading and caching.
+
+NoMeshFileError
+    Exception raised when a mesh file cannot be found.
+
+MultipleMeshFileError
+    Exception raised when multiple mesh files exist in a directory.
+
+Functions
+---------
+get_instance_filenames
+    Retrieve file paths for all instances in a dataset split.
+
+unpack_sdf_samples
+    Load and preprocess SDF samples from a .npz file.
+
+read_sdf_samples_into_ram
+    Load entire dataset into RAM for faster training.
+"""
+
+#!/usr/bin/env python3
+# Copyright 2004-present Facebook. All Rights Reserved.
+
+import glob
+import logging
+import numpy as np
+import os
+import random
+import torch
+import torch.utils.data
+
+import DeepSDFStruct.deep_sdf.workspace as ws
+
+logger = logging.getLogger(__name__)
+
+
+def get_instance_filenames(data_source, split):
+    npzfiles = []
+    for dataset in split:
+        for class_name in split[dataset]:
+            for instance_name in split[dataset][class_name]:
+                instance_filename = os.path.join(
+                    dataset, class_name, instance_name + ".npz"
+                )
+                if not os.path.isfile(
+                    os.path.join(data_source, ws.sdf_samples_subdir, instance_filename)
+                ):
+                    # raise RuntimeError(
+                    #     'Requested non-existent file "' + instance_filename + "'"
+                    # )
+                    logging.warning(
+                        "Requested non-existent file '{}'".format(instance_filename)
+                    )
+                npzfiles += [instance_filename]
+    return npzfiles
+
+
+class NoMeshFileError(RuntimeError):
+    """Raised when a mesh file is not found in a shape directory"""
+
+    pass
+
+
+class MultipleMeshFileError(RuntimeError):
+    """ "Raised when a there a multiple mesh files in a shape directory"""
+
+    pass
+
+
+def find_mesh_in_directory(shape_dir):
+    mesh_filenames = list(glob.iglob(shape_dir + "/**/*.obj")) + list(
+        glob.iglob(shape_dir + "/*.obj")
+    )
+    if len(mesh_filenames) == 0:
+        raise NoMeshFileError()
+    elif len(mesh_filenames) > 1:
+        raise MultipleMeshFileError()
+    return mesh_filenames[0]
+
+
+# and also convert to single precision
+# usually one function should only contain one functionality, but here I think
+# it makes sense, because the data is loaded at 3 different positions in the code
+def remove_nans(tensor, geom_dimension):
+    tensor_nan = torch.isnan(tensor[:, geom_dimension])
+    return tensor[~tensor_nan, :].float()
+
+
+def read_sdf_samples_into_ram(filename):
+    npz = np.load(filename)
+    pos_tensor = torch.from_numpy(npz["pos.npy"]).float()
+    neg_tensor = torch.from_numpy(npz["neg.npy"]).float()
+
+    return [pos_tensor, neg_tensor]
+
+
+def unpack_sdf_samples(filename, geom_dimension, subsample=None):
+    npz = np.load(filename)
+
+    pos_tensor = remove_nans(torch.from_numpy(npz["pos.npy"]), geom_dimension)
+    neg_tensor = remove_nans(torch.from_numpy(npz["neg.npy"]), geom_dimension)
+
+    if subsample is None:
+        return torch.cat([pos_tensor, neg_tensor], 0)
+
+    half = int(subsample / 2)
+    pos_len = len(pos_tensor)
+    neg_len = len(neg_tensor)
+    if pos_len < half:
+        neg_len = 2 * half - pos_len
+    elif neg_len < half:
+        pos_len = 2 * half - neg_len
+    else:
+        pos_len = neg_len = half
+
+    use_randperm = True
+    if use_randperm:
+        random_pos = torch.randperm(len(pos_tensor))[:pos_len]
+        random_neg = torch.randperm(len(neg_tensor))[:neg_len]
+    else:
+        # split the sample into half
+        random_pos = (torch.rand(half) * pos_tensor.shape[0]).long()
+        random_neg = (torch.rand(half) * neg_tensor.shape[0]).long()
+
+    sample_pos = torch.index_select(pos_tensor, 0, random_pos)
+    sample_neg = torch.index_select(neg_tensor, 0, random_neg)
+
+    samples = torch.cat([sample_pos, sample_neg], 0)
+
+    # if len(samples) < 50:
+    # print("less than half")
+
+    return samples
+
+
+def unpack_sdf_samples_from_ram(data, subsample=None):
+    if subsample is None:
+        return torch.cat(data, 0)
+    pos_tensor = data[0]
+    neg_tensor = data[1]
+
+    # split the sample into half
+    half = int(subsample / 2)
+
+    pos_size = pos_tensor.shape[0]
+    neg_size = neg_tensor.shape[0]
+
+    if pos_size <= half:
+        random_pos = (torch.rand(half) * pos_tensor.shape[0]).long()
+        sample_pos = torch.index_select(pos_tensor, 0, random_pos)
+    else:
+        pos_start_ind = random.randint(0, pos_size - half)
+        sample_pos = pos_tensor[pos_start_ind : (pos_start_ind + half)]
+
+    if neg_size <= half:
+        random_neg = (torch.rand(half) * neg_tensor.shape[0]).long()
+        sample_neg = torch.index_select(neg_tensor, 0, random_neg)
+    else:
+        neg_start_ind = random.randint(0, neg_size - half)
+        sample_neg = neg_tensor[neg_start_ind : (neg_start_ind + half)]
+
+    samples = torch.cat([sample_pos, sample_neg], 0)
+
+    return samples
+
+
+class SDFSamples(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        data_source,
+        split,
+        subsample,
+        geom_dimension,
+        load_ram=False,
+        print_filename=False,
+        num_files=1000000,
+    ):
+        self.subsample = subsample
+        self.geom_dimension = geom_dimension
+        self.data_source = data_source
+        self.npyfiles = get_instance_filenames(data_source, split)
+        self.filenames = []
+        logger.debug(
+            "using "
+            + str(len(self.npyfiles))
+            + " shapes from data source "
+            + data_source
+        )
+
+        self.load_ram = load_ram
+
+        self.loaded_mat_properties = []
+        if load_ram:
+            self.loaded_data = []
+            for f in self.npyfiles:
+                filename = os.path.join(self.data_source, ws.sdf_samples_subdir, f)
+                npz = np.load(filename)
+                pos_tensor = remove_nans(
+                    torch.from_numpy(npz["pos.npy"]), self.geom_dimension
+                )
+                neg_tensor = remove_nans(
+                    torch.from_numpy(npz["neg.npy"]), self.geom_dimension
+                )
+                self.loaded_data.append(
+                    [
+                        pos_tensor[torch.randperm(pos_tensor.shape[0])],
+                        neg_tensor[torch.randperm(neg_tensor.shape[0])],
+                    ]
+                )
+                if "C" in npz.keys():
+                    self.loaded_mat_properties.append(
+                        torch.from_numpy(npz["C.npy"]).to(torch.float32).reshape(-1)
+                    )
+                self.filenames.append(filename)
+
+    def __len__(self):
+        return len(self.npyfiles)
+
+    def __getitem__(self, idx):
+        filename = os.path.join(
+            self.data_source, ws.sdf_samples_subdir, self.npyfiles[idx]
+        )
+        if len(self.loaded_mat_properties) != 0:
+            mat_prop = self.loaded_mat_properties[idx]
+        else:
+            mat_prop = torch.tensor(0)
+        if self.load_ram:
+            return (
+                unpack_sdf_samples_from_ram(self.loaded_data[idx], self.subsample),
+                mat_prop.unsqueeze(0).expand(self.subsample, -1),
+                idx,
+            )
+        else:
+            return (
+                unpack_sdf_samples(filename, self.geom_dimension, self.subsample),
+                mat_prop.unsqueeze(0).expand(self.subsample, -1),
+                idx,
+            )

DeepSDFStruct/deep_sdf/metrics/__init__.py

@@ -0,0 +1,14 @@
+"""
+Geometry Evaluation Metrics
+===========================
+
+This module provides utilities for evaluating the quality of generated
+geometries by comparing them to ground truth analytical SDFs.
+
+Functions
+---------
+
+mesh_to_analytical
+    Compute SDF values at mesh vertices using an analytical SDF function.
+    Useful for quantifying reconstruction accuracy.
+"""

DeepSDFStruct/deep_sdf/metrics/mesh_to_analytical.py

@@ -0,0 +1,24 @@
+"""
+Mesh-to-Analytical SDF Evaluation
+=================================
+
+This module provides functions for evaluating mesh quality by computing
+SDF values at mesh vertices using ground truth analytical SDFs.
+
+This is useful for:
+- Quantifying reconstruction accuracy
+- Measuring surface deviation
+- Validating mesh generation algorithms
+"""
+
+import trimesh
+from DeepSDFStruct.mesh import torchSurfMesh
+from DeepSDFStruct.SDF import SDFBase
+
+
+def mesh_to_analytical(gt_sdf: SDFBase, gen_mesh: torchSurfMesh) -> float:
+    """
+    Calculates the SDF for every vertex on the mesh
+    """
+
+    return gt_sdf.forward(gen_mesh.vertices).abs().mean().item()

DeepSDFStruct/deep_sdf/models.py

@@ -0,0 +1,197 @@
+"""
+DeepSDF Model Classes
+=====================
+
+This module provides the DeepSDFModel class, which wraps a trained decoder
+network and latent vectors for evaluating learned implicit representations.
+
+The model can be used to:
+- Decode latent codes to SDF values at query points
+- Export models for deployment (TorchScript/LibTorch)
+- Integrate with gradient-based optimization
+
+Classes
+-------
+DeepSDFModel
+    Main interface for trained DeepSDF models, combining a decoder network
+    with learned latent vectors.
+"""
+
+import torch
+
+
+class DeepSDFModel:
+    """Wrapper for trained DeepSDF decoder and latent vectors.
+
+    This class provides a convenient interface for using trained DeepSDF models.
+    It combines the decoder network (which maps latent+position to SDF values)
+    with the trained latent vectors (which encode different shapes).
+
+    The model supports both constant latent codes (one code for all query points)
+    and spatially-varying codes (different codes per query point), enabling
+    flexible shape representation.
+
+    Parameters
+    ----------
+    decoder : torch.nn.Module
+        The trained decoder network. Should accept input of shape
+        (N, latent_dim + 3) and output SDF values of shape (N, 1).
+    trained_latent_vectors : torch.Tensor
+        Trained latent codes of shape (num_shapes, latent_dim).
+        Each row encodes one learned shape.
+    device : str or torch.device
+        Device for computation ('cpu' or 'cuda').
+
+    Attributes
+    ----------
+    _decoder : torch.nn.Module
+        The decoder network.
+    _trained_latent_vectors : torch.Tensor
+        The latent code library.
+    device : str or torch.device
+        Computation device.
+
+    Methods
+    -------
+    _decode_sdf(latent_vec, queries)
+        Decode SDF values from a latent code and spatial queries.
+    export_libtorch_executable(filename)
+        Export model to TorchScript for C++ deployment.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from DeepSDFStruct.deep_sdf.models import DeepSDFModel
+    >>>
+    >>> # Assume we have a trained decoder and latent vectors
+    >>> # decoder = ...
+    >>> # latents = torch.randn(10, 256)  # 10 shapes, 256-dim codes
+    >>>
+    >>> # Create model
+    >>> model = DeepSDFModel(decoder, latents, device='cuda')
+    >>>
+    >>> # Query first shape
+    >>> points = torch.rand(1000, 3, device='cuda')
+    >>> distances = model._decode_sdf(latents[0], points)
+    >>>
+    >>> # Use with SDFfromDeepSDF
+    >>> from DeepSDFStruct.SDF import SDFfromDeepSDF
+    >>> sdf = SDFfromDeepSDF(model, latent_code=latents[0])
+    >>> mesh = create_3D_mesh(sdf, N_base=64, mesh_type='surface')
+
+    Notes
+    -----
+    The decoder architecture typically consists of multiple fully-connected
+    layers with skip connections. See networks/ for architecture definitions.
+
+    References
+    ----------
+    .. [1] Park, J. J., Florence, P., Straub, J., Newcombe, R., & Lovegrove, S.
+           (2019). DeepSDF: Learning continuous signed distance functions for
+           shape representation. In CVPR.
+    """
+
+    def __init__(
+        self, decoder: torch.nn.Module, trained_latent_vectors: torch.Tensor, device
+    ):
+        self._decoder = decoder
+        self._trained_latent_vectors = trained_latent_vectors
+        self.device = device
+
+    def _decode_sdf(
+        self, latent_vec: torch.Tensor, queries: torch.Tensor
+    ) -> torch.Tensor:
+        """Decode SDF values from latent vector and xyz queries.
+
+        Combines latent codes with spatial coordinates and passes through
+        the decoder to obtain SDF values. Handles both per-query and
+        constant latent vectors.
+
+        Parameters
+        ----------
+        latent_vec : torch.Tensor
+            Latent code(s). Two formats supported:
+            - Shape (latent_dim,): Constant code used for all queries
+            - Shape (num_samples, latent_dim): Per-query codes
+        queries : torch.Tensor
+            Query point coordinates of shape (num_samples, 3).
+
+        Returns
+        -------
+        torch.Tensor
+            SDF values of shape (num_samples, 1).
+
+        Raises
+        ------
+        ValueError
+            If latent_vec shape doesn't match expected dimensions.
+
+        Notes
+        -----
+        If latent_vec is constant (1D), it's expanded to match the number
+        of query points. If latent_vec is per-query (2D), it must have
+        exactly one code per query point.
+        """
+        latent_dim = self._trained_latent_vectors[0].shape[0]
+        num_samples = queries.shape[0]
+        if latent_vec.ndim == 1:
+            if latent_vec.shape[0] != latent_dim:
+                raise ValueError(
+                    f"Latent vector shape mismatch: {latent_vec.shape} does"
+                    f"not align with latent dimension {latent_dim}."
+                )
+            latent_repeat = latent_vec.expand(-1, num_samples).T
+        elif latent_vec.ndim == 2:
+            if (latent_vec.shape[0] != num_samples) or (
+                latent_vec.shape[1] != latent_dim
+            ):
+                raise ValueError(
+                    f"Latent vector shape mismatch: {latent_vec.shape} does"
+                    f" not align with {num_samples} queries."
+                    f" Must be of shape ({num_samples}, {latent_dim})"
+                )
+            latent_repeat = latent_vec
+
+        model_input = torch.cat([latent_repeat, queries], dim=1)
+        return self._decoder(model_input)
+
+    def export_libtorch_executable(self, filename: str, use_script: bool = True):
+        """
+        Export the trained decoder model to a TorchScript file for use with LibTorch (C++).
+
+        Args:
+            filename (str): Path where the TorchScript model will be saved (e.g. "decoder.pt").
+            use_script (bool): If True, use torch.jit.script (recommended for C++ compatibility).
+                              If False, use torch.jit.trace (legacy mode, may have control flow issues).
+
+        Example:
+            >>> model.export_libtorch_executable("decoder.pt")
+            Example input:  tensor([[...]])
+            Example Output: tensor([[...]])
+            # The file "decoder.pt" is now ready for loading in LibTorch.
+        """
+        assert isinstance(
+            self._trained_latent_vectors, torch.Tensor
+        ), "trained_latent_vectors must be a tensor"
+        assert (
+            self._trained_latent_vectors.shape[0] > 0
+        ), "trained_latent_vectors must contain at least one element"
+        latent = self._trained_latent_vectors
+        example_input = torch.cat(
+            [latent[0], torch.tensor([0.0, 0.0, 0.0], device=self.device)]
+        ).unsqueeze(0)
+
+        print("Example input: ", example_input)
+        print("Example Output: ", self._decoder(example_input))
+
+        self._decoder.eval()
+
+        if use_script:
+            print("Exporting with torch.jit.script (recommended)")
+            sm = torch.jit.script(self._decoder)
+        else:
+            print("Exporting with torch.jit.trace (legacy mode)")
+            sm = torch.jit.trace(self._decoder, example_input)
+
+        sm.save(filename)
+        print(f"Saved to {filename}")