landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/synthetic/pair_generator.py

@@ -0,0 +1,208 @@
+"""Synthetic training pair generator.
+
+Creates (input, conditioning, mask, target) tuples for ControlNet fine-tuning.
+Pipeline: FFHQ image -> extract landmarks -> random FFD manipulation ->
+generate conditioning + mask -> apply clinical augmentation to input.
+
+Augmentations are applied to INPUT only, never to target (ground truth).
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from landmarkdiff.conditioning import generate_conditioning
+from landmarkdiff.landmarks import extract_landmarks, render_landmark_image
+from landmarkdiff.manipulation import (
+    PROCEDURE_LANDMARKS,
+    apply_procedure_preset,
+)
+from landmarkdiff.masking import generate_surgical_mask
+from landmarkdiff.synthetic.augmentation import apply_clinical_augmentation
+from landmarkdiff.synthetic.tps_warp import warp_image_tps
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class TrainingPair:
+    """A single training sample for ControlNet fine-tuning."""
+
+    input_image: np.ndarray  # augmented input (512x512 BGR)
+    target_image: np.ndarray  # clean target (512x512 BGR) -- TPS-warped original
+    conditioning: np.ndarray  # landmark rendering (512x512 BGR)
+    canny: np.ndarray  # canny edge map (512x512 grayscale)
+    mask: np.ndarray  # feathered surgical mask (512x512 float32)
+    procedure: str
+    intensity: float
+
+
+PROCEDURES = list(PROCEDURE_LANDMARKS.keys())
+
+
+def generate_pair(
+    image: np.ndarray,
+    procedure: str | None = None,
+    intensity: float | None = None,
+    target_size: int = 512,
+    rng: np.random.Generator | None = None,
+) -> TrainingPair | None:
+    """Generate a single training pair from a face image.
+
+    Args:
+        image: BGR input image (any size).
+        procedure: Procedure type (random if None).
+        intensity: Manipulation intensity 0-100 (random 30-90 if None).
+        target_size: Output resolution.
+        rng: Random number generator.
+
+    Returns:
+        TrainingPair or None if face detection fails.
+    """
+    rng = rng or np.random.default_rng()
+
+    # Resize to target
+    resized = cv2.resize(image, (target_size, target_size))
+
+    # Extract landmarks
+    face = extract_landmarks(resized)
+    if face is None:
+        return None
+
+    # Random procedure and intensity if not specified
+    if procedure is None:
+        procedure = rng.choice(PROCEDURES)
+    if intensity is None:
+        intensity = float(rng.uniform(30, 90))
+
+    # Manipulate landmarks
+    manipulated = apply_procedure_preset(face, procedure, intensity, target_size)
+
+    # Generate conditioning from manipulated landmarks
+    landmark_img = render_landmark_image(manipulated, target_size, target_size)
+    _, canny, _ = generate_conditioning(manipulated, target_size, target_size)
+
+    # Generate mask
+    mask = generate_surgical_mask(face, procedure, target_size, target_size)
+
+    # Generate target: TPS warp the original image to match manipulated landmarks
+    src_px = face.pixel_coords
+    dst_px = manipulated.pixel_coords
+    target = warp_image_tps(resized, src_px, dst_px)
+
+    # Apply clinical augmentation to INPUT only (never target)
+    augmented_input = apply_clinical_augmentation(resized, rng=rng)
+
+    return TrainingPair(
+        input_image=augmented_input,
+        target_image=target,
+        conditioning=landmark_img,
+        canny=canny,
+        mask=mask,
+        procedure=procedure,
+        intensity=intensity,
+    )
+
+
+def generate_pairs_from_directory(
+    image_dir: str | Path,
+    num_pairs: int = 1000,
+    target_size: int = 512,
+    seed: int = 42,
+    quality_check: bool = True,
+    min_quality: float = 45.0,
+) -> Iterator[TrainingPair]:
+    """Generate training pairs from a directory of face images.
+
+    Args:
+        image_dir: Directory containing face images.
+        num_pairs: Total number of pairs to generate.
+        target_size: Output resolution.
+        seed: Random seed.
+        quality_check: Run face verifier quality check on source images.
+        min_quality: Minimum quality score to use image (0-100).
+
+    Yields:
+        TrainingPair instances.
+    """
+    rng = np.random.default_rng(seed)
+    image_dir = Path(image_dir)
+
+    extensions = {".jpg", ".jpeg", ".png", ".webp"}
+    image_files = sorted(f for f in image_dir.iterdir() if f.suffix.lower() in extensions)
+
+    if not image_files:
+        raise FileNotFoundError(f"No images found in {image_dir}")
+
+    # Optional quality pre-filter
+    _quality_cache: dict[str, float] = {}
+    quality_rejects = 0
+
+    generated = 0
+    consecutive_failures = 0
+    idx = 0
+    while generated < num_pairs:
+        # Cycle through images
+        img_path = image_files[idx % len(image_files)]
+        idx += 1
+        image = cv2.imread(str(img_path))
+        if image is None:
+            consecutive_failures += 1
+            if consecutive_failures > len(image_files):
+                logger.warning("%d consecutive failures, stopping early", consecutive_failures)
+                break
+            continue
+
+        # Quality gate: reject low-quality source images before pair generation
+        if quality_check:
+            cache_key = str(img_path)
+            if cache_key not in _quality_cache:
+                try:
+                    from landmarkdiff.face_verifier import analyze_distortions
+
+                    resized = cv2.resize(image, (target_size, target_size))
+                    report = analyze_distortions(resized)
+                    _quality_cache[cache_key] = report.quality_score
+                except Exception:
+                    _quality_cache[cache_key] = 100.0  # Can't check -- allow through
+
+            if _quality_cache[cache_key] < min_quality:
+                quality_rejects += 1
+                if quality_rejects % 100 == 0:
+                    logger.info("  Quality filter: %d images rejected so far", quality_rejects)
+                consecutive_failures += 1
+                if consecutive_failures > len(image_files):
+                    break
+                continue
+
+        pair = generate_pair(image, target_size=target_size, rng=rng)
+        if pair is not None:
+            yield pair
+            generated += 1
+            consecutive_failures = 0
+        else:
+            consecutive_failures += 1
+            if consecutive_failures > len(image_files):
+                logger.warning("%d consecutive failures, stopping early", consecutive_failures)
+                break
+
+    if quality_rejects > 0:
+        logger.info("Quality filter: rejected %d low-quality source images", quality_rejects)
+
+
+def save_pair(pair: TrainingPair, output_dir: Path, index: int) -> None:
+    """Save a training pair to disk."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    prefix = f"{index:06d}"
+
+    cv2.imwrite(str(output_dir / f"{prefix}_input.png"), pair.input_image)
+    cv2.imwrite(str(output_dir / f"{prefix}_target.png"), pair.target_image)
+    cv2.imwrite(str(output_dir / f"{prefix}_conditioning.png"), pair.conditioning)
+    cv2.imwrite(str(output_dir / f"{prefix}_canny.png"), pair.canny)
+    cv2.imwrite(str(output_dir / f"{prefix}_mask.png"), (pair.mask * 255).astype(np.uint8))

landmarkdiff/synthetic/tps_warp.py

@@ -0,0 +1,273 @@
+"""TPS warping for synthetic pair generation.
+
+Only warps deformable tissue - rigid structures (teeth, sclera) get
+rigid translation instead. Prevents "rubber teeth" from naive TPS.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+
+def compute_tps_transform(
+    src_pts: np.ndarray,
+    dst_pts: np.ndarray,
+) -> cv2.ThinPlateSplineShapeTransformer:
+    """Fit a TPS transform from src to dst points."""
+    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))]
+
+    tps = cv2.createThinPlateSplineShapeTransformer()
+    tps.estimateTransformation(dst, src, matches)
+    return tps
+
+
+def _subsample_control_points(
+    src: np.ndarray,
+    dst: np.ndarray,
+    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."""
+    displacements = np.linalg.norm(dst - src, axis=1)
+    displaced_mask = displacements > 0.5  # moved by > 0.5px
+    displaced_idx = np.where(displaced_mask)[0]
+
+    # Add sparse anchors from non-displaced landmarks
+    non_displaced_idx = np.where(~displaced_mask)[0]
+    anchor_idx = non_displaced_idx[::anchor_stride]
+
+    selected = np.concatenate([displaced_idx, anchor_idx])
+
+    # If still too many, subsample anchors more aggressively
+    if len(selected) > max_points:
+        n_anchors = max_points - len(displaced_idx)
+        if n_anchors > 0:
+            step = max(1, len(non_displaced_idx) // n_anchors)
+            anchor_idx = non_displaced_idx[::step][:n_anchors]
+            selected = np.concatenate([displaced_idx, anchor_idx])
+        else:
+            selected = displaced_idx[:max_points]
+
+    selected = np.unique(selected)
+    return src[selected], dst[selected]
+
+
+def warp_image_tps(
+    image: np.ndarray,
+    src_landmarks: np.ndarray,
+    dst_landmarks: np.ndarray,
+    rigid_mask: np.ndarray | None = None,
+) -> np.ndarray:
+    """Apply TPS warp to an image with optional rigid region preservation."""
+    h, w = image.shape[:2]
+
+    src_pts = src_landmarks.astype(np.float32)
+    dst_pts = dst_landmarks.astype(np.float32)
+
+    # Subsample control points for speed (478 -> ~80)
+    src_sub, dst_sub = _subsample_control_points(src_pts, dst_pts)
+
+    # Compute TPS coefficients on subsampled points
+    map_x, map_y = _compute_tps_map(src_sub, dst_sub, w, h)
+
+    # Warp the image
+    warped = cv2.remap(
+        image,
+        map_x.astype(np.float32),
+        map_y.astype(np.float32),
+        interpolation=cv2.INTER_LINEAR,
+        borderMode=cv2.BORDER_REFLECT_101,
+    )
+
+    if rigid_mask is not None:
+        # For rigid regions, compute mean translation and apply rigidly
+        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 = 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
+        warped = (rigid_warped * mask_f + warped * (1 - mask_f)).astype(np.uint8)
+
+    return warped
+
+
+def _compute_tps_map(
+    src: np.ndarray,
+    dst: np.ndarray,
+    width: int,
+    height: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build remap arrays from TPS control points via RBF interpolation."""
+    # Displacement at control points
+    dx = dst[:, 0] - src[:, 0]
+    dy = dst[:, 1] - src[:, 1]
+
+    # Create grid
+    grid_x, grid_y = np.meshgrid(np.arange(width), np.arange(height))
+    grid_x = grid_x.astype(np.float64)
+    grid_y = grid_y.astype(np.float64)
+
+    # RBF interpolation using TPS kernel: r^2 * log(r)
+    map_x = grid_x.copy()
+    map_y = grid_y.copy()
+
+    n = len(src)
+    if n == 0:
+        return map_x, map_y
+
+    # Solve TPS system for x and y displacements
+    weights_x = _solve_tps_weights(src, dx)
+    weights_y = _solve_tps_weights(src, dy)
+
+    # Evaluate on grid (vectorized for speed)
+    flat_x = grid_x.ravel()
+    flat_y = grid_y.ravel()
+    pts = np.stack([flat_x, flat_y], axis=1)
+
+    disp_x = _evaluate_tps(pts, src, weights_x)
+    disp_y = _evaluate_tps(pts, src, weights_y)
+
+    map_x = (flat_x - disp_x).reshape(height, width)
+    map_y = (flat_y - disp_y).reshape(height, width)
+
+    return map_x, map_y
+
+
+def _tps_kernel(r: np.ndarray) -> np.ndarray:
+    """TPS radial basis function: r^2 * log(r), with r=0 -> 0."""
+    result = np.zeros_like(r)
+    mask = r > 0
+    result[mask] = r[mask] ** 2 * np.log(r[mask])
+    return result
+
+
+def _solve_tps_weights(
+    control_pts: np.ndarray,
+    values: np.ndarray,
+) -> np.ndarray:
+    """Solve TPS system -> weight vector [w1..wn, a0, a1, a2]."""
+    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)
+    K = np.zeros((n, n))
+    nz = r_mat > 0
+    K[nz] = r_mat[nz] ** 2 * np.log(r_mat[nz])
+
+    # Build system matrix [K P; P^T 0]
+    P = np.hstack([np.ones((n, 1)), control_pts])  # (n, 3)
+
+    L = np.zeros((n + 3, n + 3))
+    L[:n, :n] = K
+    L[:n, n:] = P
+    L[n:, :n] = P.T
+
+    # Regularization for numerical stability
+    L[:n, :n] += np.eye(n) * 1e-6
+
+    rhs = np.zeros(n + 3)
+    rhs[:n] = values
+
+    try:
+        weights = np.linalg.solve(L, rhs)
+    except np.linalg.LinAlgError:
+        weights = np.linalg.lstsq(L, rhs, rcond=None)[0]
+
+    return weights
+
+
+def _evaluate_tps(
+    points: np.ndarray,
+    control_pts: np.ndarray,
+    weights: np.ndarray,
+) -> np.ndarray:
+    """Evaluate TPS at arbitrary points (vectorized)."""
+    n = len(control_pts)
+    w = weights[:n]
+    a = weights[n:]  # affine: a0 + a1*x + a2*y
+
+    # Affine component
+    result = a[0] + a[1] * points[:, 0] + a[2] * points[:, 1]
+
+    # Vectorized RBF evaluation in batches to limit memory
+    batch_size = 50000
+    for start in range(0, len(points), batch_size):
+        end = min(start + batch_size, len(points))
+        batch = points[start:end]  # (M, 2)
+
+        # 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)
+
+        # TPS kernel: r^2 * log(r), with r=0 -> 0
+        kernel = np.zeros_like(r)
+        mask = r > 0
+        kernel[mask] = r[mask] ** 2 * np.log(r[mask])
+
+        # Weighted sum across all control points
+        result[start:end] += kernel @ w
+
+    return result
+
+
+def _compute_rigid_translation(
+    src: np.ndarray,
+    dst: np.ndarray,
+    mask: np.ndarray,
+    width: int,
+    height: int,
+) -> np.ndarray:
+    """Compute mean translation for rigid regions."""
+    # Find control points inside rigid mask
+    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 not inside:
+        return np.array([0.0, 0.0])
+
+    dx = np.mean(dst[inside, 0] - src[inside, 0])
+    dy = np.mean(dst[inside, 1] - src[inside, 1])
+    return np.array([dx, dy])
+
+
+def _apply_rigid_translation(
+    image: np.ndarray,
+    translation: np.ndarray,
+) -> np.ndarray:
+    """Apply rigid translation to an image."""
+    h, w = image.shape[:2]
+    M = np.float32([[1, 0, translation[0]], [0, 1, translation[1]]])
+    return cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REFLECT_101)
+
+
+def generate_random_warp(
+    landmarks: np.ndarray,
+    procedure_indices: list[int],
+    max_displacement: float = 15.0,
+    rng: np.random.Generator | None = None,
+) -> np.ndarray:
+    """Generate randomly warped landmarks for synthetic data."""
+    rng = rng or np.random.default_rng()
+    result = landmarks.copy()
+
+    for idx in procedure_indices:
+        if idx < len(landmarks):
+            dx = rng.uniform(-max_displacement, max_displacement)
+            dy = rng.uniform(-max_displacement, max_displacement)
+            result[idx, 0] += dx
+            result[idx, 1] += dy
+
+    return result