Add neck-based CT prealignment before VoxelMorph
This commit is contained in:
admin
297
prealign.py
Normal file
297
prealign.py
Normal file
@@ -0,0 +1,297 @@
|
||||
"""颈部基准的 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()
|
||||
Reference in New Issue
Block a user