added 2D mesh extraction and spline features

Author: mkofler96

DeepSDFStruct/SDF.py

@@ -119,6 +119,7 @@ def __init__(
         parametrization: torch.nn.Module | None = None,
         cap_border_dict: CapBorderDict = None,
         cap_outside_of_unitcube=False,
+        geometric_dim=3,
     ):
         self._deformation_spline = deformation_spline
 
@@ -130,6 +131,7 @@ def __init__(
 
         self._cap_border_dict = cap_border_dict
         self.cap_outside_of_unitcube = cap_outside_of_unitcube
+        self.geometric_dim = geometric_dim
 
     @property
     def deformation_spline(self):
@@ -186,7 +188,9 @@ def __call__(self, queries: torch.Tensor) -> torch.Tensor:
             # cap everything outside of the unit cube
             # k and d are y = k*(x-dx) + dy
         if self.cap_outside_of_unitcube:
-            sdf_values = _cap_outside_of_unitcube(queries, sdf_values)
+            sdf_values = _cap_outside_of_unitcube(
+                queries, sdf_values, max_dim=self.geometric_dim
+            )
         return sdf_values
 
     def _validate_input(self, queries: torch.Tensor):
@@ -216,6 +220,46 @@ def plot_slice(self, *args, **kwargs):
     def __add__(self, other):
         return SummedSDF(self, other)
 
+    def to2D(self, axes: list[int], offset=0.0):
+        """
+        Converts SDF to 2D
+
+        :param axis: list of axes that will be used for the 2D
+        """
+        sdf2D = SDF2D(self, axes, offset=offset)
+        sdf2D.deformation_spline = self.deformation_spline
+        sdf2D.parametrization = self.parametrization
+        return sdf2D
+
+
+class SDF2D(SDFBase):
+    def __init__(self, obj: SDFBase, axes: list[int], offset=0.0):
+        super().__init__()
+        self.obj = obj
+        assert (
+            len(axes) == 2
+        ), "List of axes must be of size 2 and needs to correspond to the 2D plane"
+        self.axes = axes
+        self.offset = offset
+
+    def _compute(self, queries):
+        queries_3D = (
+            torch.zeros(
+                (queries.shape[0], 3), dtype=queries.dtype, device=queries.device
+            )
+            + self.offset
+        )
+        queries_3D[:, self.axes[0]] = queries[:, 0]
+        queries_3D[:, self.axes[1]] = queries[:, 1]
+        result = self.obj._compute(queries_3D)
+        return result
+
+    def _get_domain_bounds(self):
+        return self.obj._get_domain_bounds()
+
+    def _set_param(self, parameter):
+        return self.obj._set_param(parameter)
+
 
 class SummedSDF(SDFBase):
     def __init__(self, obj1: SDFBase, obj2: SDFBase):
@@ -267,7 +311,26 @@ def _compute(self, queries: torch.tensor) -> torch.tensor:
         return output.reshape(-1, 1)
 
 
-def union(D, k=0):
+def union_torch(D, k=0):
+    """
+    D: np.array of shape (num_points, num_geometries)
+    k: smoothness parameter
+    """
+    if k == 0:
+        return torch.min(D, axis=1)[0]
+    # Start with the first column as d1
+    d1 = D[:, 0].copy()
+
+    # Loop over remaining columns
+    for i in range(1, D.shape[1]):
+        d2 = D[:, i]
+        h = torch.clip(0.5 + 0.5 * (d2 - d1) / k, 0, 1)
+        d1 = d2 + (d1 - d2) * h - k * h * (1 - h)
+
+    return d1
+
+
+def union_numpy(D, k=0):
     """
     D: np.array of shape (num_points, num_geometries)
     k: smoothness parameter
@@ -428,7 +491,9 @@ def _compute(self, queries: torch.Tensor | np.ndarray):
         sdf = point_segment_distance(lines[:, 0], lines[:, 1], queries_np) - self.t / 2
         if is_tensor:
             sdf = torch.tensor(sdf, dtype=orig_dtype, device=orig_device)
-        return union(sdf, k=self.smoothness)
+            return union_torch(sdf, k=self.smoothness)
+        else:
+            return union_numpy(sdf, k=self.smoothness)
 
 
 class SDFfromDeepSDF(SDFBase):
@@ -507,13 +572,12 @@ def _compute(self, queries: torch.Tensor) -> torch.Tensor:
         return sdf_values.to(orig_device).reshape(-1, 1)
 
 
-def _cap_outside_of_unitcube(samples, sdf_values):
-    dy = 0
-    for dim, k, dx in zip(
-        [0, 0, 1, 1, 2, 2], [1, -1, 1, -1, 1, -1], [0, 1, 0, 1, 0, 1]
-    ):
-        x = samples[:, dim]
-        border_sdf = k * (x - dx) + dy
+def _cap_outside_of_unitcube(samples, sdf_values, max_dim=3):
+
+    for dim in range(max_dim):
+        border_sdf = samples[:, dim]
+        sdf_values = torch.maximum(sdf_values, -border_sdf.reshape(-1, 1))
+        border_sdf = 1 - samples[:, dim]
         sdf_values = torch.maximum(sdf_values, -border_sdf.reshape(-1, 1))
     return sdf_values
 

DeepSDFStruct/flexicubes/flexicubes.py

@@ -137,7 +137,9 @@ def __init__(self, device="cuda", qef_reg_scale=1e-3, weight_scale=0.99):
         self.qef_reg_scale = qef_reg_scale
         self.weight_scale = weight_scale
 
-    def construct_voxel_grid(self, resolution) -> tuple[torch.tensor, torch.tensor]:
+    def construct_voxel_grid(
+        self, resolution, bounds=None
+    ) -> tuple[torch.tensor, torch.tensor]:
         """
         Generates a voxel grid based on the specified resolution.
 
@@ -176,8 +178,36 @@ def construct_voxel_grid(self, resolution) -> tuple[torch.tensor, torch.tensor]:
             verts_rounded, dim=0, return_inverse=True
         )
         cubes = inverse_indices[cubes.reshape(-1)].reshape(-1, 8)
-
-        return verts_unique - 0.5, cubes
+        # changes mkofler: added option to generate points in arbitrary bounds
+        if bounds is None:
+            bounds = torch.tensor(
+                [[-0.05, -0.05, -0.05], [1.05, 1.05, 1.05]],
+                device=verts_unique.device,
+                dtype=verts_unique.dtype,
+            )
+        else:
+            bounds = torch.as_tensor(
+                bounds, device=verts_unique.device, dtype=verts_unique.dtype
+            )
+            assert bounds.shape == (2, 2), "bounds must have shape [2, 2]"
+
+        # Scale samples from [0, 1] to the given bounds
+        # samples = samples * 1.1 + _torch.tensor([0.5, 0.5, 0.5], device=device)
+        samples = verts_unique - 0.5
+        verts_scaled = (
+            bounds[0]
+            + 0.5 * (bounds[1] - bounds[0])
+            + (bounds[1] - bounds[0]) * samples
+        )
+        tolerance = 1e-6
+        torch._assert(
+            torch.all(
+                verts_scaled.ge(bounds[0] - tolerance)
+                & verts_scaled.le(bounds[1] + tolerance)
+            ),
+            "Samples are out of specified bounds",
+        )
+        return verts_scaled, cubes
 
     # changes mkofler: adapted function signature to match the kaolin package
     # function signature on kaolin package: