landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/__init__.py

@@ -0,0 +1,40 @@
+"""LandmarkDiff: Anatomically-conditioned latent diffusion for facial surgery simulation."""
+
+__version__ = "0.2.3"
+
+__all__ = [
+    "api_client",
+    "arcface_torch",
+    "audit",
+    "augmentation",
+    "benchmark",
+    "checkpoint_manager",
+    "cli",
+    "clinical",
+    "conditioning",
+    "config",
+    "curriculum",
+    "data",
+    "data_version",
+    "displacement_model",
+    "ensemble",
+    "evaluation",
+    "experiment_tracker",
+    "face_verifier",
+    "fid",
+    "hyperparam",
+    "inference",
+    "landmarks",
+    "log",
+    "losses",
+    "manipulation",
+    "masking",
+    "metrics_agg",
+    "metrics_viz",
+    "model_registry",
+    "morphometry",
+    "postprocess",
+    "safety",
+    "synthetic",
+    "validation",
+]

landmarkdiff/__main__.py

@@ -0,0 +1,207 @@
+"""CLI entry point for python -m landmarkdiff."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import NoReturn
+
+
+def _error(msg: str) -> NoReturn:
+    """Print error to stderr and exit."""
+    print(f"error: {msg}", file=sys.stderr)
+    sys.exit(1)
+
+
+def _validate_image_path(path_str: str) -> Path:
+    """Validate that the image path exists and looks like an image file."""
+    p = Path(path_str)
+    if not p.exists():
+        _error(f"file not found: {path_str}")
+    if not p.is_file():
+        _error(f"not a file: {path_str}")
+    return p
+
+
+def main() -> None:
+    from landmarkdiff import __version__
+
+    parser = argparse.ArgumentParser(
+        prog="landmarkdiff",
+        description="Facial surgery outcome prediction from clinical photography",
+    )
+    parser.add_argument("--version", action="version", version=f"landmarkdiff {__version__}")
+
+    subparsers = parser.add_subparsers(dest="command")
+
+    # inference
+    infer = subparsers.add_parser("infer", help="Run inference on an image")
+    infer.add_argument("image", type=str, help="Path to input face image")
+    infer.add_argument(
+        "--procedure",
+        type=str,
+        default="rhinoplasty",
+        choices=[
+            "rhinoplasty",
+            "blepharoplasty",
+            "rhytidectomy",
+            "orthognathic",
+            "brow_lift",
+            "mentoplasty",
+        ],
+        help="Surgical procedure to simulate (default: rhinoplasty)",
+    )
+    infer.add_argument(
+        "--intensity",
+        type=float,
+        default=60.0,
+        help="Deformation intensity, 0-100 (default: 60)",
+    )
+    infer.add_argument(
+        "--mode",
+        type=str,
+        default="tps",
+        choices=["tps", "controlnet", "img2img", "controlnet_ip"],
+        help="Inference mode (default: tps, others require GPU)",
+    )
+    infer.add_argument(
+        "--output",
+        type=str,
+        default="output/",
+        help="Output directory (default: output/)",
+    )
+    infer.add_argument(
+        "--steps",
+        type=int,
+        default=30,
+        help="Number of diffusion steps (default: 30)",
+    )
+    infer.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+
+    # landmarks
+    lm = subparsers.add_parser("landmarks", help="Extract and visualize landmarks")
+    lm.add_argument("image", type=str, help="Path to input face image")
+    lm.add_argument(
+        "--output",
+        type=str,
+        default="output/landmarks.png",
+        help="Output path for landmark visualization (default: output/landmarks.png)",
+    )
+
+    # demo
+    subparsers.add_parser("demo", help="Launch Gradio web demo")
+
+    args = parser.parse_args()
+
+    if args.command is None:
+        parser.print_help()
+        return
+
+    try:
+        if args.command == "infer":
+            _run_inference(args)
+        elif args.command == "landmarks":
+            _run_landmarks(args)
+        elif args.command == "demo":
+            _run_demo()
+    except KeyboardInterrupt:
+        sys.exit(130)
+    except Exception as exc:
+        _error(str(exc))
+
+
+def _run_inference(args: argparse.Namespace) -> None:
+    import numpy as np
+    from PIL import Image
+
+    from landmarkdiff.landmarks import extract_landmarks
+    from landmarkdiff.manipulation import apply_procedure_preset
+
+    if not (0 <= args.intensity <= 100):
+        _error(f"intensity must be between 0 and 100, got {args.intensity}")
+
+    image_path = _validate_image_path(args.image)
+
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    img = Image.open(image_path).convert("RGB").resize((512, 512))
+    img_array = np.array(img)
+
+    landmarks = extract_landmarks(img_array)
+    if landmarks is None:
+        _error("no face detected in image")
+
+    deformed = apply_procedure_preset(landmarks, args.procedure, intensity=args.intensity)
+
+    if args.mode == "tps":
+        from landmarkdiff.synthetic.tps_warp import warp_image_tps
+
+        src = landmarks.pixel_coords[:, :2].copy()
+        dst = deformed.pixel_coords[:, :2].copy()
+        src[:, 0] *= 512 / landmarks.image_width
+        src[:, 1] *= 512 / landmarks.image_height
+        dst[:, 0] *= 512 / deformed.image_width
+        dst[:, 1] *= 512 / deformed.image_height
+        warped = warp_image_tps(img_array, src, dst)
+        Image.fromarray(warped).save(str(output_dir / "prediction.png"))
+        print(f"saved tps result to {output_dir / 'prediction.png'}")
+    else:
+        import torch
+
+        from landmarkdiff.inference import LandmarkDiffPipeline
+
+        pipeline = LandmarkDiffPipeline(mode=args.mode, device=torch.device("cuda"))
+        pipeline.load()
+        result = pipeline.generate(
+            img_array,
+            procedure=args.procedure,
+            intensity=args.intensity,
+            num_inference_steps=args.steps,
+            seed=args.seed,
+        )
+        result["output"].save(str(output_dir / "prediction.png"))
+        print(f"saved result to {output_dir / 'prediction.png'}")
+
+
+def _run_landmarks(args: argparse.Namespace) -> None:
+    import numpy as np
+    from PIL import Image
+
+    from landmarkdiff.landmarks import extract_landmarks, render_landmark_image
+
+    image_path = _validate_image_path(args.image)
+
+    img = np.array(Image.open(image_path).convert("RGB").resize((512, 512)))
+    landmarks = extract_landmarks(img)
+    if landmarks is None:
+        _error("no face detected in image")
+
+    mesh = render_landmark_image(landmarks, 512, 512)
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    Image.fromarray(mesh).save(str(output_path))
+    print(f"saved landmark mesh to {output_path}")
+    print(f"detected {len(landmarks.landmarks)} landmarks, confidence {landmarks.confidence:.2f}")
+
+
+def _run_demo() -> None:
+    try:
+        from scripts.app import build_app
+
+        demo = build_app()
+        demo.launch()
+    except ImportError:
+        _error("gradio not installed - run: pip install landmarkdiff[app]")
+
+
+if __name__ == "__main__":
+    main()

landmarkdiff/api_client.py

@@ -0,0 +1,316 @@
+"""Python client for the LandmarkDiff REST API.
+
+Provides a clean interface for interacting with the FastAPI server,
+handling image encoding/decoding, error handling, and session management.
+
+Usage:
+    from landmarkdiff.api_client import LandmarkDiffClient
+
+    client = LandmarkDiffClient("http://localhost:8000")
+
+    # Single prediction
+    result = client.predict("patient.png", procedure="rhinoplasty", intensity=65)
+    result.save("output.png")
+
+    # Face analysis
+    analysis = client.analyze("patient.png")
+    print(f"Fitzpatrick type: {analysis['fitzpatrick_type']}")
+
+    # Batch processing
+    results = client.batch_predict(
+        ["patient1.png", "patient2.png"],
+        procedure="blepharoplasty",
+    )
+"""
+
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+
+
+class LandmarkDiffAPIError(Exception):
+    """Base exception for LandmarkDiff API errors."""
+
+
+@dataclass
+class PredictionResult:
+    """Result from a single prediction."""
+
+    output_image: np.ndarray
+    procedure: str
+    intensity: float
+    confidence: float = 0.0
+    landmarks_before: list[Any] | None = None
+    landmarks_after: list[Any] | None = None
+    metrics: dict[str, float] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def save(self, path: str | Path, fmt: str = ".png") -> None:
+        """Save the output image to a file."""
+        cv2.imwrite(str(path), self.output_image)
+
+    def show(self) -> None:
+        """Display the output image (requires GUI)."""
+        cv2.imshow("LandmarkDiff Prediction", self.output_image)
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()
+
+
+class LandmarkDiffClient:
+    """Client for the LandmarkDiff REST API.
+
+    Args:
+        base_url: Server URL (e.g. "http://localhost:8000").
+        timeout: Request timeout in seconds.
+    """
+
+    def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 60.0) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self._session: Any = None
+
+    def _get_session(self) -> Any:
+        """Lazy-initialize requests session."""
+        if self._session is None:
+            try:
+                import requests
+            except ImportError:
+                raise ImportError("requests required. Install with: pip install requests") from None
+            self._session = requests.Session()
+            self._session.timeout = self.timeout
+        return self._session
+
+    def _read_image(self, image_path: str | Path) -> bytes:
+        """Read image file as bytes."""
+        path = Path(image_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Image not found: {path}")
+        return path.read_bytes()
+
+    def _decode_base64_image(self, b64_string: str) -> np.ndarray:
+        """Decode a base64-encoded image to numpy array."""
+        img_bytes = base64.b64decode(b64_string)
+        arr = np.frombuffer(img_bytes, np.uint8)
+        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+        if img is None:
+            raise ValueError("Failed to decode base64 image")
+        return img
+
+    # ------------------------------------------------------------------
+    # API methods
+    # ------------------------------------------------------------------
+
+    def health(self) -> dict[str, Any]:
+        """Check server health.
+
+        Returns:
+            Dict with status and version info.
+
+        Raises:
+            LandmarkDiffAPIError: If server is unreachable or returns an error.
+        """
+        session = self._get_session()
+        try:
+            resp = session.get(f"{self.base_url}/health")
+            resp.raise_for_status()
+            return resp.json()
+        except Exception as e:
+            import requests
+
+            if isinstance(e, requests.ConnectionError):
+                raise LandmarkDiffAPIError(
+                    f"Cannot connect to LandmarkDiff server at {self.base_url}. "
+                    f"Make sure the server is running (python -m landmarkdiff serve)."
+                ) from None
+            elif isinstance(e, requests.HTTPError):
+                raise LandmarkDiffAPIError(
+                    f"Server returned error {e.response.status_code}: {e.response.text[:200]}"
+                ) from None
+            else:
+                raise
+
+    def procedures(self) -> list[str]:
+        """List available surgical procedures.
+
+        Returns:
+            List of procedure names.
+
+        Raises:
+            LandmarkDiffAPIError: If server is unreachable or returns an error.
+        """
+        session = self._get_session()
+        try:
+            resp = session.get(f"{self.base_url}/procedures")
+            resp.raise_for_status()
+            return resp.json().get("procedures", [])
+        except Exception as e:
+            import requests
+
+            if isinstance(e, requests.ConnectionError):
+                raise LandmarkDiffAPIError(
+                    f"Cannot connect to LandmarkDiff server at {self.base_url}. "
+                    f"Make sure the server is running (python -m landmarkdiff serve)."
+                ) from None
+            elif isinstance(e, requests.HTTPError):
+                raise LandmarkDiffAPIError(
+                    f"Server returned error {e.response.status_code}: {e.response.text[:200]}"
+                ) from None
+            else:
+                raise
+
+    def predict(
+        self,
+        image_path: str | Path,
+        procedure: str = "rhinoplasty",
+        intensity: float = 65.0,
+        seed: int = 42,
+    ) -> PredictionResult:
+        """Run surgical outcome prediction.
+
+        Args:
+            image_path: Path to input face image.
+            procedure: Surgical procedure type.
+            intensity: Intensity of the modification (0-100).
+            seed: Random seed for reproducibility.
+
+        Returns:
+            PredictionResult with output image and metadata.
+        """
+        session = self._get_session()
+        image_bytes = self._read_image(image_path)
+
+        files = {"image": ("image.png", image_bytes, "image/png")}
+        data = {
+            "procedure": procedure,
+            "intensity": str(intensity),
+            "seed": str(seed),
+        }
+
+        resp = session.post(f"{self.base_url}/predict", files=files, data=data)
+        try:
+            resp.raise_for_status()
+            result = resp.json()
+
+            # Decode output image
+            output_img = self._decode_base64_image(result["output_image"])
+
+            return PredictionResult(
+                output_image=output_img,
+                procedure=procedure,
+                intensity=intensity,
+                confidence=result.get("confidence", 0.0),
+                metrics=result.get("metrics", {}),
+                metadata=result.get("metadata", {}),
+            )
+        except Exception as e:
+            import requests
+
+            if isinstance(e, requests.ConnectionError):
+                raise LandmarkDiffAPIError(
+                    f"Cannot connect to LandmarkDiff server at {self.base_url}. "
+                    f"Make sure the server is running (python -m landmarkdiff serve)."
+                ) from None
+            elif isinstance(e, requests.HTTPError):
+                raise LandmarkDiffAPIError(
+                    f"Server returned error {e.response.status_code}: {e.response.text[:200]}"
+                ) from None
+            else:
+                raise
+
+    def analyze(self, image_path: str | Path) -> dict[str, Any]:
+        """Analyze a face image without generating a prediction.
+
+        Returns face landmarks, Fitzpatrick type, pose estimation, etc.
+
+        Args:
+            image_path: Path to input face image.
+
+        Returns:
+            Dict with analysis results.
+
+        Raises:
+            LandmarkDiffAPIError: If server is unreachable or returns an error.
+        """
+        session = self._get_session()
+        image_bytes = self._read_image(image_path)
+
+        files = {"image": ("image.png", image_bytes, "image/png")}
+        try:
+            resp = session.post(f"{self.base_url}/analyze", files=files)
+            resp.raise_for_status()
+            return resp.json()
+        except Exception as e:
+            import requests
+
+            if isinstance(e, requests.ConnectionError):
+                raise LandmarkDiffAPIError(
+                    f"Cannot connect to LandmarkDiff server at {self.base_url}. "
+                    f"Make sure the server is running (python -m landmarkdiff serve)."
+                ) from None
+            elif isinstance(e, requests.HTTPError):
+                raise LandmarkDiffAPIError(
+                    f"Server returned error {e.response.status_code}: {e.response.text[:200]}"
+                ) from None
+            else:
+                raise
+
+    def batch_predict(
+        self,
+        image_paths: list[str | Path],
+        procedure: str = "rhinoplasty",
+        intensity: float = 65.0,
+        seed: int = 42,
+    ) -> list[PredictionResult]:
+        """Run batch prediction on multiple images.
+
+        Args:
+            image_paths: List of image file paths.
+            procedure: Procedure to apply to all images.
+            intensity: Intensity for all images.
+            seed: Base random seed.
+
+        Returns:
+            List of PredictionResult objects.
+        """
+        results = []
+        for i, path in enumerate(image_paths):
+            try:
+                result = self.predict(
+                    path,
+                    procedure=procedure,
+                    intensity=intensity,
+                    seed=seed + i,
+                )
+                results.append(result)
+            except Exception as e:
+                # Create a failed result
+                results.append(
+                    PredictionResult(
+                        output_image=np.zeros((512, 512, 3), dtype=np.uint8),
+                        procedure=procedure,
+                        intensity=intensity,
+                        metadata={"error": str(e), "path": str(path)},
+                    )
+                )
+        return results
+
+    def close(self) -> None:
+        """Close the HTTP session."""
+        if self._session is not None:
+            self._session.close()
+            self._session = None
+
+    def __enter__(self) -> LandmarkDiffClient:
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+    def __repr__(self) -> str:
+        return f"LandmarkDiffClient(base_url='{self.base_url}')"