24552038416

Committed 17 Apr 2026 06:51AM UTC coverage: 81.449% (-0.006%) from 81.455%

Build # 24552038416

Build Type

Pull #50

github

Committed by

web-flow

Commit Message

Merge 7339285d0 into 62696ed2d

Pull Request Pull Request #50: changed surface sampling to be in normal direction

Coverage Stats

368 of 440 branches covered (83.64%)

Branch coverage included in aggregate %.

3083 of 3797 relevant lines covered (81.2%)

0.81 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.98

DeepSDFStruct/sampling.py

"""
SDF Sampling and Dataset Generation
====================================

This module provides tools for sampling points from SDF representations and
generating datasets for training neural networks (DeepSDF models). It supports
various sampling strategies to create well-distributed training data.

Key Features
------------

Sampling Strategies
    - Uniform sampling in a bounding box
    - Surface-focused sampling near the zero level set
    - Importance sampling based on SDF gradients
    - Sphere-based sampling patterns
    - Combined strategies for balanced datasets

Dataset Generation
    - Batch processing of multiple geometries
    - Automatic data normalization and standardization
    - Support for multiple geometry classes
    - Metadata tracking (version, sampling parameters)
    - Export to formats compatible with DeepSDF training

The module is designed to generate high-quality training data for implicit
neural representations, with careful attention to sampling near surfaces
where accurate reconstruction is most critical.

Classes
-------
SampledSDF
    Container for sampled points and their SDF values, with utilities
    for splitting by sign and visualization.

DataSetInfo
    TypedDict for dataset metadata (name, classes, sampling strategy, etc.).

Functions
---------
generate_dataset
    Batch process multiple geometries to create a complete dataset.
sample_sdf_*
    Various sampling strategies for different use cases.
"""

import os
from functools import partial
from itertools import starmap
import multiprocessing

import vtk
import numpy as np
import json
import pathlib
import typing
import gustaf as gus
import trimesh
from DeepSDFStruct.SDF import SDFfromMesh, SDFBase
from DeepSDFStruct.mesh import torchSurfMesh
import DeepSDFStruct

# from analysis.problems.homogenization import computeHomogenizedMaterialProperties
import splinepy
import torch
from collections import defaultdict
from tqdm import tqdm
import logging
import datetime
from importlib.metadata import version

logger = logging.getLogger(DeepSDFStruct.__name__)


class DataSetInfo(typing.TypedDict):
    """Metadata for a generated SDF dataset.

    Attributes
    ----------
    dataset_name : str
        Unique identifier for the dataset.
    class_names : list of str
        Names of geometry classes in the dataset.
    sampling_strategy : str
        Description of how points were sampled.
    date_created : str
        ISO format timestamp of dataset creation.
    stds : list of float
        Standard deviations used for normalization.
    n_samples : int
        Number of sample points per geometry.
    add_surface_samples : bool
        Whether surface points were included.
    sdf_struct_version : str
        Version of DeepSDFStruct used to create the dataset.
    """

    dataset_name: str
    class_names: list[str]
    sampling_strategy: str
    date_created: str
    stds: list[float]
    n_samples: int
    add_surface_samples: bool
    sdf_struct_version: str


class SphereParameters(typing.TypedDict):
    """Parameters defining a sampling sphere."""

    cx: float
    cy: float
    cz: float
    r: float


class SampledSDF:
    """Container for sampled SDF points and their distance values.

    This class stores point samples and their corresponding SDF values,
    providing utilities for data manipulation, splitting, and visualization.

    Parameters
    ----------
    samples : torch.Tensor
        Point coordinates of shape (N, 3).
    distances : torch.Tensor
        SDF values at sample points of shape (N, 1).

    Attributes
    ----------
    samples : torch.Tensor
        The sampled point coordinates.
    distances : torch.Tensor
        The SDF distance values.

    Methods
    -------
    split_pos_neg()
        Split into separate datasets for inside (negative) and outside
        (positive) points.
    create_gus_plottable()
        Convert to gustaf Vertices for visualization.
    stacked
        Property returning concatenated samples and distances.

    Examples
    --------
    >>> import torch
    >>> from DeepSDFStruct.sampling import SampledSDF
    >>>
    >>> points = torch.rand(100, 3)
    >>> distances = torch.randn(100, 1)
    >>> sampled = SampledSDF(points, distances)
    >>>
    >>> # Split by sign
    >>> inside, outside = sampled.split_pos_neg()
    >>> print(f"Inside points: {inside.samples.shape[0]}")
    >>> print(f"Outside points: {outside.samples.shape[0]}")
    """

    samples: torch.Tensor
    distances: torch.Tensor

    def split_pos_neg(self):
        """Split samples into inside (negative) and outside (positive) points.

        Returns
        -------
        pos : SampledSDF
            Samples with non-negative distances (outside or on surface).
        neg : SampledSDF
            Samples with negative distances (inside geometry).
        """
        pos_mask = torch.where(self.distances >= 0.0)[0]
        neg_mask = torch.where(self.distances < 0.0)[0]
        pos = SampledSDF(
            samples=self.samples[pos_mask], distances=self.distances[pos_mask]
        )
        neg = SampledSDF(
            samples=self.samples[neg_mask], distances=self.distances[neg_mask]
        )
        return pos, neg

    def create_gus_plottable(self):
        """Create a gustaf Vertices object for visualization.

        Returns
        -------
        gustaf.Vertices
            Vertices with distance values stored as vertex data.
        """
        vp = gus.Vertices(vertices=self.samples)
        vp.vertex_data["distance"] = self.distances
        return vp

    @property
    def stacked(self):
        """Concatenate samples and distances into a single tensor.

        Returns
        -------
        torch.Tensor
            Tensor of shape (N, 4) with [x, y, z, distance] per row.
        """
        return torch.hstack((self.samples, self.distances))

    def __init__(self, samples, distances):
        self.samples = samples
        self.distances = distances

    def __add__(self, other):
        """Concatenate two SampledSDF objects.

        Parameters
        ----------
        other : SampledSDF
            Another SampledSDF to concatenate.

        Returns
        -------
        SampledSDF
            Combined dataset with all samples from both objects.
        """
        return SampledSDF(
            samples=torch.vstack((self.samples, other.samples)),
            distances=torch.vstack((self.distances, other.distances)),
        )


def _process_single_geometry_instance(
    geometry,
    file_name: str,
    folder_name: pathlib.Path,
    scale,
    n_samples,
    sampling_strategy,
    add_surface_samples,
    stds,
    also_save_vtk,
    also_save_mesh,
):
    fname = folder_name / file_name
    if not os.path.exists(folder_name):
        os.makedirs(folder_name, exist_ok=True)
    if os.path.isfile(fname):
        logger.warning(f"File {fname} already exists")
        return
    mesh = None
    if isinstance(geometry, SDFBase):
        sdf = geometry
    elif isinstance(geometry, trimesh.Trimesh):
        mesh = geometry
        sdf = SDFfromMesh(mesh, scale=scale)
        if also_save_mesh:
            mesh.export(fname.with_suffix(".stl"))
    else:
        raise NotImplementedError(
            f"Geometry must be either trimesh or SDFBase, but not {type(geometry)}."
        )
    sampled_sdf = random_sample_sdf(
        sdf, bounds=(-1, 1), n_samples=int(n_samples), type=sampling_strategy
    )
    if add_surface_samples:
        if not isinstance(mesh, trimesh.Trimesh):
            logger.warning(
                "Add surface samples was specified, but geometry"
                f"is not given as a mesh but as {type(geometry)}"
            )
        else:
            surf_samples = sample_mesh_surface(
                sdf, mesh, int(n_samples // 2), stds, device="cpu", dtype=torch.float32
            )
            sampled_sdf += surf_samples
    pos, neg = sampled_sdf.split_pos_neg()

    np.savez(fname, neg=neg.stacked, pos=pos.stacked)
    if also_save_vtk:
        save_points_to_vtp(fname.with_suffix(".vtp"), neg=neg.stacked, pos=pos.stacked)


class SDFSampler:
    def __init__(
        self,
        outdir,
        splitdir,
        dataset_name,
        stds=[0.05, 0.025],
        overwrite_existing=False,
    ) -> None:
        self.outdir = outdir
        self.splitdir = splitdir
        self.dataset_name = dataset_name
        self.geometries = {}
        self.stds = stds
        folder_name = pathlib.Path(outdir) / dataset_name
        if os.path.exists(folder_name):
            if not overwrite_existing:
                raise IsADirectoryError(
                    f"Dataset {folder_name} already exists. "
                    "Set overwrite_existing to true to overwrite."
                )
        else:
            os.makedirs(folder_name)

    def add_class(self, geom_list: list, class_name: str, n_faces=100) -> None:
        """
        Adds a geometry to the sampler object. Tries to transform inputs to
        trimesh data. In case the geometry is a spline object, the n_faces
        parameter determines the accuracy of the extracted mesh
        """
        instances = {}
        for i, geom in enumerate(geom_list):
            instance_name = f"{class_name}_{i:05}"
            if isinstance(geom, splinepy.Multipatch | splinepy.spline.Spline):
                geom_gus: gus.Faces = geom.extract.faces(n_faces)
                tris = gus.create.faces.to_simplex(geom_gus)
                geom = trimesh.Trimesh(vertices=tris.vertices, faces=tris.faces)

            elif isinstance(geom, torchSurfMesh):
                geom = geom.to_trimesh()

            instances[instance_name] = geom
        self.geometries[class_name] = instances

    def process_geometries(
        self,
        sampling_strategy="uniform",
        n_faces=100,
        n_samples: int = 100000,
        add_surface_samples=True,
        also_save_vtk=False,
        also_save_mesh=True,
        scale=True,
        n_workers=0,
    ):
        tasks = []

        for class_name, instance_list in self.geometries.items():
            logger.info(
                f"processing geometry list {class_name} with {len(instance_list)} items."
            )

            folder_name = pathlib.Path(self.outdir) / self.dataset_name / class_name

            func = partial(
                _process_single_geometry_instance,
                scale=scale,
                n_samples=n_samples,
                sampling_strategy=sampling_strategy,
                add_surface_samples=add_surface_samples,
                stds=self.stds,
                also_save_vtk=also_save_vtk,
                also_save_mesh=also_save_mesh,
            )

            tasks = [
                (geometry, f"{instance_id}.npz", folder_name)
                for instance_id, geometry in instance_list.items()
            ]
            if n_workers > 0:
                logger.info(f"starting multiprocessing with {n_workers} workers")
                with multiprocessing.Pool(processes=n_workers) as pool:
                    pool.starmap(func, tasks)
            else:
                logger.info("starting serial processing of geometries")
                list(starmap(func, tasks))

            logger.info(f"done processing geometry list {class_name}")

        summary = DataSetInfo(
            dataset_name=self.dataset_name,
            class_names=list(self.geometries.keys()),
            sampling_strategy=sampling_strategy,
            date_created=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            stds=self.stds,
            n_samples=n_samples,
            add_surface_samples=add_surface_samples,
            sdf_struct_version=version("DeepSDFStruct"),
        )
        with open(
            str(pathlib.Path(self.outdir) / self.dataset_name / "summary.json"), "w"
        ) as f:
            json.dump(summary, f, indent=4)

    def get_meshs_from_folder(self, foldername, mesh_type) -> list:
        """
        Reads all mesh files of a given type (extension) from a folder using meshio.

        Parameters
        ----------
        foldername : str
            Path to the folder containing the mesh files.
        mesh_type : str
            Mesh file extension (e.g., 'vtk', 'obj', 'stl', 'msh', 'xdmf').

        Returns
        -------
        list[trimesh.Trimesh]
            A list of trimesh.Trimesh objects loaded from the folder.
        """
        meshes = []

        # Normalize extension (remove dot if present)
        mesh_type = mesh_type.lstrip(".")

        # Iterate through all files in the folder

        for filename in tqdm(os.listdir(foldername), desc="Loading meshs"):
            if filename.lower().endswith("." + mesh_type.lower()):
                filepath = os.path.join(foldername, filename)
                try:
                    faces = gus.io.meshio.load(filepath)
                    trim = trimesh.Trimesh(faces.vertices, faces.elements)
                    meshes.append(trim)
                    logger.info(f"Loaded mesh: {filename}")
                except ValueError as e:
                    logger.warning(f"Could not read {filename}: {e}")

        if not meshes:
            print(f"No .{mesh_type} meshes found in {foldername}.")

        return meshes

    def write_json(self, json_fname):
        json_content = defaultdict(lambda: defaultdict(list))
        for class_name, instance_list in self.geometries.items():
            for instance_id, geometry in instance_list.items():
                file_name = f"{instance_id}"
                json_content[self.dataset_name][class_name].append(file_name)
        # json_content = {
        #     data_info["dataset_name"]: {data_info["class_name"]: split_files}
        # }
        json_fname = pathlib.Path(f"{self.splitdir}/{json_fname}")
        if not json_fname.parent.is_dir():
            os.makedirs(json_fname.parent)
        with open(json_fname, "w", encoding="utf-8") as f:
            json.dump(json_content, f, indent=4)


def move(t_mesh, new_center):
    t_mesh.vertices += new_center - t_mesh.bounding_box.centroid


def noisy_sample(t_mesh, std, count):
    return t_mesh.sample(int(count)) + torch.random.normal(
        scale=std, size=(int(count), 3)
    )


def random_points(count):
    """random points in a unit sphere centered at (0, 0, 0)"""
    points = torch.random.uniform(-1, 1, (int(count * 3), 3))
    points = points[torch.linalg.norm(points, axis=1) <= 1]
    if points.shape[0] < count:
        print("Too little random sampling points. Resampling.......")
        random_points(count=count, boundary="unit_sphere")
    elif points.shape[0] > count:
        return points[torch.random.choice(points.shape[0], count)]
    else:
        return points


def random_points_cube(count, box_size):
    """random points in a cube with size box_size centered at (0, 0, 0)"""
    points = torch.random.uniform(-box_size / 2, box_size / 2, (int(count), 3))
    return points


def random_sample_sdf(
    sdf, bounds, n_samples, type="uniform", device="cpu", dtype=torch.float32
):
    bounds = torch.tensor(bounds, dtype=dtype, device=device)
    if type == "plane":
        samples = torch.random.uniform(
            bounds[0], bounds[1], (n_samples, 2), device=device, dtype=dtype
        )
        samples = torch.hstack((samples, torch.zeros((n_samples, 1))))
    elif type == "spherical_gaussian":
        samples = torch.random.randn(n_samples, 3, device=device, dtype=dtype)
        samples /= torch.linalg.norm(samples, axis=1).reshape(-1, 1)
        # samples += torch.random.uniform(bounds[0], bounds[1], (n_samples, 3))
        samples = samples + torch.random.normal(0, 0.01, (n_samples, 3))
    elif type == "uniform":
        samples = (
            torch.rand((n_samples, 3), device=device, dtype=dtype)
            * (bounds[1] - bounds[0])
            + bounds[0]
        )
    else:
        raise ValueError(f"Unsupported sampling strategy {type}")
    distances = sdf(samples)
    return SampledSDF(samples=samples, distances=distances)


def sample_mesh_surface(
    sdf: SDFBase,
    mesh: trimesh.Trimesh,
    n_samples: int,
    stds: list[float],
    device="cpu",
    dtype=torch.float32,
) -> SampledSDF:
    """
    Sample noisy points around a mesh surface and evaluate them with a signed distance function (SDF).

    This function uses trimesh.sample to generate surface samples
    and perturbs them with Gaussian noise of varying standard deviations,
    and queries the SDF at those points.

    Args:
        sdf (SDFBase): A callable SDF object that takes 3D points and returns signed distances.
        mesh (trimesh.Trimesh): A mesh object containing the vertices and faces.
        n_samples (int): Number of mesh vertices to sample
        stds (list[float]): Standard deviations for Gaussian noise added to sampled vertices.
            - Typical values: [0.05, 0.0015].
            - Larger values spread samples farther from the surface; smaller values keep them closer.
        device (str, optional): Torch device to place tensors on (e.g., "cpu" or "cuda").
        dtype (torch.dtype, optional): Data type for generated tensors (default: torch.float32).

    Returns:
        SampledSDF: An object containing:
            - samples (torch.Tensor): The perturbed sample points of shape (n_samples * len(stds), 3).
            - distances (torch.Tensor): The corresponding SDF values at those sample points.
    """
    samples = []

    points, face_idx = trimesh.sample.sample_surface(mesh, n_samples)

    surface_points = torch.tensor(points, dtype=dtype, device=device)

    face_normals = torch.tensor(
        mesh.face_normals[face_idx], dtype=dtype, device=device
    )

    for std in stds:

        t = torch.randn(n_samples, 1, device=device, dtype=dtype) * std

        noisy = surface_points + t * face_normals
        samples.append(noisy)

    queries = torch.vstack(samples)
    distances = sdf(queries)

    return SampledSDF(samples=queries, distances=distances)


def save_points_to_vtp(filename, neg, pos):
    """
    Save pos/neg SDF sample points as a VTU point cloud using vtkPolyData.
    Each point has an SDF scalar value.
    """
    # Combine points
    all_points = np.vstack((pos, neg))
    coords = all_points[:, :3]
    sdf_vals = all_points[:, 3]

    # --- Create vtkPoints ---
    vtk_points = vtk.vtkPoints()
    for pt in coords:
        vtk_points.InsertNextPoint(pt)

    # --- Create PolyData ---
    polydata = vtk.vtkPolyData()
    polydata.SetPoints(vtk_points)

    # Add vertex cells (required for points in PolyData)
    verts = vtk.vtkCellArray()
    for i in range(len(coords)):
        verts.InsertNextCell(1)
        verts.InsertCellPoint(i)
    polydata.SetVerts(verts)

    # --- Add SDF scalar values ---
    vtk_array = vtk.vtkDoubleArray()
    vtk_array.SetName("SDF")
    vtk_array.SetNumberOfValues(len(sdf_vals))
    for i, val in enumerate(sdf_vals):
        vtk_array.SetValue(i, val)
    polydata.GetPointData().SetScalars(vtk_array)

    # --- Write to VTU ---
    writer = vtk.vtkXMLPolyDataWriter()
    writer.SetFileName(filename)
    writer.SetInputData(polydata)
    writer.Write()

    logger.debug(f"Saved {len(coords)} points with SDF to '{filename}'")


def augment_by_FFD(
    meshs: list[trimesh.Trimesh],
    n_control_points: int = 5,
    std_dev_fraction: float | None = 0.05,
    n_transformations: int = 10,
    save_meshs=False,
) -> list[trimesh.Trimesh]:
    """
    Takes list of meshs and augments the meshs by applying a freeform deformation
    """
    new_meshs = []

    for i_mesh, mesh in enumerate(tqdm(meshs, desc="Augmenting meshs")):
        bbox = mesh.bounds  # shape (2, 3)
        # Compute approximate spacing between control points along each axis
        spacing = (bbox[1] - bbox[0]) / (n_control_points - 1)
        # Use a fraction of spacing (e.g., 15%) as std_dev
        std_dev_local = std_dev_fraction * spacing

        for i_FFD in range(n_transformations):
            ffd = splinepy.FFD()
            ffd.mesh = gus.Faces(mesh.vertices, mesh.faces)
            ffd.spline.insert_knots(0, np.linspace(0, 1, n_control_points)[1:-1])
            ffd.spline.insert_knots(1, np.linspace(0, 1, n_control_points)[1:-1])
            ffd.spline.insert_knots(2, np.linspace(0, 1, n_control_points)[1:-1])
            ffd.spline.elevate_degrees([0, 1, 2])
            ffd.spline.control_points += np.random.normal(
                loc=0.0, scale=std_dev_local, size=ffd.spline.control_points.shape
            )
            new_meshs.append(trimesh.Trimesh(ffd.mesh.vertices, ffd.mesh.faces))
            if save_meshs:
                save_meshs = True

                # Make sure the directory exists
                os.makedirs("tmp", exist_ok=True)
                gus.io.meshio.export(f"tmp/mesh_{i_mesh}_{i_FFD}.obj", ffd.mesh)

    return new_meshs

1	"""
2	SDF Sampling and Dataset Generation
3	====================================
4
5	This module provides tools for sampling points from SDF representations and
6	generating datasets for training neural networks (DeepSDF models). It supports
7	various sampling strategies to create well-distributed training data.
8
9	Key Features
10	------------
11
12	Sampling Strategies
13	- Uniform sampling in a bounding box
14	- Surface-focused sampling near the zero level set
15	- Importance sampling based on SDF gradients
16	- Sphere-based sampling patterns
17	- Combined strategies for balanced datasets
18
19	Dataset Generation
20	- Batch processing of multiple geometries
21	- Automatic data normalization and standardization
22	- Support for multiple geometry classes
23	- Metadata tracking (version, sampling parameters)
24	- Export to formats compatible with DeepSDF training
25
26	The module is designed to generate high-quality training data for implicit
27	neural representations, with careful attention to sampling near surfaces
28	where accurate reconstruction is most critical.
29
30	Classes
31	-------
32	SampledSDF
33	Container for sampled points and their SDF values, with utilities
34	for splitting by sign and visualization.
35
36	DataSetInfo
37	TypedDict for dataset metadata (name, classes, sampling strategy, etc.).
38
39	Functions
40	---------
41	generate_dataset
42	Batch process multiple geometries to create a complete dataset.
43	sample_sdf_*
44	Various sampling strategies for different use cases.
45	"""
46
47	import os	1✔
48	from functools import partial	1✔
49	from itertools import starmap	1✔
50	import multiprocessing	1✔
51
52	import vtk	1✔
53	import numpy as np	1✔
54	import json	1✔
55	import pathlib	1✔
56	import typing	1✔
57	import gustaf as gus	1✔
58	import trimesh	1✔
59	from DeepSDFStruct.SDF import SDFfromMesh, SDFBase	1✔
60	from DeepSDFStruct.mesh import torchSurfMesh	1✔
61	import DeepSDFStruct	1✔
62
63	# from analysis.problems.homogenization import computeHomogenizedMaterialProperties
64	import splinepy	1✔
65	import torch	1✔
66	from collections import defaultdict	1✔
67	from tqdm import tqdm	1✔
68	import logging	1✔
69	import datetime	1✔
70	from importlib.metadata import version	1✔
71
72	logger = logging.getLogger(DeepSDFStruct.__name__)	1✔
73
74
75	class DataSetInfo(typing.TypedDict):	1✔
76	"""Metadata for a generated SDF dataset.
77
78	Attributes
79	----------
80	dataset_name : str
81	Unique identifier for the dataset.
82	class_names : list of str
83	Names of geometry classes in the dataset.
84	sampling_strategy : str
85	Description of how points were sampled.
86	date_created : str
87	ISO format timestamp of dataset creation.
88	stds : list of float
89	Standard deviations used for normalization.
90	n_samples : int
91	Number of sample points per geometry.
92	add_surface_samples : bool
93	Whether surface points were included.
94	sdf_struct_version : str
95	Version of DeepSDFStruct used to create the dataset.
96	"""
97
98	dataset_name: str	1✔
99	class_names: list[str]	1✔
100	sampling_strategy: str	1✔
101	date_created: str	1✔
102	stds: list[float]	1✔
103	n_samples: int	1✔
104	add_surface_samples: bool	1✔
105	sdf_struct_version: str	1✔
106
107
108	class SphereParameters(typing.TypedDict):	1✔
109	"""Parameters defining a sampling sphere."""
110
111	cx: float	1✔
112	cy: float	1✔
113	cz: float	1✔
114	r: float	1✔
115
116
117	class SampledSDF:	1✔
118	"""Container for sampled SDF points and their distance values.
119
120	This class stores point samples and their corresponding SDF values,
121	providing utilities for data manipulation, splitting, and visualization.
122
123	Parameters
124	----------
125	samples : torch.Tensor
126	Point coordinates of shape (N, 3).
127	distances : torch.Tensor
128	SDF values at sample points of shape (N, 1).
129
130	Attributes
131	----------
132	samples : torch.Tensor
133	The sampled point coordinates.
134	distances : torch.Tensor
135	The SDF distance values.
136
137	Methods
138	-------
139	split_pos_neg()
140	Split into separate datasets for inside (negative) and outside
141	(positive) points.
142	create_gus_plottable()
143	Convert to gustaf Vertices for visualization.
144	stacked
145	Property returning concatenated samples and distances.
146
147	Examples
148	--------
149	>>> import torch
150	>>> from DeepSDFStruct.sampling import SampledSDF
151	>>>
152	>>> points = torch.rand(100, 3)
153	>>> distances = torch.randn(100, 1)
154	>>> sampled = SampledSDF(points, distances)
155	>>>
156	>>> # Split by sign
157	>>> inside, outside = sampled.split_pos_neg()
158	>>> print(f"Inside points: {inside.samples.shape[0]}")
159	>>> print(f"Outside points: {outside.samples.shape[0]}")
160	"""
161
162	samples: torch.Tensor	1✔
163	distances: torch.Tensor	1✔
164
165	def split_pos_neg(self):	1✔
166	"""Split samples into inside (negative) and outside (positive) points.
167
168	Returns
169	-------
170	pos : SampledSDF
171	Samples with non-negative distances (outside or on surface).
172	neg : SampledSDF
173	Samples with negative distances (inside geometry).
174	"""
175	pos_mask = torch.where(self.distances >= 0.0)[0]	1✔
176	neg_mask = torch.where(self.distances < 0.0)[0]	1✔
177	pos = SampledSDF(	1✔
178	samples=self.samples[pos_mask], distances=self.distances[pos_mask]
179	)
180	neg = SampledSDF(	1✔
181	samples=self.samples[neg_mask], distances=self.distances[neg_mask]
182	)
183	return pos, neg	1✔
184
185	def create_gus_plottable(self):	1✔
186	"""Create a gustaf Vertices object for visualization.
187
188	Returns
189	-------
190	gustaf.Vertices
191	Vertices with distance values stored as vertex data.
192	"""
193	vp = gus.Vertices(vertices=self.samples)	×
194	vp.vertex_data["distance"] = self.distances	×
195	return vp	×
196
197	@property	1✔
198	def stacked(self):	1✔
199	"""Concatenate samples and distances into a single tensor.
200
201	Returns
202	-------
203	torch.Tensor
204	Tensor of shape (N, 4) with [x, y, z, distance] per row.
205	"""
206	return torch.hstack((self.samples, self.distances))	1✔
207
208	def __init__(self, samples, distances):	1✔
209	self.samples = samples	1✔
210	self.distances = distances	1✔
211
212	def __add__(self, other):	1✔
213	"""Concatenate two SampledSDF objects.
214
215	Parameters
216	----------
217	other : SampledSDF
218	Another SampledSDF to concatenate.
219
220	Returns
221	-------
222	SampledSDF
223	Combined dataset with all samples from both objects.
224	"""
225	return SampledSDF(	1✔
226	samples=torch.vstack((self.samples, other.samples)),
227	distances=torch.vstack((self.distances, other.distances)),
228	)
229
230
231	def _process_single_geometry_instance(	1✔
232	geometry,
233	file_name: str,
234	folder_name: pathlib.Path,
235	scale,
236	n_samples,
237	sampling_strategy,
238	add_surface_samples,
239	stds,
240	also_save_vtk,
241	also_save_mesh,
242	):
243	fname = folder_name / file_name	1✔
244	if not os.path.exists(folder_name):	1✔
245	os.makedirs(folder_name, exist_ok=True)	1✔
246	if os.path.isfile(fname):	1✔
247	logger.warning(f"File {fname} already exists")	×
248	return	×
249	mesh = None	1✔
250	if isinstance(geometry, SDFBase):	1✔
251	sdf = geometry	×
252	elif isinstance(geometry, trimesh.Trimesh):	1✔
253	mesh = geometry	1✔
254	sdf = SDFfromMesh(mesh, scale=scale)	1✔
255	if also_save_mesh:	1✔
256	mesh.export(fname.with_suffix(".stl"))	1✔
257	else:
258	raise NotImplementedError(	×
259	f"Geometry must be either trimesh or SDFBase, but not {type(geometry)}."
260	)
261	sampled_sdf = random_sample_sdf(	1✔
262	sdf, bounds=(-1, 1), n_samples=int(n_samples), type=sampling_strategy
263	)
264	if add_surface_samples:	1✔
265	if not isinstance(mesh, trimesh.Trimesh):	1✔
266	logger.warning(	×
267	"Add surface samples was specified, but geometry"
268	f"is not given as a mesh but as {type(geometry)}"
269	)
270	else:
271	surf_samples = sample_mesh_surface(	1✔
272	sdf, mesh, int(n_samples // 2), stds, device="cpu", dtype=torch.float32
273	)
274	sampled_sdf += surf_samples	1✔
275	pos, neg = sampled_sdf.split_pos_neg()	1✔
276
277	np.savez(fname, neg=neg.stacked, pos=pos.stacked)	1✔
278	if also_save_vtk:	1✔
279	save_points_to_vtp(fname.with_suffix(".vtp"), neg=neg.stacked, pos=pos.stacked)	1✔
280
281
282	class SDFSampler:	1✔
283	def __init__(	1✔
284	self,
285	outdir,
286	splitdir,
287	dataset_name,
288	stds=[0.05, 0.025],
289	overwrite_existing=False,
290	) -> None:
291	self.outdir = outdir	1✔
292	self.splitdir = splitdir	1✔
293	self.dataset_name = dataset_name	1✔
294	self.geometries = {}	1✔
295	self.stds = stds	1✔
296	folder_name = pathlib.Path(outdir) / dataset_name	1✔
297	if os.path.exists(folder_name):	1✔
298	if not overwrite_existing:	×
299	raise IsADirectoryError(	×
300	f"Dataset {folder_name} already exists. "
301	"Set overwrite_existing to true to overwrite."
302	)
303	else:
304	os.makedirs(folder_name)	1✔
305
306	def add_class(self, geom_list: list, class_name: str, n_faces=100) -> None:	1✔
307	"""
308	Adds a geometry to the sampler object. Tries to transform inputs to
309	trimesh data. In case the geometry is a spline object, the n_faces
310	parameter determines the accuracy of the extracted mesh
311	"""
312	instances = {}	1✔
313	for i, geom in enumerate(geom_list):	1✔
314	instance_name = f"{class_name}_{i:05}"	1✔
315	if isinstance(geom, splinepy.Multipatch \| splinepy.spline.Spline):	1✔
316	geom_gus: gus.Faces = geom.extract.faces(n_faces)	1✔
317	tris = gus.create.faces.to_simplex(geom_gus)	1✔
318	geom = trimesh.Trimesh(vertices=tris.vertices, faces=tris.faces)	1✔
319
320	elif isinstance(geom, torchSurfMesh):	1✔
321	geom = geom.to_trimesh()	×
322
323	instances[instance_name] = geom	1✔
324	self.geometries[class_name] = instances	1✔
325
326	def process_geometries(	1✔
327	self,
328	sampling_strategy="uniform",
329	n_faces=100,
330	n_samples: int = 100000,
331	add_surface_samples=True,
332	also_save_vtk=False,
333	also_save_mesh=True,
334	scale=True,
335	n_workers=0,
336	):
337	tasks = []	1✔
338
339	for class_name, instance_list in self.geometries.items():	1✔
340	logger.info(	1✔
341	f"processing geometry list {class_name} with {len(instance_list)} items."
342	)
343
344	folder_name = pathlib.Path(self.outdir) / self.dataset_name / class_name	1✔
345
346	func = partial(	1✔
347	_process_single_geometry_instance,
348	scale=scale,
349	n_samples=n_samples,
350	sampling_strategy=sampling_strategy,
351	add_surface_samples=add_surface_samples,
352	stds=self.stds,
353	also_save_vtk=also_save_vtk,
354	also_save_mesh=also_save_mesh,
355	)
356
357	tasks = [	1✔
358	(geometry, f"{instance_id}.npz", folder_name)
359	for instance_id, geometry in instance_list.items()
360	]
361	if n_workers > 0:	1✔
362	logger.info(f"starting multiprocessing with {n_workers} workers")	1✔
363	with multiprocessing.Pool(processes=n_workers) as pool:	1✔
364	pool.starmap(func, tasks)	1✔
365	else:
366	logger.info("starting serial processing of geometries")	1✔
367	list(starmap(func, tasks))	1✔
368
369	logger.info(f"done processing geometry list {class_name}")	1✔
370
371	summary = DataSetInfo(	1✔
372	dataset_name=self.dataset_name,
373	class_names=list(self.geometries.keys()),
374	sampling_strategy=sampling_strategy,
375	date_created=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
376	stds=self.stds,
377	n_samples=n_samples,
378	add_surface_samples=add_surface_samples,
379	sdf_struct_version=version("DeepSDFStruct"),
380	)
381	with open(	1✔
382	str(pathlib.Path(self.outdir) / self.dataset_name / "summary.json"), "w"
383	) as f:
384	json.dump(summary, f, indent=4)	1✔
385
386	def get_meshs_from_folder(self, foldername, mesh_type) -> list:	1✔
387	"""
388	Reads all mesh files of a given type (extension) from a folder using meshio.
389
390	Parameters
391	----------
392	foldername : str
393	Path to the folder containing the mesh files.
394	mesh_type : str
395	Mesh file extension (e.g., 'vtk', 'obj', 'stl', 'msh', 'xdmf').
396
397	Returns
398	-------
399	list[trimesh.Trimesh]
400	A list of trimesh.Trimesh objects loaded from the folder.
401	"""
402	meshes = []	1✔
403
404	# Normalize extension (remove dot if present)
405	mesh_type = mesh_type.lstrip(".")	1✔
406
407	# Iterate through all files in the folder
408
409	for filename in tqdm(os.listdir(foldername), desc="Loading meshs"):	1✔
410	if filename.lower().endswith("." + mesh_type.lower()):	1✔
411	filepath = os.path.join(foldername, filename)	1✔
412	try:	1✔
413	faces = gus.io.meshio.load(filepath)	1✔
414	trim = trimesh.Trimesh(faces.vertices, faces.elements)	1✔
415	meshes.append(trim)	1✔
416	logger.info(f"Loaded mesh: {filename}")	1✔
417	except ValueError as e:	×
418	logger.warning(f"Could not read {filename}: {e}")	×
419
420	if not meshes:	1✔
421	print(f"No .{mesh_type} meshes found in {foldername}.")	×
422
423	return meshes	1✔
424
425	def write_json(self, json_fname):	1✔
426	json_content = defaultdict(lambda: defaultdict(list))	1✔
427	for class_name, instance_list in self.geometries.items():	1✔
428	for instance_id, geometry in instance_list.items():	1✔
429	file_name = f"{instance_id}"	1✔
430	json_content[self.dataset_name][class_name].append(file_name)	1✔
431	# json_content = {
432	# data_info["dataset_name"]: {data_info["class_name"]: split_files}
433	# }
434	json_fname = pathlib.Path(f"{self.splitdir}/{json_fname}")	1✔
435	if not json_fname.parent.is_dir():	1✔
436	os.makedirs(json_fname.parent)	1✔
437	with open(json_fname, "w", encoding="utf-8") as f:	1✔
438	json.dump(json_content, f, indent=4)	1✔
439
440
441	def move(t_mesh, new_center):	1✔
442	t_mesh.vertices += new_center - t_mesh.bounding_box.centroid	×
443
444
445	def noisy_sample(t_mesh, std, count):	1✔
446	return t_mesh.sample(int(count)) + torch.random.normal(	×
447	scale=std, size=(int(count), 3)
448	)
449
450
451	def random_points(count):	1✔
452	"""random points in a unit sphere centered at (0, 0, 0)"""
453	points = torch.random.uniform(-1, 1, (int(count * 3), 3))	×
454	points = points[torch.linalg.norm(points, axis=1) <= 1]	×
455	if points.shape[0] < count:	×
456	print("Too little random sampling points. Resampling.......")	×
457	random_points(count=count, boundary="unit_sphere")	×
458	elif points.shape[0] > count:	×
459	return points[torch.random.choice(points.shape[0], count)]	×
460	else:
461	return points	×
462
463
464	def random_points_cube(count, box_size):	1✔
465	"""random points in a cube with size box_size centered at (0, 0, 0)"""
466	points = torch.random.uniform(-box_size / 2, box_size / 2, (int(count), 3))	×
467	return points	×
468
469
470	def random_sample_sdf(	1✔
471	sdf, bounds, n_samples, type="uniform", device="cpu", dtype=torch.float32
472	):
473	bounds = torch.tensor(bounds, dtype=dtype, device=device)	1✔
474	if type == "plane":	1✔
475	samples = torch.random.uniform(	×
476	bounds[0], bounds[1], (n_samples, 2), device=device, dtype=dtype
477	)
478	samples = torch.hstack((samples, torch.zeros((n_samples, 1))))	×
479	elif type == "spherical_gaussian":	1✔
480	samples = torch.random.randn(n_samples, 3, device=device, dtype=dtype)	×
481	samples /= torch.linalg.norm(samples, axis=1).reshape(-1, 1)	×
482	# samples += torch.random.uniform(bounds[0], bounds[1], (n_samples, 3))
483	samples = samples + torch.random.normal(0, 0.01, (n_samples, 3))	×
484	elif type == "uniform":	1✔
485	samples = (	1✔
486	torch.rand((n_samples, 3), device=device, dtype=dtype)
487	* (bounds[1] - bounds[0])
488	+ bounds[0]
489	)
490	else:
491	raise ValueError(f"Unsupported sampling strategy {type}")	×
492	distances = sdf(samples)	1✔
493	return SampledSDF(samples=samples, distances=distances)	1✔
494
495
496	def sample_mesh_surface(	1✔
497	sdf: SDFBase,
498	mesh: trimesh.Trimesh,
499	n_samples: int,
500	stds: list[float],
501	device="cpu",
502	dtype=torch.float32,
503	) -> SampledSDF:
504	"""
505	Sample noisy points around a mesh surface and evaluate them with a signed distance function (SDF).
506
507	This function uses trimesh.sample to generate surface samples
508	and perturbs them with Gaussian noise of varying standard deviations,
509	and queries the SDF at those points.
510
511	Args:
512	sdf (SDFBase): A callable SDF object that takes 3D points and returns signed distances.
513	mesh (trimesh.Trimesh): A mesh object containing the vertices and faces.
514	n_samples (int): Number of mesh vertices to sample
515	stds (list[float]): Standard deviations for Gaussian noise added to sampled vertices.
516	- Typical values: [0.05, 0.0015].
517	- Larger values spread samples farther from the surface; smaller values keep them closer.
518	device (str, optional): Torch device to place tensors on (e.g., "cpu" or "cuda").
519	dtype (torch.dtype, optional): Data type for generated tensors (default: torch.float32).
520
521	Returns:
522	SampledSDF: An object containing:
523	- samples (torch.Tensor): The perturbed sample points of shape (n_samples * len(stds), 3).
524	- distances (torch.Tensor): The corresponding SDF values at those sample points.
525	"""
526	samples = []	1✔
527
528	points, face_idx = trimesh.sample.sample_surface(mesh, n_samples)	1✔
529
530	surface_points = torch.tensor(points, dtype=dtype, device=device)	1✔
531
532	face_normals = torch.tensor(	1✔
533	mesh.face_normals[face_idx], dtype=dtype, device=device
534	)
535
536	for std in stds:	1✔
537
538	t = torch.randn(n_samples, 1, device=device, dtype=dtype) * std	1✔
539
540	noisy = surface_points + t * face_normals	1✔
541	samples.append(noisy)	1✔
542
543	queries = torch.vstack(samples)	1✔
544	distances = sdf(queries)	1✔
545
546	return SampledSDF(samples=queries, distances=distances)	1✔
547
548
549	def save_points_to_vtp(filename, neg, pos):	1✔
550	"""
551	Save pos/neg SDF sample points as a VTU point cloud using vtkPolyData.
552	Each point has an SDF scalar value.
553	"""
554	# Combine points
555	all_points = np.vstack((pos, neg))	1✔
556	coords = all_points[:, :3]	1✔
557	sdf_vals = all_points[:, 3]	1✔
558
559	# --- Create vtkPoints ---
560	vtk_points = vtk.vtkPoints()	1✔
561	for pt in coords:	1✔
562	vtk_points.InsertNextPoint(pt)	1✔
563
564	# --- Create PolyData ---
565	polydata = vtk.vtkPolyData()	1✔
566	polydata.SetPoints(vtk_points)	1✔
567
568	# Add vertex cells (required for points in PolyData)
569	verts = vtk.vtkCellArray()	1✔
570	for i in range(len(coords)):	1✔
571	verts.InsertNextCell(1)	1✔
572	verts.InsertCellPoint(i)	1✔
573	polydata.SetVerts(verts)	1✔
574
575	# --- Add SDF scalar values ---
576	vtk_array = vtk.vtkDoubleArray()	1✔
577	vtk_array.SetName("SDF")	1✔
578	vtk_array.SetNumberOfValues(len(sdf_vals))	1✔
579	for i, val in enumerate(sdf_vals):	1✔
580	vtk_array.SetValue(i, val)	1✔
581	polydata.GetPointData().SetScalars(vtk_array)	1✔
582
583	# --- Write to VTU ---
584	writer = vtk.vtkXMLPolyDataWriter()	1✔
585	writer.SetFileName(filename)	1✔
586	writer.SetInputData(polydata)	1✔
587	writer.Write()	1✔
588
589	logger.debug(f"Saved {len(coords)} points with SDF to '{filename}'")	1✔
590
591
592	def augment_by_FFD(	1✔
593	meshs: list[trimesh.Trimesh],
594	n_control_points: int = 5,
595	std_dev_fraction: float \| None = 0.05,
596	n_transformations: int = 10,
597	save_meshs=False,
598	) -> list[trimesh.Trimesh]:
599	"""
600	Takes list of meshs and augments the meshs by applying a freeform deformation
601	"""
602	new_meshs = []	1✔
603
604	for i_mesh, mesh in enumerate(tqdm(meshs, desc="Augmenting meshs")):	1✔
605	bbox = mesh.bounds # shape (2, 3)	1✔
606	# Compute approximate spacing between control points along each axis
607	spacing = (bbox[1] - bbox[0]) / (n_control_points - 1)	1✔
608	# Use a fraction of spacing (e.g., 15%) as std_dev
609	std_dev_local = std_dev_fraction * spacing	1✔
610
611	for i_FFD in range(n_transformations):	1✔
612	ffd = splinepy.FFD()	1✔
613	ffd.mesh = gus.Faces(mesh.vertices, mesh.faces)	1✔
614	ffd.spline.insert_knots(0, np.linspace(0, 1, n_control_points)[1:-1])	1✔
615	ffd.spline.insert_knots(1, np.linspace(0, 1, n_control_points)[1:-1])	1✔
616	ffd.spline.insert_knots(2, np.linspace(0, 1, n_control_points)[1:-1])	1✔
617	ffd.spline.elevate_degrees([0, 1, 2])	1✔
618	ffd.spline.control_points += np.random.normal(	1✔
619	loc=0.0, scale=std_dev_local, size=ffd.spline.control_points.shape
620	)
621	new_meshs.append(trimesh.Trimesh(ffd.mesh.vertices, ffd.mesh.faces))	1✔
622	if save_meshs:	1✔
623	save_meshs = True	×
624
625	# Make sure the directory exists
626	os.makedirs("tmp", exist_ok=True)	×
627	gus.io.meshio.export(f"tmp/mesh_{i_mesh}_{i_FFD}.obj", ffd.mesh)	×
628
629	return new_meshs	1✔

mkofler96 / DeepSDFStruct / 24552038416

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous