"""独立推理模块。 本模块加载官方 ``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()