Upload landmarkdiff/safety.py with huggingface_hub

landmarkdiff/safety.py

@@ -0,0 +1,380 @@
+"""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
+from typing import Optional
+
+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 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 "
+                    f"{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)