add loss functions (chamfer, beam-gap, normal alignment)

point2mesh/losses.py

@@ -0,0 +1,188 @@
+"""
+Loss functions for Point2Mesh optimisation.
+
+* Bidirectional Chamfer distance
+* Beam-gap loss  (pulls the mesh into narrow cavities)
+* Normal-alignment loss (cosine, handles unoriented normals)
+* Differentiable surface sampling
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from typing import Optional, Tuple
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Differentiable surface sampling
+# ──────────────────────────────────────────────────────────────────────
+def sample_surface(
+    verts: torch.Tensor,   # (N_v, 3)
+    faces: torch.Tensor,   # (N_f, 3) long
+    n_samples: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Uniformly sample points on the mesh surface (area-weighted).
+
+    Returns
+    -------
+    pts     : (n_samples, 3) — differentiable w.r.t. verts
+    face_ids: (n_samples,)   — which face each point was sampled from
+    """
+    v0 = verts[faces[:, 0]]
+    v1 = verts[faces[:, 1]]
+    v2 = verts[faces[:, 2]]
+
+    # Face areas (differentiable)
+    cross = torch.cross(v1 - v0, v2 - v0, dim=1)
+    areas = 0.5 * cross.norm(dim=1)  # (N_f,)
+    probs = areas / areas.sum().clamp(min=1e-12)
+
+    # Sample faces proportional to area
+    face_ids = torch.multinomial(probs, n_samples, replacement=True)
+
+    # Barycentric sampling: r1, r2 ~ U(0,1) with r1+r2 < 1
+    r1 = torch.rand(n_samples, device=verts.device)
+    r2 = torch.rand(n_samples, device=verts.device)
+    mask = (r1 + r2) >= 1.0
+    r1[mask] = 1.0 - r1[mask]
+    r2[mask] = 1.0 - r2[mask]
+
+    fv0 = verts[faces[face_ids, 0]]
+    fv1 = verts[faces[face_ids, 1]]
+    fv2 = verts[faces[face_ids, 2]]
+
+    pts = (1 - r1 - r2).unsqueeze(1) * fv0 + r1.unsqueeze(1) * fv1 + r2.unsqueeze(1) * fv2
+
+    return pts, face_ids
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Chamfer Distance (bidirectional, L2)
+# ──────────────────────────────────────────────────────────────────────
+def chamfer_loss(
+    X: torch.Tensor,      # (N, 3) target point cloud
+    Y: torch.Tensor,      # (M, 3) sampled mesh points
+    batch_size: int = 4096,
+) -> torch.Tensor:
+    """
+    Bidirectional Chamfer distance (mean of per-point min-dists).
+    Batched to avoid OOM on large point sets.
+    """
+    def _one_way(src, tgt):
+        """For each point in src, find min squared distance to tgt."""
+        total = torch.tensor(0.0, device=src.device)
+        n = src.shape[0]
+        for i in range(0, n, batch_size):
+            chunk = src[i : i + batch_size]
+            dists = torch.cdist(chunk, tgt)  # (chunk, M)
+            total = total + dists.min(dim=1).values.sum()
+        return total / max(n, 1)
+
+    return _one_way(X, Y) + _one_way(Y, X)
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Beam-Gap Loss
+# ──────────────────────────────────────────────────────────────────────
+@torch.no_grad()
+def _mutual_knn_mask(
+    Y: torch.Tensor, X: torch.Tensor, k: int = 3
+) -> torch.Tensor:
+    """
+    Boolean mask (M,): True where point y already has a 'good fit'
+    (mutual k-NN with X), so beam-gap should skip it.
+    """
+    # Y→X nearest k
+    d_yx = torch.cdist(Y, X)  # (M, N)
+    _, idx_yx = d_yx.topk(k, dim=1, largest=False)  # (M, k) indices into X
+
+    # X→Y nearest k
+    d_xy = torch.cdist(X, Y)  # (N, M)
+    _, idx_xy = d_xy.topk(k, dim=1, largest=False)  # (N, k) indices into Y
+
+    # For each y_i check: is y_i in the k-NN of any of its own k-NN x-targets?
+    good = torch.zeros(Y.shape[0], dtype=torch.bool, device=Y.device)
+    for yi in range(Y.shape[0]):
+        for xi in idx_yx[yi]:
+            if yi in idx_xy[xi.item()]:
+                good[yi] = True
+                break
+    return good
+
+
+def beam_gap_loss(
+    Y: torch.Tensor,           # (M, 3) sampled mesh points
+    normals: torch.Tensor,     # (M, 3) mesh face normals at each sample
+    X: torch.Tensor,           # (N, 3) target point cloud
+    epsilon: float = 0.5,
+    knn_k: int = 3,
+    max_samples: int = 2000,
+) -> torch.Tensor:
+    """
+    Beam-gap loss (Point2Mesh §3.3).
+
+    For each mesh sample ŷ, cast a beam along its normal, find the closest
+    target point inside an ε-cylinder, and penalise the gap.
+    Skips points that already have a mutual k-NN match.
+    """
+    M = Y.shape[0]
+    if M > max_samples:
+        # Subsample for efficiency
+        idx = torch.randperm(M, device=Y.device)[:max_samples]
+        Y = Y[idx]
+        normals = normals[idx]
+        M = max_samples
+
+    # Identify good-fit points to skip
+    good_mask = _mutual_knn_mask(Y, X, k=knn_k)
+
+    loss = torch.tensor(0.0, device=Y.device)
+    count = 0
+
+    # Vectorised cylinder test
+    # For each y: project (X - y) onto normal n
+    # along = dot(X - y, n);  perp = ||(X - y) - along * n||
+    diffs = X.unsqueeze(0) - Y.unsqueeze(1)  # (M, N, 3)
+    along = (diffs * normals.unsqueeze(1)).sum(dim=2)  # (M, N)
+    perp = (diffs - along.unsqueeze(2) * normals.unsqueeze(1)).norm(dim=2)  # (M, N)
+
+    # Inside cylinder: perp < epsilon AND along > 0 (ahead of surface)
+    in_cyl = (perp < epsilon) & (along > 0)
+
+    for i in range(M):
+        if good_mask[i]:
+            continue
+        cand = in_cyl[i]
+        if not cand.any():
+            continue
+        # Closest along the beam direction
+        dists_along = along[i].clone()
+        dists_along[~cand] = float("inf")
+        best_j = dists_along.argmin()
+        target = X[best_j]
+        loss = loss + (Y[i] - target).pow(2).sum()
+        count += 1
+
+    return loss / max(count, 1)
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Normal alignment loss
+# ──────────────────────────────────────────────────────────────────────
+def normal_loss(
+    Y: torch.Tensor,           # (M, 3) sampled mesh points
+    normals_mesh: torch.Tensor,  # (M, 3) mesh face normals at samples
+    X: torch.Tensor,           # (N, 3) target point cloud
+    normals_pc: torch.Tensor,  # (N, 3) target point-cloud normals
+) -> torch.Tensor:
+    """
+    Unoriented normal alignment: 1 − |n_mesh · n_pc|.
+    Pairs each mesh sample with its nearest point-cloud point.
+    """
+    dists = torch.cdist(Y, X)  # (M, N)
+    nn_idx = dists.argmin(dim=1)  # (M,)
+    nn_normals = normals_pc[nn_idx]
+    dot = (normals_mesh * nn_normals).sum(dim=1)
+    return (1 - dot.abs()).mean()