Add official VoxelMorph CT registration pipeline

preprocess.py

@@ -0,0 +1,327 @@
+"""医学图像预处理模块。
+
+预处理步骤：
+1. NIfTI 读取。
+2. 重采样到各向同性分辨率，例如 1x1x1 mm。
+3. 使用颈部软组织/气道窗口进行 HU 截断并归一化到 [0, 1]。
+4. 中心裁剪或填充到 VoxelMorph 兼容尺寸，例如 160x192x224。
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterable, Sequence, Tuple
+
+import numpy as np
+from scipy import ndimage
+
+from config import (
+    DEFAULT_TARGET_SHAPE,
+    DEFAULT_TARGET_SPACING,
+    DEFAULT_WINDOW_LEVEL,
+    DEFAULT_WINDOW_WIDTH,
+)
+
+
+@dataclass
+class PreprocessMeta:
+    input_path: str
+    output_path: str
+    original_shape_xyz: Tuple[int, int, int]
+    output_shape_xyz: Tuple[int, int, int]
+    original_spacing_xyz: Tuple[float, float, float]
+    target_spacing_xyz: Tuple[float, float, float]
+    window_width: float
+    window_level: float
+    crop_or_pad_offset_xyz: Tuple[int, int, int]
+
+
+def _require_nibabel():
+    try:
+        import nibabel as nib  # type: ignore
+    except Exception as exc:  # pragma: no cover
+        raise RuntimeError(
+            "缺少 nibabel，无法读取/保存 NIfTI。请先运行: pip install -r requirements.txt"
+        ) from exc
+    return nib
+
+
+def ensure_multiple_of_16(shape: Sequence[int]) -> Tuple[int, int, int]:
+    shape = tuple(int(v) for v in shape)
+    if len(shape) != 3:
+        raise ValueError("target_shape 必须包含 3 个维度。")
+    bad = [v for v in shape if v <= 0 or v % 16 != 0]
+    if bad:
+        raise ValueError(f"VoxelMorph 建议三维尺寸均为 16 的倍数，当前非法维度: {bad}")
+    return shape  # type: ignore[return-value]
+
+
+def load_nifti(path: str | Path) -> Tuple[np.ndarray, np.ndarray, Tuple[float, float, float]]:
+    """读取 NIfTI，返回 data_xyz、affine 和 spacing_xyz。"""
+
+    nib = _require_nibabel()
+    img = nib.load(str(path), mmap=True)
+    if len(img.shape) < 3:
+        raise ValueError(f"NIfTI 至少应为 3D: {path}")
+    data = np.asanyarray(img.dataobj, dtype=np.float32)
+    if data.ndim > 3:
+        data = data[..., 0]
+    data = np.nan_to_num(data.astype(np.float32, copy=False), copy=False)
+    spacing = tuple(float(v) for v in img.header.get_zooms()[:3])
+    return data, img.affine.copy(), spacing  # type: ignore[return-value]
+
+
+def estimate_memory_mb(shape: Sequence[int], dtype=np.float32, copies: int = 3) -> float:
+    """估算中间数组内存，copies 用于覆盖 scipy zoom 等临时副本。"""
+
+    voxels = int(np.prod(tuple(int(v) for v in shape)))
+    return voxels * np.dtype(dtype).itemsize * copies / (1024**2)
+
+
+def window_and_normalize(
+    data: np.ndarray,
+    window_width: float = DEFAULT_WINDOW_WIDTH,
+    window_level: float = DEFAULT_WINDOW_LEVEL,
+) -> np.ndarray:
+    """HU 截断并 Min-Max 归一化。
+
+    W=400/L=40 对颈部软组织和气道边界比较友好：既不会被骨窗拉爆，
+    也能保留软组织灰度差异供 NCC 训练使用。
+    """
+
+    low = float(window_level) - float(window_width) / 2.0
+    high = float(window_level) + float(window_width) / 2.0
+    clipped = np.clip(data.astype(np.float32, copy=False), low, high)
+    normalized = (clipped - low) / max(high - low, 1e-6)
+    return normalized.astype(np.float32, copy=False)
+
+
+def affine_with_spacing(affine: np.ndarray, target_spacing: Sequence[float]) -> np.ndarray:
+    """保持方向不变，只把 affine 三个轴的向量长度改成目标 spacing。"""
+
+    output = affine.copy()
+    for axis, spacing in enumerate(target_spacing):
+        vector = output[:3, axis]
+        length = float(np.linalg.norm(vector))
+        if length > 1e-8:
+            output[:3, axis] = vector / length * float(spacing)
+        else:
+            output[:3, axis] = 0
+            output[axis, axis] = float(spacing)
+    return output
+
+
+def resample_to_spacing(
+    data: np.ndarray,
+    affine: np.ndarray,
+    current_spacing: Sequence[float],
+    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
+    order: int = 1,
+    max_memory_mb: int = 4096,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """重采样到目标体素间距。
+
+    data 轴顺序为 (X, Y, Z)。zoom factor = 当前 spacing / 目标 spacing。
+    """
+
+    current_spacing = tuple(float(v) for v in current_spacing)
+    target_spacing = tuple(float(v) for v in target_spacing)
+    zoom_factors = tuple(current_spacing[i] / target_spacing[i] for i in range(3))
+    new_shape = tuple(max(1, int(round(data.shape[i] * zoom_factors[i]))) for i in range(3))
+
+    expected_mb = estimate_memory_mb(new_shape, dtype=np.float32, copies=4)
+    if expected_mb > max_memory_mb:
+        raise MemoryError(
+            f"重采样预计需要约 {expected_mb:.1f} MB，超过限制 {max_memory_mb} MB。"
+            "可提高 --max-memory-mb，或降低目标尺寸/分辨率。"
+        )
+
+    resampled = ndimage.zoom(
+        data.astype(np.float32, copy=False),
+        zoom=zoom_factors,
+        order=order,
+        mode="nearest",
+        prefilter=(order > 1),
+        output=np.float32,
+    )
+    return resampled.astype(np.float32, copy=False), affine_with_spacing(affine, target_spacing)
+
+
+def center_crop_or_pad(
+    data: np.ndarray,
+    target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE,
+    pad_value: float = 0.0,
+) -> Tuple[np.ndarray, Tuple[int, int, int]]:
+    """中心裁剪/填充到固定尺寸。
+
+    返回 offset_xyz，用于修正 affine 原点：
+    original_index = output_index + offset。
+    """
+
+    target_shape = ensure_multiple_of_16(target_shape)
+    output = np.full(target_shape, pad_value, dtype=np.float32)
+    input_slices = []
+    output_slices = []
+    offsets = []
+
+    for axis, target in enumerate(target_shape):
+        size = int(data.shape[axis])
+        if size >= target:
+            start = (size - target) // 2
+            input_slices.append(slice(start, start + target))
+            output_slices.append(slice(0, target))
+            offsets.append(start)
+        else:
+            pad_before = (target - size) // 2
+            input_slices.append(slice(0, size))
+            output_slices.append(slice(pad_before, pad_before + size))
+            offsets.append(-pad_before)
+
+    output[tuple(output_slices)] = data[tuple(input_slices)]
+    return output, tuple(int(v) for v in offsets)  # type: ignore[return-value]
+
+
+def shift_affine_origin(affine: np.ndarray, offset_xyz: Sequence[int]) -> np.ndarray:
+    """根据裁剪/填充偏移修正 NIfTI 原点。"""
+
+    output = affine.copy()
+    offset = np.asarray(offset_xyz, dtype=np.float64)
+    output[:3, 3] = output[:3, 3] + output[:3, :3] @ offset
+    return output
+
+
+def save_nifti(data_xyz: 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(data_xyz.astype(np.float32, copy=False), affine)
+    img.header.set_xyzt_units("mm")
+    nib.save(img, str(output_path))
+
+
+def preprocess_array(
+    data_xyz: np.ndarray,
+    affine: np.ndarray,
+    spacing_xyz: Sequence[float],
+    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
+    target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE,
+    window_width: float = DEFAULT_WINDOW_WIDTH,
+    window_level: float = DEFAULT_WINDOW_LEVEL,
+    normalize_mode: str = "window",
+    max_memory_mb: int = 4096,
+) -> Tuple[np.ndarray, np.ndarray, Tuple[int, int, int]]:
+    """数组级预处理，供 CLI 和 infer.py 共同复用。"""
+
+    if normalize_mode == "window":
+        data_xyz = window_and_normalize(data_xyz, window_width=window_width, window_level=window_level)
+    elif normalize_mode == "auto":
+        finite = data_xyz[np.isfinite(data_xyz)]
+        if finite.size and float(np.nanmin(finite)) >= -0.1 and float(np.nanmax(finite)) <= 1.1:
+            data_xyz = np.clip(data_xyz.astype(np.float32, copy=False), 0.0, 1.0)
+        else:
+            data_xyz = window_and_normalize(data_xyz, window_width=window_width, window_level=window_level)
+    elif normalize_mode == "none":
+        data_xyz = data_xyz.astype(np.float32, copy=False)
+    else:
+        raise ValueError("normalize_mode 只能是 window/auto/none。")
+
+    resampled, resampled_affine = resample_to_spacing(
+        data_xyz,
+        affine,
+        current_spacing=spacing_xyz,
+        target_spacing=target_spacing,
+        order=1,
+        max_memory_mb=max_memory_mb,
+    )
+    cropped, offset = center_crop_or_pad(resampled, target_shape=target_shape, pad_value=0.0)
+    output_affine = shift_affine_origin(resampled_affine, offset)
+    return cropped, output_affine, offset
+
+
+def preprocess_nifti(
+    input_path: str | Path,
+    output_path: str | Path,
+    target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING,
+    target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE,
+    window_width: float = DEFAULT_WINDOW_WIDTH,
+    window_level: float = DEFAULT_WINDOW_LEVEL,
+    normalize_mode: str = "window",
+    max_memory_mb: int = 4096,
+    metadata_json: str | Path | None = None,
+) -> PreprocessMeta:
+    """完整 NIfTI 预处理入口。"""
+
+    input_path = Path(input_path)
+    output_path = Path(output_path)
+    data, affine, spacing = load_nifti(input_path)
+    original_shape = tuple(int(v) for v in data.shape[:3])
+
+    processed, output_affine, offset = preprocess_array(
+        data,
+        affine,
+        spacing,
+        target_spacing=target_spacing,
+        target_shape=target_shape,
+        window_width=window_width,
+        window_level=window_level,
+        normalize_mode=normalize_mode,
+        max_memory_mb=max_memory_mb,
+    )
+    save_nifti(processed, output_affine, output_path)
+
+    meta = PreprocessMeta(
+        input_path=str(input_path),
+        output_path=str(output_path),
+        original_shape_xyz=original_shape,
+        output_shape_xyz=tuple(int(v) for v in processed.shape),
+        original_spacing_xyz=tuple(float(v) for v in spacing),
+        target_spacing_xyz=tuple(float(v) for v in target_spacing),
+        window_width=float(window_width),
+        window_level=float(window_level),
+        crop_or_pad_offset_xyz=offset,
+    )
+
+    if metadata_json is None:
+        metadata_json = Path(str(output_path).replace(".nii.gz", ".json").replace(".nii", ".json"))
+    metadata_json = Path(metadata_json)
+    metadata_json.parent.mkdir(parents=True, exist_ok=True)
+    metadata_json.write_text(json.dumps(asdict(meta), ensure_ascii=False, indent=2), encoding="utf-8")
+    return meta
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="NIfTI 重采样、窗宽窗位归一化、裁剪/填充。")
+    parser.add_argument("--input", required=True, help="输入 .nii.gz。")
+    parser.add_argument("--output", required=True, help="输出预处理后的 .nii.gz。")
+    parser.add_argument("--target-spacing", type=float, nargs=3, default=DEFAULT_TARGET_SPACING, metavar=("X", "Y", "Z"))
+    parser.add_argument("--target-shape", type=int, nargs=3, default=DEFAULT_TARGET_SHAPE, 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="window")
+    parser.add_argument("--metadata-json", default=None)
+    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)
+    meta = preprocess_nifti(
+        input_path=args.input,
+        output_path=args.output,
+        target_spacing=args.target_spacing,
+        target_shape=args.target_shape,
+        window_width=args.window_width,
+        window_level=args.window_level,
+        normalize_mode=args.normalize_mode,
+        metadata_json=args.metadata_json,
+        max_memory_mb=args.max_memory_mb,
+    )
+    print(json.dumps(asdict(meta), ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()
+