Upload losses.py

losses.py

@@ -0,0 +1,154 @@
+"""
+Loss functions for SDF and VG stages.
+All pure PyTorch, no compiled extensions needed.
+"""
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+def gaussian_kernel(points, queries):
+    """
+    points:  (n, k, 3) or (n, 3)
+    queries: (n, 3)
+    Returns: (n, k) or (n,) Gaussian weights.
+    """
+    if points.dim() == 2:
+        points = points.unsqueeze(1)  # (n, 1, 3)
+    pts_dist = torch.linalg.norm((points - queries.unsqueeze(1)), ord=2, dim=-1)  # (n, k)
+    h = pts_dist.mean(dim=-1, keepdim=True) + 1e-8
+    dist_exp = torch.exp(-pts_dist ** 2 / h ** 2)
+    gaussian_weight = dist_exp / (dist_exp.sum(dim=-1, keepdim=True) + 1e-8)
+    return gaussian_weight
+
+
+def pull_knn_loss(points, samples_moved, samples):
+    """KNN-weighted pull loss."""
+    # points: (n, k, 3); samples_moved: (n, 3); samples: (n, 3)
+    if points.dim() == 2:
+        points = points.unsqueeze(1)
+    loss_pull = torch.linalg.norm((points - samples_moved.unsqueeze(1)), ord=2, dim=-1)
+    g_weight = gaussian_kernel(points, samples).detach()
+    loss_pull = (loss_pull * g_weight).sum(dim=-1).mean()
+    return loss_pull
+
+
+def grad_consis_knn_loss(gradient_points_norm, gradients_samples_norm, points, samples):
+    if points.dim() == 2:
+        points = points.unsqueeze(1)
+    loss_consis = (1 - F.cosine_similarity(gradient_points_norm, gradients_samples_norm.unsqueeze(1), dim=-1))
+    g_weight = gaussian_kernel(points, samples).detach()
+    loss_consis = (loss_consis * g_weight).sum(dim=-1).mean()
+    return loss_consis
+
+
+def eikonal_loss(gradients_samples):
+    loss_eikonal = ((gradients_samples.norm(2, dim=-1) - 1).square()).mean()
+    return loss_eikonal
+
+
+def div_loss(samples, gradients_samples):
+    """Divergence of gradient field."""
+    def gradient(inputs, outputs):
+        inputs.requires_grad_(True)
+        d_points = torch.ones_like(outputs, requires_grad=False, device=outputs.device)
+        points_grad = torch.autograd.grad(
+            outputs=outputs,
+            inputs=inputs,
+            grad_outputs=d_points,
+            create_graph=True,
+            retain_graph=True,
+            only_inputs=True)[0]
+        return points_grad
+
+    div_dx = gradient(samples, gradients_samples[:, 0])
+    div_dy = gradient(samples, gradients_samples[:, 1])
+    div_dz = gradient(samples, gradients_samples[:, 2])
+    divergence = div_dx[:, 0] + div_dy[:, 1] + div_dz[:, 2]
+    loss_div = (torch.clamp(torch.square(divergence), 0.1, 50)).mean()
+    return loss_div
+
+
+def sdf_loss(sample_sdf):
+    return torch.abs(sample_sdf).mean()
+
+
+def cal_curvature_with_normal(pts, normals, knn=16):
+    """
+    Estimate curvature from normal variation over k-NN.
+    pts: (n, 3)    normals: (n, 3)
+    Returns: (n, 1) curvature in [0, 1].
+    """
+    from scipy.spatial import KDTree
+    pts_np = pts.detach().cpu().numpy()
+    normals_np = normals.detach().cpu().numpy()
+    tree = KDTree(pts_np)
+    _, idx = tree.query(pts_np, k=min(knn + 1, len(pts_np)))
+    idx = torch.from_numpy(idx[:, 1:]).long().to(pts.device)  # (n, k)
+    neigh_pts = pts[idx]            # (n, k, 3)
+    neigh_normals = normals[idx]    # (n, k, 3)
+    neigh_curvature = 1.0 - F.cosine_similarity(normals.unsqueeze(1), neigh_normals, dim=-1)  # (n, k)
+    g_weight = gaussian_kernel(neigh_pts, pts).detach()
+    pts_curvature_ave = (neigh_curvature * g_weight).sum(dim=-1, keepdim=True)  # (n, 1)
+
+    pts_curvature_ave = torch.sigmoid(pts_curvature_ave - torch.mean(pts_curvature_ave))
+    pts_curvature_ave = (pts_curvature_ave - torch.min(pts_curvature_ave) + 1e-6) / (
+        torch.max(pts_curvature_ave) - torch.min(pts_curvature_ave) + 1e-6)
+    return pts_curvature_ave
+
+
+def cal_nc_loss(surface_normals, sample_normals, sur_neigh_idx):
+    """Normal consistency loss."""
+    neigh_normals = sample_normals[sur_neigh_idx]
+    nc_loss = 1.0 - F.cosine_similarity(surface_normals, neigh_normals, dim=-1).mean()
+    return nc_loss
+
+
+def cal_chamfer_loss(sur_pts, sample_pts, curvature_surface, loss_w, nearest_clamp):
+    """
+    Bidirectional Chamfer with curvature weighting + repulsion.
+    loss_w: [w_curv, w_plain, w_reverse, w_repulsion]
+    """
+    from scipy.spatial import KDTree
+
+    # ---- nearest from surface to samples ----
+    sample_np = sample_pts.detach().cpu().numpy()
+    sur_np = sur_pts.detach().cpu().numpy()
+    tree = KDTree(sample_np)
+    _, sur_neigh_idx = tree.query(sur_np, k=1)
+    sur_neigh_idx = torch.from_numpy(sur_neigh_idx).long().to(sur_pts.device)
+    sur_neigh_pts = sample_pts[sur_neigh_idx]
+    dist_1 = torch.linalg.norm((sur_pts - sur_neigh_pts), ord=2, dim=-1).unsqueeze(-1) ** 2
+    weight_sur_curvature = curvature_surface
+    loss_part_1 = loss_w[0] * (dist_1 * weight_sur_curvature).mean() + loss_w[1] * dist_1.mean()
+
+    # ---- nearest from samples to surface ----
+    tree2 = KDTree(sur_np)
+    _, sample_neigh_idx = tree2.query(sample_np, k=1)
+    sample_neigh_idx = torch.from_numpy(sample_neigh_idx).long().to(sample_pts.device)
+    sample_neigh_pts = sur_pts[sample_neigh_idx]
+    dist_2 = torch.linalg.norm((sample_pts - sample_neigh_pts), ord=2, dim=-1) ** 2
+    loss_part_2 = loss_w[2] * dist_2.mean()
+
+    # ---- repulsion (self-nearest-neighbor distance) ----
+    tree3 = KDTree(sample_np)
+    _, sample_neigh_idx_self = tree3.query(sample_np, k=2)
+    sample_neigh_idx_self = torch.from_numpy(sample_neigh_idx_self[:, 1]).long().to(sample_pts.device)
+    sample_neigh_pts_self = sample_pts[sample_neigh_idx_self]
+    relative_dist = sample_neigh_pts_self - sample_pts
+    norm_dist = torch.linalg.norm(relative_dist, ord=2, dim=-1) ** 2
+    norm_dist_ = torch.clamp(norm_dist, max=nearest_clamp * norm_dist.mean())
+    loss_part_3 = -norm_dist_.mean() * loss_w[3]
+
+    chamfer_loss = loss_part_1 + loss_part_2 + loss_part_3
+    return chamfer_loss, sur_neigh_idx, sample_neigh_idx
+
+
+def cal_vg_loss(sur_pts, sur_normals, curvature_ave, sample_pts, sample_grad, loss_w, nearest_clamp):
+    """Total VG stage loss."""
+    sample_normal = F.normalize(sample_grad.detach(), dim=-1)
+    chamfer_loss, sur_neigh_idx, _ = cal_chamfer_loss(
+        sur_pts, sample_pts, curvature_ave, loss_w[:4], nearest_clamp)
+    normal_loss = cal_nc_loss(sur_normals, sample_normal, sur_neigh_idx)
+    loss = chamfer_loss + loss_w[4] * normal_loss
+    return loss