landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/safety.py

@@ -0,0 +1,395 @@
+"""Clinical safety validation for responsible deployment.
+
+Implements safety checks for surgical outcome predictions:
+1. Identity preservation: verify output preserves patient identity
+2. Anatomical plausibility: check landmark displacements are realistic
+3. Out-of-distribution detection: flag unusual inputs
+4. Watermarking: mark AI-generated images
+5. Consent metadata: embed provenance information
+
+Usage:
+    from landmarkdiff.safety import SafetyValidator
+
+    validator = SafetyValidator()
+    result = validator.validate(
+        input_image=image,
+        output_image=generated,
+        landmarks_original=face.landmarks,
+        landmarks_manipulated=manip.landmarks,
+        procedure="rhinoplasty",
+    )
+
+    if not result.passed:
+        print(f"Safety check failed: {result.failures}")
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+
+
+@dataclass
+class SafetyResult:
+    """Result of safety validation checks."""
+
+    passed: bool = True
+    failures: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+    checks: dict[str, bool] = field(default_factory=dict)
+    details: dict[str, object] = field(default_factory=dict)
+
+    def __repr__(self) -> str:
+        return (
+            f"SafetyResult(passed={self.passed}, "
+            f"failures={self.failures}, "
+            f"warnings={self.warnings}, "
+            f"checks={self.checks}, "
+            f"details={self.details})"
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SafetyResult):
+            return NotImplemented
+        return (
+            self.passed == other.passed
+            and self.failures == other.failures
+            and self.warnings == other.warnings
+            and self.checks == other.checks
+            and self.details == other.details
+        )
+
+    def add_failure(self, name: str, message: str) -> None:
+        self.passed = False
+        self.failures.append(message)
+        self.checks[name] = False
+
+    def add_warning(self, name: str, message: str) -> None:
+        self.warnings.append(message)
+
+    def add_pass(self, name: str) -> None:
+        self.checks[name] = True
+
+    def summary(self) -> str:
+        lines = [f"Safety: {'PASS' if self.passed else 'FAIL'}"]
+        for name, ok in self.checks.items():
+            lines.append(f"  [{'OK' if ok else 'FAIL'}] {name}")
+        for w in self.warnings:
+            lines.append(f"  [WARN] {w}")
+        return "\n".join(lines)
+
+
+class SafetyValidator:
+    """Clinical safety validation for surgical predictions."""
+
+    def __init__(
+        self,
+        identity_threshold: float = 0.6,
+        max_displacement_fraction: float = 0.05,
+        min_face_confidence: float = 0.5,
+        max_yaw_degrees: float = 45.0,
+        watermark_enabled: bool = True,
+        watermark_text: str = "AI-GENERATED PREDICTION",
+    ):
+        self.identity_threshold = identity_threshold
+        self.max_displacement_fraction = max_displacement_fraction
+        self.min_face_confidence = min_face_confidence
+        self.max_yaw_degrees = max_yaw_degrees
+        self.watermark_enabled = watermark_enabled
+        self.watermark_text = watermark_text
+
+    def validate(
+        self,
+        input_image: np.ndarray,
+        output_image: np.ndarray,
+        landmarks_original: np.ndarray | None = None,
+        landmarks_manipulated: np.ndarray | None = None,
+        procedure: str | None = None,
+        face_confidence: float = 1.0,
+    ) -> SafetyResult:
+        """Run all safety checks on a prediction.
+
+        Args:
+            input_image: Original patient image (BGR, uint8).
+            output_image: Generated prediction (BGR, uint8).
+            landmarks_original: Original landmarks (N, 2-3), normalized [0, 1].
+            landmarks_manipulated: Manipulated landmarks (N, 2-3), normalized [0, 1].
+            procedure: Surgical procedure name.
+            face_confidence: MediaPipe face detection confidence.
+
+        Returns:
+            SafetyResult with all check results.
+        """
+        result = SafetyResult()
+
+        # 1. Face detection confidence
+        self._check_face_confidence(result, face_confidence)
+
+        # 2. Identity preservation
+        self._check_identity(result, input_image, output_image)
+
+        # 3. Anatomical plausibility
+        if landmarks_original is not None and landmarks_manipulated is not None:
+            self._check_anatomical_plausibility(
+                result, landmarks_original, landmarks_manipulated, procedure
+            )
+
+        # 4. Output quality
+        self._check_output_quality(result, output_image)
+
+        # 5. OOD detection (basic)
+        self._check_ood(result, input_image)
+
+        return result
+
+    def _check_face_confidence(self, result: SafetyResult, confidence: float) -> None:
+        """Check face detection confidence."""
+        if confidence < self.min_face_confidence:
+            result.add_failure(
+                "face_confidence",
+                f"Face detection confidence {confidence:.2f} below threshold "
+                f"{self.min_face_confidence}",
+            )
+        else:
+            result.add_pass("face_confidence")
+            result.details["face_confidence"] = confidence
+
+    def _check_identity(
+        self,
+        result: SafetyResult,
+        input_image: np.ndarray,
+        output_image: np.ndarray,
+    ) -> None:
+        """Check identity preservation using ArcFace similarity."""
+        try:
+            from landmarkdiff.evaluation import compute_identity_similarity
+
+            sim = compute_identity_similarity(output_image, input_image)
+            result.details["identity_similarity"] = float(sim)
+
+            if sim < self.identity_threshold:
+                result.add_failure(
+                    "identity",
+                    f"Identity similarity {sim:.3f} below threshold {self.identity_threshold}",
+                )
+            else:
+                result.add_pass("identity")
+        except Exception as e:
+            result.add_warning("identity", f"Identity check failed: {e}")
+
+    def _check_anatomical_plausibility(
+        self,
+        result: SafetyResult,
+        landmarks_orig: np.ndarray,
+        landmarks_manip: np.ndarray,
+        procedure: str | None,
+    ) -> None:
+        """Check that landmark displacements are anatomically plausible."""
+        if len(landmarks_orig) != len(landmarks_manip):
+            result.add_failure(
+                "anatomical",
+                f"Landmark count mismatch: {len(landmarks_orig)} vs {len(landmarks_manip)}",
+            )
+            return
+
+        # Compute displacement magnitudes
+        n = min(len(landmarks_orig), len(landmarks_manip))
+        orig = landmarks_orig[:n, :2]  # (N, 2), normalized [0, 1]
+        manip = landmarks_manip[:n, :2]
+        displacements = np.linalg.norm(manip - orig, axis=1)
+
+        max_disp = float(displacements.max())
+        mean_disp = float(displacements.mean())
+        result.details["max_displacement"] = max_disp
+        result.details["mean_displacement"] = mean_disp
+
+        # Check maximum displacement
+        if max_disp > self.max_displacement_fraction:
+            result.add_failure(
+                "anatomical_magnitude",
+                f"Maximum displacement {max_disp:.4f} exceeds threshold "
+                f"{self.max_displacement_fraction}",
+            )
+        else:
+            result.add_pass("anatomical_magnitude")
+
+        # Check procedure-specific regions
+        if procedure:
+            self._check_procedure_regions(result, orig, manip, displacements, procedure)
+
+    def _check_procedure_regions(
+        self,
+        result: SafetyResult,
+        orig: np.ndarray,
+        manip: np.ndarray,
+        displacements: np.ndarray,
+        procedure: str,
+    ) -> None:
+        """Verify displacement is concentrated in expected anatomical regions."""
+        from landmarkdiff.landmarks import LANDMARK_REGIONS
+
+        # Expected regions by procedure
+        expected_regions = {
+            "rhinoplasty": ["nose"],
+            "blepharoplasty": ["eye_left", "eye_right"],
+            "rhytidectomy": ["jawline"],
+            "orthognathic": ["jawline", "lips"],
+        }
+
+        expected = expected_regions.get(procedure, [])
+        if not expected:
+            result.add_pass("procedure_region")
+            return
+
+        # Get expected region indices
+        expected_indices = set()
+        for region in expected:
+            if region in LANDMARK_REGIONS:
+                expected_indices.update(LANDMARK_REGIONS[region])
+
+        if not expected_indices:
+            result.add_pass("procedure_region")
+            return
+
+        # Check: is most displacement in expected regions?
+        n = min(len(displacements), len(orig))
+        expected_mask = np.array([i in expected_indices for i in range(n)])
+
+        if expected_mask.sum() > 0 and (~expected_mask).sum() > 0:
+            expected_disp = displacements[expected_mask].mean()
+            unexpected_disp = displacements[~expected_mask].mean()
+            result.details["expected_region_disp"] = float(expected_disp)
+            result.details["unexpected_region_disp"] = float(unexpected_disp)
+
+            # Expected regions should have more displacement
+            if unexpected_disp > expected_disp * 2 and unexpected_disp > 0.005:
+                result.add_warning(
+                    "procedure_region",
+                    f"{procedure}: unexpected regions displaced more than expected "
+                    f"({unexpected_disp:.4f} vs {expected_disp:.4f})",
+                )
+            else:
+                result.add_pass("procedure_region")
+        else:
+            result.add_pass("procedure_region")
+
+    def _check_output_quality(self, result: SafetyResult, output: np.ndarray) -> None:
+        """Check output image quality (not blank, not corrupted)."""
+        if output is None or output.size == 0:
+            result.add_failure("output_quality", "Output image is empty")
+            return
+
+        # Check for blank/black images
+        mean_val = output.mean()
+        if mean_val < 5:
+            result.add_failure("output_quality", f"Output is nearly black (mean={mean_val:.1f})")
+            return
+        if mean_val > 250:
+            result.add_failure("output_quality", f"Output is nearly white (mean={mean_val:.1f})")
+            return
+
+        # Check for artifacts (extreme variance)
+        std_val = output.std()
+        if std_val < 10:
+            result.add_warning(
+                "output_quality",
+                f"Output has very low variance (std={std_val:.1f}), may be uniform",
+            )
+
+        result.add_pass("output_quality")
+        result.details["output_mean"] = float(mean_val)
+        result.details["output_std"] = float(std_val)
+
+    def _check_ood(self, result: SafetyResult, image: np.ndarray) -> None:
+        """Basic out-of-distribution detection.
+
+        Checks image properties against expected ranges for face photos.
+        """
+        h, w = image.shape[:2]
+
+        # Resolution check
+        if min(h, w) < 128:
+            result.add_warning("ood", f"Image resolution too low: {w}x{h}")
+
+        # Aspect ratio (faces should be roughly square after preprocessing)
+        aspect = max(h, w) / max(min(h, w), 1)
+        if aspect > 3.0:
+            result.add_warning("ood", f"Unusual aspect ratio: {aspect:.1f}")
+
+        # Color distribution (face photos should have some skin tones)
+        if len(image.shape) == 3 and image.shape[2] == 3:
+            mean_b, mean_g, mean_r = image.mean(axis=(0, 1))
+            # Face images typically have red channel > blue channel
+            if mean_b > mean_r * 1.5:
+                result.add_warning("ood", "Image appears very blue (not typical face photo)")
+
+        result.add_pass("ood_basic")
+
+    def apply_watermark(
+        self,
+        image: np.ndarray,
+        text: str | None = None,
+        opacity: float = 0.3,
+    ) -> np.ndarray:
+        """Apply a text watermark to the output image.
+
+        Places semi-transparent text at the bottom of the image to indicate
+        it is AI-generated.
+        """
+        if not self.watermark_enabled:
+            return image
+
+        text = text or self.watermark_text
+        result = image.copy()
+        h, w = result.shape[:2]
+
+        # Create text overlay
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        font_scale = max(0.3, w / 1500)
+        thickness = max(1, int(w / 500))
+
+        text_size = cv2.getTextSize(text, font, font_scale, thickness)[0]
+        x = (w - text_size[0]) // 2
+        y = h - 10
+
+        # Semi-transparent background bar
+        bar_y1 = y - text_size[1] - 10
+        bar_y2 = h
+        overlay = result.copy()
+        cv2.rectangle(overlay, (0, bar_y1), (w, bar_y2), (0, 0, 0), -1)
+        cv2.addWeighted(overlay, opacity, result, 1 - opacity, 0, result)
+
+        # White text
+        cv2.putText(result, text, (x, y), font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA)
+
+        return result
+
+    def embed_metadata(
+        self,
+        image_path: str,
+        procedure: str,
+        intensity: float,
+        model_version: str = "0.3.0",
+    ) -> None:
+        """Embed provenance metadata in the output image.
+
+        Writes EXIF/PNG metadata with generation parameters for traceability.
+        """
+        import json
+        from pathlib import Path
+
+        meta = {
+            "generator": "LandmarkDiff",
+            "version": model_version,
+            "procedure": procedure,
+            "intensity": intensity,
+            "disclaimer": "AI-generated surgical prediction for visualization only. "
+            "Not a guarantee of surgical outcome.",
+        }
+
+        # Save as sidecar JSON (PNG doesn't have easy EXIF support)
+        meta_path = Path(image_path).with_suffix(".meta.json")
+        with open(meta_path, "w") as f:
+            json.dump(meta, f, indent=2)

landmarkdiff/synthetic/__init__.py

@@ -0,0 +1,23 @@
+"""Synthetic data generation for ControlNet fine-tuning.
+
+Modules:
+  - pair_generator: Generate training pairs from face images
+  - augmentation: Clinical degradation augmentations
+  - tps_warp: TPS warping with rigid region preservation
+"""
+
+from landmarkdiff.synthetic.augmentation import apply_clinical_augmentation
+from landmarkdiff.synthetic.pair_generator import (
+    TrainingPair,
+    generate_pair,
+    generate_pairs_from_directory,
+)
+from landmarkdiff.synthetic.tps_warp import warp_image_tps
+
+__all__ = [
+    "TrainingPair",
+    "apply_clinical_augmentation",
+    "generate_pair",
+    "generate_pairs_from_directory",
+    "warp_image_tps",
+]

landmarkdiff/synthetic/augmentation.py

@@ -0,0 +1,188 @@
+"""Clinical degradation augmentations.
+
+Degrades clean FFHQ/CelebA-HQ to match real clinical photo distribution.
+Applied from day 1 - domain gap prevention, not afterthought.
+3-5 random augmentations per sample.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+
+
+@dataclass(frozen=True)
+class AugmentationConfig:
+    """Configuration for a single augmentation."""
+
+    name: str
+    fn: Callable[[np.ndarray, np.random.Generator], np.ndarray]
+    probability: float
+
+
+def point_source_lighting(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Simulate point-source clinical lighting from a random direction."""
+    h, w = image.shape[:2]
+
+    # Random light source position
+    lx = rng.uniform(0, w)
+    ly = rng.uniform(0, h)
+    intensity = rng.uniform(0.3, 0.7)
+
+    # Distance-based falloff
+    y_grid, x_grid = np.mgrid[0:h, 0:w].astype(np.float32)
+    dist = np.sqrt((x_grid - lx) ** 2 + (y_grid - ly) ** 2)
+    max_dist = np.sqrt(w**2 + h**2)
+    light_map = 1.0 - (dist / max_dist) * intensity
+
+    light_map = np.clip(light_map, 0.3, 1.0)
+    light_3ch = np.stack([light_map] * 3, axis=-1)
+
+    return np.clip(image.astype(np.float32) * light_3ch, 0, 255).astype(np.uint8)
+
+
+def color_temperature_jitter(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Jitter color temperature +/- 2000K equivalent."""
+    shift = rng.uniform(-0.15, 0.15)
+
+    result = image.astype(np.float32)
+    if shift > 0:
+        # Warmer: boost red, reduce blue
+        result[:, :, 2] *= 1 + shift  # red (BGR)
+        result[:, :, 0] *= 1 - shift * 0.5  # blue
+    else:
+        # Cooler: boost blue, reduce red
+        result[:, :, 0] *= 1 + abs(shift)
+        result[:, :, 2] *= 1 - abs(shift) * 0.5
+
+    return np.clip(result, 0, 255).astype(np.uint8)
+
+
+def green_fluorescent_cast(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Add green fluorescent lighting cast (common in clinical settings)."""
+    intensity = rng.uniform(0.05, 0.15)
+    result = image.astype(np.float32)
+    result[:, :, 1] *= 1 + intensity  # green channel boost
+    result[:, :, 0] *= 1 - intensity * 0.3  # slight blue reduction
+    result[:, :, 2] *= 1 - intensity * 0.3  # slight red reduction
+    return np.clip(result, 0, 255).astype(np.uint8)
+
+
+def jpeg_compression(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Simulate JPEG compression artifacts (quality 40-85)."""
+    quality = int(rng.uniform(40, 85))
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
+    _, encoded = cv2.imencode(".jpg", image, encode_param)
+    return cv2.imdecode(encoded, cv2.IMREAD_COLOR)
+
+
+def gaussian_sensor_noise(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Add Gaussian sensor noise (sigma 5-25)."""
+    sigma = rng.uniform(5, 25)
+    noise = rng.normal(0, sigma, image.shape).astype(np.float32)
+    return np.clip(image.astype(np.float32) + noise, 0, 255).astype(np.uint8)
+
+
+def barrel_distortion(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Apply barrel/pincushion distortion simulating phone camera lens."""
+    h, w = image.shape[:2]
+    k1 = rng.uniform(-0.2, 0.2)
+
+    fx = fy = max(w, h)
+    cx, cy = w / 2, h / 2
+
+    camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
+    dist_coeffs = np.array([k1, 0, 0, 0, 0], dtype=np.float64)
+
+    map1, map2 = cv2.initUndistortRectifyMap(
+        camera_matrix, dist_coeffs, None, camera_matrix, (w, h), cv2.CV_32FC1
+    )
+    return cv2.remap(image, map1, map2, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
+
+
+def motion_blur(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Slight motion blur (common in handheld clinical photos)."""
+    size = int(rng.uniform(3, 7))
+    angle = rng.uniform(0, 180)
+
+    kernel = np.zeros((size, size))
+    kernel[size // 2, :] = 1.0 / size
+
+    M = cv2.getRotationMatrix2D((size / 2, size / 2), angle, 1)
+    kernel = cv2.warpAffine(kernel, M, (size, size))
+    ksum = kernel.sum()
+    if ksum > 0:
+        kernel = kernel / ksum
+    else:
+        # rotation can zero out the kernel - fall back to identity
+        kernel = np.zeros_like(kernel)
+        kernel[size // 2, size // 2] = 1.0
+
+    return cv2.filter2D(image, -1, kernel)
+
+
+def vignette(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Add lens vignetting (darkened corners)."""
+    h, w = image.shape[:2]
+    strength = rng.uniform(0.3, 0.7)
+
+    y, x = np.mgrid[0:h, 0:w].astype(np.float32)
+    cx, cy = w / 2, h / 2
+    dist = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
+    max_dist = np.sqrt(cx**2 + cy**2)
+
+    mask = 1 - strength * (dist / max_dist) ** 2
+    mask = np.clip(mask, 0.3, 1.0)
+    mask_3ch = np.stack([mask] * 3, axis=-1)
+
+    return np.clip(image.astype(np.float32) * mask_3ch, 0, 255).astype(np.uint8)
+
+
+# Augmentation pool with probabilities from the spec
+AUGMENTATION_POOL: list[AugmentationConfig] = [
+    AugmentationConfig("point_source_lighting", point_source_lighting, 0.40),
+    AugmentationConfig("color_temperature", color_temperature_jitter, 0.60),
+    AugmentationConfig("green_fluorescent", green_fluorescent_cast, 0.25),
+    AugmentationConfig("jpeg_compression", jpeg_compression, 0.30),
+    AugmentationConfig("sensor_noise", gaussian_sensor_noise, 0.40),
+    AugmentationConfig("barrel_distortion", barrel_distortion, 0.30),
+    AugmentationConfig("motion_blur", motion_blur, 0.20),
+    AugmentationConfig("vignette", vignette, 0.25),
+]
+
+
+def apply_clinical_augmentation(
+    image: np.ndarray,
+    min_augmentations: int = 3,
+    max_augmentations: int = 5,
+    rng: np.random.Generator | None = None,
+) -> np.ndarray:
+    """Apply random clinical degradation augmentations to an image."""
+    rng = rng or np.random.default_rng()
+
+    # Select augmentations by probability
+    selected = []
+    for aug in AUGMENTATION_POOL:
+        if rng.random() < aug.probability:
+            selected.append(aug)
+
+    # Ensure min/max bounds
+    if len(selected) < min_augmentations:
+        remaining = [a for a in AUGMENTATION_POOL if a not in selected]
+        rng.shuffle(remaining)
+        selected.extend(remaining[: min_augmentations - len(selected)])
+
+    if len(selected) > max_augmentations:
+        rng.shuffle(selected)
+        selected = selected[:max_augmentations]
+
+    # Apply in random order
+    rng.shuffle(selected)
+    result = image.copy()
+    for aug in selected:
+        result = aug.fn(result, rng)
+
+    return result