diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 7571720..0e7b0bc 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -15,6 +15,7 @@ Thanks to everyone who has contributed to LandmarkDiff. | [lshariprasad](https://github.com/lshariprasad) | histogram_match_skin tests ([#263](https://github.com/dreamlessx/LandmarkDiff-public/pull/263)) | | [dagangtj](https://github.com/dagangtj) | SafetyResult dataclass improvements ([#235](https://github.com/dreamlessx/LandmarkDiff-public/pull/235)), file logging handler ([#236](https://github.com/dreamlessx/LandmarkDiff-public/pull/236)), api_client error messages ([#237](https://github.com/dreamlessx/LandmarkDiff-public/pull/237)) | | [PredictiveManish](https://github.com/PredictiveManish) | OpenAPI specification ([#340](https://github.com/dreamlessx/LandmarkDiff-public/pull/340)), batch processing notebook ([#334](https://github.com/dreamlessx/LandmarkDiff-public/pull/334)) | +| [mimosapudical](https://github.com/mimosapudical) | Canonical 3D surgical deformation bridge, FLAME integration, and held-out multi-view validation ([#418](https://github.com/dreamlessx/LandmarkDiff-public/issues/418)) | | [srikar117](https://github.com/srikar117) | Dark mode for Gradio interface ([#204](https://github.com/dreamlessx/LandmarkDiff-public/issues/204), in progress) | | [passionworkeer](https://github.com/passionworkeer) | Docker architecture feedback ([#5](https://github.com/dreamlessx/LandmarkDiff-public/issues/5)) | diff --git a/docs/SHARED_3D_SURGERY.md b/docs/SHARED_3D_SURGERY.md new file mode 100644 index 0000000..71b7ed5 --- /dev/null +++ b/docs/SHARED_3D_SURGERY.md @@ -0,0 +1,204 @@ +# Shared 3D Surgical Deformation + +This module adds a geometry-first alternative to LandmarkDiff's independent +per-view 2D procedure presets. + +## What changes + +Independent 2D mode applies a separate landmark preset to every uploaded view. +Shared 3D mode instead: + +1. fits one patient-specific canonical face geometry from the available views; +2. estimates a full 3D scaled-orthographic camera for each view; +3. applies one surgical deformation in canonical 3D; +4. reprojects that same deformation into every camera; +5. reuses the existing TPS/image pipeline for rendering. + +The goal is not to make frontal and profile images look identical. The goal is +for both images to be projections of the same underlying anatomical edit. + +## Geometry backends + +### Dependency-free canonical fallback + +The default backend uses LandmarkDiff's 478-point MediaPipe topology. It is +intended for tests, demos, and environments where FLAME assets are unavailable. +It must not be interpreted as the official FLAME statistical head model. + +### Official FLAME backend + +LandmarkDiff can load a user-supplied official FLAME model asset: + +```python +from landmarkdiff.flame_model import FlameModel + +model = FlameModel.from_flame_assets( + "/path/to/flame2023.pkl", + "/path/to/mediapipe_to_flame.json", +) +``` + +The correspondence may be either: + +- JSON: a 478-entry direct vertex list, or a possibly partial object mapping + MediaPipe ids to FLAME vertex ids; +- NPZ direct map: a 478-entry `vertex_indices` array; or +- standard barycentric NPZ embedding with `landmark_indices`, + `lmk_face_idx`, and `lmk_b_coords`. + +In the standard format, `landmark_indices` are MediaPipe landmark ids. The +FLAME surface location is represented by a triangle id plus barycentric +coordinates. LandmarkDiff evaluates those barycentric locations directly for +camera and shape fitting. Partial embeddings, including the commonly used +selected-landmark form, are supported: mapped points are reprojected from the +3D edit while unmapped MediaPipe points retain their original coordinates for +the downstream 478-point TPS pipeline. + +FLAME model files and third-party correspondence assets are not redistributed +by this repository. Some public trackers publish MediaPipe-to-FLAME embeddings +under licenses that prohibit redistribution, so callers must provide an +embedding they are entitled to use. + +The first FLAME backend uses the neutral FLAME template/topology and identity +shape basis. Global head pose is represented by the per-view camera. Expression +and jaw articulation remain neutral in this first surgical-geometry version. + +## Supported 3D procedures + +The current canonical handle sets cover: + +- rhinoplasty +- alarplasty +- mentoplasty +- genioplasty + +Shared 3D rejects unsupported procedures explicitly instead of silently falling +back to independent 2D. + +## Held-out-view benchmark + +The benchmark is designed to avoid circular evaluation. Ground-truth synthetic +surgery is generated by an analytic Gaussian 3D field, while Shared 3D uses +sparse surgical handles plus Laplacian propagation. + +The training views are 0°, 30°, and 60°. The default held-out view is 80°. +The held-out pre-op image may be used to estimate its camera, but its post-op +target is never used by the fit or deformation. + +Fallback backend: + +```bash +python scripts/benchmark_surgery3d.py \ + --procedure mentoplasty \ + --out artifacts/surgery3d_benchmark +``` + +Official FLAME backend: + +```bash +python scripts/benchmark_surgery3d.py \ + --flame-model /path/to/flame2023.pkl \ + --correspondence /path/to/mediapipe_to_flame.npz \ + --procedure mentoplasty \ + --out artifacts/surgery3d_benchmark_flame +``` + +Primary metric: + +- held-out reprojection NME + +Secondary diagnostics: + +- canonical vertex RMSE, converted to millimetres using fitted IPD; +- held-out camera rotation error; +- front/profile chin displacement ratio for a qualitative sanity check. + +The benchmark also writes a front / 45° / profile comparison grid and a +canonical before/after 3D deformation visualization. + +### Synthetic validation results + +Using the dependency-free canonical backend, the sweep-selected regularization +(`lambda_smooth=30`, `lambda_locality=0.3`) was evaluated on 27 held-out-view +conditions spanning three procedures, three intensities (40/60/100), and three +held-out yaw angles (20°/45°/85°). + +| Procedure | Wins vs. independent 2D | Mean relative NME change | +| --- | ---: | ---: | +| mentoplasty | 9 / 9 | -12.0% | +| genioplasty | 9 / 9 | -47.2% | +| rhinoplasty | 6 / 9 | -1.1% | +| **Overall** | **24 / 27** | **-20.1%** | + +Median relative NME change across all 27 cases was -18.1%. Maximum synthetic +camera-rotation error was approximately 0.0001°. + +The three losses are all near-frontal rhinoplasty cases at 20° yaw, where Shared +3D is about 23% worse in NME than the existing 2D preset. This is an important +trade-off rather than a hidden failure: the existing rhinoplasty preset is +explicitly hand-tuned for frontal-image alar/dorsum narrowing, while the Shared +3D path is designed to preserve one anatomical edit across changing viewpoints. +A separate regularization sweep did not remove this near-frontal regression +without changing the deformation model itself. + +These numbers are synthetic geometry validation, not clinical outcome accuracy. +They demonstrate the cross-view geometry behavior of the method and should not +be interpreted as evidence of patient-specific surgical prediction quality. + +## Numerical implementation + +The 3D surgical deformation solves a Laplacian objective with a small locality +penalty: + +[ +\sum_h w_h\|\Delta v_h-d_h\|^2 ++ \lambda_s\|L\Delta V\|^2 ++ \lambda_l\|\Delta V\|^2. +] + +The locality term prevents the constant-vector nullspace of the graph Laplacian +from turning a local procedure into an almost global translation. + +For FLAME-scale meshes the system is solved matrix-free with preconditioned +conjugate gradient rather than constructing and factorizing a dense 5K-by-5K +matrix. + +## Validation results + +Using the dependency-free canonical backend, held-out-view validation was run +across 27 combinations of: + +- procedures: mentoplasty, genioplasty, rhinoplasty; +- intensities: 40, 60, 100; +- held-out yaw: 20°, 45°, 85°. + +With the sweep-selected regularization `lambda_smooth=30` and +`lambda_locality=0.3`: + +- Shared 3D wins 24 / 27 held-out cases; +- mean relative held-out NME change is -20.1%; +- genioplasty wins 9 / 9 cases, mean -47.2%; +- mentoplasty wins 9 / 9 cases, mean -12.0%; +- rhinoplasty wins 6 / 9 cases, mean -1.1%; +- maximum synthetic camera-rotation error is about 0.0001°. + +The three losses are all rhinoplasty at 20° yaw, where Shared 3D is about 23% +worse than the independent 2D preset. The same rhinoplasty deformation improves +at 45° and 85° yaw. This is treated as a method limitation rather than hidden by +procedure-specific tuning. + +A focused 70° / 80° held-out matrix gives 6 / 6 wins with about 23.7% mean +relative NME improvement. + +## Current limitations + +- The real FLAME path requires an externally supplied MediaPipe-to-FLAME + correspondence map. +- A projected vertex-color PLY preview is available; UV-unwrapped or multi-view + photometric texture fitting is not part of the current geometry bridge. +- The 3D handle magnitudes are method parameters and should be clinically + calibrated before being interpreted as patient-specific surgical predictions. +- Real clinical multi-view validation remains separate from the synthetic + held-out geometry benchmark. +- Near-frontal rhinoplasty remains a documented trade-off: the current 3D + deformation wins at 45°/85° in the synthetic validation but loses at 20°. diff --git a/landmarkdiff/__init__.py b/landmarkdiff/__init__.py index 20f427e..36143c8 100644 --- a/landmarkdiff/__init__.py +++ b/landmarkdiff/__init__.py @@ -16,12 +16,15 @@ "curriculum", "data", "data_version", + "deformation3d", "displacement_model", "ensemble", "evaluation", "experiment_tracker", "face_verifier", "fid", + "flame_fitting", + "flame_model", "hyperparam", "inference", "landmarks", @@ -35,6 +38,7 @@ "morphometry", "postprocess", "safety", + "surgery3d", "synthetic", "validation", ] diff --git a/landmarkdiff/deformation3d.py b/landmarkdiff/deformation3d.py new file mode 100644 index 0000000..d77107c --- /dev/null +++ b/landmarkdiff/deformation3d.py @@ -0,0 +1,339 @@ +"""Canonical 3D surgical deformation via Laplacian mesh editing. + +Surgical edits are defined once on a shared patient-specific mesh using +a small set of anatomical handles, then propagated smoothly:: + + min_ΔV Σ_h w_h ||Δv_h - d_h||² + λ ||L ΔV||² + +where L is the mesh graph Laplacian. This avoids concave dents from +naively moving isolated vertices. + +First-version procedures: rhinoplasty, mentoplasty / genioplasty, +alarplasty. Unsupported procedures raise explicitly. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import numpy as np + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class SurgicalHandle3D: + """One anatomical handle in canonical 3D space.""" + + vertex_index: int + displacement: np.ndarray # (3,) in canonical mm units + weight: float = 1.0 + + +# Anatomical MediaPipe indices used as surgical handles. +_NOSE_TIP = (1, 2, 4, 5, 94, 19) +_NOSE_DORSUM = (6, 168, 195, 197) +_LEFT_ALAR = (98, 97, 99, 64, 60, 240, 236, 141) # include common alar set +_RIGHT_ALAR = (327, 326, 328, 294, 290, 460, 456, 279) +# Prefer the manipulation.py alar sets when present on the mesh +_LEFT_ALAR_PRIMARY = (240, 236, 141, 363, 370) +_RIGHT_ALAR_PRIMARY = (460, 456, 274, 275, 278, 279) +_CHIN_TIP = (152, 175) +_CHIN_CONTOUR = (148, 149, 150, 176, 377, 400, 378) +_JAW_ANGLES = (172, 397, 58, 288) + + +def get_surgical_handles( + procedure: str, + intensity: float = 50.0, + face_scale_mm: float = 63.0, +) -> list[SurgicalHandle3D]: + """Build canonical 3D handles for a procedure. + + Intensity is the same 0-100 UI scale used by ``apply_procedure_preset``. + Displacements are expressed in millimetres relative to a ~63 mm IPD. + """ + scale = float(np.clip(intensity, 0.0, 100.0)) / 100.0 + # Scale displacements mildly with face size (IPD proxy) + mm = face_scale_mm / 63.0 + procedure = procedure.lower().strip() + handles: list[SurgicalHandle3D] = [] + + if procedure == "rhinoplasty": + # Alar wings inward + for idx in _LEFT_ALAR_PRIMARY: + handles.append( + SurgicalHandle3D(idx, np.array([2.5 * scale * mm, 0.0, 0.0]), weight=1.0) + ) + for idx in _RIGHT_ALAR_PRIMARY: + handles.append( + SurgicalHandle3D(idx, np.array([-2.5 * scale * mm, 0.0, 0.0]), weight=1.0) + ) + # Tip: slightly anterior + superior + for idx in _NOSE_TIP: + handles.append( + SurgicalHandle3D( + idx, + np.array([0.0, 1.5 * scale * mm, 2.0 * scale * mm]), + weight=1.2, + ) + ) + # Dorsum: flatten / reduce projection slightly + for idx in _NOSE_DORSUM: + handles.append( + SurgicalHandle3D( + idx, + np.array([0.0, 0.0, -1.2 * scale * mm]), + weight=0.8, + ) + ) + + elif procedure in ("mentoplasty", "genioplasty"): + # Chin advancement: primarily anterior (+Z). Front view sees little + # lateral change; profile sees a clear projection change. + chin_mm = 6.0 * scale * mm # ~6 mm equivalent at intensity=100 + for idx in _CHIN_TIP: + handles.append( + SurgicalHandle3D(idx, np.array([0.0, -0.5 * scale * mm, chin_mm]), weight=1.5) + ) + for idx in _CHIN_CONTOUR: + handles.append( + SurgicalHandle3D( + idx, + np.array([0.0, -0.3 * scale * mm, 0.7 * chin_mm]), + weight=1.0, + ) + ) + if procedure == "genioplasty": + for idx in _JAW_ANGLES: + handles.append( + SurgicalHandle3D( + idx, + np.array([0.0, 0.0, 0.35 * chin_mm]), + weight=0.6, + ) + ) + + elif procedure == "alarplasty": + for idx in _LEFT_ALAR_PRIMARY: + handles.append( + SurgicalHandle3D(idx, np.array([3.0 * scale * mm, 0.0, 0.0]), weight=1.0) + ) + for idx in _RIGHT_ALAR_PRIMARY: + handles.append( + SurgicalHandle3D(idx, np.array([-3.0 * scale * mm, 0.0, 0.0]), weight=1.0) + ) + + else: + raise ValueError( + f"3D deformation not yet defined for procedure '{procedure}'. " + "Supported: rhinoplasty, mentoplasty, genioplasty, alarplasty." + ) + + # Drop out-of-range indices defensively + return [h for h in handles if h.vertex_index >= 0] + + +def _mesh_adjacency( + faces: np.ndarray, + n_vertices: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return directed unique graph edges and vertex degrees.""" + tri = np.asarray(faces, dtype=np.int64) + undirected = np.concatenate( + [ + tri[:, [0, 1]], + tri[:, [1, 2]], + tri[:, [2, 0]], + ], + axis=0, + ) + undirected = np.sort(undirected, axis=1) + undirected = np.unique(undirected, axis=0) + + src = np.concatenate([undirected[:, 0], undirected[:, 1]]) + dst = np.concatenate([undirected[:, 1], undirected[:, 0]]) + degree = np.bincount(src, minlength=n_vertices).astype(np.float64) + return src, dst, degree + + +def _laplacian_matvec( + x: np.ndarray, + src: np.ndarray, + dst: np.ndarray, + degree: np.ndarray, +) -> np.ndarray: + out = degree * x + np.add.at(out, src, -x[dst]) + return out + + +def _pcg( + matvec, + b: np.ndarray, + preconditioner_inv: np.ndarray, + max_iter: int = 500, + tol: float = 1e-8, +) -> np.ndarray: + """Preconditioned conjugate gradient for a symmetric positive system.""" + x = np.zeros_like(b, dtype=np.float64) + r = b - matvec(x) + if float(np.linalg.norm(r)) <= tol: + return x + + z = preconditioner_inv * r + p = z.copy() + rz_old = float(np.dot(r, z)) + + for _ in range(max_iter): + ap = matvec(p) + denom = float(np.dot(p, ap)) + if abs(denom) < 1e-20: + break + alpha = rz_old / denom + x += alpha * p + r -= alpha * ap + if float(np.linalg.norm(r)) <= tol: + break + z = preconditioner_inv * r + rz_new = float(np.dot(r, z)) + if abs(rz_old) < 1e-20: + break + p = z + (rz_new / rz_old) * p + rz_old = rz_new + + return x + + +def laplacian_deform( + vertices: np.ndarray, + faces: np.ndarray, + handles: list[SurgicalHandle3D], + lambda_smooth: float = 30.0, + lambda_locality: float = 0.3, +) -> np.ndarray: + """Solve soft-handle Laplacian deformation without a dense N x N matrix. + + The objective is identical to the original dense formulation: + + sum_h w_h ||delta_h - d_h||^2 + lambda ||L delta||^2 + + plus a small locality term lambda_locality * ||delta||^2. The locality term + prevents the Laplacian nullspace from turning a local surgical edit into a + near-global rigid translation. The system is evaluated from mesh adjacency + and solved with preconditioned conjugate gradient, so it remains practical + for FLAME-scale meshes with roughly 5K vertices. + """ + if not handles: + return vertices.copy() + + n = vertices.shape[0] + src, dst, degree = _mesh_adjacency(faces, n) + + weights = np.zeros(n, dtype=np.float64) + target = np.zeros((n, 3), dtype=np.float64) + for handle in handles: + i = handle.vertex_index + if not 0 <= i < n: + continue + weights[i] += float(handle.weight) + target[i] += float(handle.weight) * np.asarray( + handle.displacement, + dtype=np.float64, + ) + + active = weights > 0 + target[active] /= weights[active, None] + + ridge = 1e-4 + + def matvec(x: np.ndarray) -> np.ndarray: + lx = _laplacian_matvec(x, src, dst, degree) + ltlx = _laplacian_matvec(lx, src, dst, degree) + return weights * x + lambda_smooth * ltlx + (lambda_locality + ridge) * x + + # diag(L^T L) = degree^2 + degree for an unweighted simple graph. + diagonal = weights + lambda_smooth * (degree**2 + degree) + lambda_locality + ridge + preconditioner_inv = 1.0 / np.maximum(diagonal, 1e-12) + + delta = np.zeros((n, 3), dtype=np.float64) + for axis in range(3): + rhs = weights * target[:, axis] + delta[:, axis] = _pcg( + matvec, + rhs, + preconditioner_inv, + ) + + return vertices + delta + + +def apply_canonical_deformation( + mesh_vertices: np.ndarray, + faces: np.ndarray, + procedure: str, + intensity: float = 50.0, + lambda_smooth: float = 30.0, + lambda_locality: float = 0.3, + face_scale_mm: float = 63.0, + landmark_vertex_indices: np.ndarray | None = None, +) -> tuple[np.ndarray, list[SurgicalHandle3D]]: + """Apply anatomy-aware 3D surgical deformation on a canonical mesh. + + Args: + mesh_vertices: (N, 3) shared canonical vertices. + faces: (F, 3) triangle indices. + procedure: Surgical procedure name. + intensity: 0-100 UI intensity. + lambda_smooth: Laplacian smoothness weight. + lambda_locality: Penalty keeping non-surgical regions near zero displacement. + face_scale_mm: Face scale proxy (IPD in mm). + landmark_vertex_indices: Optional (478,) MediaPipe-id to mesh-vertex map. + + Returns: + (deformed_vertices, handles_used) + """ + handles = get_surgical_handles(procedure, intensity=intensity, face_scale_mm=face_scale_mm) + + if landmark_vertex_indices is not None: + mapping = np.asarray(landmark_vertex_indices, dtype=np.int64) + if mapping.shape != (478,): + raise ValueError("landmark_vertex_indices must have shape (478,)") + handles = [ + SurgicalHandle3D( + vertex_index=int(mapping[h.vertex_index]), + displacement=h.displacement.copy(), + weight=h.weight, + ) + for h in handles + ] + + n = mesh_vertices.shape[0] + handles = [h for h in handles if 0 <= h.vertex_index < n] + deformed = laplacian_deform( + mesh_vertices, + faces, + handles, + lambda_smooth=lambda_smooth, + lambda_locality=lambda_locality, + ) + logger.info( + "Applied 3D %s (intensity=%.1f) with %d handles; max |Δ|=%.2f mm", + procedure, + intensity, + len(handles), + float(np.linalg.norm(deformed - mesh_vertices, axis=1).max()), + ) + return deformed, handles + + +def deformation_vectors( + before: np.ndarray, + after: np.ndarray, + stride: int = 1, +) -> tuple[np.ndarray, np.ndarray]: + """Return origins and displacement vectors for visualization arrows.""" + origins = before[::stride] + vectors = (after - before)[::stride] + return origins, vectors diff --git a/landmarkdiff/flame_fitting.py b/landmarkdiff/flame_fitting.py new file mode 100644 index 0000000..b38dfa3 --- /dev/null +++ b/landmarkdiff/flame_fitting.py @@ -0,0 +1,713 @@ +"""Canonical 3D face fitting from MediaPipe landmarks. + +Fits a shared patient-specific 3D face mesh and per-view weak-perspective +cameras from one or more MediaPipe Face Mesh observations:: + + min_{V, {R_v, t_v, s_v}} Σ_v Σ_i w_vi robust(||pi_v(V_i) - l_vi||^2) + + λ_shape ||V - V_mean||² + + λ_smooth ||L V||² + +The built-in backend uses the MediaPipe 478-vertex tessellation as a lightweight +canonical mesh. It is deliberately *not* presented as the official FLAME statistical +head model. ``FlameModel`` remains a compatibility container so an official FLAME +backend plus an explicit MediaPipe-to-FLAME correspondence map can be plugged in +without changing the shared-surgery API. + +Mathematical role matches issue #418 + the 3D surgical bridge: + +- shared identity geometry V across views +- independent camera / pose per view +- surgical edits applied once in canonical 3D, then reprojected +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np + +from landmarkdiff.flame_model import FlameModel +from landmarkdiff.landmarks import FaceLandmarks + +logger = logging.getLogger(__name__) + +# Stable semantic correspondences for fitting (subset of MediaPipe 478). +# Prefer high-confidence anatomical anchors over dense coverage. +FIT_LANDMARK_INDICES: tuple[int, ...] = ( + # Eyes + 33, + 133, + 362, + 263, + 159, + 386, + # Nose + 1, + 2, + 4, + 5, + 6, + 19, + 94, + 168, + 195, + 197, + # Mouth + 61, + 291, + 13, + 14, + 78, + 308, + # Chin / jaw + 152, + 175, + 377, + 148, + 172, + 397, + # Contour anchors + 234, + 454, + 10, + 338, + 109, +) + +# Interpupillary distance used as metric scale reference (approx. adult mean). +_IPD_MM = 63.0 + + +@dataclass +class CameraParams: + """Weak-perspective camera for one view.""" + + rotation: np.ndarray # (3, 3) + translation: np.ndarray # (2,) image-plane translation in pixels + scale: float # pixels per canonical unit + image_width: int = 512 + image_height: int = 512 + + def project(self, vertices: np.ndarray) -> np.ndarray: + """Project canonical 3D vertices to 2D pixel coordinates. + + Args: + vertices: (N, 3) canonical mesh vertices. + + Returns: + (N, 2) pixel coordinates. + """ + rotated = vertices @ self.rotation.T + xy = np.empty((len(rotated), 2), dtype=np.float64) + xy[:, 0] = rotated[:, 0] * self.scale + self.translation[0] + # Canonical/anatomical +Y is superior (up), while image +y is down. + xy[:, 1] = -rotated[:, 1] * self.scale + self.translation[1] + return xy + + +@dataclass +class FlameFitResult: + """Result of canonical face fitting.""" + + vertices: np.ndarray # (N, 3) shared canonical mesh + faces: np.ndarray # (F, 3) triangle indices + cameras: list[CameraParams] + shape_params: np.ndarray # low-dim residual coeffs (may be empty) + mean_vertices: np.ndarray # (N, 3) template used as prior + landmark_vertex_indices: np.ndarray # (478,) representative vertex id or -1 + landmark_face_indices: np.ndarray | None = None + landmark_bary_coords: np.ndarray | None = None + reprojection_errors: list[float] = field(default_factory=list) + n_iterations: int = 0 + + @property + def n_vertices(self) -> int: + return int(self.vertices.shape[0]) + + def project_view(self, view_index: int) -> np.ndarray: + """Project fitted mesh into a specific view.""" + return self.cameras[view_index].project(self.vertices) + + def to_obj(self, path: str | Path, scale: float = 1.0) -> Path: + """Write the fitted canonical mesh as Wavefront OBJ.""" + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as f: + f.write("# LandmarkDiff canonical face mesh\n") + f.write(f"# {self.n_vertices} vertices, {len(self.faces)} faces\n\n") + for x, y, z in self.vertices: + f.write(f"v {x * scale:.6f} {y * scale:.6f} {z * scale:.6f}\n") + f.write("\n") + for v0, v1, v2 in self.faces: + f.write(f"f {int(v0) + 1} {int(v1) + 1} {int(v2) + 1}\n") + return out + + def to_colored_ply( + self, + path: str | Path, + image_bgr: np.ndarray, + view_index: int = 0, + scale: float = 1.0, + ) -> Path: + """Export the fitted mesh as ASCII PLY with projected vertex colors. + + Colors are sampled from a selected fitted view by projecting every mesh + vertex through that camera. This is a lightweight texture preview rather + than a UV-unwrapped photometric texture model. + """ + if image_bgr.ndim != 3 or image_bgr.shape[2] != 3: + raise ValueError("image_bgr must have shape (H, W, 3)") + if not 0 <= view_index < len(self.cameras): + raise IndexError("view_index out of range") + + h, w = image_bgr.shape[:2] + xy = self.cameras[view_index].project(self.vertices) + px = np.clip(np.rint(xy[:, 0]).astype(np.int64), 0, w - 1) + py = np.clip(np.rint(xy[:, 1]).astype(np.int64), 0, h - 1) + bgr = image_bgr[py, px] + rgb = bgr[:, ::-1] + + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w") as fh: + fh.write("ply\n") + fh.write("format ascii 1.0\n") + fh.write("comment LandmarkDiff projected vertex-color texture preview\n") + fh.write(f"element vertex {self.n_vertices}\n") + fh.write("property float x\n") + fh.write("property float y\n") + fh.write("property float z\n") + fh.write("property uchar red\n") + fh.write("property uchar green\n") + fh.write("property uchar blue\n") + fh.write(f"element face {len(self.faces)}\n") + fh.write("property list uchar int vertex_indices\n") + fh.write("end_header\n") + + for vertex, color in zip(self.vertices, rgb, strict=True): + x, y, z = vertex * scale + r, g, b = (int(color[0]), int(color[1]), int(color[2])) + fh.write(f"{x:.6f} {y:.6f} {z:.6f} {r} {g} {b}\n") + for v0, v1, v2 in self.faces: + fh.write(f"3 {int(v0)} {int(v1)} {int(v2)}\n") + + return out + + +def landmarks_to_metric_vertices(face: FaceLandmarks) -> np.ndarray: + """Lift a single MediaPipe observation into canonical-ish 3D (mm). + + Uses interpupillary distance for XY scale and MediaPipe z relative + depth, flipped into (+X right, +Y up, +Z anterior). + """ + lm = face.landmarks.astype(np.float64) + # Pixel space + px = lm[:, 0] * face.image_width + py = lm[:, 1] * face.image_height + left_eye = np.array([px[33], py[33]]) + right_eye = np.array([px[263], py[263]]) + ipd_px = float(np.linalg.norm(right_eye - left_eye)) + if ipd_px < 1e-6: + ipd_px = 0.25 * face.image_width + scale = _IPD_MM / ipd_px + + # Center at nose tip in image, flip Y, scale z relative to face + origin = np.array([px[1], py[1]]) + x = (px - origin[0]) * scale + y = -(py - origin[1]) * scale # flip to +Y up + z_raw = lm[:, 2] + z = -(z_raw - np.median(z_raw)) * (_IPD_MM * 2.0) # anterior positive + return np.column_stack([x, y, z]) + + +def estimate_weak_perspective_camera( + vertices: np.ndarray, + landmarks_px: np.ndarray, + weights: np.ndarray | None = None, + image_width: int = 512, + image_height: int = 512, + indices: Sequence[int] | np.ndarray | None = None, + vertex_indices: Sequence[int] | np.ndarray | None = None, +) -> CameraParams: + """Estimate a full 3D scaled-orthographic camera from 3D-to-2D correspondences. + + The previous implementation solved only an in-plane 2D Procrustes problem and + embedded that rotation in a 3x3 identity matrix. That cannot recover yaw or + pitch, which are precisely the quantities needed for front/three-quarter/profile + consistency. + + Here we first solve the weighted unconstrained 2x3 affine map from centered 3D + points to centered 2D points, then project its two rows onto the closest pair of + orthonormal rows with a shared scale. The third row is their cross product, + yielding a proper SO(3) rotation. + + This is the classic scaled-orthographic/weak-perspective Procrustes camera used + in many 3D morphable-model fitting pipelines. + """ + mp_idx = np.asarray( + indices if indices is not None else FIT_LANDMARK_INDICES, + dtype=np.int64, + ) + vert_idx = np.asarray( + vertex_indices if vertex_indices is not None else mp_idx, + dtype=np.int64, + ) + if mp_idx.shape != vert_idx.shape: + raise ValueError("indices and vertex_indices must have the same shape") + + src = np.asarray(vertices[vert_idx], dtype=np.float64) + dst = np.asarray(landmarks_px[mp_idx], dtype=np.float64) + + if weights is None: + w = np.ones(len(mp_idx), dtype=np.float64) + else: + w = np.asarray(weights, dtype=np.float64)[mp_idx] + w = np.clip(w, 1e-6, None) + + w = w / (float(w.sum()) + 1e-12) + mu_s = np.sum(src * w[:, None], axis=0) + mu_d = np.sum(dst * w[:, None], axis=0) + x_centered = src - mu_s + y_centered = dst - mu_d + + # Weighted least-squares unconstrained 2x3 camera map. + sw = np.sqrt(w)[:, None] + affine_t, *_ = np.linalg.lstsq( + x_centered * sw, + y_centered * sw, + rcond=None, + ) + affine = affine_t.T + + # The second observed image row points down, while the canonical camera + # frame uses +Y up. Convert the fitted 2D affine map back to physical + # camera rows before projecting onto SO(3). + physical_affine = affine.copy() + physical_affine[1] *= -1.0 + + # Closest scaled row-orthonormal 2x3 physical camera matrix. + u_mat, singular, vt_mat = np.linalg.svd( + physical_affine, + full_matrices=True, + ) + rotation_rows = u_mat @ vt_mat[:2, :] + scale = float(np.mean(singular[:2])) + if not np.isfinite(scale) or scale <= 1e-8: + scale = 1e-6 + + r1 = rotation_rows[0] / (np.linalg.norm(rotation_rows[0]) + 1e-12) + r2 = rotation_rows[1] - np.dot(rotation_rows[1], r1) * r1 + r2 = r2 / (np.linalg.norm(r2) + 1e-12) + r3 = np.cross(r1, r2) + r3 = r3 / (np.linalg.norm(r3) + 1e-12) + rotation = np.stack([r1, r2, r3], axis=0) + + # Ensure a right-handed proper rotation. + if np.linalg.det(rotation) < 0: + rotation[2] *= -1.0 + + projected_mu = np.array( + [ + scale * float(mu_s @ rotation[0]), + -scale * float(mu_s @ rotation[1]), + ], + dtype=np.float64, + ) + translation = mu_d - projected_mu + + return CameraParams( + rotation=rotation.astype(np.float64), + translation=translation.astype(np.float64), + scale=scale, + image_width=image_width, + image_height=image_height, + ) + + +def _huber_weights(residuals: np.ndarray, delta: float = 4.0) -> np.ndarray: + """Per-point Huber weights from 2D residual magnitudes (pixels).""" + r = np.asarray(residuals, dtype=np.float64) + w = np.ones_like(r) + mask = r > delta + w[mask] = delta / np.maximum(r[mask], 1e-8) + return w + + +def fit_flame_from_landmarks( + faces: Sequence[FaceLandmarks], + model: FlameModel | None = None, + n_iters: int = 8, + lambda_shape: float = 1e-2, + lambda_smooth: float = 1e-3, + fit_indices: Sequence[int] | None = None, +) -> FlameFitResult: + """Fit one shared canonical mesh and a full 3D camera for each view. + + Official FLAME fitting uses whichever MediaPipe semantics are present in the + supplied correspondence embedding. Standard partial barycentric embeddings + therefore work without pretending to cover all 478 MediaPipe points. + """ + if not faces: + raise ValueError("At least one FaceLandmarks observation is required") + + model = model or FlameModel.mediapipe_template() + + if fit_indices is None: + requested = ( + np.asarray(FIT_LANDMARK_INDICES, dtype=np.int64) + if model.backend == "flame" + else np.arange(478, dtype=np.int64) + ) + else: + requested = np.asarray(fit_indices, dtype=np.int64) + + mp_indices = model.available_mediapipe_indices(requested) + if len(mp_indices) < 4: + raise ValueError("At least four MediaPipe correspondences are required for 3D fitting") + local_indices = np.arange(len(mp_indices), dtype=np.int64) + + all_landmark_vertices = model.vertex_indices_for_mediapipe( + np.arange(478), + allow_missing=True, + ) + + betas = np.zeros( + model.shape_basis.shape[-1] if model.shape_basis is not None else 0, + dtype=np.float64, + ) + verts = model.apply_shape(betas) + + cameras: list[CameraParams] = [] + pixel_targets: list[np.ndarray] = [] + conf_weights: list[np.ndarray] = [] + + for face in faces: + pixel_targets.append(face.pixel_coords.astype(np.float64)) + conf_weights.append(face.landmark_confidence.astype(np.float64)) + source_points = model.sample_mediapipe_points(verts, mp_indices) + cameras.append( + estimate_weak_perspective_camera( + source_points, + pixel_targets[-1], + weights=conf_weights[-1], + image_width=face.image_width, + image_height=face.image_height, + indices=mp_indices, + vertex_indices=local_indices, + ) + ) + + for _ in range(n_iters): + source_points = model.sample_mediapipe_points(verts, mp_indices) + for v, face in enumerate(faces): + cameras[v] = estimate_weak_perspective_camera( + source_points, + pixel_targets[v], + weights=conf_weights[v], + image_width=face.image_width, + image_height=face.image_height, + indices=mp_indices, + vertex_indices=local_indices, + ) + + if model.backend == "flame": + betas = _solve_flame_shape_betas( + model, + verts, + cameras, + pixel_targets, + conf_weights, + mp_indices, + lambda_shape=lambda_shape, + ) + verts = model.apply_shape(betas) + else: + verts = model.mean_vertices.copy() + + errors: list[float] = [] + source_points = model.sample_mediapipe_points(verts, mp_indices) + for v, face in enumerate(faces): + cameras[v] = estimate_weak_perspective_camera( + source_points, + pixel_targets[v], + weights=conf_weights[v], + image_width=face.image_width, + image_height=face.image_height, + indices=mp_indices, + vertex_indices=local_indices, + ) + proj = cameras[v].project(source_points) + tgt = pixel_targets[v][mp_indices] + w = conf_weights[v][mp_indices] + err = float( + np.sqrt( + np.average( + np.sum((proj - tgt) ** 2, axis=1), + weights=w, + ) + ) + ) + errors.append(err) + + return FlameFitResult( + vertices=verts.astype(np.float64), + faces=model.faces.copy(), + cameras=cameras, + shape_params=betas.astype(np.float64), + mean_vertices=model.mean_vertices.copy(), + landmark_vertex_indices=all_landmark_vertices.astype(np.int64), + landmark_face_indices=( + None if model.landmark_face_indices is None else model.landmark_face_indices.copy() + ), + landmark_bary_coords=( + None if model.landmark_bary_coords is None else model.landmark_bary_coords.copy() + ), + reprojection_errors=errors, + n_iterations=n_iters, + ) + + +def _solve_flame_shape_betas( + model: FlameModel, + current_vertices: np.ndarray, + cameras: Sequence[CameraParams], + pixel_targets: Sequence[np.ndarray], + conf_weights: Sequence[np.ndarray], + mp_indices: np.ndarray, + lambda_shape: float, +) -> np.ndarray: + """Solve FLAME identity coefficients under fixed cameras.""" + if model.shape_basis is None or model.shape_basis.shape[-1] == 0: + return np.zeros(0, dtype=np.float64) + + k = model.shape_basis.shape[-1] + basis = model.sample_mediapipe_basis(model.shape_basis, mp_indices) + mean_points = model.sample_mediapipe_points(model.mean_vertices, mp_indices) + current_points = model.sample_mediapipe_points(current_vertices, mp_indices) + + rows: list[np.ndarray] = [] + targets: list[np.ndarray] = [] + + for cam, target, confidence in zip( + cameras, + pixel_targets, + conf_weights, + strict=True, + ): + mean_xy = cam.project(mean_points) + residual = target[mp_indices] - mean_xy + + projected_basis = np.einsum( + "mck,rc->mrk", + basis, + cam.rotation[:2, :], + ) + projected_basis[:, 1, :] *= -1.0 + projected_basis *= cam.scale + + current_xy = cam.project(current_points) + residual_mag = np.linalg.norm( + current_xy - target[mp_indices], + axis=1, + ) + robust = _huber_weights(residual_mag) * confidence[mp_indices] + sqrt_w = np.repeat(np.sqrt(np.clip(robust, 1e-8, None)), 2) + + design = projected_basis.reshape(-1, k) + rhs = residual.reshape(-1) + rows.append(design * sqrt_w[:, None]) + targets.append(rhs * sqrt_w) + + design_all = np.vstack(rows) + target_all = np.concatenate(targets) + ridge = max(float(lambda_shape), 1e-6) + system = design_all.T @ design_all + ridge * np.eye(k) + betas = np.linalg.solve(system, design_all.T @ target_all) + return np.clip(betas, -3.0, 3.0) + + +def _procrustes_align( + source: np.ndarray, + target: np.ndarray, + indices: np.ndarray, +) -> np.ndarray: + """Similarity-align source to target on a landmark subset, apply to all.""" + s = source[indices] + t = target[indices] + mu_s = s.mean(axis=0) + mu_t = t.mean(axis=0) + x_centered = s - mu_s + y_centered = t - mu_t + covariance = x_centered.T @ y_centered + u_mat, singular, vt_mat = np.linalg.svd(covariance) + rotation = vt_mat.T @ u_mat.T + if np.linalg.det(rotation) < 0: + vt_mat[-1, :] *= -1 + rotation = vt_mat.T @ u_mat.T + scale = float(np.sum(singular) / (np.sum(x_centered**2) + 1e-8)) + return (source - mu_s) @ rotation.T * scale + mu_t + + +def build_mesh_laplacian(faces: np.ndarray, n_vertices: int) -> np.ndarray: + """Build a combinatorial graph Laplacian from triangle faces.""" + adjacency = np.zeros((n_vertices, n_vertices), dtype=np.float64) + for tri in faces: + i, j, k = int(tri[0]), int(tri[1]), int(tri[2]) + for a, b in ((i, j), (j, k), (k, i)): + adjacency[a, b] = 1.0 + adjacency[b, a] = 1.0 + degree = adjacency.sum(axis=1) + laplacian = np.diag(degree) - adjacency + return laplacian + + +def _refine_shared_vertices( + verts: np.ndarray, + prior: np.ndarray, + cameras: Sequence[CameraParams], + pixel_targets: Sequence[np.ndarray], + conf_weights: Sequence[np.ndarray], + mp_indices: np.ndarray, + vertex_indices: np.ndarray, + laplacian: np.ndarray, + lambda_shape: float, + lambda_smooth: float, +) -> np.ndarray: + """One robust shared-geometry update under mapped 2D observations.""" + n = verts.shape[0] + accum = np.zeros((n, 3), dtype=np.float64) + weight = np.zeros(n, dtype=np.float64) + + for cam, tgt, conf in zip(cameras, pixel_targets, conf_weights, strict=True): + rotated = verts @ cam.rotation.T + current = rotated[vertex_indices] + depth = current[:, 2:3] + desired_xy = (tgt[mp_indices] - cam.translation) / max(cam.scale, 1e-6) + desired_xy[:, 1] *= -1.0 + desired_rot = np.concatenate([desired_xy, depth], axis=1) + desired_can = desired_rot @ cam.rotation + + proj = cam.project(verts)[vertex_indices] + resid = np.linalg.norm(proj - tgt[mp_indices], axis=1) + robust = _huber_weights(resid) * conf[mp_indices] + + for local_i, mesh_i in enumerate(vertex_indices): + wi = float(robust[local_i]) + accum[mesh_i] += wi * desired_can[local_i] + weight[mesh_i] += wi + + weight_matrix = np.diag(weight) + smooth = laplacian.T @ laplacian + system = weight_matrix + lambda_shape * np.eye(n) + lambda_smooth * smooth + 1e-6 * np.eye(n) + updated = np.zeros_like(verts) + + mask = weight > 0 + for axis in range(3): + data = np.zeros(n, dtype=np.float64) + data[mask] = accum[mask, axis] / np.maximum(weight[mask], 1e-8) + data[~mask] = verts[~mask, axis] + rhs = weight * data + lambda_shape * prior[:, axis] + updated[:, axis] = np.linalg.solve(system, rhs) + + return updated + + +def project_model_landmarks( + vertices: np.ndarray, + camera: CameraParams, + model: FlameModel, + source_face: FaceLandmarks | None = None, + image_width: int | None = None, + image_height: int | None = None, +) -> FaceLandmarks: + """Project model geometry through its MediaPipe correspondence definition. + + Partial FLAME embeddings are supported. When source_face is provided, + MediaPipe points without a FLAME correspondence retain their original + coordinates, while mapped points are replaced by the projected 3D edit. + """ + w = ( + source_face.image_width + if source_face is not None + else (image_width if image_width is not None else camera.image_width) + ) + h = ( + source_face.image_height + if source_face is not None + else (image_height if image_height is not None else camera.image_height) + ) + + available = model.available_mediapipe_indices() + if len(available) == 0: + raise ValueError("Model has no MediaPipe correspondences") + + sampled = model.sample_mediapipe_points(vertices, available) + xy = camera.project(sampled) + + if source_face is not None: + lm = source_face.landmarks.copy().astype(np.float32) + confidence = source_face.confidence + else: + lm = np.zeros((478, 3), dtype=np.float32) + confidence = 1.0 + + lm[available, 0] = np.clip(xy[:, 0] / max(w, 1), 0.0, 1.0) + lm[available, 1] = np.clip(xy[:, 1] / max(h, 1), 0.0, 1.0) + + if source_face is None: + rotated_z = (sampled @ camera.rotation.T)[:, 2] + lm[available, 2] = (-(rotated_z - np.median(rotated_z)) / (_IPD_MM * 2.0)).astype( + np.float32 + ) + + return FaceLandmarks( + landmarks=lm, + image_width=w, + image_height=h, + confidence=confidence, + ) + + +def project_landmarks( + vertices: np.ndarray, + camera: CameraParams, + image_width: int | None = None, + image_height: int | None = None, + source_z: np.ndarray | None = None, + landmark_vertex_indices: np.ndarray | None = None, +) -> FaceLandmarks: + """Project a canonical mesh into MediaPipe landmark observation space.""" + w = image_width if image_width is not None else camera.image_width + h = image_height if image_height is not None else camera.image_height + + if landmark_vertex_indices is None: + sampled = vertices + else: + idx = np.asarray(landmark_vertex_indices, dtype=np.int64) + if idx.shape != (478,): + raise ValueError("landmark_vertex_indices must have shape (478,)") + sampled = vertices[idx] + + xy = camera.project(sampled) + lm = np.zeros((sampled.shape[0], 3), dtype=np.float32) + lm[:, 0] = np.clip(xy[:, 0] / max(w, 1), 0.0, 1.0) + lm[:, 1] = np.clip(xy[:, 1] / max(h, 1), 0.0, 1.0) + + if source_z is not None: + z = np.asarray(source_z, dtype=np.float32) + if z.shape[0] != sampled.shape[0]: + raise ValueError("source_z length must match projected landmark count") + lm[:, 2] = z + else: + rotated_z = (sampled @ camera.rotation.T)[:, 2] + lm[:, 2] = (-(rotated_z - np.median(rotated_z)) / (_IPD_MM * 2.0)).astype(np.float32) + + return FaceLandmarks( + landmarks=lm, + image_width=w, + image_height=h, + confidence=1.0, + ) diff --git a/landmarkdiff/flame_model.py b/landmarkdiff/flame_model.py new file mode 100644 index 0000000..9107bc5 --- /dev/null +++ b/landmarkdiff/flame_model.py @@ -0,0 +1,418 @@ +"""Canonical face-model backends for the Shared 3D surgical bridge. + +The lightweight MediaPipe-topology model keeps tests and demos dependency-free. +The optional FLAME backend loads official user-supplied FLAME assets plus an +explicit MediaPipe-to-FLAME correspondence map. LandmarkDiff does not redistribute +the FLAME model files. +""" + +from __future__ import annotations + +import json +import pickle +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + + +@dataclass +class FlameModel: + """Canonical parametric face model used by the shared-3D bridge. + + ``mediapipe_template()`` is a lightweight canonical fallback for tests and + dependency-free demos; it is not the official FLAME statistical model. + ``model_path`` is reserved for an external FLAME backend. + """ + + mean_vertices: np.ndarray # (N, 3) + faces: np.ndarray # (F, 3) + shape_basis: np.ndarray | None = None # (N, 3, K) or None + model_path: Path | None = None + landmark_vertex_indices: np.ndarray | None = None + landmark_face_indices: np.ndarray | None = None + landmark_bary_coords: np.ndarray | None = None + backend: str = "mediapipe_canonical" + + @classmethod + def mediapipe_template(cls, n_shape: int = 10) -> FlameModel: + """Build the lightweight MediaPipe-topology canonical fallback.""" + from landmarkdiff.export import _get_tessellation_triangles + + mean = _default_mean_face(478) + faces = np.asarray(_get_tessellation_triangles(), dtype=np.int32) + basis = _random_orthonormal_basis(mean, n_shape) + return cls( + mean_vertices=mean, + faces=faces, + shape_basis=basis, + landmark_vertex_indices=np.arange(478, dtype=np.int64), + backend="mediapipe_canonical", + ) + + @classmethod + def from_flame_assets( + cls, + model_path: str | Path, + correspondence_path: str | Path, + n_shape: int = 100, + ) -> FlameModel: + """Load official FLAME neutral identity geometry from local assets. + + The FLAME model is not redistributed here. The caller supplies an + official FLAME pickle plus either a direct JSON/NPZ vertex map or a + standard MediaPipe-to-FLAME barycentric NPZ embedding containing + landmark_indices, lmk_face_idx, and lmk_b_coords. Partial embeddings + are supported. Rigid head pose is handled by the per-view camera; + expression and jaw pose remain neutral in this first backend. + """ + model_path = Path(model_path) + correspondence_path = Path(correspondence_path) + if not model_path.exists(): + raise FileNotFoundError(f"FLAME model asset not found: {model_path}") + if not correspondence_path.exists(): + raise FileNotFoundError( + f"MediaPipe-to-FLAME correspondence not found: {correspondence_path}" + ) + + with open(model_path, "rb") as fh: + raw = pickle.load(fh, encoding="latin1") + + def field(name: str): + return raw[name] if isinstance(raw, dict) else getattr(raw, name) + + def array(value, dtype): + if hasattr(value, "r"): + value = value.r + if hasattr(value, "toarray"): + value = value.toarray() + return np.asarray(value, dtype=dtype) + + mean = array(field("v_template"), np.float64) + faces = array(field("f"), np.int32) + shapedirs = array(field("shapedirs"), np.float64) + if shapedirs.ndim != 3 or shapedirs.shape[:2] != mean.shape: + raise ValueError( + f"Unexpected FLAME shapedirs shape {shapedirs.shape}; expected (N, 3, K)" + ) + basis = shapedirs[:, :, : min(int(n_shape), shapedirs.shape[-1])].copy() + + landmark_face_indices: np.ndarray | None = None + landmark_bary_coords: np.ndarray | None = None + + if correspondence_path.suffix.lower() == ".npz": + packed = np.load(correspondence_path, allow_pickle=False) + keys = set(packed.files) + standard_keys = {"landmark_indices", "lmk_face_idx", "lmk_b_coords"} + + if standard_keys.issubset(keys): + mp_ids = np.asarray(packed["landmark_indices"], dtype=np.int64) + face_ids = np.asarray(packed["lmk_face_idx"], dtype=np.int64) + bary = np.asarray(packed["lmk_b_coords"], dtype=np.float64) + + if mp_ids.ndim != 1 or face_ids.shape != mp_ids.shape: + raise ValueError( + "Standard NPZ embedding requires 1D landmark_indices " + "and matching lmk_face_idx" + ) + if bary.shape != (len(mp_ids), 3): + raise ValueError( + "Standard NPZ embedding requires lmk_b_coords with shape (K, 3)" + ) + if np.any(mp_ids < 0) or np.any(mp_ids >= 478): + raise ValueError("MediaPipe landmark ids must be in [0, 477]") + if np.any(face_ids < 0) or np.any(face_ids >= len(faces)): + raise ValueError("Embedding contains FLAME face ids outside the mesh") + + mapping = np.full(478, -1, dtype=np.int64) + landmark_face_indices = np.full(478, -1, dtype=np.int64) + landmark_bary_coords = np.zeros((478, 3), dtype=np.float64) + + for mp_id, face_id, weights in zip( + mp_ids, + face_ids, + bary, + strict=True, + ): + face_vertices = faces[int(face_id)] + dominant = int(face_vertices[int(np.argmax(weights))]) + mapping[int(mp_id)] = dominant + landmark_face_indices[int(mp_id)] = int(face_id) + landmark_bary_coords[int(mp_id)] = weights + + elif "vertex_indices" in keys: + mapping = np.asarray(packed["vertex_indices"], dtype=np.int64) + if mapping.shape != (478,): + raise ValueError("NPZ vertex_indices correspondence must have shape (478,)") + else: + raise ValueError( + "NPZ correspondence must contain either " + "{landmark_indices, lmk_face_idx, lmk_b_coords} or " + "a 478-entry vertex_indices array" + ) + else: + data = json.loads(correspondence_path.read_text()) + if isinstance(data, dict): + mapping = np.full(478, -1, dtype=np.int64) + for mp_id, vertex_id in data.items(): + i = int(mp_id) + if 0 <= i < 478: + mapping[i] = int(vertex_id) + elif isinstance(data, list): + mapping = np.asarray(data, dtype=np.int64) + else: + raise ValueError("Correspondence JSON must be an object or a 478-entry list") + + if mapping.shape != (478,): + raise ValueError("Correspondence vertex map must have shape (478,)") + valid_vertices = mapping[mapping >= 0] + if len(valid_vertices) == 0: + raise ValueError("Correspondence does not define any MediaPipe landmarks") + if np.any(valid_vertices >= mean.shape[0]): + raise ValueError("Correspondence map contains vertex ids outside the FLAME mesh") + + return cls( + mean_vertices=mean, + faces=faces, + shape_basis=basis, + model_path=model_path, + landmark_vertex_indices=mapping, + landmark_face_indices=landmark_face_indices, + landmark_bary_coords=landmark_bary_coords, + backend="flame", + ) + + def available_mediapipe_indices( + self, + indices: Sequence[int] | np.ndarray | None = None, + ) -> np.ndarray: + """Return MediaPipe ids with a usable mesh correspondence.""" + mp = ( + np.arange(478, dtype=np.int64) + if indices is None + else np.asarray(indices, dtype=np.int64) + ) + if self.landmark_vertex_indices is None: + if self.mean_vertices.shape[0] == 478: + return mp + return np.empty(0, dtype=np.int64) + mapping = np.asarray(self.landmark_vertex_indices, dtype=np.int64) + if mapping.shape != (478,): + raise ValueError("landmark_vertex_indices must have shape (478,)") + return mp[mapping[mp] >= 0] + + def vertex_indices_for_mediapipe( + self, + indices: Sequence[int] | np.ndarray, + *, + allow_missing: bool = False, + ) -> np.ndarray: + """Map MediaPipe ids to representative mesh vertices. + + Standard barycentric FLAME embeddings are reduced to the dominant + triangle vertex only for operations that require a concrete vertex id, + such as sparse surgical handles. Fitting and projection use the exact + barycentric location through sample_mediapipe_points. + """ + mp = np.asarray(indices, dtype=np.int64) + if self.landmark_vertex_indices is None: + if self.mean_vertices.shape[0] != 478: + raise ValueError("A MediaPipe-to-mesh correspondence map is required") + return mp + if self.landmark_vertex_indices.shape != (478,): + raise ValueError("landmark_vertex_indices must have shape (478,)") + mapped = self.landmark_vertex_indices[mp] + if not allow_missing and np.any(mapped < 0): + missing = mp[mapped < 0].tolist() + raise ValueError(f"No mesh correspondence for MediaPipe landmarks {missing}") + return mapped + + def sample_mediapipe_points( + self, + vertices: np.ndarray, + indices: Sequence[int] | np.ndarray, + ) -> np.ndarray: + """Sample MediaPipe semantic points on a mesh.""" + mp = np.asarray(indices, dtype=np.int64) + out = np.empty((len(mp), 3), dtype=np.float64) + face_indices = self.landmark_face_indices + bary_coords = self.landmark_bary_coords + + for j, mp_id in enumerate(mp): + i = int(mp_id) + if face_indices is not None and bary_coords is not None: + face_id = int(face_indices[i]) + if face_id >= 0: + tri = self.faces[face_id] + weights = bary_coords[i] + out[j] = np.sum( + vertices[tri] * weights[:, None], + axis=0, + ) + continue + + vertex_id = int(self.vertex_indices_for_mediapipe([i])[0]) + out[j] = vertices[vertex_id] + + return out + + def sample_mediapipe_basis( + self, + basis: np.ndarray, + indices: Sequence[int] | np.ndarray, + ) -> np.ndarray: + """Sample a vertex basis of shape (V, 3, K) at MediaPipe points.""" + mp = np.asarray(indices, dtype=np.int64) + k = basis.shape[-1] + out = np.empty((len(mp), 3, k), dtype=np.float64) + face_indices = self.landmark_face_indices + bary_coords = self.landmark_bary_coords + + for j, mp_id in enumerate(mp): + i = int(mp_id) + if face_indices is not None and bary_coords is not None: + face_id = int(face_indices[i]) + if face_id >= 0: + tri = self.faces[face_id] + weights = bary_coords[i] + out[j] = np.sum( + basis[tri] * weights[:, None, None], + axis=0, + ) + continue + + vertex_id = int(self.vertex_indices_for_mediapipe([i])[0]) + out[j] = basis[vertex_id] + + return out + + def apply_shape(self, betas: np.ndarray) -> np.ndarray: + """Return mean + shape_basis @ betas.""" + verts = self.mean_vertices.copy() + if self.shape_basis is None or len(betas) == 0: + return verts + k = min(len(betas), self.shape_basis.shape[-1]) + for i in range(k): + verts = verts + float(betas[i]) * self.shape_basis[:, :, i] + return verts + + +def _default_mean_face(n: int = 478) -> np.ndarray: + """Construct a smooth neutral face-like mean mesh in canonical coords. + + Coordinate convention (canonical / anatomical): + +X = subject right, +Y = up, +Z = anterior (out of face). + Units are approximately millimetres after IPD normalisation. + + Key MediaPipe indices are placed at anatomically plausible positions so + surgical handles (nose tip, alae, chin, eyes) land on the right regions + even before a patient-specific fit. + """ + rng = np.random.default_rng(0) + verts = np.zeros((n, 3), dtype=np.float64) + + # Explicit anatomical anchors (mm). Subject-right = +X. + anchors: dict[int, tuple[float, float, float]] = { + # Eyes + 33: (-31.5, 25.0, 8.0), + 133: (-12.0, 25.0, 10.0), + 362: (12.0, 25.0, 10.0), + 263: (31.5, 25.0, 8.0), + 159: (-22.0, 28.0, 9.0), + 386: (22.0, 28.0, 9.0), + # Nose + 1: (0.0, -5.0, 28.0), # tip + 2: (0.0, -2.0, 26.0), + 4: (0.0, 2.0, 24.0), + 5: (0.0, 8.0, 22.0), + 6: (0.0, 15.0, 18.0), # bridge + 19: (0.0, -8.0, 24.0), + 94: (0.0, -10.0, 20.0), + 168: (0.0, 20.0, 14.0), + 195: (-4.0, 12.0, 16.0), + 197: (4.0, 12.0, 16.0), + # Alae + 240: (-18.0, -6.0, 16.0), + 236: (-14.0, -2.0, 18.0), + 141: (-12.0, -8.0, 14.0), + 363: (-10.0, -4.0, 15.0), + 370: (-16.0, -10.0, 12.0), + 460: (18.0, -6.0, 16.0), + 456: (14.0, -2.0, 18.0), + 274: (12.0, -8.0, 14.0), + 275: (10.0, -4.0, 15.0), + 278: (16.0, -10.0, 12.0), + 279: (15.0, -8.0, 13.0), + # Mouth + 61: (-20.0, -30.0, 12.0), + 291: (20.0, -30.0, 12.0), + 13: (0.0, -28.0, 14.0), + 14: (0.0, -34.0, 12.0), + 78: (-8.0, -30.0, 13.0), + 308: (8.0, -30.0, 13.0), + # Chin / jaw + 152: (0.0, -70.0, 18.0), + 175: (0.0, -65.0, 16.0), + 148: (-12.0, -66.0, 12.0), + 149: (-20.0, -62.0, 10.0), + 150: (-28.0, -55.0, 8.0), + 176: (-8.0, -68.0, 14.0), + 377: (12.0, -66.0, 12.0), + 400: (20.0, -62.0, 10.0), + 378: (28.0, -55.0, 8.0), + 172: (-45.0, -40.0, 0.0), + 397: (45.0, -40.0, 0.0), + 58: (-42.0, -35.0, 2.0), + 288: (42.0, -35.0, 2.0), + # Contour / forehead + 234: (-55.0, 10.0, -5.0), + 454: (55.0, 10.0, -5.0), + 10: (0.0, 75.0, 5.0), + 338: (25.0, 70.0, 3.0), + 109: (-25.0, 70.0, 3.0), + } + + for idx, xyz in anchors.items(): + if idx < n: + verts[idx] = xyz + + anchor_idx = np.array(sorted(i for i in anchors if i < n), dtype=np.int64) + anchor_pos = verts[anchor_idx] + + # Fill remaining vertices by RBF-ish blend of anchors + mild ellipsoid prior. + for i in range(n): + if i in anchors: + continue + # Deterministic prior on an ellipsoid + t = (i + 0.5) / n + yaw = 2.0 * np.pi * i * 0.61803398875 + pitch = np.pi * (0.35 * (t - 0.5)) + prior = np.array( + [ + 55.0 * np.sin(yaw) * np.cos(pitch), + 75.0 * np.sin(pitch) + 5.0, + 25.0 * np.cos(yaw) * np.cos(pitch) + 5.0, + ] + ) + d2 = np.sum((anchor_pos - prior) ** 2, axis=1) + w = np.exp(-d2 / (2 * 35.0**2)) + w_sum = float(w.sum()) + 1e-8 + verts[i] = (w[:, None] * anchor_pos).sum(axis=0) / w_sum + verts[i] = 0.65 * verts[i] + 0.35 * prior + + verts += rng.normal(0.0, 0.02, size=verts.shape) + return verts + + +def _random_orthonormal_basis(mean: np.ndarray, k: int) -> np.ndarray: + """Low-rank shape basis around the mean (deterministic).""" + rng = np.random.default_rng(1) + n = mean.shape[0] + flat = rng.normal(0.0, 1.0, size=(n * 3, k)) + # Orthonormalize columns + q, _ = np.linalg.qr(flat) + basis = q[:, :k].reshape(n, 3, k) + # Scale so β≈1 is a few mm of variation + basis *= 3.0 + return basis diff --git a/landmarkdiff/surgery3d.py b/landmarkdiff/surgery3d.py new file mode 100644 index 0000000..cef4c70 --- /dev/null +++ b/landmarkdiff/surgery3d.py @@ -0,0 +1,192 @@ +"""Shared 3D surgical deformation bridge. + +Replaces independent per-view 2D procedure presets with:: + + views → shared canonical mesh → one 3D surgery → reproject → TPS + +so front / 45° / profile show projections of the *same* Δ³D. +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np + +from landmarkdiff.deformation3d import apply_canonical_deformation +from landmarkdiff.flame_fitting import ( + FlameFitResult, + FlameModel, + fit_flame_from_landmarks, + project_model_landmarks, +) +from landmarkdiff.landmarks import FaceLandmarks +from landmarkdiff.manipulation import apply_procedure_preset + +logger = logging.getLogger(__name__) + +# Procedures with a dedicated 3D handle set. +SUPPORTED_3D_PROCEDURES = frozenset({"rhinoplasty", "mentoplasty", "genioplasty", "alarplasty"}) + + +@dataclass +class Shared3DSurgeryResult: + """Outputs of a shared-3D multi-view surgery.""" + + fit: FlameFitResult + vertices_before: np.ndarray + vertices_after: np.ndarray + faces: np.ndarray + manipulated_faces: list[FaceLandmarks] + mode: str = "shared_3d" + + +def apply_shared_3d_surgery( + faces: Sequence[FaceLandmarks], + procedure: str, + intensity: float = 50.0, + model: FlameModel | None = None, + lambda_smooth: float = 30.0, + fit_iters: int = 6, +) -> Shared3DSurgeryResult: + """Fit once, deform once, reproject to every view. + + Args: + faces: Multi-view FaceLandmarks for the same person (order preserved). + procedure: Procedure name. + intensity: 0-100 intensity (same scale as 2D presets). + model: Optional FlameModel override. + lambda_smooth: Laplacian smoothness for 3D deformation. + fit_iters: Canonical fitting iterations. + + Returns: + Shared3DSurgeryResult with per-view manipulated FaceLandmarks. + """ + if not faces: + raise ValueError("faces must be non-empty") + + procedure = procedure.lower().strip() + if procedure not in SUPPORTED_3D_PROCEDURES: + raise ValueError( + f"Shared 3D mode supports {sorted(SUPPORTED_3D_PROCEDURES)}, got '{procedure}'" + ) + + active_model = model or FlameModel.mediapipe_template() + fit = fit_flame_from_landmarks( + faces, + model=active_model, + n_iters=fit_iters, + ) + + eye_ids = active_model.available_mediapipe_indices([33, 263]) + if len(eye_ids) == 2: + eye_points = active_model.sample_mediapipe_points( + fit.vertices, + np.array([33, 263], dtype=np.int64), + ) + ipd = float(np.linalg.norm(eye_points[1] - eye_points[0])) + else: + ipd = 63.0 + face_scale = ipd if ipd > 1e-3 else 63.0 + + deformed, _handles = apply_canonical_deformation( + fit.vertices, + fit.faces, + procedure=procedure, + intensity=intensity, + lambda_smooth=lambda_smooth, + face_scale_mm=face_scale, + landmark_vertex_indices=fit.landmark_vertex_indices, + ) + + if not _handles: + raise ValueError( + "The supplied MediaPipe-to-FLAME embedding does not cover any " + f"surgical handles required for '{procedure}'" + ) + + manipulated: list[FaceLandmarks] = [] + for i, face in enumerate(faces): + projected = project_model_landmarks( + deformed, + fit.cameras[i], + active_model, + source_face=face, + ) + manipulated.append(projected) + + return Shared3DSurgeryResult( + fit=fit, + vertices_before=fit.vertices.copy(), + vertices_after=deformed, + faces=fit.faces.copy(), + manipulated_faces=manipulated, + mode="shared_3d", + ) + + +def apply_surgery_landmarks( + faces: Sequence[FaceLandmarks] | FaceLandmarks, + procedure: str, + intensity: float = 50.0, + mode: str = "shared_3d", + **kwargs, +) -> list[FaceLandmarks]: + """Unified entry point for Independent 2D vs Shared 3D landmark surgery. + + Args: + faces: Single FaceLandmarks or a sequence of views. + procedure: Procedure name. + intensity: 0-100. + mode: ``"independent_2d"`` or ``"shared_3d"``. + **kwargs: Forwarded to the underlying implementation. + + Returns: + List of manipulated FaceLandmarks (one per input view). + """ + if isinstance(faces, FaceLandmarks): + face_list: list[FaceLandmarks] = [faces] + else: + face_list = list(faces) + + mode = mode.lower().strip() + if mode in ("independent_2d", "2d", "independent"): + return [ + apply_procedure_preset( + f, + procedure, + intensity, + **{ + k: v + for k, v in kwargs.items() + if k + in { + "image_size", + "clinical_flags", + "displacement_model_path", + "noise_scale", + "regional_intensity", + } + }, + ) + for f in face_list + ] + + if mode in ("shared_3d", "3d", "shared"): + if procedure.lower().strip() not in SUPPORTED_3D_PROCEDURES: + raise ValueError( + f"Shared 3D mode supports {sorted(SUPPORTED_3D_PROCEDURES)}, got '{procedure}'" + ) + result = apply_shared_3d_surgery( + face_list, + procedure=procedure, + intensity=intensity, + model=kwargs.get("model"), + lambda_smooth=float(kwargs.get("lambda_smooth", 30.0)), + fit_iters=int(kwargs.get("fit_iters", 6)), + ) + return result.manipulated_faces + + raise ValueError(f"Unknown surgery mode '{mode}'. Use 'independent_2d' or 'shared_3d'.") diff --git a/scripts/app.py b/scripts/app.py index 620c3b7..dc308b4 100644 --- a/scripts/app.py +++ b/scripts/app.py @@ -22,6 +22,7 @@ ) from landmarkdiff.manipulation import apply_procedure_preset from landmarkdiff.masking import generate_surgical_mask +from landmarkdiff.surgery3d import SUPPORTED_3D_PROCEDURES, apply_surgery_landmarks from landmarkdiff.synthetic.tps_warp import warp_image_tps @@ -551,8 +552,9 @@ def analyze_face(image_rgb): "Upload photos from multiple angles for more accurate 3D-aware prediction.\n" "Follow the guide: **front -> left 45 -> right 45" " -> left profile -> right profile**.\n\n" - "The system validates each angle using landmark-based pose estimation " - "and combines all views for the final result." + "Choose **Shared 3D** to fit one patient-specific canonical mesh, apply a " + "single surgical deformation, and reproject it to every view — instead of " + "running an independent 2D preset per photo." ) with gr.Row(): @@ -581,6 +583,17 @@ def analyze_face(image_rgb): value="rhinoplasty", label="Procedure", ) + ma_mode = gr.Radio( + choices=["Independent 2D", "Shared 3D"], + value="Shared 3D", + label="Surgery mode", + info=( + "Independent 2D: each view gets its own 2D preset. " + "Shared 3D: one canonical mesh deformation, reprojected " + "to every view (supports: " + f"{', '.join(sorted(SUPPORTED_3D_PROCEDURES))})." + ), + ) ma_intensity = gr.Slider( minimum=0, maximum=100, @@ -670,7 +683,14 @@ def validate_all_angles(front, left45, right45, left_prof, right_prof): return "\n".join(lines) def generate_multi_angle( - front, left45, right45, left_prof, right_prof, procedure, intensity + front, + left45, + right45, + left_prof, + right_prof, + procedure, + intensity, + mode_label, ): """Generate surgical outcome prediction for each uploaded angle.""" images = [ @@ -680,32 +700,69 @@ def generate_multi_angle( (left_prof, "Left Profile"), (right_prof, "Right Profile"), ] - results = [] - gallery_items = [] + mode = ( + "shared_3d" + if str(mode_label).lower().startswith("shared") + else "independent_2d" + ) + + if ( + mode == "shared_3d" + and str(procedure).lower().strip() not in SUPPORTED_3D_PROCEDURES + ): + raise gr.Error( + "Shared 3D currently supports: " + + ", ".join(sorted(SUPPORTED_3D_PROCEDURES)) + + ". Switch to Independent 2D for this procedure." + ) + # Extract landmarks for all uploaded views first so Shared 3D + # can fit one canonical mesh across them. + prepared: list[tuple[np.ndarray | None, object | None, str]] = [] + face_list = [] for img_rgb, label in images: if img_rgb is None: - results.append(None) + prepared.append((None, None, label)) continue image_bgr = img_rgb[:, :, ::-1].copy() image_bgr = cv2.resize(image_bgr, (512, 512)) face = extract_landmarks(image_bgr) + prepared.append((image_bgr, face, label)) + if face is not None: + face_list.append(face) + + manip_by_slot: dict[int, object] = {} + if face_list: + manipulated = apply_surgery_landmarks( + face_list, + procedure, + float(intensity), + mode=mode, + ) + manip_iter = iter(manipulated) + for prep_i, (_img, face, _label) in enumerate(prepared): + if face is not None: + manip_by_slot[prep_i] = next(manip_iter) + + results = [] + gallery_items = [] + for prep_i, (image_bgr, face, label) in enumerate(prepared): + if image_bgr is None: + results.append(None) + continue if face is None: results.append(bgr_to_rgb(image_bgr)) continue - manip = apply_procedure_preset( - face, procedure, float(intensity), image_size=512 - ) + manip = manip_by_slot[prep_i] mask = generate_surgical_mask(face, procedure, 512, 512) warped = warp_image_tps(image_bgr, face.pixel_coords, manip.pixel_coords) comp = mask_composite(warped, image_bgr, mask) result_rgb = bgr_to_rgb(comp) results.append(result_rgb) gallery_items.append((bgr_to_rgb(image_bgr), f"{label} (Before)")) - gallery_items.append((result_rgb, f"{label} (After)")) + gallery_items.append((result_rgb, f"{label} ({mode})")) - # Pad results to 5 entries while len(results) < 5: results.append(None) @@ -720,7 +777,7 @@ def generate_multi_angle( ) ma_generate_btn.click( fn=generate_multi_angle, - inputs=ma_angle_inputs + [ma_procedure, ma_intensity], + inputs=ma_angle_inputs + [ma_procedure, ma_intensity, ma_mode], outputs=[ ma_out_front, ma_out_left45, diff --git a/scripts/benchmark_surgery3d.py b/scripts/benchmark_surgery3d.py new file mode 100644 index 0000000..d5d9b32 --- /dev/null +++ b/scripts/benchmark_surgery3d.py @@ -0,0 +1,637 @@ +#!/usr/bin/env python3 +"""Held-out-view benchmark: Independent 2D vs Shared 3D surgery. + +This benchmark deliberately avoids generating ground truth with the same +Laplacian operator used by the Shared 3D method. + +Protocol: +1. Build a canonical synthetic face M. +2. Apply an independent analytic 3D deformation field to obtain M_gt. +3. Project the *pre-op* face at multiple known yaw angles. +4. Fit Shared 3D using only the training views. +5. Estimate the held-out camera from the held-out *pre-op* view. +6. Predict the held-out post-op geometry without using held-out post-op GT. +7. Compare against Independent 2D on the exact same held-out input. + +Primary metric: +- held-out reprojection NME (lower is better) + +Secondary diagnostics: +- canonical vertex RMSE against independently generated GT +- camera rotation error on the held-out pre-op view +- profile/front displacement ratio for the chin sanity check + +Usage: + python scripts/benchmark_surgery3d.py --out artifacts/surgery3d_benchmark +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from landmarkdiff.deformation3d import apply_canonical_deformation +from landmarkdiff.flame_fitting import ( + FIT_LANDMARK_INDICES, + CameraParams, + FlameModel, + estimate_weak_perspective_camera, + project_landmarks, +) +from landmarkdiff.manipulation import apply_procedure_preset +from landmarkdiff.surgery3d import apply_shared_3d_surgery + + +def _yaw_rotation(degrees: float) -> np.ndarray: + rad = np.deg2rad(degrees) + c, s = np.cos(rad), np.sin(rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def _rotation_error_deg(pred: np.ndarray, gt: np.ndarray) -> float: + rel = pred @ gt.T + cos_angle = np.clip((np.trace(rel) - 1.0) / 2.0, -1.0, 1.0) + return float(np.degrees(np.arccos(cos_angle))) + + +def project_mesh( + vertices: np.ndarray, + yaw_deg: float, + scale: float = 3.0, + size: int = 512, + landmark_vertex_indices: np.ndarray | None = None, +): + cam = CameraParams( + rotation=_yaw_rotation(yaw_deg), + translation=np.array([size / 2.0, size / 2.0]), + scale=scale, + image_width=size, + image_height=size, + ) + face = project_landmarks( + vertices, + cam, + image_width=size, + image_height=size, + landmark_vertex_indices=landmark_vertex_indices, + ) + return face, cam + + +def nme(pred: np.ndarray, gt: np.ndarray, norm: float) -> float: + return float(np.mean(np.linalg.norm(pred - gt, axis=1)) / max(norm, 1e-6)) + + +def vertex_rmse(a: np.ndarray, b: np.ndarray) -> float: + return float(np.sqrt(np.mean(np.sum((a - b) ** 2, axis=1)))) + + +def _gaussian_field( + vertices: np.ndarray, + center: np.ndarray, + sigma_mm: float, +) -> np.ndarray: + d2 = np.sum((vertices - center[None, :]) ** 2, axis=1) + return np.exp(-d2 / (2.0 * sigma_mm**2)) + + +def independent_ground_truth_deformation( + vertices: np.ndarray, + procedure: str, + intensity: float, + landmark_vertex_indices: np.ndarray | None = None, +) -> np.ndarray: + """Create independent analytic 3D ground truth on any supported topology.""" + mapping = ( + np.arange(478, dtype=np.int64) + if landmark_vertex_indices is None + else np.asarray(landmark_vertex_indices, dtype=np.int64) + ) + if mapping.shape != (478,): + raise ValueError("landmark_vertex_indices must have shape (478,)") + + def vertex(mp_id: int) -> np.ndarray: + return vertices[int(mapping[mp_id])] + + ipd = float(np.linalg.norm(vertex(263) - vertex(33))) + unit_per_mm = ipd / 63.0 if ipd > 1e-8 else 1.0 + + out = vertices.copy() + s = float(np.clip(intensity, 0.0, 100.0)) / 100.0 + proc = procedure.lower().strip() + + if proc in {"mentoplasty", "genioplasty"}: + center = 0.5 * (vertex(152) + vertex(175)) + w = _gaussian_field(vertices, center, sigma_mm=26.0 * unit_per_mm) + out[:, 2] += 6.0 * unit_per_mm * s * w + out[:, 1] -= 0.5 * unit_per_mm * s * w + if proc == "genioplasty": + jaw_center = 0.5 * (vertex(172) + vertex(397)) + w_jaw = _gaussian_field( + vertices, + jaw_center, + sigma_mm=34.0 * unit_per_mm, + ) + out[:, 2] += 1.5 * unit_per_mm * s * w_jaw + return out + + if proc == "rhinoplasty": + tip = vertex(1) + w_tip = _gaussian_field( + vertices, + tip, + sigma_mm=18.0 * unit_per_mm, + ) + out[:, 2] += 2.0 * unit_per_mm * s * w_tip + out[:, 1] += 1.5 * unit_per_mm * s * w_tip + + left = vertex(240) + right = vertex(460) + w_left = _gaussian_field( + vertices, + left, + sigma_mm=14.0 * unit_per_mm, + ) + w_right = _gaussian_field( + vertices, + right, + sigma_mm=14.0 * unit_per_mm, + ) + out[:, 0] += 2.5 * unit_per_mm * s * w_left + out[:, 0] -= 2.5 * unit_per_mm * s * w_right + + dorsum = vertex(6) + w_dorsum = _gaussian_field( + vertices, + dorsum, + sigma_mm=16.0 * unit_per_mm, + ) + out[:, 2] -= 1.0 * unit_per_mm * s * w_dorsum + return out + + raise ValueError( + "Held-out benchmark GT is defined for mentoplasty, genioplasty, and " + f"rhinoplasty; got {procedure}" + ) + + +def run_benchmark( + procedure: str = "mentoplasty", + intensity: float = 80.0, + heldout_yaw: float = 80.0, + model: FlameModel | None = None, +) -> dict: + model = model or FlameModel.mediapipe_template(n_shape=5) + mapping = model.vertex_indices_for_mediapipe(np.arange(478)) + mesh_before = model.mean_vertices.copy() + mesh_gt = independent_ground_truth_deformation( + mesh_before, + procedure, + intensity, + landmark_vertex_indices=mapping, + ) + + train_yaws = [0.0, 30.0, 60.0] + train_faces = [ + project_mesh( + mesh_before, + yaw, + landmark_vertex_indices=mapping, + )[0] + for yaw in train_yaws + ] + heldout_before, heldout_cam_gt = project_mesh( + mesh_before, + heldout_yaw, + landmark_vertex_indices=mapping, + ) + heldout_after_gt, _ = project_mesh( + mesh_gt, + heldout_yaw, + landmark_vertex_indices=mapping, + ) + + # Isolate the surgery operator from the fitting problem: apply the same + # canonical 3D surgery directly to the known pre-op mesh, then project with + # the known held-out camera. This is not the final method score; it tells us + # whether a failure comes from canonical fitting or from the deformation + # operator itself. + reference_ipd = float( + np.linalg.norm(mesh_before[int(mapping[263])] - mesh_before[int(mapping[33])]) + ) + reference_unit_per_mm = reference_ipd / 63.0 if reference_ipd > 1e-8 else 1.0 + direct_mesh_after, _ = apply_canonical_deformation( + mesh_before, + model.faces, + procedure, + intensity, + face_scale_mm=reference_ipd if reference_ipd > 1e-8 else 63.0, + landmark_vertex_indices=mapping, + ) + direct_heldout = project_landmarks( + direct_mesh_after, + heldout_cam_gt, + image_width=heldout_before.image_width, + image_height=heldout_before.image_height, + source_z=heldout_before.landmarks[:, 2], + landmark_vertex_indices=mapping, + ) + + # Fit and deform using training views only. + shared = apply_shared_3d_surgery( + train_faces, + procedure=procedure, + intensity=intensity, + model=model, + fit_iters=8, + ) + + # The held-out *pre-op* image may be used to estimate its camera. No held-out + # post-op landmark or target is used in fitting or deformation. + heldout_fit_indices = ( + np.asarray(FIT_LANDMARK_INDICES, dtype=np.int64) + if model.backend == "flame" + else np.arange(478, dtype=np.int64) + ) + heldout_cam_pred = estimate_weak_perspective_camera( + shared.vertices_before, + heldout_before.pixel_coords, + weights=heldout_before.landmark_confidence, + image_width=heldout_before.image_width, + image_height=heldout_before.image_height, + indices=heldout_fit_indices, + vertex_indices=model.vertex_indices_for_mediapipe(heldout_fit_indices), + ) + shared_heldout = project_landmarks( + shared.vertices_after, + heldout_cam_pred, + image_width=heldout_before.image_width, + image_height=heldout_before.image_height, + source_z=heldout_before.landmarks[:, 2], + landmark_vertex_indices=mapping, + ) + + independent_heldout = apply_procedure_preset( + heldout_before, + procedure, + intensity, + ) + + ref_pts = train_faces[0].pixel_coords + face_diag = float(np.linalg.norm(ref_pts.max(axis=0) - ref_pts.min(axis=0))) + + shared_nme = nme( + shared_heldout.pixel_coords, + heldout_after_gt.pixel_coords, + face_diag, + ) + independent_nme = nme( + independent_heldout.pixel_coords, + heldout_after_gt.pixel_coords, + face_diag, + ) + + direct_deformation_nme = nme( + direct_heldout.pixel_coords, + heldout_after_gt.pixel_coords, + face_diag, + ) + fit_before_rmse_mm = vertex_rmse(shared.vertices_before, mesh_before) / reference_unit_per_mm + deformation_delta_rmse_mm = ( + vertex_rmse( + shared.vertices_after - shared.vertices_before, + mesh_gt - mesh_before, + ) + / reference_unit_per_mm + ) + direct_deformation_rmse_mm = vertex_rmse(direct_mesh_after, mesh_gt) / reference_unit_per_mm + + canonical_rmse_units = vertex_rmse(shared.vertices_after, mesh_gt) + fitted_ipd = float( + np.linalg.norm( + shared.vertices_before[int(mapping[263])] - shared.vertices_before[int(mapping[33])] + ) + ) + unit_per_mm = fitted_ipd / 63.0 if fitted_ipd > 1e-8 else 1.0 + canonical_rmse_mm = canonical_rmse_units / unit_per_mm + camera_error = _rotation_error_deg(heldout_cam_pred.rotation, heldout_cam_gt.rotation) + + chin = 152 + front_before, front_cam = project_mesh( + mesh_before, + 0.0, + landmark_vertex_indices=mapping, + ) + front_shared = project_landmarks( + shared.vertices_after, + front_cam, + image_width=512, + image_height=512, + source_z=front_before.landmarks[:, 2], + landmark_vertex_indices=mapping, + ) + front_delta = float( + np.linalg.norm(front_shared.pixel_coords[chin] - front_before.pixel_coords[chin]) + ) + profile_delta = float( + np.linalg.norm(shared_heldout.pixel_coords[chin] - heldout_before.pixel_coords[chin]) + ) + + return { + "backend": model.backend, + "protocol": { + "train_yaws_deg": train_yaws, + "heldout_yaw_deg": heldout_yaw, + "heldout_postop_used_for_fit": False, + "gt_operator": "independent_gaussian_3d_field", + }, + "procedure": procedure, + "intensity": intensity, + "shared_3d": { + "heldout_reprojection_nme": shared_nme, + "canonical_vertex_rmse_mm": canonical_rmse_mm, + "fit_before_rmse_mm": fit_before_rmse_mm, + "deformation_delta_rmse_mm": deformation_delta_rmse_mm, + "heldout_camera_rotation_error_deg": camera_error, + "chin_front_px": front_delta, + "chin_heldout_px": profile_delta, + "chin_heldout_over_front": profile_delta / max(front_delta, 1e-6), + }, + "independent_2d": { + "heldout_reprojection_nme": independent_nme, + }, + "deformation_only_oracle_geometry": { + "heldout_reprojection_nme": direct_deformation_nme, + "canonical_vertex_rmse_mm": direct_deformation_rmse_mm, + }, + "delta": { + "heldout_nme_shared_minus_independent": shared_nme - independent_nme, + "heldout_nme_relative_change": ( + (shared_nme - independent_nme) / max(independent_nme, 1e-12) + ), + }, + } + + +def render_demo_grid( + out_dir: Path, + procedure: str = "mentoplasty", + intensity: float = 80.0, + model: FlameModel | None = None, +) -> Path: + """Save a qualitative front/45/profile comparison and canonical 3D edit.""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + model = model or FlameModel.mediapipe_template(n_shape=5) + mapping = model.vertex_indices_for_mediapipe(np.arange(478)) + mesh = model.mean_vertices.copy() + ipd = float(np.linalg.norm(mesh[int(mapping[263])] - mesh[int(mapping[33])])) + mesh_after, _ = apply_canonical_deformation( + mesh, + model.faces, + procedure, + intensity, + face_scale_mm=ipd if ipd > 1e-8 else 63.0, + landmark_vertex_indices=mapping, + ) + + yaws = [0.0, 45.0, 80.0] + labels = ["FRONT", "45°", "PROFILE"] + faces_before = [ + project_mesh( + mesh, + y, + landmark_vertex_indices=mapping, + )[0] + for y in yaws + ] + + shared = apply_shared_3d_surgery( + faces_before, + procedure=procedure, + intensity=intensity, + model=model, + fit_iters=8, + ) + indep = [apply_procedure_preset(f, procedure, intensity) for f in faces_before] + + fig, axes = plt.subplots(3, 3, figsize=(10, 10), constrained_layout=True) + fig.suptitle( + f"3D Surgical Bridge — {procedure} @ {intensity:.0f}\n" + "Row1: input · Row2: Independent 2D · Row3: Shared 3D", + fontsize=12, + ) + for col, (face, label) in enumerate(zip(faces_before, labels, strict=True)): + rows = [ + face.pixel_coords, + indep[col].pixel_coords, + shared.manipulated_faces[col].pixel_coords, + ] + for row, pts in enumerate(rows): + ax = axes[row, col] + ax.scatter(pts[:, 0], pts[:, 1], s=2, c="0.3") + ax.scatter(pts[152, 0], pts[152, 1], s=40, c="C3") + ax.scatter(pts[1, 0], pts[1, 1], s=40, c="C0") + ax.set_xlim(0, 512) + ax.set_ylim(512, 0) + ax.set_aspect("equal") + ax.set_xticks([]) + ax.set_yticks([]) + if row == 0: + ax.set_title(label) + if col == 0: + ax.set_ylabel(["Input", "Indep. 2D", "Shared 3D"][row]) + + grid_path = out_dir / f"demo_grid_{procedure}.png" + fig.savefig(grid_path, dpi=140) + plt.close(fig) + + fig = plt.figure(figsize=(10, 4.5)) + for i, (verts, title) in enumerate([(mesh, "Before"), (mesh_after, "After")]): + ax = fig.add_subplot(1, 2, i + 1, projection="3d") + ax.scatter(verts[:, 0], verts[:, 2], verts[:, 1], s=2, c=verts[:, 2], cmap="coolwarm") + if i == 1: + delta = mesh_after - mesh + mag = np.linalg.norm(delta, axis=1) + idx = np.where(mag > 0.4)[0][::3] + ax.quiver( + mesh[idx, 0], + mesh[idx, 2], + mesh[idx, 1], + delta[idx, 0], + delta[idx, 2], + delta[idx, 1], + length=1.0, + normalize=False, + color="k", + linewidth=0.6, + arrow_length_ratio=0.2, + ) + ax.set_title(title) + ax.set_xlabel("X") + ax.set_ylabel("Z (ant.)") + ax.set_zlabel("Y") + + mesh_path = out_dir / f"demo_mesh_{procedure}.png" + fig.tight_layout() + fig.savefig(mesh_path, dpi=140) + plt.close(fig) + return grid_path + + +def render_displacement_panel( + out_dir: Path, + procedure: str = "mentoplasty", + intensity: float = 80.0, + model: FlameModel | None = None, +) -> Path: + """Visualize the actual view-dependent projection of one shared 3D edit.""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + model = model or FlameModel.mediapipe_template(n_shape=5) + mapping = model.vertex_indices_for_mediapipe(np.arange(478)) + mesh = model.mean_vertices.copy() + yaws = [0.0, 45.0, 80.0] + labels = ["FRONT", "45°", "PROFILE"] + faces_before = [ + project_mesh( + mesh, + yaw, + landmark_vertex_indices=mapping, + )[0] + for yaw in yaws + ] + shared = apply_shared_3d_surgery( + faces_before, + procedure=procedure, + intensity=intensity, + model=model, + fit_iters=8, + ) + + fig, axes = plt.subplots(1, 3, figsize=(12, 4), constrained_layout=True) + fig.suptitle( + f"One canonical 3D edit, reprojected by view — {procedure} @ {intensity:.0f}", + fontsize=12, + ) + + for ax, label, before, after in zip( + axes, + labels, + faces_before, + shared.manipulated_faces, + strict=True, + ): + p0 = before.pixel_coords + p1 = after.pixel_coords + delta = p1 - p0 + magnitude = np.linalg.norm(delta, axis=1) + threshold = max(float(np.quantile(magnitude, 0.85)), 0.05) + moving = magnitude >= threshold + + ax.scatter( + p0[:, 0], + p0[:, 1], + s=5, + alpha=0.30, + label="before", + ) + ax.scatter( + p1[moving, 0], + p1[moving, 1], + s=12, + alpha=0.85, + label="after (moving region)", + ) + ax.quiver( + p0[moving, 0], + p0[moving, 1], + delta[moving, 0], + delta[moving, 1], + angles="xy", + scale_units="xy", + scale=1.0, + width=0.004, + ) + ax.set_title(f"{label}\nmax displacement={float(magnitude.max()):.1f}px") + ax.set_xlim(0, 512) + ax.set_ylim(512, 0) + ax.set_aspect("equal") + ax.set_xticks([]) + ax.set_yticks([]) + + axes[0].legend(loc="upper left", fontsize=8) + path = out_dir / f"demo_displacement_{procedure}.png" + fig.savefig(path, dpi=160) + plt.close(fig) + return path + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--out", + type=Path, + default=Path("artifacts/surgery3d_benchmark"), + ) + parser.add_argument("--procedure", default="mentoplasty") + parser.add_argument("--intensity", type=float, default=80.0) + parser.add_argument("--heldout-yaw", type=float, default=80.0) + parser.add_argument("--flame-model", type=Path) + parser.add_argument("--correspondence", type=Path) + args = parser.parse_args() + args.out.mkdir(parents=True, exist_ok=True) + + if bool(args.flame_model) != bool(args.correspondence): + parser.error("--flame-model and --correspondence must be provided together") + + model = ( + FlameModel.from_flame_assets( + args.flame_model, + args.correspondence, + ) + if args.flame_model is not None + else FlameModel.mediapipe_template(n_shape=5) + ) + + results = run_benchmark( + procedure=args.procedure, + intensity=args.intensity, + heldout_yaw=args.heldout_yaw, + model=model, + ) + json_path = args.out / "metrics.json" + json_path.write_text(json.dumps(results, indent=2)) + print(json.dumps(results, indent=2)) + + grid = render_demo_grid( + args.out, + procedure=args.procedure, + intensity=args.intensity, + model=model, + ) + displacement = render_displacement_panel( + args.out, + procedure=args.procedure, + intensity=args.intensity, + model=model, + ) + print(f"Wrote {json_path}") + print(f"Wrote {grid}") + print(f"Wrote {displacement}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_surgery3d.py b/tests/test_surgery3d.py new file mode 100644 index 0000000..1bdb385 --- /dev/null +++ b/tests/test_surgery3d.py @@ -0,0 +1,485 @@ +"""Tests for canonical/FLAME fitting, 3D deformation, and the Shared 3D bridge.""" + +from __future__ import annotations + +import json +import pickle + +import numpy as np +import pytest + +from landmarkdiff.deformation3d import ( + SurgicalHandle3D, + apply_canonical_deformation, + get_surgical_handles, + laplacian_deform, +) +from landmarkdiff.flame_fitting import ( + CameraParams, + FlameModel, + build_mesh_laplacian, + estimate_weak_perspective_camera, + fit_flame_from_landmarks, + landmarks_to_metric_vertices, + project_landmarks, + project_model_landmarks, +) +from landmarkdiff.landmarks import FaceLandmarks +from landmarkdiff.surgery3d import ( + SUPPORTED_3D_PROCEDURES, + apply_shared_3d_surgery, + apply_surgery_landmarks, +) + + +def _yaw_rotation(degrees: float) -> np.ndarray: + rad = np.deg2rad(degrees) + c, s = np.cos(rad), np.sin(rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def _synthetic_face_from_mesh( + vertices: np.ndarray, + yaw_deg: float = 0.0, + image_size: int = 512, + scale: float = 3.0, +) -> FaceLandmarks: + """Project a canonical mesh with a yawed weak-perspective camera.""" + rotation = _yaw_rotation(yaw_deg) + cam = CameraParams( + rotation=rotation, + translation=np.array([image_size / 2.0, image_size / 2.0]), + scale=scale, + image_width=image_size, + image_height=image_size, + ) + return project_landmarks(vertices, cam, image_width=image_size, image_height=image_size) + + +@pytest.fixture +def flame_model(): + return FlameModel.mediapipe_template(n_shape=5) + + +@pytest.fixture +def multi_view_faces(flame_model): + verts = flame_model.mean_vertices + yaws = [0.0, 30.0, 60.0, 90.0] + return [_synthetic_face_from_mesh(verts, yaw) for yaw in yaws] + + +class TestFlameFitting: + def test_official_flame_asset_loader_uses_real_shape_basis(self, tmp_path): + """External backend must preserve supplied FLAME topology and shapedirs.""" + n_vertices = 478 + model_path = tmp_path / "flame.pkl" + corr_path = tmp_path / "mediapipe_to_flame.json" + + payload = { + "v_template": np.zeros((n_vertices, 3), dtype=np.float64), + "f": np.array([[0, 1, 2], [2, 3, 4]], dtype=np.int32), + "shapedirs": np.ones((n_vertices, 3, 6), dtype=np.float64), + } + with open(model_path, "wb") as fh: + pickle.dump(payload, fh) + corr_path.write_text(json.dumps({str(i): i for i in range(478)})) + + model = FlameModel.from_flame_assets( + model_path, + corr_path, + n_shape=4, + ) + assert model.backend == "flame" + assert model.shape_basis is not None + assert model.shape_basis.shape == (478, 3, 4) + assert np.array_equal( + model.vertex_indices_for_mediapipe([1, 152, 263]), + np.array([1, 152, 263]), + ) + + def test_official_flame_asset_loader_accepts_standard_barycentric_npz( + self, + tmp_path, + ): + model_path = tmp_path / "flame.pkl" + corr_path = tmp_path / "mediapipe_to_flame.npz" + + mean = np.array( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [2.0, 0.0, 0.0], + [2.0, 1.0, 0.0], + [3.0, 0.0, 0.0], + ], + dtype=np.float64, + ) + faces = np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int32) + payload = { + "v_template": mean, + "f": faces, + "shapedirs": np.ones((6, 3, 3), dtype=np.float64), + } + with open(model_path, "wb") as fh: + pickle.dump(payload, fh) + + mp_ids = np.array([1, 33, 152, 263], dtype=np.int64) + face_ids = np.array([0, 0, 1, 1], dtype=np.int64) + bary = np.array( + [ + [0.2, 0.7, 0.1], + [0.1, 0.2, 0.7], + [0.1, 0.8, 0.1], + [0.2, 0.3, 0.5], + ], + dtype=np.float64, + ) + np.savez( + corr_path, + landmark_indices=mp_ids, + lmk_face_idx=face_ids, + lmk_b_coords=bary, + ) + + model = FlameModel.from_flame_assets( + model_path, + corr_path, + n_shape=2, + ) + assert model.backend == "flame" + assert np.array_equal( + model.available_mediapipe_indices([1, 2, 33, 152, 263]), + np.array([1, 33, 152, 263]), + ) + assert model.vertex_indices_for_mediapipe([1])[0] == 1 + + expected = np.sum(mean[faces[0]] * bary[0, :, None], axis=0) + sampled = model.sample_mediapipe_points(mean, [1])[0] + assert np.allclose(sampled, expected) + + camera = CameraParams( + rotation=np.eye(3), + translation=np.array([256.0, 256.0]), + scale=50.0, + image_width=512, + image_height=512, + ) + observed = np.full((478, 3), 0.5, dtype=np.float32) + sampled_all = model.sample_mediapipe_points(mean, mp_ids) + pixels = camera.project(sampled_all) + observed[mp_ids, 0] = pixels[:, 0] / 512.0 + observed[mp_ids, 1] = pixels[:, 1] / 512.0 + face = FaceLandmarks( + landmarks=observed, + image_width=512, + image_height=512, + confidence=1.0, + ) + fit = fit_flame_from_landmarks( + [face], + model=model, + n_iters=1, + fit_indices=mp_ids.tolist(), + ) + assert np.isfinite(fit.reprojection_errors[0]) + assert fit.reprojection_errors[0] < 1e-5 + + with pytest.raises(ValueError, match="No mesh correspondence"): + model.vertex_indices_for_mediapipe([2]) + + def test_partial_flame_projection_preserves_unmapped_mediapipe_points( + self, + tmp_path, + ): + model_path = tmp_path / "flame.pkl" + corr_path = tmp_path / "mediapipe_to_flame.npz" + + mean = np.array( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [2.0, 0.0, 0.0], + [2.0, 1.0, 0.0], + [3.0, 0.0, 0.0], + ], + dtype=np.float64, + ) + payload = { + "v_template": mean, + "f": np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int32), + "shapedirs": np.zeros((6, 3, 2), dtype=np.float64), + } + with open(model_path, "wb") as fh: + pickle.dump(payload, fh) + + np.savez( + corr_path, + landmark_indices=np.array([1, 33, 152, 263], dtype=np.int64), + lmk_face_idx=np.array([0, 0, 1, 1], dtype=np.int64), + lmk_b_coords=np.array( + [ + [0.2, 0.7, 0.1], + [0.1, 0.2, 0.7], + [0.1, 0.8, 0.1], + [0.2, 0.3, 0.5], + ], + dtype=np.float64, + ), + ) + model = FlameModel.from_flame_assets(model_path, corr_path, n_shape=2) + + source_landmarks = np.full((478, 3), 0.25, dtype=np.float32) + source = FaceLandmarks( + landmarks=source_landmarks, + image_width=512, + image_height=512, + confidence=0.9, + ) + camera = CameraParams( + rotation=np.eye(3), + translation=np.array([256.0, 256.0]), + scale=50.0, + image_width=512, + image_height=512, + ) + + projected = project_model_landmarks( + mean, + camera, + model, + source_face=source, + ) + + assert np.array_equal(projected.landmarks[2], source.landmarks[2]) + assert not np.array_equal(projected.landmarks[1, :2], source.landmarks[1, :2]) + assert projected.landmarks[1, 2] == source.landmarks[1, 2] + assert projected.confidence == source.confidence + + def test_mediapipe_template_shapes(self, flame_model): + assert flame_model.mean_vertices.shape == (478, 3) + assert flame_model.faces.shape[1] == 3 + assert flame_model.faces.shape[0] > 100 + + @pytest.mark.parametrize("yaw_deg", [0.0, 30.0, 60.0, 80.0]) + def test_camera_recovers_out_of_plane_yaw(self, flame_model, yaw_deg): + """Known synthetic yaw must be recovered as a true 3D rotation.""" + verts = flame_model.mean_vertices + face = _synthetic_face_from_mesh(verts, yaw_deg=yaw_deg) + cam = estimate_weak_perspective_camera( + verts, + face.pixel_coords, + image_width=face.image_width, + image_height=face.image_height, + ) + gt = _yaw_rotation(yaw_deg) + rel = cam.rotation @ gt.T + cos_angle = np.clip((np.trace(rel) - 1.0) / 2.0, -1.0, 1.0) + error_deg = float(np.degrees(np.arccos(cos_angle))) + assert error_deg < 2.0 + + def test_projection_preserves_anatomical_vertical_orientation(self, flame_model): + """Canonical forehead (+Y) must appear above the chin in image coordinates.""" + face = _synthetic_face_from_mesh( + flame_model.mean_vertices, + yaw_deg=0.0, + ) + forehead_y = float(face.pixel_coords[10, 1]) + chin_y = float(face.pixel_coords[152, 1]) + assert forehead_y < chin_y + + def test_single_view_fit(self, flame_model, multi_view_faces): + fit = fit_flame_from_landmarks([multi_view_faces[0]], model=flame_model, n_iters=4) + assert fit.vertices.shape == (478, 3) + assert len(fit.cameras) == 1 + assert fit.reprojection_errors[0] < 25.0 # pixels, synthetic + + def test_multi_view_shared_vertices(self, flame_model, multi_view_faces): + fit = fit_flame_from_landmarks(multi_view_faces[:3], model=flame_model, n_iters=5) + assert len(fit.cameras) == 3 + assert fit.vertices.shape == (478, 3) + # All views should get finite reprojection error + assert all(np.isfinite(e) for e in fit.reprojection_errors) + + def test_obj_export(self, flame_model, multi_view_faces, tmp_path): + fit = fit_flame_from_landmarks([multi_view_faces[0]], model=flame_model, n_iters=2) + out = fit.to_obj(tmp_path / "face.obj") + text = out.read_text() + assert text.count("\nv ") + (1 if "\nv " in text else 0) >= 1 + assert "v " in text and "f " in text + + def test_colored_ply_export_samples_input_view(self, flame_model, multi_view_faces, tmp_path): + fit = fit_flame_from_landmarks( + [multi_view_faces[0]], + model=flame_model, + n_iters=2, + ) + image = np.zeros((512, 512, 3), dtype=np.uint8) + image[:, :] = np.array([10, 20, 30], dtype=np.uint8) # BGR + out = fit.to_colored_ply( + tmp_path / "face_colored.ply", + image, + ) + text = out.read_text() + assert "property uchar red" in text + assert "property uchar green" in text + assert "property uchar blue" in text + assert "30 20 10" in text + + def test_landmarks_to_metric_vertices(self, multi_view_faces): + verts = landmarks_to_metric_vertices(multi_view_faces[0]) + assert verts.shape == (478, 3) + # IPD should be near 63 mm + ipd = np.linalg.norm(verts[263] - verts[33]) + assert 40.0 < ipd < 90.0 + + +class TestDeformation3D: + def test_rhinoplasty_handles_nonzero(self): + handles = get_surgical_handles("rhinoplasty", intensity=70.0) + assert len(handles) > 0 + assert any(np.linalg.norm(h.displacement) > 0 for h in handles) + + def test_mentoplasty_primarily_anterior(self): + handles = get_surgical_handles("mentoplasty", intensity=100.0) + # Chin tip handle should have dominant +Z + chin = next(h for h in handles if h.vertex_index == 152) + assert chin.displacement[2] > abs(chin.displacement[0]) + assert chin.displacement[2] > abs(chin.displacement[1]) + + def test_laplacian_smooth_propagation(self, flame_model): + verts = flame_model.mean_vertices + handles = get_surgical_handles("mentoplasty", intensity=80.0) + deformed = laplacian_deform(verts, flame_model.faces, handles, lambda_smooth=10.0) + delta = deformed - verts + # Handle moves + assert np.linalg.norm(delta[152]) > 1.0 + # Neighbor of chin should move some, far forehead less + assert np.linalg.norm(delta[175]) > 0.1 + assert np.linalg.norm(delta[10]) < np.linalg.norm(delta[152]) + + def test_surgical_handles_respect_mesh_correspondence(self, flame_model): + mapping = np.arange(478, dtype=np.int64) + mapping[152], mapping[151] = mapping[151], mapping[152] + _, handles = apply_canonical_deformation( + flame_model.mean_vertices, + flame_model.faces, + procedure="mentoplasty", + intensity=80.0, + landmark_vertex_indices=mapping, + ) + mapped_chin = int(mapping[152]) + assert any(h.vertex_index == mapped_chin for h in handles) + + def test_matrix_free_laplacian_matches_dense_reference(self): + vertices = np.zeros((4, 3), dtype=np.float64) + faces = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.int32) + handles = [ + SurgicalHandle3D( + vertex_index=0, + displacement=np.array([1.0, 0.0, 0.0]), + weight=1.0, + ) + ] + + lambda_smooth = 2.0 + lambda_locality = 1e-2 + actual = laplacian_deform( + vertices, + faces, + handles, + lambda_smooth=lambda_smooth, + lambda_locality=lambda_locality, + ) + + laplacian = build_mesh_laplacian(faces, n_vertices=4) + weights = np.diag([1.0, 0.0, 0.0, 0.0]) + system = ( + weights + + lambda_smooth * (laplacian.T @ laplacian) + + (lambda_locality + 1e-4) * np.eye(4) + ) + rhs = np.array([1.0, 0.0, 0.0, 0.0]) + expected_x = np.linalg.solve(system, rhs) + + assert np.allclose(actual[:, 0], expected_x, atol=1e-7) + assert np.allclose(actual[:, 1:], 0.0, atol=1e-9) + + def test_apply_canonical_deformation(self, flame_model): + deformed, handles = apply_canonical_deformation( + flame_model.mean_vertices, + flame_model.faces, + procedure="rhinoplasty", + intensity=60.0, + ) + assert deformed.shape == flame_model.mean_vertices.shape + assert len(handles) > 0 + assert not np.allclose(deformed, flame_model.mean_vertices) + + def test_unsupported_procedure_raises(self, flame_model): + with pytest.raises(ValueError, match="not yet defined"): + apply_canonical_deformation( + flame_model.mean_vertices, + flame_model.faces, + procedure="blepharoplasty", + intensity=50.0, + ) + + +class TestSurgery3DBridge: + def test_shared_3d_returns_one_per_view(self, multi_view_faces): + result = apply_shared_3d_surgery( + multi_view_faces[:3], + procedure="mentoplasty", + intensity=70.0, + fit_iters=4, + ) + assert len(result.manipulated_faces) == 3 + assert result.vertices_after.shape == result.vertices_before.shape + assert not np.allclose(result.vertices_after, result.vertices_before) + + def test_mentoplasty_profile_moves_more_than_front(self, flame_model): + """Eye-test metric: same Δ³D → small front Δ, large profile Δ.""" + verts = flame_model.mean_vertices + front = _synthetic_face_from_mesh(verts, yaw_deg=0.0) + profile = _synthetic_face_from_mesh(verts, yaw_deg=90.0) + result = apply_shared_3d_surgery( + [front, profile], + procedure="mentoplasty", + intensity=100.0, + model=flame_model, + fit_iters=5, + ) + front_before = front.pixel_coords + front_after = result.manipulated_faces[0].pixel_coords + prof_before = profile.pixel_coords + prof_after = result.manipulated_faces[1].pixel_coords + + # Chin landmark 152 displacement magnitude in image space + front_d = float(np.linalg.norm(front_after[152] - front_before[152])) + prof_d = float(np.linalg.norm(prof_after[152] - prof_before[152])) + assert prof_d > front_d + + def test_independent_2d_mode(self, multi_view_faces): + out = apply_surgery_landmarks( + multi_view_faces[0], + procedure="rhinoplasty", + intensity=40.0, + mode="independent_2d", + ) + assert len(out) == 1 + assert out[0].landmarks.shape == (478, 3) + + def test_shared_mode_rejects_unsupported_procedure(self, multi_view_faces): + with pytest.raises(ValueError, match="Shared 3D mode supports"): + apply_surgery_landmarks( + multi_view_faces[0], + procedure="blepharoplasty", + intensity=40.0, + mode="shared_3d", + ) + + def test_supported_procedures_set(self): + assert "rhinoplasty" in SUPPORTED_3D_PROCEDURES + assert "mentoplasty" in SUPPORTED_3D_PROCEDURES