Skip to content

feat(manipulation): modularize NumPy RBF pre-warp stage for TPS pipeline PR#4 #423

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Flames4fun wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(manipulation): modularize NumPy RBF pre-warp stage for TPS pipeline PR#4 #423

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 39 additions & 1 deletion docs/api/manipulation.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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. |
Expand Down
150 changes: 98 additions & 52 deletions landmarkdiff/manipulation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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,
Expand Down
Loading
Loading