Add official VoxelMorph CT registration pipeline

infer.py

@@ -0,0 +1,242 @@
+"""独立推理模块。
+
+本模块加载官方 ``voxelmorph.nn.models.VxmPairwise`` 训练出的权重，输出：
+- warped_moving.nii.gz：形变后的 Moving。
+- ddf_mm.nii.gz：Dense Displacement Field，shape = (X, Y, Z, 3)，单位 mm。
+- metrics.json：配准前后 NCC/MSE/MAE 与 DDF 统计。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Dict, Iterable, Sequence, Tuple
+
+import numpy as np
+
+from config import (
+    DEFAULT_CHECKPOINT,
+    DEFAULT_TARGET_SPACING,
+    DEFAULT_WINDOW_LEVEL,
+    DEFAULT_WINDOW_WIDTH,
+    INFERENCE_DIR,
+)
+from metrics import ddf_summary, registration_metrics
+from model_and_train import build_vxm_model, require_official_voxelmorph, resolve_device
+from preprocess import load_nifti, preprocess_array, save_nifti
+
+
+def _require_nibabel():
+    try:
+        import nibabel as nib  # type: ignore
+    except Exception as exc:  # pragma: no cover
+        raise RuntimeError("缺少 nibabel，请先安装项目依赖。") from exc
+    return nib
+
+
+def _require_torch():
+    try:
+        import torch
+    except Exception as exc:  # pragma: no cover
+        raise RuntimeError("缺少 PyTorch，请先创建/激活 CUDA 环境。") from exc
+    return torch
+
+
+def _checkpoint_model_shape_xyz(checkpoint: Dict) -> Tuple[int, int, int]:
+    """读取模型训练时的输入尺寸。
+
+    兼容早期草稿里的 input_shape_dhw，但正规版本使用 input_shape_xyz。
+    """
+
+    if "input_shape_xyz" in checkpoint:
+        return tuple(int(v) for v in checkpoint["input_shape_xyz"])
+    if "input_shape_dhw" in checkpoint:
+        d, h, w = tuple(int(v) for v in checkpoint["input_shape_dhw"])
+        return (w, h, d)
+    raise KeyError("checkpoint 缺少 input_shape_xyz，无法确定模型输入尺寸。")
+
+
+def prepare_nifti_for_model(
+    nifti_path: str | Path,
+    target_shape_xyz: Sequence[int],
+    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
+    window_width: float = DEFAULT_WINDOW_WIDTH,
+    window_level: float = DEFAULT_WINDOW_LEVEL,
+    normalize_mode: str = "auto",
+    max_memory_mb: int = 4096,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """把任意 NIfTI 预处理成模型输入尺寸。"""
+
+    data, affine, spacing = load_nifti(nifti_path)
+    processed, output_affine, _ = preprocess_array(
+        data,
+        affine,
+        spacing,
+        target_spacing=target_spacing,
+        target_shape=target_shape_xyz,
+        window_width=window_width,
+        window_level=window_level,
+        normalize_mode=normalize_mode,
+        max_memory_mb=max_memory_mb,
+    )
+    return processed, output_affine
+
+
+def xyz_to_model_tensor(data_xyz: np.ndarray, torch, device) -> "torch.Tensor":
+    """NIfTI (X,Y,Z) -> 官方 VoxelMorph (B,C,X,Y,Z)。"""
+
+    return torch.from_numpy(data_xyz.astype(np.float32, copy=True))[None, None].to(device=device, dtype=torch.float32)
+
+
+def model_tensor_to_xyz(tensor) -> np.ndarray:
+    """官方 VoxelMorph (B,C,X,Y,Z) -> NIfTI (X,Y,Z)。"""
+
+    return tensor.detach().cpu().numpy()[0, 0].astype(np.float32, copy=False)
+
+
+def field_to_ddf_xyz_mm(field_tensor, spacing_xyz: Sequence[float]) -> np.ndarray:
+    """官方 displacement field (B,3,X,Y,Z) voxel -> DDF (X,Y,Z,3) mm。"""
+
+    field_cxyz = field_tensor.detach().cpu().numpy()[0].astype(np.float32, copy=False)
+    ddf_xyz = np.moveaxis(field_cxyz, 0, -1)
+    spacing = np.asarray(spacing_xyz, dtype=np.float32)
+    return (ddf_xyz * spacing[None, None, None, :]).astype(np.float32, copy=False)
+
+
+def save_ddf_nifti(ddf_xyz_mm: np.ndarray, affine: np.ndarray, output_path: str | Path) -> None:
+    nib = _require_nibabel()
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    img = nib.Nifti1Image(ddf_xyz_mm.astype(np.float32, copy=False), affine)
+    img.header.set_xyzt_units("mm")
+    img.header.set_intent("vector")
+    nib.save(img, str(output_path))
+
+
+def run_inference(
+    moving_path: str | Path,
+    fixed_path: str | Path,
+    checkpoint_path: str | Path = DEFAULT_CHECKPOINT,
+    out_dir: str | Path = INFERENCE_DIR,
+    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
+    window_width: float = DEFAULT_WINDOW_WIDTH,
+    window_level: float = DEFAULT_WINDOW_LEVEL,
+    normalize_mode: str = "auto",
+    device: str = "auto",
+    max_memory_mb: int = 4096,
+) -> Dict[str, str | float]:
+    """执行官方 VoxelMorph 推理并保存结果。"""
+
+    require_official_voxelmorph()
+    torch = _require_torch()
+    device_obj = resolve_device(device)
+    checkpoint_path = Path(checkpoint_path)
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"模型权重不存在: {checkpoint_path}")
+
+    try:
+        checkpoint = torch.load(str(checkpoint_path), map_location=device_obj, weights_only=True)
+    except TypeError:  # PyTorch < 2.0
+        checkpoint = torch.load(str(checkpoint_path), map_location=device_obj)
+    target_shape_xyz = _checkpoint_model_shape_xyz(checkpoint)
+    nb_features = checkpoint.get("nb_features", [16, 16, 16, 16, 16])
+    integration_steps = int(checkpoint.get("integration_steps", 0))
+
+    moving_xyz, _ = prepare_nifti_for_model(
+        moving_path,
+        target_shape_xyz=target_shape_xyz,
+        target_spacing=target_spacing,
+        window_width=window_width,
+        window_level=window_level,
+        normalize_mode=normalize_mode,
+        max_memory_mb=max_memory_mb,
+    )
+    fixed_xyz, fixed_affine = prepare_nifti_for_model(
+        fixed_path,
+        target_shape_xyz=target_shape_xyz,
+        target_spacing=target_spacing,
+        window_width=window_width,
+        window_level=window_level,
+        normalize_mode=normalize_mode,
+        max_memory_mb=max_memory_mb,
+    )
+
+    moving_tensor = xyz_to_model_tensor(moving_xyz, torch=torch, device=device_obj)
+    fixed_tensor = xyz_to_model_tensor(fixed_xyz, torch=torch, device=device_obj)
+
+    model = build_vxm_model(nb_features=nb_features, integration_steps=integration_steps, device=device_obj)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    with torch.no_grad():
+        displacement, warped_tensor = model(
+            moving_tensor,
+            fixed_tensor,
+            return_warped_source=True,
+            return_field_type="displacement",
+        )
+
+    warped_xyz = model_tensor_to_xyz(warped_tensor)
+    ddf_xyz_mm = field_to_ddf_xyz_mm(displacement, spacing_xyz=target_spacing)
+
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    warped_path = out_dir / "warped_moving.nii.gz"
+    ddf_path = out_dir / "ddf_mm.nii.gz"
+    metrics_path = out_dir / "metrics.json"
+
+    save_nifti(warped_xyz, fixed_affine, warped_path)
+    save_ddf_nifti(ddf_xyz_mm, fixed_affine, ddf_path)
+
+    metrics = registration_metrics(fixed_xyz, moving_xyz, warped_xyz)
+    metrics.update(ddf_summary(ddf_xyz_mm))
+    metrics.update(
+        {
+            "moving_path": str(moving_path),
+            "fixed_path": str(fixed_path),
+            "checkpoint_path": str(checkpoint_path),
+            "warped_path": str(warped_path),
+            "ddf_path": str(ddf_path),
+            "core_library": "voxelmorph.nn.models.VxmPairwise",
+        }
+    )
+    metrics_path.write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")
+    metrics["metrics_path"] = str(metrics_path)
+    return metrics
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="官方 VoxelMorph 3D 配准推理。")
+    parser.add_argument("--moving", required=True, help="Moving NIfTI 路径。")
+    parser.add_argument("--fixed", required=True, help="Fixed NIfTI 路径。")
+    parser.add_argument("--checkpoint", default=str(DEFAULT_CHECKPOINT), help="模型权重路径。")
+    parser.add_argument("--out-dir", default=str(INFERENCE_DIR), help="推理输出目录。")
+    parser.add_argument("--target-spacing", type=float, nargs=3, default=DEFAULT_TARGET_SPACING, metavar=("X", "Y", "Z"))
+    parser.add_argument("--window-width", type=float, default=DEFAULT_WINDOW_WIDTH)
+    parser.add_argument("--window-level", type=float, default=DEFAULT_WINDOW_LEVEL)
+    parser.add_argument("--normalize-mode", choices=["window", "auto", "none"], default="auto")
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--max-memory-mb", type=int, default=4096)
+    return parser
+
+
+def main(argv: Iterable[str] | None = None) -> None:
+    args = build_arg_parser().parse_args(argv)
+    result = run_inference(
+        moving_path=args.moving,
+        fixed_path=args.fixed,
+        checkpoint_path=args.checkpoint,
+        out_dir=args.out_dir,
+        target_spacing=args.target_spacing,
+        window_width=args.window_width,
+        window_level=args.window_level,
+        normalize_mode=args.normalize_mode,
+        device=args.device,
+        max_memory_mb=args.max_memory_mb,
+    )
+    print(json.dumps(result, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()