"""颈部基准的 CT 预配准。 这一阶段发生在 VoxelMorph 训练/推理之前： 1. Fixed 与 Moving 原始 NIfTI 先重采样到同一 spacing。 2. 基于 CT foreground mask 估计颈部/身体区域的稳健中心。 3. 将 Moving 平移到 Fixed 的颈部中心。 4. 在同一个物理网格中裁剪/填充并归一化，得到后续 VoxelMorph 输入。 当前实现是确定性的平移预配准，目标是先解决层数不同、扫描范围不同导致的 初始空间不一致；细致非线性变形仍交给官方 VoxelMorph。 """ 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_FIXED_NIFTI, DEFAULT_MOVING_NIFTI, DEFAULT_TARGET_SHAPE, DEFAULT_TARGET_SPACING, DEFAULT_WINDOW_LEVEL, DEFAULT_WINDOW_WIDTH, PREPROCESSED_DIR, ) from preprocess import ( affine_with_spacing, ensure_multiple_of_16, load_nifti, resample_to_spacing, save_nifti, window_and_normalize, ) @dataclass class NeckAlignmentMeta: fixed_input_path: str moving_input_path: str fixed_output_path: str moving_output_path: str fixed_resampled_shape_xyz: Tuple[int, int, int] moving_resampled_shape_xyz: Tuple[int, int, int] target_shape_xyz: Tuple[int, int, int] target_spacing_xyz: Tuple[float, float, float] fixed_neck_center_index_xyz: Tuple[float, float, float] moving_neck_center_index_xyz: Tuple[float, float, float] fixed_neck_center_world_mm: Tuple[float, float, float] moving_neck_center_world_mm: Tuple[float, float, float] moving_translation_mm: Tuple[float, float, float] body_threshold_hu: float foreground_percentiles: Tuple[float, float] method: str = "neck-mask-centroid-translation" def _tuple_float(values: Sequence[float]) -> Tuple[float, float, float]: return tuple(float(v) for v in values) # type: ignore[return-value] def _tuple_int(values: Sequence[int]) -> Tuple[int, int, int]: return tuple(int(v) for v in values) # type: ignore[return-value] def index_to_world(affine: np.ndarray, index_xyz: Sequence[float]) -> np.ndarray: index = np.asarray(index_xyz, dtype=np.float64) return affine[:3, :3] @ index + affine[:3, 3] def largest_foreground_component(mask: np.ndarray) -> np.ndarray: structure = ndimage.generate_binary_structure(rank=3, connectivity=1) labeled, count = ndimage.label(mask, structure=structure) if count < 1: raise ValueError("未能从 CT 中提取颈部 foreground mask。") sizes = np.bincount(labeled.ravel()) sizes[0] = 0 return labeled == int(np.argmax(sizes)) def estimate_neck_center_index( data_hu: np.ndarray, body_threshold_hu: float = -500.0, foreground_percentiles: Sequence[float] = (2.0, 98.0), ) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """估计颈部 foreground 的稳健中心。 先用 HU 阈值去掉空气，再保留最大连通域，最后用坐标百分位框的中心而不是 简单均值，减少局部噪声、少量外部物体和扫描边界的影响。 """ data = np.asarray(data_hu, dtype=np.float32) mask = np.isfinite(data) & (data > float(body_threshold_hu)) & (data < 3000.0) mask = largest_foreground_component(mask) coords = np.where(mask) if not coords or coords[0].size == 0: raise ValueError("颈部 foreground mask 为空。") p_low, p_high = (float(v) for v in foreground_percentiles) lows = np.asarray([np.percentile(axis_coords, p_low) for axis_coords in coords], dtype=np.float64) highs = np.asarray([np.percentile(axis_coords, p_high) for axis_coords in coords], dtype=np.float64) center = (lows + highs) / 2.0 return center, (lows, highs) def target_affine_from_center( reference_affine: np.ndarray, center_world_mm: Sequence[float], target_shape: Sequence[int], target_spacing: Sequence[float], ) -> np.ndarray: output_affine = affine_with_spacing(reference_affine, target_spacing) center_index = (np.asarray(target_shape, dtype=np.float64) - 1.0) / 2.0 center_world = np.asarray(center_world_mm, dtype=np.float64) output_affine[:3, 3] = center_world - output_affine[:3, :3] @ center_index return output_affine def resample_to_reference_grid( data: np.ndarray, input_affine: np.ndarray, output_affine: np.ndarray, output_shape: Sequence[int], order: int = 1, cval: float = -1000.0, ) -> np.ndarray: input_to_output = np.linalg.inv(input_affine) @ output_affine matrix = input_to_output[:3, :3] offset = input_to_output[:3, 3] output = ndimage.affine_transform( data.astype(np.float32, copy=False), matrix=matrix, offset=offset, output_shape=tuple(int(v) for v in output_shape), order=order, mode="constant", cval=float(cval), prefilter=(order > 1), output=np.float32, ) return output.astype(np.float32, copy=False) def prealign_pair( fixed_input_path: str | Path, moving_input_path: str | Path, fixed_output_path: str | Path = DEFAULT_FIXED_NIFTI, moving_output_path: str | Path = DEFAULT_MOVING_NIFTI, 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, body_threshold_hu: float = -500.0, foreground_percentiles: Sequence[float] = (2.0, 98.0), max_memory_mb: int = 8192, metadata_json: str | Path | None = None, ) -> NeckAlignmentMeta: target_shape = ensure_multiple_of_16(target_shape) fixed_input_path = Path(fixed_input_path) moving_input_path = Path(moving_input_path) fixed_output_path = Path(fixed_output_path) moving_output_path = Path(moving_output_path) fixed_hu, fixed_affine, fixed_spacing = load_nifti(fixed_input_path) moving_hu, moving_affine, moving_spacing = load_nifti(moving_input_path) fixed_resampled, fixed_resampled_affine = resample_to_spacing( fixed_hu, fixed_affine, current_spacing=fixed_spacing, target_spacing=target_spacing, order=1, max_memory_mb=max_memory_mb, ) moving_resampled, moving_resampled_affine = resample_to_spacing( moving_hu, moving_affine, current_spacing=moving_spacing, target_spacing=target_spacing, order=1, max_memory_mb=max_memory_mb, ) fixed_center_idx, _ = estimate_neck_center_index( fixed_resampled, body_threshold_hu=body_threshold_hu, foreground_percentiles=foreground_percentiles, ) moving_center_idx, _ = estimate_neck_center_index( moving_resampled, body_threshold_hu=body_threshold_hu, foreground_percentiles=foreground_percentiles, ) fixed_center_world = index_to_world(fixed_resampled_affine, fixed_center_idx) moving_center_world = index_to_world(moving_resampled_affine, moving_center_idx) moving_translation = fixed_center_world - moving_center_world moving_aligned_affine = moving_resampled_affine.copy() moving_aligned_affine[:3, 3] = moving_aligned_affine[:3, 3] + moving_translation output_affine = target_affine_from_center( fixed_resampled_affine, center_world_mm=fixed_center_world, target_shape=target_shape, target_spacing=target_spacing, ) fixed_grid_hu = resample_to_reference_grid( fixed_resampled, input_affine=fixed_resampled_affine, output_affine=output_affine, output_shape=target_shape, order=1, cval=-1000.0, ) moving_grid_hu = resample_to_reference_grid( moving_resampled, input_affine=moving_aligned_affine, output_affine=output_affine, output_shape=target_shape, order=1, cval=-1000.0, ) fixed_norm = window_and_normalize(fixed_grid_hu, window_width=window_width, window_level=window_level) moving_norm = window_and_normalize(moving_grid_hu, window_width=window_width, window_level=window_level) save_nifti(fixed_norm, output_affine, fixed_output_path) save_nifti(moving_norm, output_affine, moving_output_path) meta = NeckAlignmentMeta( fixed_input_path=str(fixed_input_path), moving_input_path=str(moving_input_path), fixed_output_path=str(fixed_output_path), moving_output_path=str(moving_output_path), fixed_resampled_shape_xyz=_tuple_int(fixed_resampled.shape), moving_resampled_shape_xyz=_tuple_int(moving_resampled.shape), target_shape_xyz=_tuple_int(target_shape), target_spacing_xyz=_tuple_float(target_spacing), fixed_neck_center_index_xyz=_tuple_float(fixed_center_idx), moving_neck_center_index_xyz=_tuple_float(moving_center_idx), fixed_neck_center_world_mm=_tuple_float(fixed_center_world), moving_neck_center_world_mm=_tuple_float(moving_center_world), moving_translation_mm=_tuple_float(moving_translation), body_threshold_hu=float(body_threshold_hu), foreground_percentiles=(float(foreground_percentiles[0]), float(foreground_percentiles[1])), ) if metadata_json is None: metadata_json = PREPROCESSED_DIR / "patient1_neck_alignment.json" metadata_path = Path(metadata_json) metadata_path.parent.mkdir(parents=True, exist_ok=True) metadata_path.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="以颈部 foreground 为基准的 CT 平移预配准。") parser.add_argument("--fixed", required=True, help="Fixed 原始 NIfTI。") parser.add_argument("--moving", required=True, help="Moving 原始 NIfTI。") parser.add_argument("--fixed-output", default=str(DEFAULT_FIXED_NIFTI)) parser.add_argument("--moving-output", default=str(DEFAULT_MOVING_NIFTI)) 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("--body-threshold-hu", type=float, default=-500.0) parser.add_argument("--foreground-percentiles", type=float, nargs=2, default=(2.0, 98.0)) parser.add_argument("--metadata-json", default=str(PREPROCESSED_DIR / "patient1_neck_alignment.json")) parser.add_argument("--max-memory-mb", type=int, default=8192) return parser def main(argv: Iterable[str] | None = None) -> None: args = build_arg_parser().parse_args(argv) meta = prealign_pair( fixed_input_path=args.fixed, moving_input_path=args.moving, fixed_output_path=args.fixed_output, moving_output_path=args.moving_output, target_spacing=args.target_spacing, target_shape=args.target_shape, window_width=args.window_width, window_level=args.window_level, body_threshold_hu=args.body_threshold_hu, foreground_percentiles=args.foreground_percentiles, 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()