DeepSDFStruct 1.7.2__py3-none-any.whl

DeepSDFStruct/deep_sdf/nn_utils.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+
+activations = {
+    "relu": nn.ReLU(),
+    "tanh": nn.Tanh(),
+    "gelu": nn.GELU(),
+    "silu": nn.SiLU(),
+    "leaky_relu": nn.LeakyReLU(),
+    "elu": nn.ELU(),
+    "selu": nn.SELU(),
+    "sigmoid": nn.Sigmoid(),
+    "softplus": nn.Softplus(),
+}
+
+
+class ClampedL1Loss(torch.nn.Module):
+    def __init__(self, clamp_val=0.1):
+        super().__init__()
+        self.clamp_val = clamp_val
+        self.loss = torch.nn.L1Loss()
+
+    def forward(self, input, target):
+        input_clamped = input.clamp(-self.clamp_val, self.clamp_val)
+        target_clamped = target.clamp(-self.clamp_val, self.clamp_val)
+        return self.loss(input_clamped, target_clamped)
+
+
+class LeakyClampedL1Loss(torch.nn.Module):
+    """L1 loss that clamps values within [-clamp_val, clamp_val], but applies a
+    leaky linear continuation outside that range instead of hard clamping.
+
+    This behaves like a `clamp` in the center and a `LeakyReLU`-style slope
+    outside. Both input and target are transformed the same way before L1.
+    """
+
+    def __init__(self, clamp_val=0.1, leak=0.01):
+        super().__init__()
+        self.clamp_val = clamp_val
+        self.leak = leak
+        self.loss = torch.nn.L1Loss()
+
+    def _leaky_clamp(self, x: torch.Tensor) -> torch.Tensor:
+        cv = self.clamp_val
+        lk = self.leak
+        # If x > c -> c + l*(x-c); if x < -c -> -c + l*(x + c); else x
+        pos = torch.where(x > cv, cv + lk * (x - cv), x)
+        out = torch.where(pos < -cv, -cv + lk * (pos + cv), pos)
+        return out
+
+    def forward(self, input, target):
+        input_trans = self._leaky_clamp(input)
+        target_trans = self._leaky_clamp(target)
+        return self.loss(input_trans, target_trans)
+
+
+def get_loss_function(loss_fun_spec: str):
+    """Return a loss function instance for the given spec string."""
+    if loss_fun_spec == "clampedL1":
+        return ClampedL1Loss()
+    elif loss_fun_spec == "leakyClampedL1":
+        return LeakyClampedL1Loss()
+    elif loss_fun_spec == "L1":
+        return torch.nn.L1Loss()
+    elif loss_fun_spec == "MSE":
+        return torch.nn.MSELoss()
+    elif loss_fun_spec == "huber":
+        return torch.nn.HuberLoss()
+    else:
+        raise ValueError(f"Unsupported loss function: {loss_fun_spec}")

DeepSDFStruct/deep_sdf/plotting.py

@@ -0,0 +1,143 @@
+"""
+Training Visualization and Logging
+==================================
+
+This module provides utilities for visualizing DeepSDF training progress
+and generating plots of training metrics.
+
+Functions
+---------
+
+plot_logs
+    Plot training loss curves with smoothing and optional learning rate overlay.
+    Useful for monitoring training progress and diagnosing convergence issues.
+
+plot_reconstruction_loss
+    Visualize loss during shape reconstruction, showing how well the model
+    fits to target geometries.
+
+extract_paths
+    Helper for extracting file paths from nested data structures.
+
+running_mean
+    Compute running average for smoothing noisy loss curves.
+
+The module integrates with the workspace utilities to load training logs
+and generate publication-quality plots for analysis and presentation.
+"""
+
+import numpy as np
+import os
+import logging
+import torch
+import matplotlib.pyplot as plt
+
+import DeepSDFStruct.deep_sdf.workspace as ws
+
+logger = logging.getLogger(__name__)
+
+
+def extract_paths(data, current_path=""):
+    paths = []
+
+    if isinstance(data, dict):
+        for key, value in data.items():
+            new_path = f"{current_path}/{key}" if current_path else key
+            paths.extend(extract_paths(value, new_path))
+
+    elif isinstance(data, list):
+        for item in data:
+            paths.extend(extract_paths(item, current_path))
+
+    else:
+        paths.append(f"{current_path}/{data}")
+
+    return paths
+
+
+def running_mean(x, N):
+    cumsum = np.cumsum(np.insert(x, 0, 0))
+    return (cumsum[N:] - cumsum[:-N]) / float(N)
+
+
+def plot_logs(experiment_directory, show_lr=False, ax=None, filename=None):
+
+    logs = torch.load(os.path.join(experiment_directory, ws.logs_filename))
+
+    num_iters = len(logs["loss"])
+    iters_per_epoch = num_iters / logs["epoch"]
+
+    smoothed_loss_41 = running_mean(logs["loss"], 41)
+
+    show_plt = False
+
+    if show_lr:
+        if ax is None:
+            fig, ax = plt.subplots(2, 1)
+            fig.tight_layout()
+            show_plt = True
+    else:
+        if ax is None:
+            fig, ax = plt.subplots()
+            show_plt = True
+        ax = [ax]
+
+    ax[0].plot(
+        np.arange(num_iters) / iters_per_epoch,
+        logs["loss"],
+        "#82c6eb",
+        np.arange(20, num_iters - 20) / iters_per_epoch,
+        smoothed_loss_41,
+        "#2a9edd",
+    )
+    ax[0].set_yscale("log")
+
+    ax[0].set(xlabel="Epoch", ylabel="Loss")
+    ax[0].legend(["Loss", "Loss (Running Mean)", "Loss (Running Mean 41)"])
+
+    if show_lr:
+        combined_lrs = np.array(logs["learning_rate"])
+        ax[1].plot(
+            np.arange(combined_lrs.shape[0]),
+            combined_lrs[:, 0],
+            np.arange(combined_lrs.shape[0]),
+            combined_lrs[:, 1],
+        )
+        ax[1].set(xlabel="Epoch", ylabel="Learning Rate")
+        ax[1].legend(["Decoder", "Latent Vector"])
+
+    for axis in ax:
+        axis.grid()
+    if filename is not None:
+        plt.savefig(filename, bbox_inches="tight")
+        plt.close()
+    elif show_plt:
+        plt.show()
+
+
+def plot_reconstruction_loss(loss_history, iters_per_epoch, filename=None):
+
+    losses = np.array(loss_history)
+    num_iters = len(losses)
+
+    latest_epoch = num_iters / iters_per_epoch
+    logger.info("latest epoch is {}".format(latest_epoch))
+    logger.info("{} iters per epoch".format(iters_per_epoch))
+
+    smoothed_loss_41 = running_mean(losses, 41)
+
+    fig, ax = plt.subplots()
+
+    ax.plot(
+        np.arange(num_iters) / iters_per_epoch,
+        losses,
+        "#82c6eb",
+        np.arange(20, num_iters - 20) / iters_per_epoch,
+        smoothed_loss_41,
+        "#2a9edd",
+    )
+    ax.set_yscale("log")
+    ax.set(xlabel="Epoch", ylabel="Loss")
+    ax.legend(["Loss", "Loss (Running Mean 41)"])
+    ax.grid()
+    plt.savefig(filename, bbox_inches="tight")

DeepSDFStruct/deep_sdf/reconstruction.py

@@ -0,0 +1,165 @@
+"""
+Shape Reconstruction and Latent Code Optimization
+=================================================
+
+This module provides methods for reconstructing shapes from SDF samples
+by optimizing latent codes. This is a key capability of DeepSDF models,
+enabling fitting of new geometries not seen during training.
+
+Key Functions
+-------------
+
+reconstruct_from_samples
+    Optimize a latent code to match given SDF samples. Uses gradient
+    descent to find the latent representation that best reproduces
+    the target geometry.
+
+The reconstruction process:
+1. Start with initial latent code (often zeros or random)
+2. Forward pass through decoder to get SDF predictions
+3. Compute loss between predictions and target samples
+4. Backpropagate and update latent code
+5. Repeat until convergence
+
+This enables:
+- Fitting trained models to new geometries
+- Shape interpolation by blending latent codes
+- Shape editing by modifying latent codes
+- Compression by storing only latent codes
+
+Examples
+--------
+Reconstruct a shape from samples::
+
+    from DeepSDFStruct.deep_sdf.reconstruction import reconstruct_from_samples
+    from DeepSDFStruct.sampling import sample_sdf_near_surface
+
+    # Sample target geometry
+    target_samples = sample_sdf_near_surface(target_sdf, n_samples=100000)
+
+    # Reconstruct using trained model
+    reconstructed_sdf = reconstruct_from_samples(
+        model,
+        target_samples,
+        num_iterations=1000,
+        lr=5e-4
+    )
+"""
+
+from DeepSDFStruct.SDF import SDFBase
+from DeepSDFStruct.deep_sdf.nn_utils import ClampedL1Loss
+from torch.utils.data import TensorDataset, DataLoader
+from DeepSDFStruct.sampling import SampledSDF
+from DeepSDFStruct.deep_sdf.plotting import plot_reconstruction_loss
+from DeepSDFStruct.torch_spline import TorchSpline, TorchScaling
+from tqdm import trange
+import torch
+import torch.nn as nn
+
+
+def reconstruct_from_samples(
+    sdf: SDFBase,
+    sdfSample: SampledSDF,
+    num_iterations=1000,
+    lr=5e-4,
+    loss_fn="ClampedL1",
+    batch_size=512,
+    drop_last=True,
+    use_tanh_on_gt=False,
+    loss_plot_path=None,
+    optimizer_name="adam",
+    deformation_function=None | TorchSpline | TorchScaling,
+):
+    if optimizer_name == "adam":
+        optimizer = torch.optim.Adam(sdf.parameters(), lr=lr)
+    elif optimizer_name == "lbfgs":
+        optimizer = torch.optim.LBFGS(
+            sdf.parameters(),
+            lr=lr,
+            max_iter=20,  # inner Newton iterations
+            history_size=100,  # curvature memory
+            line_search_fn="strong_wolfe",
+        )
+    else:
+        raise NotImplementedError(f"Optimizer {optimizer_name} not available.")
+
+    verts_min = sdfSample.samples.min(axis=0)
+    verts_max = sdfSample.samples.max(axis=0)
+
+    print("Min/Max in PHYSICAL space:\n")
+    for name, mn, mx in zip(["x", "y", "z"], verts_min.values, verts_max.values):
+        print(f"{name}: min={mn:.6f}, max={mx:.6f}")
+
+    if deformation_function is not None:
+        queries_parameter_space = deformation_function.inverse_target_points(
+            sdfSample.samples
+        ).detach()
+    else:
+        queries_parameter_space = sdfSample.samples.detach()
+
+    queries_min = queries_parameter_space.min(dim=0).values
+    queries_max = queries_parameter_space.max(dim=0).values
+
+    print("\nMin/Max in QUERY space:\n")
+    for name, mn, mx in zip(
+        ["x", "y", "z"], queries_min.tolist(), queries_max.tolist()
+    ):
+        print(f"{name}: min={mn:.6f}, max={mx:.6f}")
+
+    gt_dist = sdfSample.distances
+    if use_tanh_on_gt:
+        gt_dist = torch.tanh(gt_dist)
+
+    pbar = trange(num_iterations, desc="Reconstructing SDF from mesh", leave=True)
+
+    if loss_fn == "L1":
+        Loss = torch.nn.L1Loss()
+    elif loss_fn == "ClampedL1":
+        Loss = ClampedL1Loss(clamp_val=0.1)
+    elif loss_fn == "MSE":
+        Loss = torch.nn.MSELoss()
+    else:
+        raise NotImplementedError(f"Loss function {loss_fn} not available.")
+
+    dataset = TensorDataset(queries_parameter_space, gt_dist)
+    if drop_last and (batch_size > len(dataset)):
+        print(
+            "Warning: drop_last was set to true, "
+            f"but batch size ({batch_size}) is larger "
+            f"than the size of the dataset ({len(dataset)})."
+            " setting drop_last=False"
+        )
+        drop_last = False
+    dataloader = DataLoader(
+        dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last
+    )
+
+    loss_history = []
+    for e in pbar:
+        for querie_batch, gt_batch in dataloader:
+
+            def closure() -> torch.Tensor:
+                optimizer.zero_grad()
+                pred_dist = sdf(querie_batch)
+                loss = Loss(pred_dist, gt_batch)
+                loss.backward()
+                return loss
+
+            if optimizer_name == "adam":
+                loss = closure()
+                optimizer.step()
+            elif optimizer_name == "lbfgs":
+                loss = optimizer.step(closure)
+            loss_num = loss.detach().item()
+            pbar.set_postfix({"loss": f"{loss_num:.5f}"})
+            loss_history.append(loss_num)
+
+    if loss_plot_path is not None:
+        plot_reconstruction_loss(
+            loss_history, iters_per_epoch=len(dataloader), filename=loss_plot_path
+        )
+
+    params = list(sdf.parameters())
+    print(params)
+
+    return params