landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/conditioning.py

@@ -0,0 +1,278 @@
+"""Conditioning signal generation: static adjacency wireframe + auto-Canny.
+
+Uses a pre-defined anatomical adjacency matrix (NOT dynamic Delaunay) to prevent
+triangle inversion on drastic landmark displacements. Auto-Canny adapts thresholds
+to skin tone (Fitzpatrick I-VI safe).
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+
+from landmarkdiff.landmarks import FaceLandmarks
+
+# Static anatomical adjacency for MediaPipe 478 landmarks.
+# Connects landmarks along anatomically meaningful contours:
+# jawline, nasal dorsum, orbital rim, lip vermilion, eyebrow arch.
+# This is invariant to landmark displacement (unlike Delaunay).
+
+JAWLINE_CONTOUR = [
+    10,
+    338,
+    297,
+    332,
+    284,
+    251,
+    389,
+    356,
+    454,
+    323,
+    361,
+    288,
+    397,
+    365,
+    379,
+    378,
+    400,
+    377,
+    152,
+    148,
+    176,
+    149,
+    150,
+    136,
+    172,
+    58,
+    132,
+    93,
+    234,
+    127,
+    162,
+    21,
+    54,
+    103,
+    67,
+    109,
+    10,
+]
+
+LEFT_EYE_CONTOUR = [
+    33,
+    7,
+    163,
+    144,
+    145,
+    153,
+    154,
+    155,
+    133,
+    173,
+    157,
+    158,
+    159,
+    160,
+    161,
+    246,
+    33,
+]
+
+RIGHT_EYE_CONTOUR = [
+    362,
+    382,
+    381,
+    380,
+    374,
+    373,
+    390,
+    249,
+    263,
+    466,
+    388,
+    387,
+    386,
+    385,
+    384,
+    398,
+    362,
+]
+
+LEFT_EYEBROW = [70, 63, 105, 66, 107, 55, 65, 52, 53, 46]
+RIGHT_EYEBROW = [300, 293, 334, 296, 336, 285, 295, 282, 283, 276]
+
+NOSE_BRIDGE = [168, 6, 197, 195, 5, 4, 1]
+NOSE_TIP = [94, 2, 326, 327, 294, 278, 279, 275, 274, 460, 456, 363, 370]
+NOSE_BOTTOM = [19, 1, 274, 275, 440, 344, 278, 294, 460, 305, 289, 392]
+
+OUTER_LIPS = [
+    61,
+    146,
+    91,
+    181,
+    84,
+    17,
+    314,
+    405,
+    321,
+    375,
+    291,
+    308,
+    324,
+    318,
+    402,
+    317,
+    14,
+    87,
+    178,
+    88,
+    95,
+    78,
+    61,
+]
+
+INNER_LIPS = [
+    78,
+    191,
+    80,
+    81,
+    82,
+    13,
+    312,
+    311,
+    310,
+    415,
+    308,
+    324,
+    318,
+    402,
+    317,
+    14,
+    87,
+    178,
+    88,
+    95,
+    78,
+]
+
+# Auto-Canny threshold factors (median-relative)
+_CANNY_LOW_FACTOR = 0.66
+_CANNY_HIGH_FACTOR = 1.33
+_CANNY_DEFAULT_MEDIAN = 128.0  # fallback when no non-zero pixels exist
+
+ALL_CONTOURS = [
+    JAWLINE_CONTOUR,
+    LEFT_EYE_CONTOUR,
+    RIGHT_EYE_CONTOUR,
+    LEFT_EYEBROW,
+    RIGHT_EYEBROW,
+    NOSE_BRIDGE,
+    NOSE_TIP,
+    NOSE_BOTTOM,
+    OUTER_LIPS,
+    INNER_LIPS,
+]
+
+
+def render_wireframe(
+    face: FaceLandmarks,
+    width: int | None = None,
+    height: int | None = None,
+    thickness: int = 1,
+) -> np.ndarray:
+    """Render static anatomical adjacency wireframe on black canvas.
+
+    Args:
+        face: Facial landmarks (normalized coordinates).
+        width: Canvas width.
+        height: Canvas height.
+        thickness: Line thickness in pixels.
+
+    Returns:
+        Grayscale wireframe image.
+    """
+    w = width or face.image_width
+    h = height or face.image_height
+    canvas = np.zeros((h, w), dtype=np.uint8)
+
+    coords = face.landmarks[:, :2].copy()
+    coords[:, 0] *= w
+    coords[:, 1] *= h
+    pts = coords.astype(np.int32)
+
+    for contour in ALL_CONTOURS:
+        for i in range(len(contour) - 1):
+            p1 = tuple(pts[contour[i]])
+            p2 = tuple(pts[contour[i + 1]])
+            cv2.line(canvas, p1, p2, 255, thickness)
+
+    return canvas
+
+
+def auto_canny(image: np.ndarray) -> np.ndarray:
+    """Auto-Canny edge detection with adaptive thresholds.
+
+    Uses median-based thresholds (0.66*median, 1.33*median) instead of
+    hardcoded 50/150 to handle all Fitzpatrick skin types.
+    Post-processes with morphological skeletonization for 1-pixel edges.
+
+    Args:
+        image: Grayscale input image.
+
+    Returns:
+        Binary edge map (uint8, 0 or 255).
+    """
+    median = np.median(image[image > 0]) if np.any(image > 0) else _CANNY_DEFAULT_MEDIAN
+    low = int(max(0, _CANNY_LOW_FACTOR * median))
+    high = int(min(255, _CANNY_HIGH_FACTOR * median))
+
+    edges = cv2.Canny(image, low, high)
+
+    # Morphological skeletonization for guaranteed 1-pixel thickness
+    # ControlNet blurs on 2+ pixel edges
+    skeleton = np.zeros_like(edges)
+    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
+    temp = edges.copy()
+
+    max_iterations = max(edges.shape[0], edges.shape[1])
+    for _ in range(max_iterations):
+        eroded = cv2.erode(temp, element)
+        dilated = cv2.dilate(eroded, element)
+        diff = cv2.subtract(temp, dilated)
+        skeleton = cv2.bitwise_or(skeleton, diff)
+        temp = eroded.copy()
+        if cv2.countNonZero(temp) == 0:
+            break
+
+    return skeleton
+
+
+def generate_conditioning(
+    face: FaceLandmarks,
+    width: int | None = None,
+    height: int | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Generate full conditioning signal for ControlNet.
+
+    Returns three channels per the spec:
+    1. Rendered landmark dots (colored, BGR)
+    2. Canny edge map from static wireframe (grayscale)
+    3. Wireframe rendering (grayscale)
+
+    Args:
+        face: Extracted facial landmarks.
+        width: Output width.
+        height: Output height.
+
+    Returns:
+        Tuple of (landmark_image, canny_edges, wireframe).
+    """
+    from landmarkdiff.landmarks import render_landmark_image
+
+    w = width or face.image_width
+    h = height or face.image_height
+
+    landmark_img = render_landmark_image(face, w, h)
+    wireframe = render_wireframe(face, w, h)
+    canny = auto_canny(wireframe)
+
+    return landmark_img, canny, wireframe

landmarkdiff/config.py

@@ -0,0 +1,358 @@
+"""YAML-based experiment configuration for reproducible training and evaluation.
+
+Provides typed dataclasses that can be loaded from YAML files, enabling
+reproducible experiments with version-tracked configs.
+
+Usage:
+    from landmarkdiff.config import ExperimentConfig
+    config = ExperimentConfig.from_yaml("configs/rhinoplasty_phaseA.yaml")
+    print(config.training.learning_rate)
+
+    # Or create programmatically
+    config = ExperimentConfig(
+        experiment_name="rhino_v1",
+        training=TrainingConfig(phase="A", learning_rate=1e-5),
+    )
+    config.to_yaml("configs/rhino_v1.yaml")
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass
+class ModelConfig:
+    """ControlNet and base model configuration."""
+
+    base_model: str = "runwayml/stable-diffusion-v1-5"
+    controlnet_conditioning_channels: int = 3
+    controlnet_conditioning_scale: float = 1.0
+    use_ema: bool = True
+    ema_decay: float = 0.9999
+    gradient_checkpointing: bool = True
+
+
+@dataclass
+class TrainingConfig:
+    """Training hyperparameters."""
+
+    phase: str = "A"  # "A" or "B"
+    learning_rate: float = 1e-5
+    batch_size: int = 4
+    gradient_accumulation_steps: int = 4
+    max_train_steps: int = 50000
+    warmup_steps: int = 500
+    mixed_precision: str = "bf16"
+    seed: int = 42
+    ema_decay: float = 0.9999
+
+    # Optimizer
+    optimizer: str = "adamw"  # "adamw", "adam8bit", "prodigy"
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.999
+    weight_decay: float = 1e-2
+    max_grad_norm: float = 1.0
+
+    # LR scheduler
+    lr_scheduler: str = "cosine"
+    lr_scheduler_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    # Logging intervals
+    log_every: int = 100
+    sample_every: int = 1000
+
+    # Phase B specific
+    identity_loss_weight: float = 0.1
+    perceptual_loss_weight: float = 0.05
+    use_differentiable_arcface: bool = False
+    arcface_weights_path: str | None = None
+
+    # Loss weights (alternative to individual weights)
+    loss_weights: dict[str, float] = field(default_factory=dict)
+
+    # Checkpointing
+    save_every_n_steps: int = 5000
+    resume_from_checkpoint: str | None = None
+    resume_phase_a: str | None = None
+
+    # Validation
+    validate_every_n_steps: int = 2500
+    num_validation_samples: int = 4
+
+
+@dataclass
+class DataConfig:
+    """Dataset configuration."""
+
+    train_dir: str = "data/training_combined"
+    val_dir: str = "data/splits/val"
+    test_dir: str = "data/splits/test"
+    image_size: int = 512
+    num_workers: int = 4
+    pin_memory: bool = True
+
+    # Augmentation
+    random_flip: bool = True
+    random_rotation: float = 5.0  # degrees
+    color_jitter: float = 0.1
+    clinical_augment: bool = False
+    geometric_augment: bool = True
+
+    # Procedure filtering
+    procedures: list[str] = field(
+        default_factory=lambda: [
+            "rhinoplasty",
+            "blepharoplasty",
+            "rhytidectomy",
+            "orthognathic",
+            "brow_lift",
+            "mentoplasty",
+        ]
+    )
+    intensity_range: tuple[float, float] = (30.0, 100.0)
+
+    # Data-driven displacement
+    displacement_model_path: str | None = None
+    noise_scale: float = 0.1
+
+
+@dataclass
+class InferenceConfig:
+    """Inference / generation configuration."""
+
+    num_inference_steps: int = 30
+    guidance_scale: float = 7.5
+    scheduler: str = "dpmsolver++"  # "ddpm", "ddim", "dpmsolver++"
+    controlnet_conditioning_scale: float = 1.0
+
+    # Post-processing
+    use_neural_postprocess: bool = False
+    restore_mode: str = "codeformer"
+    codeformer_fidelity: float = 0.7
+    use_realesrgan: bool = True
+    use_laplacian_blend: bool = True
+    sharpen_strength: float = 0.25
+
+    # Identity verification
+    verify_identity: bool = True
+    identity_threshold: float = 0.6
+
+
+@dataclass
+class EvaluationConfig:
+    """Evaluation configuration."""
+
+    compute_fid: bool = True
+    compute_lpips: bool = True
+    compute_nme: bool = True
+    compute_identity: bool = True
+    compute_ssim: bool = True
+    stratify_fitzpatrick: bool = True
+    stratify_procedure: bool = True
+    max_eval_samples: int = 0  # 0 = all
+
+
+@dataclass
+class WandbConfig:
+    """Weights & Biases logging configuration."""
+
+    enabled: bool = True
+    project: str = "landmarkdiff"
+    entity: str | None = None
+    run_name: str | None = None
+    tags: list[str] = field(default_factory=list)
+    mode: str = "online"  # "online", "offline", "disabled"
+
+
+@dataclass
+class SlurmConfig:
+    """SLURM job submission parameters."""
+
+    partition: str = "batch_gpu"
+    account: str = ""  # Set via YAML or SLURM_ACCOUNT env var
+    gpu_type: str = "nvidia_rtx_a6000"
+    num_gpus: int = 1
+    mem: str = "48G"
+    cpus_per_task: int = 8
+    time_limit: str = "48:00:00"
+    job_prefix: str = "surgery_"
+
+
+@dataclass
+class SafetyConfig:
+    """Clinical safety and responsible AI parameters."""
+
+    identity_threshold: float = 0.6
+    max_displacement_fraction: float = 0.05
+    watermark_enabled: bool = True
+    watermark_text: str = "AI-GENERATED PREDICTION"
+    ood_detection_enabled: bool = True
+    ood_confidence_threshold: float = 0.3
+    min_face_confidence: float = 0.5
+    max_yaw_degrees: float = 45.0
+
+
+@dataclass
+class ExperimentConfig:
+    """Top-level experiment configuration."""
+
+    experiment_name: str = "default"
+    description: str = ""
+    version: str = "0.3.2"
+
+    model: ModelConfig = field(default_factory=ModelConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    inference: InferenceConfig = field(default_factory=InferenceConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
+    wandb: WandbConfig = field(default_factory=WandbConfig)
+    slurm: SlurmConfig = field(default_factory=SlurmConfig)
+    safety: SafetyConfig = field(default_factory=SafetyConfig)
+
+    # Output
+    output_dir: str = "outputs"
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> ExperimentConfig:
+        """Load config from a YAML file."""
+        path = Path(path)
+        with open(path) as f:
+            raw = yaml.safe_load(f)
+
+        if raw is None:
+            return cls()
+
+        return cls(
+            experiment_name=raw.get("experiment_name", "default"),
+            description=raw.get("description", ""),
+            version=raw.get("version", "0.3.2"),
+            model=_from_dict(ModelConfig, raw.get("model", {})),
+            training=_from_dict(TrainingConfig, raw.get("training", {})),
+            data=_from_dict(DataConfig, raw.get("data", {})),
+            inference=_from_dict(InferenceConfig, raw.get("inference", {})),
+            evaluation=_from_dict(EvaluationConfig, raw.get("evaluation", {})),
+            wandb=_from_dict(WandbConfig, raw.get("wandb", {})),
+            slurm=_from_dict(SlurmConfig, raw.get("slurm", {})),
+            safety=_from_dict(SafetyConfig, raw.get("safety", {})),
+            output_dir=raw.get("output_dir", "outputs"),
+        )
+
+    def to_yaml(self, path: str | Path) -> None:
+        """Save config to a YAML file."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        d = _convert_tuples(asdict(self))
+        with open(path, "w") as f:
+            yaml.dump(d, f, default_flow_style=False, sort_keys=False)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        return asdict(self)
+
+
+_FIELD_ALIASES: dict[str, str] = {
+    # YAML name -> dataclass field name
+    "max_steps": "max_train_steps",
+    "save_interval": "save_every_n_steps",
+    "sample_interval": "sample_every",
+    "log_interval": "log_every",
+    "adam_weight_decay": "weight_decay",
+    "lr_warmup_steps": "warmup_steps",
+    "resume_from": "resume_from_checkpoint",
+}
+
+
+def _from_dict(cls: type, d: dict) -> Any:
+    """Create a dataclass from a dict, ignoring unknown keys.
+
+    Supports field aliases so YAML configs using train_controlnet.py-style
+    names (e.g. max_steps) map to dataclass fields (max_train_steps).
+    """
+    import dataclasses
+
+    field_map = {f.name: f for f in dataclasses.fields(cls)}
+    filtered = {}
+    for k, v in d.items():
+        # Resolve aliases
+        canonical = _FIELD_ALIASES.get(k, k)
+        if canonical not in field_map:
+            continue
+        # Don't overwrite if the canonical name was already set explicitly
+        if canonical in filtered:
+            continue
+        # Convert lists back to tuples where the field type is tuple
+        f = field_map[canonical]
+        if isinstance(v, list) and "tuple" in str(f.type):
+            v = tuple(v)
+        filtered[canonical] = v
+    return cls(**filtered)
+
+
+def _convert_tuples(obj: Any) -> Any:
+    """Recursively convert tuples to lists for YAML serialization."""
+    if isinstance(obj, dict):
+        return {k: _convert_tuples(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_convert_tuples(item) for item in obj]
+    return obj
+
+
+def load_config(
+    config_path: str | Path | None = None,
+    overrides: dict[str, object] | None = None,
+) -> ExperimentConfig:
+    """Load config with optional dot-notation overrides.
+
+    Args:
+        config_path: Path to YAML config. None returns defaults.
+        overrides: Dict of "section.key" -> value overrides.
+            E.g., {"training.learning_rate": 5e-6}
+
+    Returns:
+        ExperimentConfig with overrides applied.
+    """
+    config = ExperimentConfig.from_yaml(config_path) if config_path else ExperimentConfig()
+
+    if overrides:
+        for key, value in overrides.items():
+            parts = key.split(".")
+            obj = config
+            resolved = True
+            for part in parts[:-1]:
+                if hasattr(obj, part):
+                    obj = getattr(obj, part)
+                else:
+                    resolved = False
+                    break
+            if resolved and hasattr(obj, parts[-1]):
+                setattr(obj, parts[-1], value)
+
+    return config
+
+
+def validate_config(config: ExperimentConfig) -> list[str]:
+    """Validate config and return list of warnings."""
+    warnings = []
+
+    if config.training.phase == "B" and not config.training.resume_from_checkpoint:
+        warnings.append("Phase B should resume from a Phase A checkpoint")
+
+    eff_batch = config.training.batch_size * config.training.gradient_accumulation_steps
+    if eff_batch < 8:
+        warnings.append(f"Effective batch size {eff_batch} < 8 may cause instability")
+
+    if config.training.learning_rate > 1e-4:
+        warnings.append("Learning rate > 1e-4 is unusually high for fine-tuning")
+
+    if config.data.image_size != 512:
+        warnings.append(f"Image size {config.data.image_size} != 512; SD1.5 expects 512")
+
+    if config.safety.identity_threshold < 0.3:
+        warnings.append("Identity threshold < 0.3 may pass poor quality outputs")
+
+    return warnings