orient_scan.py

import argparse
import json
import sys
from pathlib import Path
import debugpy

PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT / "src"))

import numpy as np
import torch

from scannormalizer.scan_inference import load_normalizer, normalize_scan, transform_scan
from scannormalizer.postprocessing import ScanPostprocessing


def parse_args():
    parser = argparse.ArgumentParser(description="Orient one scan with a trained normalizer.")
    parser.add_argument("scan", type=Path)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--output-dir", default="data/output")
    parser.add_argument("--points", type=int, default=None)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--debug", action="store_true", help="Enable debug mode with debugpy.")
    parser.add_argument(
        "--orient-only",
        action="store_true",
        help="Apply only the predicted orientation matrix; do not center or scale the output to the unit sphere.",
    )
    parser.add_argument(
        "--center-and-orient",
        action="store_true",
        help="Translate the scan to the origin and apply the orientation matrix without scaling to the unit sphere.",
    )
    parser.add_argument(
        "--preserve-occlusion",
        action="store_true",
        help="Treat the scan argument as a patient folder, or a scan inside one, and transform sibling lower.stl and upper.stl together.",
    )
    parser.add_argument(
        "--export-pca",
        action="store_true",
        help="Export a PCA-aligned version of the scan alongside the oriented scan.",
    )
    parser.add_argument(
        "--save-matrix-npy",
        nargs="?",
        const=True,
        default=None,
        metavar="PATH",
        help=(
            "Save the transformation matrix as a .npy file. Optionally pass an output path: "
            "a file path, or a directory to write <scan>_oriented.npy into. "
            "When no path is given, it is saved alongside each transformed scan."
        ),
    )
    parser.add_argument(
        "--save-matrix-json",
        nargs="?",
        const=True,
        default=None,
        metavar="PATH",
        help=(
            "Save the transformation matrix as a JSON file. Optionally pass an output path: "
            "a file path, or a directory to write <scan>_oriented.json into. "
            "When no path is given, it is saved alongside each transformed scan."
        ),
    )
    parser.add_argument(
        "--postprocessing",
        type=Path,
        default=None,
        metavar="YAML",
        help=(
            "Optional YAML file with extra rotations to apply after alignment, in the "
            "aligned reference frame. Supports a top-level 'transforms' list and/or "
            "'lower'/'upper' sections; each transform is 'type: rotate' with 'axis' and "
            "'angle_deg'. With --preserve-occlusion the 'lower' rotations are used for "
            "both scans."
        ),
    )
    args = parser.parse_args()
    if args.orient_only and args.center_and_orient:
        parser.error("choose either --orient-only or --center-and-orient")
    return args


def main():
    args = parse_args()
    if args.orient_only and args.center_and_orient:
        raise RuntimeError("Choose either --orient-only or --center-and-orient")
    if args.debug == True:
        print("Hello, happy debugging.")
        debugpy.listen(("0.0.0.0", 5681))
        print(">>> Debugger is listening on port 5681. Waiting for client to attach...")
        debugpy.wait_for_client()
        print(">>> Debugger attached. Resuming execution.")
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    normalizer = load_normalizer(args.checkpoint, args.device, points=args.points)

    postprocessing = (
        ScanPostprocessing.load(args.postprocessing) if args.postprocessing else None
    )

    if args.preserve_occlusion:
        patient_dir = args.scan if args.scan.is_dir() else args.scan.parent
        lower_path = find_patient_scan(patient_dir, "lower")
        upper_path = find_patient_scan(patient_dir, "upper")
        lower_output_path = output_dir / f"{lower_path.stem}_oriented{lower_path.suffix}"
        upper_output_path = output_dir / f"{upper_path.stem}_oriented{upper_path.suffix}"
        if args.export_pca:
            pca_output_path = output_dir / f"{lower_path.stem}_pca{lower_path.suffix}"
        else:
            pca_output_path = None
        result = normalize_scan(
            lower_path,
            lower_output_path,
            normalizer,
            pca_output_path=pca_output_path,
            orient_only=args.orient_only,
            center_and_orient=args.center_and_orient,
            postprocess_matrix=(
                postprocessing.matrix_for("lower") if postprocessing else None
            ),
        )
        transform_scan(
            upper_path,
            upper_output_path,
            result.matrix,
            center=result.center,
            scale=result.scale,
            orient_only=args.orient_only,
            center_and_orient=args.center_and_orient,
        )
        if args.save_matrix_npy is not None or args.save_matrix_json is not None:
            save_affine(
                result.matrix,
                result.center,
                result.scale,
                lower_output_path,
                orient_only=args.orient_only,
                center_and_orient=args.center_and_orient,
                save_npy=args.save_matrix_npy,
                save_json=args.save_matrix_json,
            )
            save_affine(
                result.matrix,
                result.center,
                result.scale,
                upper_output_path,
                orient_only=args.orient_only,
                center_and_orient=args.center_and_orient,
                save_npy=args.save_matrix_npy,
                save_json=args.save_matrix_json,
            )

        print(f"rotation logits: {result.logits}")
        print(f"selected rotation index: {result.rotation_index}")
        print(f"pca saved: {result.pca_output_path}")
        print(f"lower saved: {result.output_path}")
        print(f"upper saved: {upper_output_path}")
        return

    pca_output_path = output_dir / f"{args.scan.stem}_pca{args.scan.suffix}"
    output_path = output_dir / f"{args.scan.stem}_oriented{args.scan.suffix}"
    postprocess_matrix = None
    if postprocessing is not None:
        jaw = classify_scan(args.scan)
        postprocess_matrix = postprocessing.matrix_for(jaw)
        if not postprocessing.has_rotation_for(jaw):
            print(
                f"warning: --postprocessing has no rotations for {args.scan.name!r} "
                f"(jaw={jaw!r}); output is unchanged"
            )
    result = normalize_scan(
        args.scan,
        output_path,
        normalizer,
        pca_output_path=pca_output_path,
        orient_only=args.orient_only,
        center_and_orient=args.center_and_orient,
        postprocess_matrix=postprocess_matrix,
    )
    if args.save_matrix_npy is not None or args.save_matrix_json is not None:
        save_affine(
            result.matrix,
            result.center,
            result.scale,
            output_path,
            orient_only=args.orient_only,
            center_and_orient=args.center_and_orient,
            save_npy=args.save_matrix_npy,
            save_json=args.save_matrix_json,
        )

    print(f"rotation logits: {result.logits}")
    print(f"selected rotation index: {result.rotation_index}")
    print(f"pca saved: {result.pca_output_path}")
    print(f"saved: {result.output_path}")


def _resolve_matrix_path(dest, output_scan_path, suffix):
    if dest is True:
        return output_scan_path.with_suffix(suffix)
    dest = Path(dest)
    if dest.suffix:
        return dest
    return dest / (output_scan_path.stem + suffix)


def save_affine(
    matrix,
    center,
    scale,
    output_scan_path,
    orient_only=False,
    center_and_orient=False,
    save_npy=None,
    save_json=None,
):
    matrix = np.asarray(matrix, dtype=np.float32)
    center = None if center is None else np.asarray(center, dtype=np.float32)

    affine = np.eye(4, dtype=np.float32)

    if orient_only:
        affine[:3, :3] = matrix.T
    elif center_and_orient:
        affine[:3, :3] = matrix.T
        affine[:3, 3] = -(matrix.T @ center)
    else:
        affine[:3, :3] = matrix.T / scale
        affine[:3, 3] = -(matrix.T @ (center / scale))

    if save_npy is not None:
        npy_path = _resolve_matrix_path(save_npy, output_scan_path, ".npy")
        npy_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(npy_path, affine)

    if save_json is not None:
        json_path = _resolve_matrix_path(save_json, output_scan_path, ".json")
        json_path.parent.mkdir(parents=True, exist_ok=True)
        with open(json_path, "w") as handle:
            json.dump(affine.tolist(), handle, indent=2)


def find_patient_scan(patient_dir, scan_type):
    matches = []
    for path in Path(patient_dir).iterdir():
        if not path.is_file() or path.suffix.lower() != ".stl":
            continue
        classified_type = classify_scan(path)
        if classified_type == scan_type:
            matches.append(path)
    if len(matches) != 1:
        raise RuntimeError(f"Expected exactly one {scan_type} scan under {patient_dir}")
    return matches[0]


def classify_scan(scan_path):
    name = scan_path.name.lower()
    is_lower = "lower" in name or "mandibular" in name
    is_upper = "upper" in name or "maxillary" in name
    if is_lower and is_upper:
        raise RuntimeError(f"Ambiguous lower/upper scan name: {scan_path}")
    if is_lower:
        return "lower"
    if is_upper:
        return "upper"
    return None


if __name__ == "__main__":
    main()