landmarkdiff/__main__.py

"""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()