landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/cli.py

@@ -0,0 +1,252 @@
+"""Unified CLI for LandmarkDiff.
+
+Usage:
+    landmarkdiff infer IMAGE --procedure rhinoplasty --intensity 65
+    landmarkdiff evaluate --test-dir data/test --checkpoint checkpoints/latest
+    landmarkdiff train --config configs/phaseA.yaml
+    landmarkdiff demo IMAGE --output demo_report.png
+    landmarkdiff config --show
+    landmarkdiff validate IMAGE --output validated.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+
+logger = logging.getLogger(__name__)
+
+
+def cmd_infer(args: argparse.Namespace) -> None:
+    """Run single-image inference."""
+    from pathlib import Path
+
+    import cv2
+
+    from landmarkdiff.inference import LandmarkDiffPipeline
+
+    image = cv2.imread(args.image)
+    if image is None:
+        logger.error("Cannot read image: %s", args.image)
+        sys.exit(1)
+
+    image = cv2.resize(image, (512, 512))
+
+    pipeline = LandmarkDiffPipeline(
+        mode=args.mode,
+        controlnet_checkpoint=args.checkpoint,
+        displacement_model_path=args.displacement_model,
+    )
+    pipeline.load()
+
+    result = pipeline.generate(
+        image,
+        procedure=args.procedure,
+        intensity=args.intensity,
+        seed=args.seed,
+    )
+
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    cv2.imwrite(str(out_path), result["output"])
+    print(f"Output saved: {out_path}")
+
+    if args.watermark:
+        from landmarkdiff.safety import SafetyValidator
+
+        validator = SafetyValidator()
+        watermarked = validator.apply_watermark(result["output"])
+        wm_path = out_path.with_stem(out_path.stem + "_watermarked")
+        cv2.imwrite(str(wm_path), watermarked)
+        print(f"Watermarked: {wm_path}")
+
+
+def cmd_ensemble(args: argparse.Namespace) -> None:
+    """Run ensemble inference."""
+    from landmarkdiff.ensemble import ensemble_inference
+
+    ensemble_inference(
+        image_path=args.image,
+        procedure=args.procedure,
+        intensity=args.intensity,
+        output_dir=args.output,
+        n_samples=args.n_samples,
+        strategy=args.strategy,
+        mode=args.mode,
+        controlnet_checkpoint=args.checkpoint,
+        displacement_model_path=args.displacement_model,
+        seed=args.seed,
+    )
+
+
+def cmd_evaluate(args: argparse.Namespace) -> None:
+    """Run evaluation on test set."""
+    from pathlib import Path
+
+    # Import evaluation functions
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+    from scripts.run_evaluation import run_evaluation
+
+    run_evaluation(
+        test_dir=args.test_dir,
+        output_dir=args.output,
+        checkpoint=args.checkpoint,
+        max_samples=args.max_samples,
+    )
+
+
+def cmd_config(args: argparse.Namespace) -> None:
+    """Show or validate configuration."""
+    from landmarkdiff.config import ExperimentConfig, load_config, validate_config
+
+    config = load_config(args.file) if args.file else ExperimentConfig()
+
+    if args.validate:
+        warnings = validate_config(config)
+        if warnings:
+            print("Validation warnings:")
+            for w in warnings:
+                print(f"  - {w}")
+        else:
+            print("Configuration valid (no warnings).")
+    else:
+        from dataclasses import asdict
+
+        import yaml
+
+        print(yaml.dump(asdict(config), default_flow_style=False, sort_keys=False))
+
+
+def cmd_validate(args: argparse.Namespace) -> None:
+    """Run safety validation on an output image."""
+    import cv2
+
+    from landmarkdiff.safety import SafetyValidator
+
+    input_img = cv2.imread(args.input)
+    output_img = cv2.imread(args.output_image)
+
+    if input_img is None or output_img is None:
+        logger.error("Cannot read input or output image.")
+        sys.exit(1)
+
+    validator = SafetyValidator(
+        watermark_enabled=args.watermark,
+    )
+
+    result = validator.validate(
+        input_image=input_img,
+        output_image=output_img,
+        face_confidence=args.face_confidence,
+    )
+
+    print(result.summary())
+
+    if not result.passed:
+        sys.exit(1)
+
+
+def cmd_version(args: argparse.Namespace) -> None:
+    """Print version info."""
+    from landmarkdiff import __version__
+
+    print(f"LandmarkDiff v{__version__}")
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(
+        prog="landmarkdiff",
+        description="LandmarkDiff: Facial surgery outcome prediction via latent diffusion",
+    )
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # --- infer ---
+    p_infer = subparsers.add_parser("infer", help="Run single-image inference")
+    p_infer.add_argument("image", help="Input face image path")
+    p_infer.add_argument(
+        "--procedure",
+        default="rhinoplasty",
+        choices=[
+            "rhinoplasty",
+            "blepharoplasty",
+            "rhytidectomy",
+            "orthognathic",
+            "brow_lift",
+            "mentoplasty",
+        ],
+    )
+    p_infer.add_argument("--intensity", type=float, default=65.0)
+    p_infer.add_argument("--output", default="output.png")
+    p_infer.add_argument(
+        "--mode",
+        default="tps",
+        choices=["controlnet", "controlnet_ip", "controlnet_fast", "img2img", "tps"],
+    )
+    p_infer.add_argument("--checkpoint", default=None)
+    p_infer.add_argument("--displacement-model", default=None)
+    p_infer.add_argument("--seed", type=int, default=42)
+    p_infer.add_argument("--watermark", action="store_true")
+    p_infer.set_defaults(func=cmd_infer)
+
+    # --- ensemble ---
+    p_ensemble = subparsers.add_parser("ensemble", help="Run ensemble inference")
+    p_ensemble.add_argument("image", help="Input face image path")
+    p_ensemble.add_argument("--procedure", default="rhinoplasty")
+    p_ensemble.add_argument("--intensity", type=float, default=65.0)
+    p_ensemble.add_argument("--output", default="ensemble_output")
+    p_ensemble.add_argument("--n-samples", type=int, default=5)
+    p_ensemble.add_argument(
+        "--strategy",
+        default="best_of_n",
+        choices=["pixel_average", "weighted_average", "best_of_n", "median"],
+    )
+    p_ensemble.add_argument(
+        "--mode",
+        default="tps",
+        choices=["controlnet", "controlnet_ip", "controlnet_fast", "img2img", "tps"],
+    )
+    p_ensemble.add_argument("--checkpoint", default=None)
+    p_ensemble.add_argument("--displacement-model", default=None)
+    p_ensemble.add_argument("--seed", type=int, default=42)
+    p_ensemble.set_defaults(func=cmd_ensemble)
+
+    # --- evaluate ---
+    p_eval = subparsers.add_parser("evaluate", help="Evaluate on test set")
+    p_eval.add_argument("--test-dir", required=True)
+    p_eval.add_argument("--output", default="eval_results")
+    p_eval.add_argument("--mode", default="tps")
+    p_eval.add_argument("--checkpoint", default=None)
+    p_eval.add_argument("--displacement-model", default=None)
+    p_eval.add_argument("--max-samples", type=int, default=0)
+    p_eval.set_defaults(func=cmd_evaluate)
+
+    # --- config ---
+    p_config = subparsers.add_parser("config", help="Show or validate configuration")
+    p_config.add_argument("--file", default=None, help="YAML config file")
+    p_config.add_argument("--validate", action="store_true")
+    p_config.set_defaults(func=cmd_config)
+
+    # --- validate ---
+    p_validate = subparsers.add_parser("validate", help="Run safety validation")
+    p_validate.add_argument("input", help="Original input image")
+    p_validate.add_argument("output_image", help="Generated output image")
+    p_validate.add_argument("--watermark", action="store_true")
+    p_validate.add_argument("--face-confidence", type=float, default=1.0)
+    p_validate.set_defaults(func=cmd_validate)
+
+    # --- version ---
+    p_version = subparsers.add_parser("version", help="Print version")
+    p_version.set_defaults(func=cmd_version)
+
+    args = parser.parse_args(argv)
+    if not hasattr(args, "func"):
+        parser.print_help()
+        sys.exit(1)
+
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

landmarkdiff/clinical.py

@@ -0,0 +1,223 @@
+"""Clinical edge case handling for pathological conditions.
+
+Implements special-case logic for:
+- Vitiligo: preserve depigmented patches (don't blend over them)
+- Bell's palsy: disable bilateral symmetry in deformation vectors
+- Keloid: flag keloid-prone areas to reduce aggressive compositing
+- Ehlers-Danlos: wider influence radii for hypermobile tissue
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+
+from landmarkdiff.landmarks import FaceLandmarks
+
+
+@dataclass
+class ClinicalFlags:
+    """Clinical condition flags that modify pipeline behavior.
+
+    Set flags to True to enable condition-specific handling.
+    """
+
+    vitiligo: bool = False
+    bells_palsy: bool = False
+    bells_palsy_side: str = "left"  # affected side: "left" or "right"
+    keloid_prone: bool = False
+    keloid_regions: list[str] = field(default_factory=list)  # e.g. ["jawline", "nose"]
+    ehlers_danlos: bool = False
+
+    def has_any(self) -> bool:
+        return self.vitiligo or self.bells_palsy or self.keloid_prone or self.ehlers_danlos
+
+
+def detect_vitiligo_patches(
+    image: np.ndarray,
+    face: FaceLandmarks,
+    l_threshold: float = 85.0,
+    min_patch_area: int = 200,
+) -> np.ndarray:
+    """Detect depigmented (vitiligo) patches on face using LAB luminance.
+
+    Vitiligo patches appear as high-L, low-saturation regions that deviate
+    significantly from surrounding skin tone.
+
+    Args:
+        image: BGR face image.
+        face: Extracted landmarks for face ROI.
+        l_threshold: Luminance threshold (patches brighter than surrounding skin).
+        min_patch_area: Minimum contour area in pixels to count as a patch.
+
+    Returns:
+        Binary mask (uint8, 0/255) of detected vitiligo patches.
+    """
+    h, w = image.shape[:2]
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB).astype(np.float32)
+
+    # Create face ROI mask from landmarks
+    coords = face.pixel_coords.astype(np.int32)
+    hull = cv2.convexHull(coords)
+    face_mask = np.zeros((h, w), dtype=np.uint8)
+    cv2.fillConvexPoly(face_mask, hull, 255)
+
+    # Get face-region luminance statistics
+    l_channel = lab[:, :, 0]
+    face_pixels = l_channel[face_mask > 0]
+    if len(face_pixels) == 0:
+        return np.zeros((h, w), dtype=np.uint8)
+
+    l_mean = np.mean(face_pixels)
+    l_std = np.std(face_pixels)
+
+    # Vitiligo patches: significantly brighter than mean skin
+    threshold = min(l_threshold, l_mean + 2.0 * l_std)
+    bright_mask = ((l_channel > threshold) & (face_mask > 0)).astype(np.uint8) * 255
+
+    # Also check for low saturation (a,b channels close to 128)
+    a_channel = lab[:, :, 1]
+    b_channel = lab[:, :, 2]
+    low_sat = ((np.abs(a_channel - 128) < 15) & (np.abs(b_channel - 128) < 15)).astype(
+        np.uint8
+    ) * 255
+
+    # Combined: bright AND low-saturation within face
+    vitiligo_raw = cv2.bitwise_and(bright_mask, low_sat)
+
+    # Filter small noise patches
+    contours, _ = cv2.findContours(vitiligo_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    result = np.zeros((h, w), dtype=np.uint8)
+    for cnt in contours:
+        if cv2.contourArea(cnt) >= min_patch_area:
+            cv2.fillPoly(result, [cnt], 255)
+
+    return result
+
+
+def adjust_mask_for_vitiligo(
+    mask: np.ndarray,
+    vitiligo_patches: np.ndarray,
+    preservation_factor: float = 0.3,
+) -> np.ndarray:
+    """Reduce mask intensity over vitiligo patches to preserve them.
+
+    Instead of full blending over depigmented patches, we reduce the
+    mask weight so the original vitiligo pattern shows through.
+
+    Args:
+        mask: Float32 surgical mask [0-1].
+        vitiligo_patches: Binary mask of vitiligo regions (0/255 uint8).
+        preservation_factor: How much to reduce blending (0=full blend, 1=fully preserve).
+
+    Returns:
+        Modified mask with reduced intensity over vitiligo patches.
+    """
+    patches_f = vitiligo_patches.astype(np.float32) / 255.0
+    reduction = patches_f * preservation_factor
+    return np.clip(mask - reduction, 0.0, 1.0)
+
+
+def get_bells_palsy_side_indices(
+    side: str,
+) -> dict[str, list[int]]:
+    """Get landmark indices for the affected side in Bell's palsy.
+
+    In Bell's palsy, one side of the face is paralyzed. We should NOT
+    apply bilateral symmetric deformations — only deform the healthy side.
+
+    Returns:
+        Dict mapping region names to landmark indices on the affected side.
+    """
+    if side == "left":
+        return {
+            "eye": [33, 7, 163, 144, 145, 153, 154, 155, 133, 173, 157, 158, 159, 160, 161, 246],
+            "eyebrow": [70, 63, 105, 66, 107, 55, 65, 52, 53, 46],
+            "mouth_corner": [61, 146, 91, 181, 84],
+            "jawline": [132, 136, 172, 58, 150, 176, 148, 149],
+        }
+    else:
+        return {
+            "eye": [362, 382, 381, 380, 374, 373, 390, 249, 263, 466, 388, 387, 386, 385, 384, 398],
+            "eyebrow": [300, 293, 334, 296, 336, 285, 295, 282, 283, 276],
+            "mouth_corner": [291, 308, 324, 318, 402],
+            "jawline": [361, 365, 397, 288, 379, 400, 377, 378],
+        }
+
+
+def get_keloid_exclusion_mask(
+    face: FaceLandmarks,
+    regions: list[str],
+    width: int,
+    height: int,
+    margin_px: int = 10,
+) -> np.ndarray:
+    """Generate mask of keloid-prone regions to exclude from aggressive compositing.
+
+    Keloid patients should have reduced blending intensity and no sharp
+    boundary transitions in prone areas (typically jawline, ears, chest).
+
+    Args:
+        face: Extracted landmarks.
+        regions: List of region names prone to keloids.
+        width: Image width.
+        height: Image height.
+        margin_px: Extra margin around keloid regions.
+
+    Returns:
+        Float32 mask [0-1] where 1 = keloid-prone area.
+    """
+    from landmarkdiff.landmarks import LANDMARK_REGIONS
+
+    mask = np.zeros((height, width), dtype=np.float32)
+    coords = face.pixel_coords.astype(np.int32)
+
+    for region in regions:
+        indices = LANDMARK_REGIONS.get(region, [])
+        if not indices:
+            continue
+        pts = coords[indices]
+        hull = cv2.convexHull(pts)
+        cv2.fillConvexPoly(mask, hull, 1.0)
+
+    # Dilate by margin
+    if margin_px > 0:
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, (2 * margin_px + 1, 2 * margin_px + 1)
+        )
+        mask = cv2.dilate(mask, kernel)
+
+    return np.clip(mask, 0.0, 1.0)
+
+
+def adjust_mask_for_keloid(
+    mask: np.ndarray,
+    keloid_mask: np.ndarray,
+    reduction_factor: float = 0.5,
+) -> np.ndarray:
+    """Soften mask transitions in keloid-prone areas.
+
+    Reduces the mask gradient steepness to prevent hard boundaries
+    that could trigger keloid formation in real surgical planning.
+
+    Args:
+        mask: Float32 surgical mask [0-1].
+        keloid_mask: Float32 keloid region mask [0-1].
+        reduction_factor: How much to reduce mask intensity in keloid areas.
+
+    Returns:
+        Modified mask with gentler transitions in keloid regions.
+    """
+    # Reduce mask intensity in keloid-prone areas
+    keloid_reduction = keloid_mask * reduction_factor
+    modified = mask * (1.0 - keloid_reduction)
+
+    # Extra Gaussian blur in keloid regions for softer transitions
+    blur_kernel = 31
+    blurred = cv2.GaussianBlur(modified, (blur_kernel, blur_kernel), 10.0)
+
+    # Use blurred version only in keloid regions
+    result = modified * (1.0 - keloid_mask) + blurred * keloid_mask
+    return np.clip(result, 0.0, 1.0)