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