landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/model_registry.py

@@ -0,0 +1,362 @@
+"""Model registry for checkpoint discovery and management.
+
+Provides a unified interface for finding, loading, and comparing model
+checkpoints across local directories and remote sources.
+
+Usage:
+    from landmarkdiff.model_registry import ModelRegistry
+
+    registry = ModelRegistry("checkpoints/")
+
+    # Discover all checkpoints
+    models = registry.list_models()
+
+    # Get best checkpoint by metric
+    best = registry.get_best("loss")
+
+    # Load a specific checkpoint
+    state = registry.load("checkpoint-5000")
+
+    # Compare multiple checkpoints
+    comparison = registry.compare(["checkpoint-1000", "checkpoint-5000"])
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+
+
+@dataclass
+class ModelEntry:
+    """Metadata for a registered model checkpoint."""
+
+    name: str
+    path: Path
+    step: int = 0
+    phase: str = ""
+    metrics: dict[str, float] = field(default_factory=dict)
+    size_mb: float = 0.0
+    has_ema: bool = False
+    has_training_state: bool = False
+
+    @property
+    def inference_path(self) -> Path | None:
+        """Path to inference-ready weights (EMA preferred)."""
+        ema_dir = self.path / "controlnet_ema"
+        if ema_dir.exists():
+            return ema_dir
+        # Fallback to training state
+        state_path = self.path / "training_state.pt"
+        if state_path.exists():
+            return state_path
+        return None
+
+
+class ModelRegistry:
+    """Central registry for discovering and managing model checkpoints.
+
+    Args:
+        checkpoint_dirs: One or more directories to scan for checkpoints.
+        scan_on_init: Whether to scan directories immediately on creation.
+    """
+
+    def __init__(
+        self,
+        *checkpoint_dirs: str | Path,
+        scan_on_init: bool = True,
+    ) -> None:
+        self.checkpoint_dirs = [Path(d) for d in checkpoint_dirs]
+        self._models: dict[str, ModelEntry] = {}
+
+        if scan_on_init:
+            self.scan()
+
+    def scan(self) -> int:
+        """Scan checkpoint directories and register all found models.
+
+        Returns:
+            Number of models found.
+        """
+        self._models.clear()
+        for base_dir in self.checkpoint_dirs:
+            if not base_dir.exists():
+                continue
+            self._scan_directory(base_dir)
+        return len(self._models)
+
+    def _scan_directory(self, base_dir: Path) -> None:
+        """Scan a single directory for checkpoint subdirectories."""
+        # Look for checkpoint-* directories
+        for ckpt_dir in sorted(base_dir.glob("checkpoint-*")):
+            if not ckpt_dir.is_dir():
+                continue
+            entry = self._load_entry(ckpt_dir)
+            if entry is not None:
+                self._models[entry.name] = entry
+
+        # Also check for "final" and "best" directories/symlinks
+        for special in ["final", "best", "latest"]:
+            special_dir = base_dir / special
+            if special_dir.exists() and special_dir.is_dir():
+                entry = self._load_entry(special_dir)
+                if entry is not None:
+                    entry.name = f"{base_dir.name}/{special}"
+                    self._models[entry.name] = entry
+
+    def _load_entry(self, ckpt_dir: Path) -> ModelEntry | None:
+        """Load metadata for a single checkpoint directory."""
+        has_training = (ckpt_dir / "training_state.pt").exists()
+        has_ema = (ckpt_dir / "controlnet_ema").exists()
+
+        if not has_training and not has_ema:
+            return None
+
+        # Try to load metadata.json (from CheckpointManager)
+        meta_path = ckpt_dir / "metadata.json"
+        if meta_path.exists():
+            with open(meta_path) as f:
+                meta = json.load(f)
+            return ModelEntry(
+                name=ckpt_dir.name,
+                path=ckpt_dir,
+                step=meta.get("step", 0),
+                phase=meta.get("phase", ""),
+                metrics=meta.get("metrics", {}),
+                size_mb=meta.get("size_mb", 0.0),
+                has_ema=has_ema,
+                has_training_state=has_training,
+            )
+
+        # Fallback: extract step from directory name
+        step = 0
+        parts = ckpt_dir.name.split("-")
+        if len(parts) >= 2 and parts[-1].isdigit():
+            step = int(parts[-1])
+
+        # Compute size
+        size_mb = sum(f.stat().st_size for f in ckpt_dir.rglob("*") if f.is_file()) / (1024 * 1024)
+
+        return ModelEntry(
+            name=ckpt_dir.name,
+            path=ckpt_dir,
+            step=step,
+            size_mb=round(size_mb, 1),
+            has_ema=has_ema,
+            has_training_state=has_training,
+        )
+
+    # ------------------------------------------------------------------
+    # Queries
+    # ------------------------------------------------------------------
+
+    def list_models(self, sort_by: str = "step") -> list[ModelEntry]:
+        """List all registered models.
+
+        Args:
+            sort_by: Sort key — "step", "name", or a metric name.
+
+        Returns:
+            Sorted list of ModelEntry objects.
+        """
+        models = list(self._models.values())
+        if sort_by == "step":
+            models.sort(key=lambda m: m.step)
+        elif sort_by == "name":
+            models.sort(key=lambda m: m.name)
+        else:
+            # Sort by metric value
+            models.sort(
+                key=lambda m: m.metrics.get(sort_by, float("inf")),
+            )
+        return models
+
+    def get(self, name: str) -> ModelEntry | None:
+        """Get a model entry by name."""
+        return self._models.get(name)
+
+    def get_best(
+        self,
+        metric: str = "loss",
+        lower_is_better: bool = True,
+    ) -> ModelEntry | None:
+        """Get the best model by a specific metric.
+
+        Args:
+            metric: Metric name to rank by.
+            lower_is_better: If True, lower values are better.
+
+        Returns:
+            Best ModelEntry, or None if no models have the metric.
+        """
+        candidates = [m for m in self._models.values() if metric in m.metrics]
+        if not candidates:
+            return None
+
+        return (
+            min(candidates, key=lambda m: m.metrics[metric])
+            if lower_is_better
+            else max(candidates, key=lambda m: m.metrics[metric])
+        )
+
+    def get_by_step(self, step: int) -> ModelEntry | None:
+        """Get a model by its training step."""
+        for model in self._models.values():
+            if model.step == step:
+                return model
+        return None
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+
+    def load(
+        self,
+        name: str,
+        map_location: str = "cpu",
+    ) -> dict[str, Any]:
+        """Load training state from a checkpoint.
+
+        Args:
+            name: Checkpoint name (e.g. "checkpoint-5000").
+            map_location: Device to load tensors to.
+
+        Returns:
+            State dict containing controlnet, ema_controlnet, optimizer, etc.
+
+        Raises:
+            KeyError: If checkpoint not found.
+            FileNotFoundError: If training_state.pt missing.
+        """
+        entry = self._models.get(name)
+        if entry is None:
+            raise KeyError(f"Checkpoint '{name}' not found in registry")
+
+        state_path = entry.path / "training_state.pt"
+        if not state_path.exists():
+            raise FileNotFoundError(f"No training_state.pt in {entry.path}")
+
+        return torch.load(state_path, map_location=map_location, weights_only=True)
+
+    def load_controlnet(
+        self,
+        name: str,
+        use_ema: bool = True,
+    ) -> Any:
+        """Load a ControlNet model from checkpoint.
+
+        Args:
+            name: Checkpoint name.
+            use_ema: If True, load EMA weights (preferred for inference).
+
+        Returns:
+            ControlNetModel instance.
+        """
+        from diffusers import ControlNetModel
+
+        entry = self._models.get(name)
+        if entry is None:
+            raise KeyError(f"Checkpoint '{name}' not found in registry")
+
+        if use_ema and entry.has_ema:
+            return ControlNetModel.from_pretrained(str(entry.path / "controlnet_ema"))
+
+        # Fallback: load from training state
+        state = self.load(name)
+        model = ControlNetModel.from_pretrained(
+            "lllyasviel/control_v11p_sd15_openpose",
+            subfolder="diffusion_sd15",
+        )
+        key = "ema_controlnet" if use_ema else "controlnet"
+        model.load_state_dict(state[key])
+        return model
+
+    # ------------------------------------------------------------------
+    # Comparison
+    # ------------------------------------------------------------------
+
+    def compare(
+        self,
+        names: list[str],
+        metrics: list[str] | None = None,
+    ) -> dict[str, Any]:
+        """Compare multiple checkpoints side-by-side.
+
+        Args:
+            names: List of checkpoint names to compare.
+            metrics: Specific metrics to include. None = all available.
+
+        Returns:
+            Dict with comparison table data.
+        """
+        entries = []
+        for name in names:
+            entry = self._models.get(name)
+            if entry is not None:
+                entries.append(entry)
+
+        if not entries:
+            return {"error": "No valid checkpoints found"}
+
+        # Collect all available metrics
+        if metrics is None:
+            all_metrics: set[str] = set()
+            for e in entries:
+                all_metrics.update(e.metrics.keys())
+            metrics = sorted(all_metrics)
+
+        rows = []
+        for e in entries:
+            row = {
+                "name": e.name,
+                "step": e.step,
+                "phase": e.phase,
+                "size_mb": e.size_mb,
+            }
+            for m in metrics:
+                row[m] = e.metrics.get(m)
+            rows.append(row)
+
+        return {
+            "metrics": metrics,
+            "rows": rows,
+            "count": len(rows),
+        }
+
+    # ------------------------------------------------------------------
+    # Summary
+    # ------------------------------------------------------------------
+
+    def summary(self) -> str:
+        """Return a human-readable summary."""
+        models = self.list_models()
+        if not models:
+            return "No models registered."
+
+        total_size = sum(m.size_mb for m in models)
+        lines = [
+            f"Model Registry: {len(models)} checkpoints ({total_size:.0f} MB)",
+            f"  Steps: {models[0].step} — {models[-1].step}",
+        ]
+
+        # Show metrics ranges
+        all_metrics: set[str] = set()
+        for m in models:
+            all_metrics.update(m.metrics.keys())
+
+        for metric in sorted(all_metrics):
+            values = [m.metrics[metric] for m in models if metric in m.metrics]
+            if values:
+                lines.append(f"  {metric}: {min(values):.4f} — {max(values):.4f}")
+
+        return "\n".join(lines)
+
+    def __len__(self) -> int:
+        return len(self._models)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._models

landmarkdiff/morphometry.py

@@ -0,0 +1,342 @@
+"""Nasal morphometry and facial symmetry evaluation.
+
+Geometric evaluation metrics derived from Varghaei et al. (2025),
+adapted for evaluating surgical prediction outputs.
+
+Computes five nasal ratios plus bilateral facial symmetry from
+MediaPipe 478-point landmarks, enabling interpretable clinical
+quality assessment beyond perceptual metrics (LPIPS, FID).
+
+Usage::
+
+    from landmarkdiff.morphometry import NasalMorphometry, FacialSymmetry
+
+    morph = NasalMorphometry()
+    ratios = morph.compute(landmarks_478)
+
+    sym = FacialSymmetry()
+    score = sym.compute(landmarks_478)
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# MediaPipe landmark indices (478-point mesh)
+# Reference: https://github.com/google/mediapipe/blob/master/mediapipe/modules/face_geometry/data/canonical_face_model_uv_visualization.png
+NOSE_TIP = 1
+LEFT_NOSTRIL = 98
+RIGHT_NOSTRIL = 327
+LEFT_INNER_EYE = 133
+RIGHT_INNER_EYE = 362
+LEFT_OUTER_EYE = 33
+RIGHT_OUTER_EYE = 263
+LEFT_CHEEK = 234
+RIGHT_CHEEK = 454
+CHIN = 152
+FOREHEAD = 10
+GLABELLA = 168
+
+
+@dataclass
+class NasalRatios:
+    """Five nasal morphometric ratios from Varghaei et al. (2025).
+
+    Attributes:
+        alar_intercanthal: Alar width / intercanthal distance.
+            Ideal ~1.0 (nose width equals eye spacing).
+        alar_face_width: Alar width / face width.
+            Ideal ~0.20 (nose is 1/5 of face width).
+        nose_length_face_height: Nose length / face height.
+            Proportional measure of nose vertical extent.
+        tip_midline_deviation: Horizontal offset of nose tip from
+            facial midline, normalized by face width. Lower is better.
+        nostril_vertical_asymmetry: Vertical height difference between
+            nostrils, normalized by face height. Lower is better.
+    """
+
+    alar_intercanthal: float = 0.0
+    alar_face_width: float = 0.0
+    nose_length_face_height: float = 0.0
+    tip_midline_deviation: float = 0.0
+    nostril_vertical_asymmetry: float = 0.0
+
+    def improvement_score(self, reference: NasalRatios) -> dict[str, bool]:
+        """Check which ratios improved relative to reference (pre-op).
+
+        A ratio 'improved' if the prediction moved it closer to the
+        anthropometric ideal compared to the reference.
+        """
+        ideals = {
+            "alar_intercanthal": 1.0,
+            "alar_face_width": 0.20,
+        }
+        results = {}
+        for name, ideal in ideals.items():
+            pred_val = getattr(self, name)
+            ref_val = getattr(reference, name)
+            results[name] = abs(pred_val - ideal) < abs(ref_val - ideal)
+
+        # For deviation/asymmetry, lower is always better
+        results["tip_midline_deviation"] = (
+            self.tip_midline_deviation < reference.tip_midline_deviation
+        )
+        results["nostril_vertical_asymmetry"] = (
+            self.nostril_vertical_asymmetry < reference.nostril_vertical_asymmetry
+        )
+        return results
+
+    def to_dict(self) -> dict[str, float]:
+        return {
+            "alar_intercanthal": self.alar_intercanthal,
+            "alar_face_width": self.alar_face_width,
+            "nose_length_face_height": self.nose_length_face_height,
+            "tip_midline_deviation": self.tip_midline_deviation,
+            "nostril_vertical_asymmetry": self.nostril_vertical_asymmetry,
+        }
+
+
+class NasalMorphometry:
+    """Compute nasal morphometric ratios from MediaPipe landmarks.
+
+    Five geometric features following Varghaei et al. (2025):
+    1. Alar width / intercanthal distance (ideal ~1.0)
+    2. Alar width / face width (ideal ~0.20)
+    3. Nose length / face height
+    4. Tip midline deviation (normalized)
+    5. Nostril vertical asymmetry (normalized)
+    """
+
+    def compute(self, landmarks: np.ndarray) -> NasalRatios:
+        """Compute all five nasal ratios.
+
+        Args:
+            landmarks: (N, 2) or (N, 3) array of MediaPipe landmarks.
+                Must have at least 478 points. Uses only x, y.
+
+        Returns:
+            NasalRatios dataclass with computed values.
+        """
+        pts = landmarks[:, :2]  # use only x, y
+
+        # Key points
+        nose_tip = pts[NOSE_TIP]
+        left_nostril = pts[LEFT_NOSTRIL]
+        right_nostril = pts[RIGHT_NOSTRIL]
+        left_inner_eye = pts[LEFT_INNER_EYE]
+        right_inner_eye = pts[RIGHT_INNER_EYE]
+        left_cheek = pts[LEFT_CHEEK]
+        right_cheek = pts[RIGHT_CHEEK]
+        forehead = pts[FOREHEAD]
+        chin = pts[CHIN]
+        glabella = pts[GLABELLA]
+
+        # Distances (cast to float for mypy compatibility)
+        alar_width: float = float(np.linalg.norm(left_nostril - right_nostril))
+        intercanthal: float = max(float(np.linalg.norm(left_inner_eye - right_inner_eye)), 1e-6)
+        face_width: float = max(float(np.linalg.norm(left_cheek - right_cheek)), 1e-6)
+        face_height: float = max(float(np.linalg.norm(forehead - chin)), 1e-6)
+        nose_length: float = float(np.linalg.norm(glabella - nose_tip))
+
+        # Facial midline (between outer eye corners)
+        midline_x = (pts[LEFT_OUTER_EYE][0] + pts[RIGHT_OUTER_EYE][0]) / 2
+
+        # Ratios
+        alar_intercanthal = float(alar_width / intercanthal)
+        alar_face = float(alar_width / face_width)
+        nose_face = float(nose_length / face_height)
+        tip_deviation = float(abs(nose_tip[0] - midline_x) / face_width)
+        nostril_asymmetry = float(abs(left_nostril[1] - right_nostril[1]) / face_height)
+
+        return NasalRatios(
+            alar_intercanthal=alar_intercanthal,
+            alar_face_width=alar_face,
+            nose_length_face_height=nose_face,
+            tip_midline_deviation=tip_deviation,
+            nostril_vertical_asymmetry=nostril_asymmetry,
+        )
+
+    def compute_from_image(self, image: np.ndarray) -> NasalRatios | None:
+        """Extract landmarks from image and compute ratios.
+
+        Args:
+            image: BGR uint8 image (H, W, 3).
+
+        Returns:
+            NasalRatios or None if landmark detection fails.
+        """
+        try:
+            import mediapipe as mp
+        except ImportError:
+            logger.warning("mediapipe required for landmark extraction")
+            return None
+
+        with mp.solutions.face_mesh.FaceMesh(
+            static_image_mode=True,
+            max_num_faces=1,
+            refine_landmarks=True,
+            min_detection_confidence=0.5,
+        ) as face_mesh:
+            import cv2
+
+            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            results = face_mesh.process(rgb)
+
+            if not results.multi_face_landmarks:
+                return None
+
+            h, w = image.shape[:2]
+            face = results.multi_face_landmarks[0]
+            landmarks = np.array([(lm.x * w, lm.y * h) for lm in face.landmark])
+            return self.compute(landmarks)
+
+
+class FacialSymmetry:
+    """Bilateral facial symmetry scoring.
+
+    Measures deviation from perfect bilateral symmetry by reflecting
+    left-side landmarks across the facial midline and computing
+    distances to nearest right-side counterparts.
+
+    Lower scores indicate greater symmetry.
+    """
+
+    def compute(
+        self,
+        landmarks: np.ndarray,
+        left_eye_idx: int = LEFT_OUTER_EYE,
+        right_eye_idx: int = RIGHT_OUTER_EYE,
+    ) -> float:
+        """Compute bilateral symmetry error.
+
+        Args:
+            landmarks: (N, 2) or (N, 3) array. Uses only x, y.
+            left_eye_idx: Landmark index for left outer eye corner.
+            right_eye_idx: Landmark index for right outer eye corner.
+
+        Returns:
+            Mean symmetry error (lower = more symmetric).
+            Normalized by inter-ocular distance.
+        """
+        pts = landmarks[:, :2].copy()
+
+        # Midline from eye corners
+        midline_x = (pts[left_eye_idx][0] + pts[right_eye_idx][0]) / 2
+        iod = abs(pts[left_eye_idx][0] - pts[right_eye_idx][0])
+        if iod < 1e-6:
+            return 0.0
+
+        # Partition into left and right
+        left_mask = pts[:, 0] < midline_x
+        right_mask = pts[:, 0] > midline_x
+
+        left_pts = pts[left_mask]
+        right_pts = pts[right_mask]
+
+        if len(left_pts) == 0 or len(right_pts) == 0:
+            return 0.0
+
+        # Reflect left across midline
+        reflected = left_pts.copy()
+        reflected[:, 0] = 2 * midline_x - reflected[:, 0]
+
+        # KDTree nearest-neighbor matching
+        try:
+            from scipy.spatial import KDTree
+
+            tree = KDTree(right_pts)
+            distances, _ = tree.query(reflected)
+            return float(np.mean(distances) / iod)
+        except ImportError:
+            # Fallback: brute force
+            total = 0.0
+            for pt in reflected:
+                dists = np.linalg.norm(right_pts - pt, axis=1)
+                total += np.min(dists)
+            return float(total / (len(reflected) * iod))
+
+    def compute_from_image(self, image: np.ndarray) -> float | None:
+        """Extract landmarks from image and compute symmetry.
+
+        Args:
+            image: BGR uint8 image (H, W, 3).
+
+        Returns:
+            Symmetry error or None if detection fails.
+        """
+        try:
+            import mediapipe as mp
+        except ImportError:
+            logger.warning("mediapipe required for landmark extraction")
+            return None
+
+        with mp.solutions.face_mesh.FaceMesh(
+            static_image_mode=True,
+            max_num_faces=1,
+            refine_landmarks=True,
+            min_detection_confidence=0.5,
+        ) as face_mesh:
+            import cv2
+
+            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            results = face_mesh.process(rgb)
+
+            if not results.multi_face_landmarks:
+                return None
+
+            h, w = image.shape[:2]
+            face = results.multi_face_landmarks[0]
+            landmarks = np.array([(lm.x * w, lm.y * h) for lm in face.landmark])
+            return self.compute(landmarks)
+
+
+def compare_morphometry(
+    pred_image: np.ndarray,
+    input_image: np.ndarray,
+    procedure: str = "rhinoplasty",
+) -> dict:
+    """Compare morphometric quality between prediction and input.
+
+    Computes nasal ratios and symmetry for both images and reports
+    which metrics improved. Useful for evaluating whether the predicted
+    surgical output shows clinically meaningful improvement.
+
+    Args:
+        pred_image: Predicted output (BGR uint8).
+        input_image: Original input (BGR uint8).
+        procedure: Procedure type (affects which metrics are relevant).
+
+    Returns:
+        Dict with 'input_ratios', 'pred_ratios', 'improvements',
+        'input_symmetry', 'pred_symmetry', 'symmetry_improved'.
+    """
+    morph = NasalMorphometry()
+    sym = FacialSymmetry()
+
+    input_ratios = morph.compute_from_image(input_image)
+    pred_ratios = morph.compute_from_image(pred_image)
+    input_sym = sym.compute_from_image(input_image)
+    pred_sym = sym.compute_from_image(pred_image)
+
+    result: dict = {
+        "procedure": procedure,
+        "input_ratios": input_ratios.to_dict() if input_ratios else None,
+        "pred_ratios": pred_ratios.to_dict() if pred_ratios else None,
+        "input_symmetry": input_sym,
+        "pred_symmetry": pred_sym,
+        "symmetry_improved": (
+            pred_sym < input_sym if pred_sym is not None and input_sym is not None else None
+        ),
+    }
+
+    if input_ratios and pred_ratios:
+        result["improvements"] = pred_ratios.improvement_score(input_ratios)
+    else:
+        result["improvements"] = None
+
+    return result