327 lines
12 KiB
Python
327 lines
12 KiB
Python
"""医学图像预处理模块。
|
||
|
||
预处理步骤：
|
||
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 % 32 != 0]
|
||
if bad:
|
||
raise ValueError(f"VoxelMorph 当前 U-Net 建议三维尺寸均为 32 的倍数，当前非法维度: {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()
|