diff --git a/docs/api/manipulation.md b/docs/api/manipulation.md index 180af73..1141546 100644 --- a/docs/api/manipulation.md +++ b/docs/api/manipulation.md @@ -58,6 +58,44 @@ A single deformation control point. Each handle moves its target landmark by `di ## Functions +### `build_rbf_prewarp_handles` + +```python +def build_rbf_prewarp_handles( + face: FaceLandmarks, + procedure: str, + intensity: float = 50.0, + clinical_flags: ClinicalFlags | None = None, + regional_intensity: RegionalIntensity | None = None, +) -> list[DeformationHandle] +``` + +Builds the explicit NumPy RBF pre-warp handles for a procedure before any landmark deformation is applied. + +This function isolates policy (procedure/intensity/clinical flags) from execution so the deformation stage can be audited and tested independently. +It is intentionally confidence-agnostic. + +--- + +### `apply_rbf_prewarp_stage` + +```python +def apply_rbf_prewarp_stage( + face: FaceLandmarks, + procedure: str, + intensity: float = 50.0, + clinical_flags: ClinicalFlags | None = None, + regional_intensity: RegionalIntensity | None = None, +) -> FaceLandmarks +``` + +Runs the modular NumPy RBF deformation stage used before TPS warping. + +Equivalent to the legacy inline RBF path previously inside `apply_procedure_preset`, now extracted as a reusable stage. +Unlike `build_rbf_prewarp_handles(...)`, this stage applies landmark-confidence weighting before deformation. + +--- + ### `gaussian_rbf_deform` ```python @@ -117,7 +155,7 @@ The `intensity` parameter uses a 0-100 scale. Internally, it is divided by 100 ( | `face` | `FaceLandmarks` | (required) | Input face landmarks | | `procedure` | `str` | (required) | One of `"rhinoplasty"`, `"blepharoplasty"`, `"rhytidectomy"`, `"orthognathic"`, `"brow_lift"`, `"mentoplasty"` | | `intensity` | `float` | `50.0` | Deformation strength on a 0 to 100 scale. Mild ~ 33, moderate ~ 66, aggressive ~ 100. | -| `image_size` | `int` | `512` | Reference image size for displacement scaling. Displacements are calibrated at 512x512. | +| `image_size` | `int` | `512` | Legacy compatibility argument. The modular NumPy RBF path scales from `face.image_width`/`face.image_height` (geometric mean) and does not use this value directly. | | `clinical_flags` | `ClinicalFlags \| None` | `None` | Clinical condition flags (see [clinical](clinical.md)). Enables condition-specific handling: Ehlers-Danlos widens radii by 1.5x, Bell's palsy removes handles on the paralyzed side. | | `displacement_model_path` | `str \| None` | `None` | Path to a fitted `DisplacementModel` (`.npz`). When provided, uses data-driven displacements from real surgery pairs instead of hand-tuned RBF vectors. | | `noise_scale` | `float` | `0.0` | Random variation added to data-driven displacements (0 = deterministic). Only used when `displacement_model_path` is set. | diff --git a/landmarkdiff/manipulation.py b/landmarkdiff/manipulation.py index 41c9b84..b204fc0 100644 --- a/landmarkdiff/manipulation.py +++ b/landmarkdiff/manipulation.py @@ -630,75 +630,39 @@ def apply_combined_procedures( ) -def apply_procedure_preset( +def build_rbf_prewarp_handles( face: FaceLandmarks, procedure: str, intensity: float = 50.0, - image_size: int = 512, clinical_flags: ClinicalFlags | None = None, - displacement_model_path: str | None = None, - noise_scale: float = 0.0, regional_intensity: RegionalIntensity | None = None, -) -> FaceLandmarks: - """Apply a surgical procedure preset to landmarks. - - Args: - face: Input face landmarks. - procedure: Procedure name (rhinoplasty, blepharoplasty, etc.). - intensity: Relative intensity 0-100 (mild=33, moderate=66, aggressive=100). - image_size: Reference image size for displacement scaling. - clinical_flags: Optional clinical condition flags. - displacement_model_path: Path to a fitted DisplacementModel (.npz). - When provided, uses data-driven displacements from real surgery pairs - instead of hand-tuned RBF vectors. - noise_scale: Variation noise scale for data-driven mode (0=deterministic). +) -> list[DeformationHandle]: + """Build confidence-agnostic RBF handles for the TPS pre-warp stage. - Returns: - New FaceLandmarks with manipulated landmarks. + This function isolates procedure/intensity/clinical policy from the + deformation application step so the pre-warp stage can be tested and + reused independently of the full pipeline. """ if procedure not in PROCEDURE_LANDMARKS: raise ValueError(f"Unknown procedure: {procedure}. Choose from {list(PROCEDURE_LANDMARKS)}") - landmarks = face.landmarks.copy() scale = intensity / 100.0 - - # Data-driven displacement mode (fall back to RBF if procedure not in model) - # Map UI intensity (0-100) to displacement model intensity (0-2): - # 50 -> 1.0x mean displacement, matching inference.py scaling - if displacement_model_path is not None: - dm_scale = intensity / 50.0 - try: - return _apply_data_driven( - face, - procedure, - dm_scale, - displacement_model_path, - noise_scale, - ) - except KeyError: - logger.warning( - "Procedure '%s' not in displacement model, falling back to RBF preset", - procedure, - ) - # Fall through to RBF-based preset below - indices = PROCEDURE_LANDMARKS[procedure] radius = PROCEDURE_RADIUS[procedure] - # Ehlers-Danlos: wider influence radii for hypermobile tissue + # Ehlers-Danlos: wider influence radii for hypermobile tissue. if clinical_flags and clinical_flags.ehlers_danlos: radius *= 1.5 - # Scale radius based on geometric mean of actual image dimensions. - # Radii are calibrated for 512x512; using geometric mean handles - # non-square inputs without asymmetric deformation. + # Radii are calibrated at 512x512. Scale by geometric mean so + # non-square inputs are handled symmetrically. geo_mean = math.sqrt(face.image_width * face.image_height) pixel_scale = geo_mean / 512.0 handles = _get_procedure_handles( procedure, indices, scale, radius * pixel_scale, regional_intensity ) - # Bell's palsy: remove handles on the affected (paralyzed) side + # Bell's palsy: remove handles on the affected side. if clinical_flags and clinical_flags.bells_palsy: from landmarkdiff.clinical import get_bells_palsy_side_indices @@ -708,14 +672,19 @@ def apply_procedure_preset( affected_indices.update(region_indices) handles = [h for h in handles if h.landmark_index not in affected_indices] - # Convert to pixel space for deformation - pixel_landmarks = landmarks.copy() + return handles + + +def _apply_rbf_handles_to_face( + face: FaceLandmarks, + handles: list[DeformationHandle], +) -> FaceLandmarks: + """Apply a prepared list of handles to a face in pixel space.""" + pixel_landmarks = face.landmarks.copy() pixel_landmarks[:, 0] *= face.image_width pixel_landmarks[:, 1] *= face.image_height - # Scale each handle's displacement by the confidence of its anchor - # landmark. Low-confidence landmarks (e.g., near face boundary on - # profile views) are deformed less aggressively. + # Confidence-weight each handle by its anchor reliability. conf = face.landmark_confidence scaled_handles = [] for handle in handles: @@ -733,7 +702,6 @@ def apply_procedure_preset( pixel_landmarks = gaussian_rbf_deform_batch(pixel_landmarks, scaled_handles) - # Convert back to normalized result = pixel_landmarks.copy() result[:, 0] /= face.image_width result[:, 1] /= face.image_height @@ -746,6 +714,84 @@ def apply_procedure_preset( ) +def apply_rbf_prewarp_stage( + face: FaceLandmarks, + procedure: str, + intensity: float = 50.0, + clinical_flags: ClinicalFlags | None = None, + regional_intensity: RegionalIntensity | None = None, +) -> FaceLandmarks: + """Run the NumPy RBF deformation stage used before TPS image warping.""" + handles = build_rbf_prewarp_handles( + face=face, + procedure=procedure, + intensity=intensity, + clinical_flags=clinical_flags, + regional_intensity=regional_intensity, + ) + return _apply_rbf_handles_to_face(face, handles) + + +def apply_procedure_preset( + face: FaceLandmarks, + procedure: str, + intensity: float = 50.0, + image_size: int = 512, + clinical_flags: ClinicalFlags | None = None, + displacement_model_path: str | None = None, + noise_scale: float = 0.0, + regional_intensity: RegionalIntensity | None = None, +) -> FaceLandmarks: + """Apply a surgical procedure preset to landmarks. + + Args: + face: Input face landmarks. + procedure: Procedure name (rhinoplasty, blepharoplasty, etc.). + intensity: Relative intensity 0-100 (mild=33, moderate=66, aggressive=100). + image_size: Legacy compatibility argument. The modular NumPy RBF path + scales by ``face.image_width``/``face.image_height`` (geometric + mean) and does not consume this value directly. + clinical_flags: Optional clinical condition flags. + displacement_model_path: Path to a fitted DisplacementModel (.npz). + When provided, uses data-driven displacements from real surgery pairs + instead of hand-tuned RBF vectors. + noise_scale: Variation noise scale for data-driven mode (0=deterministic). + + Returns: + New FaceLandmarks with manipulated landmarks. + """ + if procedure not in PROCEDURE_LANDMARKS: + raise ValueError(f"Unknown procedure: {procedure}. Choose from {list(PROCEDURE_LANDMARKS)}") + + # Data-driven displacement mode (fall back to RBF if procedure not in model) + # Map UI intensity (0-100) to displacement model intensity (0-2): + # 50 -> 1.0x mean displacement, matching inference.py scaling + if displacement_model_path is not None: + dm_scale = intensity / 50.0 + try: + return _apply_data_driven( + face, + procedure, + dm_scale, + displacement_model_path, + noise_scale, + ) + except KeyError: + logger.warning( + "Procedure '%s' not in displacement model, falling back to RBF preset", + procedure, + ) + # Fall through to RBF-based preset below + + return apply_rbf_prewarp_stage( + face=face, + procedure=procedure, + intensity=intensity, + clinical_flags=clinical_flags, + regional_intensity=regional_intensity, + ) + + def _apply_data_driven( face: FaceLandmarks, procedure: str, diff --git a/tests/test_manipulation_rbf_stage.py b/tests/test_manipulation_rbf_stage.py new file mode 100644 index 0000000..d3d21f1 --- /dev/null +++ b/tests/test_manipulation_rbf_stage.py @@ -0,0 +1,298 @@ +"""Tests for modular NumPy RBF pre-warp stage. + +These tests lock parity with the legacy inline path from +`apply_procedure_preset` while validating the new explicit stage API. +""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from landmarkdiff.clinical import ClinicalFlags +from landmarkdiff.landmarks import FaceLandmarks +from landmarkdiff.manipulation import ( + PROCEDURE_LANDMARKS, + PROCEDURE_RADIUS, + DeformationHandle, + RegionalIntensity, + _get_procedure_handles, + apply_procedure_preset, + apply_rbf_prewarp_stage, + build_rbf_prewarp_handles, + gaussian_rbf_deform_batch, +) + + +def _make_face(seed: int = 7, width: int = 512, height: int = 512) -> FaceLandmarks: + rng = np.random.default_rng(seed) + landmarks = rng.uniform(0.2, 0.8, size=(478, 3)).astype(np.float32) + return FaceLandmarks( + landmarks=landmarks, + image_width=width, + image_height=height, + confidence=0.93, + ) + + +def _legacy_build_rbf_handles( + face: FaceLandmarks, + procedure: str, + intensity: float, + clinical_flags: ClinicalFlags | None = None, + regional_intensity: RegionalIntensity | None = None, +) -> list[DeformationHandle]: + scale = intensity / 100.0 + indices = PROCEDURE_LANDMARKS[procedure] + radius = PROCEDURE_RADIUS[procedure] + + if clinical_flags and clinical_flags.ehlers_danlos: + radius *= 1.5 + + geo_mean = math.sqrt(face.image_width * face.image_height) + pixel_scale = geo_mean / 512.0 + handles = _get_procedure_handles( + procedure, + indices, + scale, + radius * pixel_scale, + regional_intensity, + ) + + if clinical_flags and clinical_flags.bells_palsy: + from landmarkdiff.clinical import get_bells_palsy_side_indices + + affected = get_bells_palsy_side_indices(clinical_flags.bells_palsy_side) + affected_indices = set() + for region_indices in affected.values(): + affected_indices.update(region_indices) + handles = [h for h in handles if h.landmark_index not in affected_indices] + + return handles + + +def _legacy_apply_rbf_prewarp( + face: FaceLandmarks, + procedure: str, + intensity: float, + clinical_flags: ClinicalFlags | None = None, + regional_intensity: RegionalIntensity | None = None, +) -> FaceLandmarks: + handles = _legacy_build_rbf_handles( + face=face, + procedure=procedure, + intensity=intensity, + clinical_flags=clinical_flags, + regional_intensity=regional_intensity, + ) + + pixel_landmarks = face.landmarks.copy() + pixel_landmarks[:, 0] *= face.image_width + pixel_landmarks[:, 1] *= face.image_height + + conf = face.landmark_confidence + scaled_handles = [] + for handle in handles: + c = float(conf[handle.landmark_index]) + if c < 1.0: + scaled_handles.append( + DeformationHandle( + landmark_index=handle.landmark_index, + displacement=handle.displacement * c, + influence_radius=handle.influence_radius, + ) + ) + else: + scaled_handles.append(handle) + + pixel_landmarks = gaussian_rbf_deform_batch(pixel_landmarks, scaled_handles) + + result = pixel_landmarks.copy() + result[:, 0] /= face.image_width + result[:, 1] /= face.image_height + + return FaceLandmarks( + landmarks=result, + image_width=face.image_width, + image_height=face.image_height, + confidence=face.confidence, + ) + + +class TestRBFPrewarpStageParity: + @pytest.mark.parametrize("procedure", sorted(PROCEDURE_LANDMARKS.keys())) + @pytest.mark.parametrize("intensity", [0.0, 50.0, 100.0]) + def test_stage_matches_legacy_behavior(self, procedure: str, intensity: float): + face = _make_face() + stage_out = apply_rbf_prewarp_stage(face, procedure, intensity=intensity) + legacy_out = _legacy_apply_rbf_prewarp(face, procedure, intensity=intensity) + + np.testing.assert_allclose(stage_out.landmarks, legacy_out.landmarks, atol=1e-7, rtol=0) + assert stage_out.image_width == legacy_out.image_width + assert stage_out.image_height == legacy_out.image_height + assert stage_out.confidence == legacy_out.confidence + + def test_stage_matches_legacy_with_clinical_flags(self): + face = _make_face() + flags = ClinicalFlags(bells_palsy=True, bells_palsy_side="left", ehlers_danlos=True) + + stage_out = apply_rbf_prewarp_stage( + face, + "rhytidectomy", + intensity=65.0, + clinical_flags=flags, + ) + legacy_out = _legacy_apply_rbf_prewarp( + face, + "rhytidectomy", + intensity=65.0, + clinical_flags=flags, + ) + + np.testing.assert_allclose(stage_out.landmarks, legacy_out.landmarks, atol=1e-7, rtol=0) + + def test_stage_matches_legacy_with_regional_intensity(self): + face = _make_face() + regional = RegionalIntensity(tip=1.2, bridge=0.85, alar=1.1) + + stage_out = apply_rbf_prewarp_stage( + face, + "rhinoplasty", + intensity=70.0, + regional_intensity=regional, + ) + legacy_out = _legacy_apply_rbf_prewarp( + face, + "rhinoplasty", + intensity=70.0, + regional_intensity=regional, + ) + + np.testing.assert_allclose(stage_out.landmarks, legacy_out.landmarks, atol=1e-7, rtol=0) + + +class TestRBFHandleBuilderParity: + def test_handle_builder_matches_legacy(self): + face = _make_face(width=640, height=768) + flags = ClinicalFlags(bells_palsy=True, bells_palsy_side="right", ehlers_danlos=True) + regional = RegionalIntensity(tip=1.1, bridge=0.9, alar=1.05) + + built = build_rbf_prewarp_handles( + face, + "rhinoplasty", + intensity=55.0, + clinical_flags=flags, + regional_intensity=regional, + ) + legacy = _legacy_build_rbf_handles( + face, + "rhinoplasty", + intensity=55.0, + clinical_flags=flags, + regional_intensity=regional, + ) + + assert len(built) == len(legacy) + for current, previous in zip(built, legacy): + assert current.landmark_index == previous.landmark_index + np.testing.assert_allclose( + current.displacement, previous.displacement, atol=1e-7, rtol=0 + ) + assert current.influence_radius == pytest.approx(previous.influence_radius) + + +class TestProcedurePresetRouting: + def test_apply_procedure_uses_rbf_stage_without_model(self, monkeypatch): + face = _make_face() + calls: dict[str, object] = {} + + expected = FaceLandmarks( + landmarks=np.zeros((478, 3), dtype=np.float32), + image_width=face.image_width, + image_height=face.image_height, + confidence=face.confidence, + ) + + def fake_stage( + face: FaceLandmarks, + procedure: str, + intensity: float = 50.0, + clinical_flags: ClinicalFlags | None = None, + regional_intensity: RegionalIntensity | None = None, + ) -> FaceLandmarks: + calls["face"] = face + calls["procedure"] = procedure + calls["intensity"] = intensity + calls["clinical_flags"] = clinical_flags + calls["regional_intensity"] = regional_intensity + return expected + + monkeypatch.setattr("landmarkdiff.manipulation.apply_rbf_prewarp_stage", fake_stage) + + result = apply_procedure_preset(face, "rhinoplasty", intensity=42.0) + + assert result is expected + assert calls["face"] is face + assert calls["procedure"] == "rhinoplasty" + assert calls["intensity"] == 42.0 + assert calls["clinical_flags"] is None + assert calls["regional_intensity"] is None + + def test_apply_procedure_prefers_data_driven_when_model_path_present(self, monkeypatch): + face = _make_face() + expected = FaceLandmarks( + landmarks=np.ones((478, 3), dtype=np.float32), + image_width=face.image_width, + image_height=face.image_height, + confidence=face.confidence, + ) + + monkeypatch.setattr( + "landmarkdiff.manipulation._apply_data_driven", + lambda *_args, **_kwargs: expected, + ) + monkeypatch.setattr( + "landmarkdiff.manipulation.apply_rbf_prewarp_stage", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("RBF stage should not run when data-driven path succeeds") + ), + ) + + result = apply_procedure_preset( + face, + "rhinoplasty", + intensity=50.0, + displacement_model_path="dummy.npz", + ) + + assert result is expected + + def test_image_size_is_compatibility_only_in_rbf_path(self): + face = _make_face(width=640, height=640) + out_512 = apply_procedure_preset(face, "rhinoplasty", intensity=60.0, image_size=512) + out_1024 = apply_procedure_preset(face, "rhinoplasty", intensity=60.0, image_size=1024) + np.testing.assert_allclose(out_512.landmarks, out_1024.landmarks, atol=1e-7, rtol=0) + + +class TestConfidenceSemantics: + def test_builder_is_confidence_agnostic_stage_applies_confidence_weighting(self): + face = _make_face(seed=21) + handles = build_rbf_prewarp_handles(face, "rhinoplasty", intensity=70.0) + + # Builder returns raw policy handles; confidence scaling is applied + # only by the explicit stage execution step. + conf = face.landmark_confidence + assert any(float(conf[h.landmark_index]) < 1.0 for h in handles) + + pixel_landmarks = face.landmarks.copy() + pixel_landmarks[:, 0] *= face.image_width + pixel_landmarks[:, 1] *= face.image_height + raw_pixels = gaussian_rbf_deform_batch(pixel_landmarks, handles) + raw_norm = raw_pixels.copy() + raw_norm[:, 0] /= face.image_width + raw_norm[:, 1] /= face.image_height + + stage = apply_rbf_prewarp_stage(face, "rhinoplasty", intensity=70.0) + assert not np.allclose(raw_norm, stage.landmarks)