Update landmarkdiff/synthetic/tps_warp.py to v0.3.2

landmarkdiff/synthetic/tps_warp.py

@@ -1,7 +1,8 @@
-"""TPS warping for synthetic pair generation.
+"""Thin-Plate Spline warping for synthetic training pair generation.
 
-Only warps deformable tissue - rigid structures (teeth, sclera) get
-rigid translation instead. Prevents "rubber teeth" from naive TPS.
+Applies TPS warp ONLY to deformable tissue regions. Rigid structures
+(teeth, sclera) are rigidly translated, not warped. This prevents
+the "rubber teeth" artifact from naive TPS.
 """
 
 from __future__ import annotations
@@ -14,7 +15,15 @@ def compute_tps_transform(
     src_pts: np.ndarray,
     dst_pts: np.ndarray,
 ) -> cv2.ThinPlateSplineShapeTransformer:
-    """Fit a TPS transform from src to dst points."""
+    """Compute a TPS transform from source to destination point pairs.
+
+    Args:
+        src_pts: (N, 2) source control points.
+        dst_pts: (N, 2) destination control points.
+
+    Returns:
+        Fitted TPS transformer.
+    """
     src = src_pts.reshape(1, -1, 2).astype(np.float32)
     dst = dst_pts.reshape(1, -1, 2).astype(np.float32)
     matches = [cv2.DMatch(i, i, 0) for i in range(len(src_pts))]
@@ -30,7 +39,12 @@ def _subsample_control_points(
     max_points: int = 80,
     anchor_stride: int = 8,
 ) -> tuple[np.ndarray, np.ndarray]:
-    """Keep all displaced points + sparse anchors. ~80 pts instead of 478, ~30x faster."""
+    """Subsample control points for faster TPS: all displaced + sparse anchors.
+
+    With 478 MediaPipe landmarks, full TPS requires solving a 481x481 system
+    and evaluating 478 RBFs at each pixel — very slow. Subsampling to ~80
+    points gives nearly identical results ~30x faster.
+    """
     displacements = np.linalg.norm(dst - src, axis=1)
     displaced_mask = displacements > 0.5  # moved by > 0.5px
     displaced_idx = np.where(displaced_mask)[0]
@@ -61,7 +75,18 @@ def warp_image_tps(
     dst_landmarks: np.ndarray,
     rigid_mask: np.ndarray | None = None,
 ) -> np.ndarray:
-    """Apply TPS warp to an image with optional rigid region preservation."""
+    """Apply TPS warp to an image with optional rigid region preservation.
+
+    Args:
+        image: BGR input image.
+        src_landmarks: (N, 2) original landmark pixel coords.
+        dst_landmarks: (N, 2) target landmark pixel coords.
+        rigid_mask: Optional binary mask of rigid regions (teeth, sclera).
+                    These regions are rigidly translated, not TPS-warped.
+
+    Returns:
+        Warped image.
+    """
     h, w = image.shape[:2]
 
     src_pts = src_landmarks.astype(np.float32)
@@ -87,8 +112,10 @@ def warp_image_tps(
         rigid_translation = _compute_rigid_translation(src_pts, dst_pts, rigid_mask, w, h)
         rigid_warped = _apply_rigid_translation(image, rigid_translation)
 
+        # Translate the mask to match the rigidly-shifted content
+        translated_mask = _apply_rigid_translation(rigid_mask, rigid_translation)
         # Composite: use rigid warp in rigid regions, TPS elsewhere
-        mask_f = rigid_mask.astype(np.float32)
+        mask_f = translated_mask.astype(np.float32)
         if len(mask_f.shape) == 2:
             mask_f = np.stack([mask_f] * 3, axis=-1)
         mask_f = mask_f / 255.0 if mask_f.max() > 1 else mask_f
@@ -103,7 +130,10 @@ def _compute_tps_map(
     width: int,
     height: int,
 ) -> tuple[np.ndarray, np.ndarray]:
-    """Build remap arrays from TPS control points via RBF interpolation."""
+    """Compute pixel displacement maps from TPS control points.
+
+    Uses RBF interpolation of control point displacements.
+    """
     # Displacement at control points
     dx = dst[:, 0] - src[:, 0]
     dy = dst[:, 1] - src[:, 1]
@@ -151,12 +181,16 @@ def _solve_tps_weights(
     control_pts: np.ndarray,
     values: np.ndarray,
 ) -> np.ndarray:
-    """Solve TPS system -> weight vector [w1..wn, a0, a1, a2]."""
+    """Solve for TPS weights given control points and target values.
+
+    Returns weight vector [w1..wn, a0, a1, a2] for n control points
+    plus affine terms.
+    """
     n = len(control_pts)
 
     # Build kernel matrix K (vectorized)
     diff = control_pts[:, np.newaxis, :] - control_pts[np.newaxis, :, :]  # (n, n, 2)
-    r_mat = np.sqrt((diff**2).sum(axis=2))  # (n, n)
+    r_mat = np.sqrt((diff ** 2).sum(axis=2))  # (n, n)
     K = np.zeros((n, n))
     nz = r_mat > 0
     K[nz] = r_mat[nz] ** 2 * np.log(r_mat[nz])
@@ -205,7 +239,7 @@ def _evaluate_tps(
         # Compute all distances at once: (M, n)
         dx = batch[:, 0:1] - control_pts[:, 0]  # (M, n) via broadcasting
         dy = batch[:, 1:2] - control_pts[:, 1]  # (M, n)
-        r = np.sqrt(dx**2 + dy**2)
+        r = np.sqrt(dx ** 2 + dy ** 2)
 
         # TPS kernel: r^2 * log(r), with r=0 -> 0
         kernel = np.zeros_like(r)
@@ -230,8 +264,9 @@ def _compute_rigid_translation(
     inside = []
     for i, (x, y) in enumerate(src):
         ix, iy = int(x), int(y)
-        if 0 <= ix < width and 0 <= iy < height and mask[iy, ix] > 0:
-            inside.append(i)
+        if 0 <= ix < width and 0 <= iy < height:
+            if mask[iy, ix] > 0:
+                inside.append(i)
 
     if not inside:
         return np.array([0.0, 0.0])
@@ -257,7 +292,17 @@ def generate_random_warp(
     max_displacement: float = 15.0,
     rng: np.random.Generator | None = None,
 ) -> np.ndarray:
-    """Generate randomly warped landmarks for synthetic data."""
+    """Generate randomly warped landmarks for synthetic data.
+
+    Args:
+        landmarks: (N, 2) pixel coordinates.
+        procedure_indices: Which landmarks to warp.
+        max_displacement: Max pixel displacement.
+        rng: Random number generator.
+
+    Returns:
+        New landmark array with random deformations.
+    """
     rng = rng or np.random.default_rng()
     result = landmarks.copy()