diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f1a2a06 --- /dev/null +++ b/.gitignore @@ -0,0 +1,23 @@ +__pycache__/ +*.py[cod] +*.pyo +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.venv/ +venv/ +env/ + +# 运行产物通常较大，保留目录占位但不纳入版本库。 +outputs/**/* +!outputs/**/ +!outputs/**/.gitkeep + +# 模型权重和医学影像中间文件建议按需单独归档。 +*.pt +*.pth +*.ckpt +*.nii +*.nii.gz +*.log + diff --git a/README.md b/README.md new file mode 100644 index 0000000..656a413 --- /dev/null +++ b/README.md @@ -0,0 +1,141 @@ +# VoxelMorph Head CT Deformable Registration + +面向“患者平扫 CT（中立位）到仰头 CT（极度后仰位）”的 3D 形变配准工程。项目包含 DICOM 转 NIfTI、医学图像预处理、官方 VoxelMorph 训练适配器、独立推理，以及 Streamlit 交互式结果查看界面。 + +核心模型使用官方仓库 `voxelmorph/voxelmorph` 的 PyTorch 实现： + +- 官方仓库：https://github.com/voxelmorph/voxelmorph +- 固定提交：`db73f34b910bcefcb520f7f40a1bc4a3e0b6401d` +- 核心类：`voxelmorph.nn.models.VxmPairwise` +- 平滑正则：`neurite.nn.modules.SpatialGradient` +- 相似度：默认使用带 epsilon 的稳定 Local NCC；可通过 `--ncc-impl neurite` 切换到 `neurite.nn.modules.NCC` + +## 工程目录 + +```text +Voxelmorph_Head_CT/ +├── Data/ +│ ├── 患者1-平扫CT/ # Moving 原始 DICOM +│ ├── 患者1-仰头CT/ # Fixed 原始 DICOM +│ └── 患者2-平扫CT/ +├── app.py # Streamlit Web 可视化界面 +├── config.py # 默认路径与预处理参数 +├── data_loader.py # DICOM 序列读取与 NIfTI 转换 +├── environment.yml # Conda CUDA 环境 +├── infer.py # 独立推理，输出 warped image 与 DDF +├── metrics.py # NCC/MSE/MAE/DDF 等量化指标 +├── model_and_train.py # 官方 VoxelMorph 训练适配器 +├── preprocess.py # 重采样、窗宽窗位、裁剪/填充 +├── requirements.txt +└── outputs/ + ├── nifti/ # DICOM 转换结果 + ├── preprocessed/ # 预处理结果 + ├── checkpoints/ # 模型权重 + └── inference/ # 推理输出 +``` + +## 环境 + +推荐使用 Conda 创建 CUDA 环境。本机如有 NVIDIA 驱动，不需要单独安装系统 CUDA toolkit，`pytorch-cuda=12.4` 会随 Conda 环境提供运行时。 + +为减少 Conda 包缓存和 Qt 依赖占用，环境采用两步安装：Conda 只安装 Python、PyTorch 与 CUDA runtime，业务依赖和官方 VoxelMorph 用 pip 安装。 + +```bash +conda env create -f environment.yml +conda activate voxelmorph-head-ct +python -m pip install --no-cache-dir -r requirements.txt +python -c "import torch; print(torch.__version__, torch.cuda.is_available())" +``` + +也可以直接运行项目脚本： + +```bash +bash scripts/setup_env.sh +``` + +如果只做 CPU 调试，也可以使用 `pip install -r requirements.txt`，但 3D 训练强烈建议使用 CUDA。 + +## 1. DICOM 转 NIfTI + +```bash +python data_loader.py \ + --dicom-dir "Data/患者1-平扫CT" \ + --output "outputs/nifti/patient1_moving.nii.gz" + +python data_loader.py \ + --dicom-dir "Data/患者1-仰头CT" \ + --output "outputs/nifti/patient1_fixed.nii.gz" +``` + +`data_loader.py` 会优先按 `InstanceNumber` 排序，其次按 `SliceLocation` 排序，并保存 spacing、层厚、排序依据等元数据 JSON。 + +## 2. 预处理 + +```bash +python preprocess.py \ + --input "outputs/nifti/patient1_moving.nii.gz" \ + --output "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \ + --target-spacing 1 1 1 \ + --target-shape 160 192 224 + +python preprocess.py \ + --input "outputs/nifti/patient1_fixed.nii.gz" \ + --output "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \ + --target-spacing 1 1 1 \ + --target-shape 160 192 224 +``` + +默认窗口为 `W=400, L=40`，适合观察颈部软组织和气道。 + +## 3. 训练 + +```bash +python model_and_train.py \ + --moving "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \ + --fixed "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \ + --checkpoint "outputs/checkpoints/vxm_head_ct.pt" \ + --epochs 200 \ + --image-loss ncc \ + --ncc-impl local \ + --smooth-weight 0.01 +``` + +训练脚本会调用官方 `vxm.nn.models.VxmPairwise`。如果显存不足，可降低 `--nb-features` 或改小 `--target-shape` 后重新预处理。 +由于 CT 预处理后常有大面积零填充背景，默认 `--ncc-impl local` 会在 NCC 方差项中加入 epsilon，避免低方差窗口导致非有限 loss。 + +## 4. 推理 + +```bash +python infer.py \ + --moving "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \ + --fixed "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \ + --checkpoint "outputs/checkpoints/vxm_head_ct.pt" \ + --out-dir "outputs/inference" +``` + +推理会输出： + +- `outputs/inference/warped_moving.nii.gz` +- `outputs/inference/ddf_mm.nii.gz` +- `outputs/inference/metrics.json` + +## 5. Web 结果展示 + +```bash +streamlit run app.py +``` + +网页提供： + +- Moving/Fixed/模型权重/输出目录输入。 +- “开始推理”按钮。 +- Axial、Coronal、Sagittal 正交三视图。 +- Fixed 与 Warped 的 Alpha 融合或棋盘格对比。 +- DDF 位移强度热力图。 +- NCC、MSE、MAE、逐切片误差曲线、DDF 位移分布等量化图。 + +## 内存注意事项 + +- DICOM 转换和重采样都有 `--max-memory-mb` 防护。 +- Web 界面对超大 NIfTI 会自动 stride 下采样，只影响浏览器展示，不改变磁盘结果。 +- 训练阶段的主要瓶颈是 3D U-Net 显存；`160x192x224` 是较重的 3D 输入，建议优先使用 CUDA GPU。 diff --git a/app.py b/app.py new file mode 100644 index 0000000..79942c6 --- /dev/null +++ b/app.py @@ -0,0 +1,430 @@ +"""Streamlit 交互式结果展示界面。 + +运行： + streamlit run app.py +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Dict, Iterable, List, Sequence, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import streamlit as st + +from config import ( + DEFAULT_CHECKPOINT, + DEFAULT_FIXED_NIFTI, + DEFAULT_MOVING_NIFTI, + INFERENCE_DIR, + OUTPUT_ROOT, + PROJECT_ROOT, +) +from metrics import crop_to_common_shape, ddf_summary, registration_metrics, slice_metric_curve + + +st.set_page_config( + page_title="VoxelMorph 颈部 CT 配准工作台", + layout="wide", + initial_sidebar_state="expanded", +) + + +CUSTOM_CSS = """ + +""" + + +def _require_nibabel(): + try: + import nibabel as nib # type: ignore + except Exception as exc: + raise RuntimeError("缺少 nibabel，请先运行: pip install -r requirements.txt") from exc + return nib + + +def discover_nifti_files() -> List[str]: + roots = [OUTPUT_ROOT, PROJECT_ROOT] + paths: List[Path] = [] + for root in roots: + if root.exists(): + paths.extend(root.rglob("*.nii")) + paths.extend(root.rglob("*.nii.gz")) + unique = sorted({str(path) for path in paths}) + return unique + + +def choose_path(label: str, default_path: Path, candidates: Sequence[str]) -> str: + options = ["手动输入"] + list(candidates) + default_str = str(default_path) + selected_index = options.index(default_str) if default_str in options else 0 + selected = st.selectbox(label, options=options, index=selected_index) + if selected == "手动输入": + return st.text_input(f"{label}路径", value=default_str) + return selected + + +@st.cache_data(max_entries=8, show_spinner=False) +def load_nifti_cached(path: str, max_voxels: int = 14_000_000) -> Tuple[np.ndarray, Tuple[float, float, float], int]: + """为 Web 展示读取 NIfTI。 + + Streamlit/浏览器端不适合一次渲染超大体数据。如果体素数过大，按等比例 + stride 下采样，只影响网页查看，不改变磁盘上的推理结果。 + """ + + if not path: + raise ValueError("路径为空。") + nib = _require_nibabel() + img = nib.load(path, mmap=True) + data = np.asanyarray(img.dataobj, dtype=np.float32) + if data.ndim > 4: + raise ValueError(f"暂不支持超过 4D 的 NIfTI: {path}") + + spatial_shape = data.shape[:3] + voxels = int(np.prod(spatial_shape)) + stride = max(1, int(np.ceil((voxels / max_voxels) ** (1.0 / 3.0)))) if voxels > max_voxels else 1 + + if data.ndim == 4: + data = data[::stride, ::stride, ::stride, :] + else: + data = data[::stride, ::stride, ::stride] + data = np.nan_to_num(data.astype(np.float32, copy=False), copy=False) + spacing = tuple(float(v) * stride for v in img.header.get_zooms()[:3]) + return data, spacing, stride # type: ignore[return-value] + + +def normalize_slice(image: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray: + image = image.astype(np.float32, copy=False) + low, high = np.percentile(image, [p_low, p_high]) + if high <= low: + return np.zeros_like(image, dtype=np.float32) + return np.clip((image - low) / (high - low), 0.0, 1.0) + + +def orient_for_display(image: np.ndarray) -> np.ndarray: + """统一显示方向，让三视图在网页中更接近医学浏览器观感。""" + + return np.rot90(image) + + +def get_slice(volume_xyz: np.ndarray, plane: str, index: int) -> np.ndarray: + if plane == "Axial": + return volume_xyz[:, :, index] + if plane == "Coronal": + return volume_xyz[:, index, :] + if plane == "Sagittal": + return volume_xyz[index, :, :] + raise ValueError(f"未知平面: {plane}") + + +def plane_axis(plane: str) -> int: + return {"Sagittal": 0, "Coronal": 1, "Axial": 2}[plane] + + +def checkerboard(fixed_slice: np.ndarray, warped_slice: np.ndarray, tile: int = 24) -> np.ndarray: + fixed_norm = normalize_slice(fixed_slice) + warped_norm = normalize_slice(warped_slice) + yy, xx = np.indices(fixed_norm.shape) + mask = ((yy // tile + xx // tile) % 2).astype(bool) + return np.where(mask, fixed_norm, warped_norm) + + +def alpha_overlay(fixed_slice: np.ndarray, warped_slice: np.ndarray, alpha: float = 0.45) -> np.ndarray: + fixed_norm = normalize_slice(fixed_slice) + warped_norm = normalize_slice(warped_slice) + rgb = np.zeros((*fixed_norm.shape, 3), dtype=np.float32) + rgb[..., 0] = warped_norm + rgb[..., 1] = fixed_norm * (1.0 - alpha) + warped_norm * alpha + rgb[..., 2] = fixed_norm + return np.clip(rgb, 0.0, 1.0) + + +def ddf_magnitude(ddf_xyz: np.ndarray) -> np.ndarray: + if ddf_xyz.ndim != 4 or ddf_xyz.shape[-1] != 3: + raise ValueError("DDF 应为 (X, Y, Z, 3)。") + return np.linalg.norm(ddf_xyz.astype(np.float32, copy=False), axis=-1) + + +def render_image(image: np.ndarray, caption: str, cmap: str = "gray") -> None: + fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120) + ax.imshow(orient_for_display(image), cmap=cmap, interpolation="nearest") + ax.set_title(caption, fontsize=10) + ax.axis("off") + st.pyplot(fig, use_container_width=True) + plt.close(fig) + + +def render_rgb(image: np.ndarray, caption: str) -> None: + fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120) + ax.imshow(orient_for_display(image), interpolation="nearest") + ax.set_title(caption, fontsize=10) + ax.axis("off") + st.pyplot(fig, use_container_width=True) + plt.close(fig) + + +def render_heatmap_overlay(base_slice: np.ndarray, mag_slice: np.ndarray, alpha: float, caption: str) -> None: + fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120) + ax.imshow(orient_for_display(normalize_slice(base_slice)), cmap="gray", interpolation="nearest") + heat = orient_for_display(mag_slice) + im = ax.imshow(heat, cmap="inferno", alpha=alpha, interpolation="nearest") + ax.set_title(caption, fontsize=10) + ax.axis("off") + fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02) + st.pyplot(fig, use_container_width=True) + plt.close(fig) + + +def render_three_views(volume_xyz: np.ndarray, title: str) -> None: + st.subheader(title) + columns = st.columns(3) + planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")] + for col, (plane, cn_name) in zip(columns, planes): + axis = plane_axis(plane) + with col: + index = st.slider(f"{cn_name}", 0, volume_xyz.shape[axis] - 1, volume_xyz.shape[axis] // 2, key=f"{title}-{plane}") + render_image(get_slice(volume_xyz, plane, index), f"{cn_name} #{index}") + + +def render_overlay_views(fixed_xyz: np.ndarray, warped_xyz: np.ndarray) -> None: + fixed_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, warped_xyz) + mode = st.radio("对比模式", ["Alpha 融合", "棋盘格"], horizontal=True) + alpha = st.slider("透明度", 0.0, 1.0, 0.45, 0.05) + tile = st.slider("棋盘格尺寸", 8, 64, 24, 4) + + columns = st.columns(3) + planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")] + for col, (plane, cn_name) in zip(columns, planes): + axis = plane_axis(plane) + with col: + index = st.slider(f"{cn_name}切片", 0, fixed_xyz.shape[axis] - 1, fixed_xyz.shape[axis] // 2, key=f"overlay-{plane}") + fixed_slice = get_slice(fixed_xyz, plane, index) + warped_slice = get_slice(warped_xyz, plane, index) + if mode == "棋盘格": + render_image(checkerboard(fixed_slice, warped_slice, tile=tile), f"{cn_name} 棋盘格 #{index}") + else: + render_rgb(alpha_overlay(fixed_slice, warped_slice, alpha=alpha), f"{cn_name} 融合 #{index}") + + +def render_ddf_views(fixed_xyz: np.ndarray, ddf_xyz: np.ndarray) -> None: + fixed_xyz, ddf_xyz = crop_to_common_shape(fixed_xyz, ddf_xyz) + mag = ddf_magnitude(ddf_xyz) + alpha = st.slider("热力图透明度", 0.0, 1.0, 0.55, 0.05) + + stats = ddf_summary(ddf_xyz) + c1, c2, c3, c4 = st.columns(4) + c1.metric("平均位移 mm", f"{stats['ddf_mean']:.3f}") + c2.metric("P95 mm", f"{stats['ddf_p95']:.3f}") + c3.metric("最大 mm", f"{stats['ddf_max']:.3f}") + c4.metric("标准差", f"{stats['ddf_std']:.3f}") + + columns = st.columns(3) + planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")] + for col, (plane, cn_name) in zip(columns, planes): + axis = plane_axis(plane) + with col: + index = st.slider(f"{cn_name}DDF", 0, fixed_xyz.shape[axis] - 1, fixed_xyz.shape[axis] // 2, key=f"ddf-{plane}") + render_heatmap_overlay( + get_slice(fixed_xyz, plane, index), + get_slice(mag, plane, index), + alpha=alpha, + caption=f"{cn_name} 形变强度 #{index}", + ) + + +def render_metric_charts(fixed_xyz: np.ndarray, moving_xyz: np.ndarray, warped_xyz: np.ndarray, ddf_xyz: np.ndarray | None) -> None: + fixed_xyz, moving_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, moving_xyz, warped_xyz) + metrics = registration_metrics(fixed_xyz, moving_xyz, warped_xyz) + + c1, c2, c3 = st.columns(3) + c1.metric("NCC", f"{metrics['after_ncc']:.4f}", delta=f"{metrics['ncc_improvement']:+.4f}") + c2.metric("MSE", f"{metrics['after_mse']:.5f}", delta=f"{-metrics['mse_improvement']:+.5f}") + c3.metric("MAE", f"{metrics['after_mae']:.5f}", delta=f"{-metrics['mae_improvement']:+.5f}") + + labels = ["NCC", "MSE", "MAE"] + before = [metrics["before_ncc"], metrics["before_mse"], metrics["before_mae"]] + after = [metrics["after_ncc"], metrics["after_mse"], metrics["after_mae"]] + + fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), dpi=130) + x = np.arange(len(labels)) + width = 0.36 + axes[0].bar(x - width / 2, before, width, label="配准前", color="#64748b") + axes[0].bar(x + width / 2, after, width, label="配准后", color="#0f766e") + axes[0].set_xticks(x, labels) + axes[0].set_title("配准前后指标对比") + axes[0].legend(frameon=False) + axes[0].grid(axis="y", alpha=0.22) + + curve = slice_metric_curve(fixed_xyz, moving_xyz, warped_xyz, axis=2) + axes[1].plot(curve["slice_index"], curve["before_mse"], label="配准前 MSE", color="#64748b", linewidth=1.8) + axes[1].plot(curve["slice_index"], curve["after_mse"], label="配准后 MSE", color="#b45309", linewidth=1.8) + axes[1].set_title("轴状面逐切片误差") + axes[1].set_xlabel("Slice") + axes[1].grid(alpha=0.22) + axes[1].legend(frameon=False) + st.pyplot(fig, use_container_width=True) + plt.close(fig) + + if ddf_xyz is not None: + mag = ddf_magnitude(ddf_xyz) + fig2, ax = plt.subplots(figsize=(10, 3.2), dpi=130) + ax.hist(mag.ravel(), bins=80, color="#7c3f00", alpha=0.82) + ax.set_title("DDF 位移强度分布") + ax.set_xlabel("mm") + ax.set_ylabel("Voxel count") + ax.grid(axis="y", alpha=0.2) + st.pyplot(fig2, use_container_width=True) + plt.close(fig2) + + +def load_result_paths(out_dir: str) -> Dict[str, str]: + out = Path(out_dir) + return { + "warped": str(out / "warped_moving.nii.gz"), + "ddf": str(out / "ddf_mm.nii.gz"), + "metrics": str(out / "metrics.json"), + } + + +def run_inference_from_ui(moving_path: str, fixed_path: str, checkpoint_path: str, out_dir: str) -> Dict: + from infer import run_inference + + return run_inference( + moving_path=moving_path, + fixed_path=fixed_path, + checkpoint_path=checkpoint_path, + out_dir=out_dir, + ) + + +def main() -> None: + st.markdown(CUSTOM_CSS, unsafe_allow_html=True) + st.title("VoxelMorph 颈部 CT 配准工作台") + + candidates = discover_nifti_files() + with st.sidebar: + st.header("输入") + moving_path = choose_path("Moving", DEFAULT_MOVING_NIFTI, candidates) + fixed_path = choose_path("Fixed", DEFAULT_FIXED_NIFTI, candidates) + checkpoint_path = st.text_input("模型权重", value=str(DEFAULT_CHECKPOINT)) + out_dir = st.text_input("输出目录", value=str(INFERENCE_DIR)) + + start = st.button("开始推理", type="primary", use_container_width=True) + if start: + with st.spinner("推理运行中"): + try: + result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir) + st.session_state["last_result"] = result + st.success("推理完成") + except Exception as exc: + st.error(str(exc)) + + result_paths = load_result_paths(out_dir) + warped_path = str(st.session_state.get("last_result", {}).get("warped_path", result_paths["warped"])) + ddf_path = str(st.session_state.get("last_result", {}).get("ddf_path", result_paths["ddf"])) + + status_cols = st.columns(4) + status_cols[0].metric("Moving", "存在" if Path(moving_path).exists() else "缺失") + status_cols[1].metric("Fixed", "存在" if Path(fixed_path).exists() else "缺失") + status_cols[2].metric("Warped", "存在" if Path(warped_path).exists() else "缺失") + status_cols[3].metric("DDF", "存在" if Path(ddf_path).exists() else "缺失") + + try: + moving_xyz, moving_spacing, moving_stride = load_nifti_cached(moving_path) + fixed_xyz, fixed_spacing, fixed_stride = load_nifti_cached(fixed_path) + except Exception as exc: + st.warning(f"待显示数据尚未就绪: {exc}") + return + + warped_xyz = None + ddf_xyz = None + if Path(warped_path).exists(): + try: + warped_xyz, _, _ = load_nifti_cached(warped_path) + except Exception as exc: + st.warning(f"Warped 读取失败: {exc}") + if Path(ddf_path).exists(): + try: + ddf_xyz, _, _ = load_nifti_cached(ddf_path) + except Exception as exc: + st.warning(f"DDF 读取失败: {exc}") + + if moving_stride > 1 or fixed_stride > 1: + st.info(f"网页显示已自动下采样：Moving stride={moving_stride}, Fixed stride={fixed_stride}") + + tab_views, tab_overlay, tab_ddf, tab_metrics = st.tabs(["正交三视图", "重叠对比", "形变场", "量化图"]) + with tab_views: + view_target = st.radio("显示对象", ["Fixed", "Moving", "Warped"], horizontal=True) + if view_target == "Fixed": + render_three_views(fixed_xyz, "Fixed Image") + elif view_target == "Moving": + render_three_views(moving_xyz, "Moving Image") + elif warped_xyz is not None: + render_three_views(warped_xyz, "Warped Image") + else: + st.warning("尚未生成 Warped Image。") + + with tab_overlay: + if warped_xyz is None: + st.warning("尚未生成 Warped Image。") + else: + render_overlay_views(fixed_xyz, warped_xyz) + + with tab_ddf: + if ddf_xyz is None: + st.warning("尚未生成 DDF。") + else: + render_ddf_views(fixed_xyz, ddf_xyz) + + with tab_metrics: + if warped_xyz is None: + st.warning("尚未生成 Warped Image。") + else: + render_metric_charts(fixed_xyz, moving_xyz, warped_xyz, ddf_xyz) + metrics_file = Path(result_paths["metrics"]) + if metrics_file.exists(): + try: + st.json(json.loads(metrics_file.read_text(encoding="utf-8"))) + except Exception: + pass + + +if __name__ == "__main__": + main() diff --git a/config.py b/config.py new file mode 100644 index 0000000..f7faa94 --- /dev/null +++ b/config.py @@ -0,0 +1,32 @@ +"""项目级默认配置。 + +这些值只作为命令行和 Web 界面的默认项。真正运行时仍可通过参数覆盖， +便于把同一套代码迁移到其他患者或其他机器。 +""" + +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parent + +DATA_ROOT = PROJECT_ROOT / "Data" +DEFAULT_MOVING_DICOM_DIR = DATA_ROOT / "患者1-平扫CT" +DEFAULT_FIXED_DICOM_DIR = DATA_ROOT / "患者1-仰头CT" + +OUTPUT_ROOT = PROJECT_ROOT / "outputs" +NIFTI_DIR = OUTPUT_ROOT / "nifti" +PREPROCESSED_DIR = OUTPUT_ROOT / "preprocessed" +CHECKPOINT_DIR = OUTPUT_ROOT / "checkpoints" +INFERENCE_DIR = OUTPUT_ROOT / "inference" + +DEFAULT_MOVING_NIFTI = PREPROCESSED_DIR / "patient1_moving_preprocessed.nii.gz" +DEFAULT_FIXED_NIFTI = PREPROCESSED_DIR / "patient1_fixed_preprocessed.nii.gz" +DEFAULT_CHECKPOINT = CHECKPOINT_DIR / "vxm_head_ct.pt" + +# VoxelMorph 的 3D U-Net 多次下采样，三维尺寸建议均为 16 的倍数。 +DEFAULT_TARGET_SHAPE = (160, 192, 224) # NIfTI 轴顺序: X, Y, Z +DEFAULT_TARGET_SPACING = (1.0, 1.0, 1.0) # mm, X/Y/Z + +# 颈部软组织/气道观察常用窗口：W=400, L=40。 +DEFAULT_WINDOW_WIDTH = 400.0 +DEFAULT_WINDOW_LEVEL = 40.0 diff --git a/data_loader.py b/data_loader.py new file mode 100644 index 0000000..0b94ad2 --- /dev/null +++ b/data_loader.py @@ -0,0 +1,354 @@ +"""DICOM 序列读取与 NIfTI 转换模块。 + +核心职责： +1. 读取同一序列目录下的 DICOM 文件。 +2. 严格按照 InstanceNumber 或 SliceLocation 排序，防止 Z 轴切片错乱。 +3. 提取 PixelSpacing、SliceThickness 等空间信息。 +4. 应用 RescaleSlope/RescaleIntercept，把原始像素转换为 HU。 +5. 保存为 NIfTI，供预处理、训练和 Web 可视化使用。 +""" + +from __future__ import annotations + +import argparse +import json +import math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable, List, Sequence, Tuple + +import numpy as np +import pydicom + + +@dataclass +class DicomSeriesMeta: + """记录 DICOM 到 NIfTI 转换时的重要元数据。""" + + dicom_dir: str + output_path: str + num_slices: int + rows: int + columns: int + pixel_spacing_row_col: Tuple[float, float] + slice_thickness: float + spacing_xyz: Tuple[float, float, float] + sort_key: str + intensity_unit: str = "HU" + + +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 _to_float(value, default: float = math.nan) -> float: + try: + return float(value) + except Exception: + return default + + +def _candidate_dicom_files(dicom_dir: Path) -> List[Path]: + """列出可能的 DICOM 文件。 + + 有些设备导出的 DICOM 没有 .dcm 后缀，所以这里不只依赖扩展名，而是 + 后续通过 pydicom 读取头信息判断是否可用。 + """ + + if not dicom_dir.exists(): + raise FileNotFoundError(f"DICOM 目录不存在: {dicom_dir}") + if not dicom_dir.is_dir(): + raise NotADirectoryError(f"不是目录: {dicom_dir}") + + return sorted(path for path in dicom_dir.iterdir() if path.is_file() and not path.name.startswith(".")) + + +def _read_headers(dicom_dir: Path) -> List[Tuple[Path, pydicom.dataset.FileDataset]]: + headers: List[Tuple[Path, pydicom.dataset.FileDataset]] = [] + skipped: List[str] = [] + + for path in _candidate_dicom_files(dicom_dir): + try: + ds = pydicom.dcmread(path, stop_before_pixels=True, force=True) + if not hasattr(ds, "Rows") or not hasattr(ds, "Columns"): + skipped.append(path.name) + continue + headers.append((path, ds)) + except Exception: + skipped.append(path.name) + + if not headers: + raise RuntimeError(f"目录中没有可读取的 DICOM 切片: {dicom_dir}") + + if skipped: + print(f"[data_loader] 跳过 {len(skipped)} 个非 DICOM/无像素头文件。") + + return headers + + +def _unique_numeric_values( + headers: Sequence[Tuple[Path, pydicom.dataset.FileDataset]], + attr: str, +) -> List[float] | None: + values: List[float] = [] + for _, ds in headers: + if not hasattr(ds, attr): + return None + value = _to_float(getattr(ds, attr)) + if math.isnan(value): + return None + values.append(value) + if len(set(values)) != len(values): + return None + return values + + +def sort_dicom_headers( + headers: Sequence[Tuple[Path, pydicom.dataset.FileDataset]], +) -> Tuple[List[Tuple[Path, pydicom.dataset.FileDataset]], str]: + """按可靠的 Z 轴顺序排列切片。 + + 优先级： + 1. InstanceNumber：多数 CT 序列最稳定、最直观。 + 2. SliceLocation：用户明确要求的第二排序依据。 + 3. ImagePositionPatient 投影：当上面两个标签缺失时作为医学影像常见兜底。 + 4. 文件名：最后兜底，并给出明显标记，提醒用户检查。 + """ + + instance_numbers = _unique_numeric_values(headers, "InstanceNumber") + if instance_numbers is not None: + return sorted(headers, key=lambda item: float(item[1].InstanceNumber)), "InstanceNumber" + + slice_locations = _unique_numeric_values(headers, "SliceLocation") + if slice_locations is not None: + return sorted(headers, key=lambda item: float(item[1].SliceLocation)), "SliceLocation" + + if all(hasattr(ds, "ImagePositionPatient") for _, ds in headers): + try: + first = headers[0][1] + orientation = [float(v) for v in getattr(first, "ImageOrientationPatient", [])] + row_cos = np.array(orientation[:3], dtype=np.float64) + col_cos = np.array(orientation[3:], dtype=np.float64) + normal = np.cross(row_cos, col_cos) + if np.linalg.norm(normal) > 0: + normal = normal / np.linalg.norm(normal) + return ( + sorted( + headers, + key=lambda item: float(np.dot(np.asarray(item[1].ImagePositionPatient, dtype=np.float64), normal)), + ), + "ImagePositionPatient", + ) + except Exception: + pass + + print("[data_loader] 警告：未找到可靠 Z 轴标签，退回文件名排序，请人工确认切片顺序。") + return sorted(headers, key=lambda item: item[0].name), "Filename" + + +def _infer_slice_thickness( + sorted_headers: Sequence[Tuple[Path, pydicom.dataset.FileDataset]], +) -> float: + first = sorted_headers[0][1] + thickness = _to_float(getattr(first, "SliceThickness", math.nan)) + if not math.isnan(thickness) and thickness > 0: + return thickness + + locations = _unique_numeric_values(sorted_headers, "SliceLocation") + if locations and len(locations) > 1: + diffs = np.diff(sorted(locations)) + diffs = np.abs(diffs[diffs != 0]) + if diffs.size: + return float(np.median(diffs)) + + positions = [] + for _, ds in sorted_headers: + if hasattr(ds, "ImagePositionPatient"): + positions.append(np.asarray(ds.ImagePositionPatient, dtype=np.float64)) + if len(positions) > 1: + distances = [np.linalg.norm(positions[i + 1] - positions[i]) for i in range(len(positions) - 1)] + distances = [d for d in distances if d > 0] + if distances: + return float(np.median(distances)) + + raise RuntimeError("无法从 DICOM 中提取 SliceThickness，也无法根据切片位置估计层厚。") + + +def _make_affine( + sorted_headers: Sequence[Tuple[Path, pydicom.dataset.FileDataset]], + row_spacing: float, + col_spacing: float, + slice_thickness: float, +) -> np.ndarray: + """构建 NIfTI affine。 + + NIfTI 保存的数据轴为 (X, Y, Z)，由 DICOM 的 (Z, Rows, Columns) + 转置而来： + - X 轴对应 DICOM Columns，物理间距为 PixelSpacing[1] + - Y 轴对应 DICOM Rows，物理间距为 PixelSpacing[0] + - Z 轴对应切片方向，间距为 SliceThickness 或切片位置差 + """ + + affine = np.eye(4, dtype=np.float64) + first = sorted_headers[0][1] + + try: + orientation = [float(v) for v in first.ImageOrientationPatient] + origin = np.asarray(first.ImagePositionPatient, dtype=np.float64) + row_cos = np.asarray(orientation[:3], dtype=np.float64) + col_cos = np.asarray(orientation[3:], dtype=np.float64) + normal = np.cross(row_cos, col_cos) + + affine[:3, 0] = row_cos * col_spacing + affine[:3, 1] = col_cos * row_spacing + + if len(sorted_headers) > 1 and hasattr(sorted_headers[-1][1], "ImagePositionPatient"): + last_origin = np.asarray(sorted_headers[-1][1].ImagePositionPatient, dtype=np.float64) + step = (last_origin - origin) / max(len(sorted_headers) - 1, 1) + if np.linalg.norm(step) > 0: + affine[:3, 2] = step + else: + affine[:3, 2] = normal * slice_thickness + else: + affine[:3, 2] = normal * slice_thickness + affine[:3, 3] = origin + except Exception: + affine[0, 0] = col_spacing + affine[1, 1] = row_spacing + affine[2, 2] = slice_thickness + + return affine + + +def load_dicom_series( + dicom_dir: str | Path, + max_memory_mb: int = 4096, +) -> Tuple[np.ndarray, DicomSeriesMeta, np.ndarray]: + """读取 DICOM 序列为 3D NumPy 数组。 + + 返回： + - volume_zyx: shape = (Z, Y, X)，数值单位为 HU。 + - meta: 关键元数据。 + - affine: 与转置后的 NIfTI 数据 (X, Y, Z) 匹配的 affine。 + """ + + dicom_dir = Path(dicom_dir) + headers = _read_headers(dicom_dir) + sorted_headers, sort_key = sort_dicom_headers(headers) + + first = sorted_headers[0][1] + rows = int(first.Rows) + columns = int(first.Columns) + pixel_spacing = [float(v) for v in first.PixelSpacing] + row_spacing, col_spacing = pixel_spacing[0], pixel_spacing[1] + slice_thickness = _infer_slice_thickness(sorted_headers) + + expected_mb = len(sorted_headers) * rows * columns * np.dtype(np.float32).itemsize / (1024**2) + if expected_mb > max_memory_mb: + raise MemoryError( + f"预计载入体数据约 {expected_mb:.1f} MB，超过限制 {max_memory_mb} MB。" + "可增大 --max-memory-mb，或先减少序列范围。" + ) + + volume = np.empty((len(sorted_headers), rows, columns), dtype=np.float32) + + for index, (path, _) in enumerate(sorted_headers): + ds = pydicom.dcmread(path, force=True) + pixel_array = ds.pixel_array + if pixel_array.ndim == 3 and pixel_array.shape[0] == 1: + pixel_array = pixel_array[0] + if pixel_array.ndim != 2: + raise RuntimeError(f"暂不支持多帧或非 2D 切片: {path}") + if pixel_array.shape != (rows, columns): + raise RuntimeError(f"切片尺寸不一致: {path}, {pixel_array.shape} != {(rows, columns)}") + + slope = _to_float(getattr(ds, "RescaleSlope", 1.0), default=1.0) + intercept = _to_float(getattr(ds, "RescaleIntercept", 0.0), default=0.0) + volume[index] = pixel_array.astype(np.float32) * slope + intercept + + affine = _make_affine(sorted_headers, row_spacing, col_spacing, slice_thickness) + meta = DicomSeriesMeta( + dicom_dir=str(dicom_dir), + output_path="", + num_slices=len(sorted_headers), + rows=rows, + columns=columns, + pixel_spacing_row_col=(row_spacing, col_spacing), + slice_thickness=slice_thickness, + spacing_xyz=(col_spacing, row_spacing, float(np.linalg.norm(affine[:3, 2]))), + sort_key=sort_key, + ) + return volume, meta, affine + + +def save_volume_as_nifti( + volume_zyx: np.ndarray, + affine: np.ndarray, + output_path: str | Path, +) -> None: + """保存 NIfTI。 + + DICOM 读取后是 (Z, Y, X)，NIfTI 这里统一保存为 (X, Y, Z)。 + 这样后续 preprocess/infer/app 都能用同一套轴约定。 + """ + + nib = _require_nibabel() + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + volume_xyz = np.transpose(volume_zyx, (2, 1, 0)).astype(np.float32, copy=False) + img = nib.Nifti1Image(volume_xyz, affine) + img.header.set_xyzt_units("mm") + nib.save(img, str(output_path)) + + +def convert_dicom_series_to_nifti( + dicom_dir: str | Path, + output_path: str | Path, + metadata_json: str | Path | None = None, + max_memory_mb: int = 4096, +) -> DicomSeriesMeta: + """从 DICOM 目录生成 .nii.gz，并写出元数据 JSON。""" + + volume, meta, affine = load_dicom_series(dicom_dir, max_memory_mb=max_memory_mb) + save_volume_as_nifti(volume, affine, output_path) + + meta.output_path = str(output_path) + 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="读取 DICOM 序列并转换为 NIfTI。") + parser.add_argument("--dicom-dir", required=True, help="DICOM 序列目录。") + parser.add_argument("--output", required=True, help="输出 .nii.gz 路径。") + parser.add_argument("--metadata-json", default=None, help="可选：元数据 JSON 输出路径。") + parser.add_argument("--max-memory-mb", type=int, default=4096, help="载入体数据允许使用的最大内存估计值。") + return parser + + +def main(argv: Iterable[str] | None = None) -> None: + args = build_arg_parser().parse_args(argv) + meta = convert_dicom_series_to_nifti( + dicom_dir=args.dicom_dir, + output_path=args.output, + 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() + diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..7a7db1f --- /dev/null +++ b/environment.yml @@ -0,0 +1,10 @@ +name: voxelmorph-head-ct +channels: + - pytorch + - nvidia + - conda-forge +dependencies: + - python=3.11 + - pip + - pytorch + - pytorch-cuda=12.4 diff --git a/infer.py b/infer.py new file mode 100644 index 0000000..73feb12 --- /dev/null +++ b/infer.py @@ -0,0 +1,242 @@ +"""独立推理模块。 + +本模块加载官方 ``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() diff --git a/metrics.py b/metrics.py new file mode 100644 index 0000000..9e8c94a --- /dev/null +++ b/metrics.py @@ -0,0 +1,123 @@ +"""配准质量评估工具。 + +这里刻意只依赖 NumPy，便于在训练、推理和 Streamlit 三处复用。 +""" + +from __future__ import annotations + +from typing import Dict, Iterable, Tuple + +import numpy as np + + +def crop_to_common_shape(*arrays: np.ndarray) -> Tuple[np.ndarray, ...]: + """把多个 3D/4D 数组中心裁剪到共同尺寸，避免形状不一致导致评估失败。""" + + if not arrays: + return tuple() + + spatial_shapes = [arr.shape[:3] for arr in arrays] + common = tuple(min(shape[axis] for shape in spatial_shapes) for axis in range(3)) + cropped = [] + + for arr in arrays: + slices = [] + for axis, target in enumerate(common): + start = max((arr.shape[axis] - target) // 2, 0) + slices.append(slice(start, start + target)) + if arr.ndim > 3: + slices.append(slice(None)) + cropped.append(arr[tuple(slices)]) + + return tuple(cropped) + + +def global_ncc(a: np.ndarray, b: np.ndarray, eps: float = 1e-8) -> float: + """计算全局 NCC。训练时使用局部 NCC，这里用于快速量化展示。""" + + a = np.asarray(a, dtype=np.float32) + b = np.asarray(b, dtype=np.float32) + a = a - float(np.mean(a)) + b = b - float(np.mean(b)) + denom = float(np.sqrt(np.sum(a * a) * np.sum(b * b)) + eps) + return float(np.sum(a * b) / denom) + + +def mse(a: np.ndarray, b: np.ndarray) -> float: + diff = np.asarray(a, dtype=np.float32) - np.asarray(b, dtype=np.float32) + return float(np.mean(diff * diff)) + + +def mae(a: np.ndarray, b: np.ndarray) -> float: + diff = np.abs(np.asarray(a, dtype=np.float32) - np.asarray(b, dtype=np.float32)) + return float(np.mean(diff)) + + +def registration_metrics( + fixed: np.ndarray, + moving: np.ndarray, + warped: np.ndarray, +) -> Dict[str, float]: + """输出配准前后可比较的常用指标。""" + + fixed, moving, warped = crop_to_common_shape(fixed, moving, warped) + before_mse = mse(fixed, moving) + after_mse = mse(fixed, warped) + before_mae = mae(fixed, moving) + after_mae = mae(fixed, warped) + before_ncc = global_ncc(fixed, moving) + after_ncc = global_ncc(fixed, warped) + + return { + "before_mse": before_mse, + "after_mse": after_mse, + "before_mae": before_mae, + "after_mae": after_mae, + "before_ncc": before_ncc, + "after_ncc": after_ncc, + "mse_improvement": before_mse - after_mse, + "mae_improvement": before_mae - after_mae, + "ncc_improvement": after_ncc - before_ncc, + } + + +def slice_metric_curve( + fixed: np.ndarray, + moving: np.ndarray, + warped: np.ndarray, + axis: int = 2, +) -> Dict[str, Iterable[float]]: + """逐切片计算 MSE，适合生成“配准前后误差曲线”。""" + + fixed, moving, warped = crop_to_common_shape(fixed, moving, warped) + before = [] + after = [] + + for index in range(fixed.shape[axis]): + selector = [slice(None)] * 3 + selector[axis] = index + selector = tuple(selector) + before.append(mse(fixed[selector], moving[selector])) + after.append(mse(fixed[selector], warped[selector])) + + return { + "slice_index": list(range(fixed.shape[axis])), + "before_mse": before, + "after_mse": after, + } + + +def ddf_summary(ddf_xyz: np.ndarray) -> Dict[str, float]: + """统计形变场向量大小。输入应为 X/Y/Z/3，单位可为 voxel 或 mm。""" + + if ddf_xyz.ndim != 4 or ddf_xyz.shape[-1] != 3: + raise ValueError("DDF 必须是形状为 (X, Y, Z, 3) 的 4D 数组。") + + mag = np.linalg.norm(ddf_xyz.astype(np.float32), axis=-1) + return { + "ddf_mean": float(np.mean(mag)), + "ddf_std": float(np.std(mag)), + "ddf_p95": float(np.percentile(mag, 95)), + "ddf_max": float(np.max(mag)), + } + diff --git a/model_and_train.py b/model_and_train.py new file mode 100644 index 0000000..649e74f --- /dev/null +++ b/model_and_train.py @@ -0,0 +1,400 @@ +"""官方 VoxelMorph 训练适配器。 + +本项目以 https://github.com/voxelmorph/voxelmorph 的 PyTorch 核心为运行核心： +- 网络: ``voxelmorph.nn.models.VxmPairwise`` +- 空间变换: 官方 ``voxelmorph.nn.modules.SpatialTransformer`` +- 平滑正则: 官方推荐的 ``neurite.nn.modules.SpatialGradient`` +- 相似度: 默认使用带 epsilon 的本地稳定 NCC；可切换到 ``neurite.nn.modules.NCC`` + +输入张量约定： +- NIfTI 数组轴顺序为 (X, Y, Z)。 +- 官方 PyTorch VoxelMorph 输入为 (B, C, X, Y, Z)。 +- DDF 通道顺序与空间轴一致，即 (X, Y, Z)。 +""" + +from __future__ import annotations + +import argparse +from contextlib import nullcontext +import json +import os +from pathlib import Path +from typing import Dict, Iterable, List, Sequence, Tuple + +import numpy as np + +from config import DEFAULT_CHECKPOINT + +OFFICIAL_VXM_REPO = "https://github.com/voxelmorph/voxelmorph" +OFFICIAL_VXM_COMMIT = "db73f34b910bcefcb520f7f40a1bc4a3e0b6401d" + +try: # 让缺 torch 的环境仍能 py_compile。 + import torch + import torch.nn as nn + import torch.nn.functional as F +except Exception as _torch_exc: # pragma: no cover + torch = None # type: ignore + nn = None # type: ignore + F = None # type: ignore + _TORCH_IMPORT_ERROR = _torch_exc +else: + _TORCH_IMPORT_ERROR = None + + +def configure_voxelmorph_backend() -> None: + """显式启用官方 VoxelMorph/Neurite 的 PyTorch 后端。""" + + os.environ.setdefault("NEURITE_BACKEND", "pytorch") + os.environ.setdefault("VXM_BACKEND", "pytorch") + + +def require_torch(): + if torch is None: + raise RuntimeError( + "缺少 PyTorch，无法训练或推理。建议使用 conda 创建 CUDA 环境：" + "conda env create -f environment.yml" + ) from _TORCH_IMPORT_ERROR + return torch + + +def require_official_voxelmorph(): + """导入官方 VoxelMorph 与 Neurite，并检查关键 API。""" + + configure_voxelmorph_backend() + require_torch() + try: + import neurite as ne # type: ignore + import voxelmorph as vxm # type: ignore + except Exception as exc: # pragma: no cover + raise RuntimeError( + "缺少官方 VoxelMorph/Neurite。请运行: conda env create -f environment.yml " + "或 pip install -r requirements.txt" + ) from exc + + if not hasattr(vxm.nn.models, "VxmPairwise"): + raise RuntimeError("当前 voxelmorph 版本缺少 vxm.nn.models.VxmPairwise，请安装官方 dev 版本。") + return vxm, ne + + +def _require_nibabel(): + try: + import nibabel as nib # type: ignore + except Exception as exc: # pragma: no cover + raise RuntimeError("缺少 nibabel，请先安装项目依赖。") from exc + return nib + + +if torch is not None: + + class LocalNCCLoss(nn.Module): + """本地 NCC 兜底。 + + 官方 VoxelMorph 当前把 NCC 放在 Neurite 中；若用户安装的 Neurite + 接口短期变动，本类保证训练脚本仍可运行，但网络核心仍是官方 VoxelMorph。 + """ + + def __init__(self, window_size: int = 9, eps: float = 1e-5): + super().__init__() + self.window_size = int(window_size) + self.eps = float(eps) + + def forward(self, fixed, warped): + win = [self.window_size] * 3 + padding = self.window_size // 2 + filt = torch.ones((1, 1, *win), dtype=fixed.dtype, device=fixed.device) + win_volume = float(np.prod(win)) + + fixed_sum = F.conv3d(fixed, filt, padding=padding) + warped_sum = F.conv3d(warped, filt, padding=padding) + fixed2_sum = F.conv3d(fixed * fixed, filt, padding=padding) + warped2_sum = F.conv3d(warped * warped, filt, padding=padding) + cross_sum = F.conv3d(fixed * warped, filt, padding=padding) + + fixed_mean = fixed_sum / win_volume + warped_mean = warped_sum / win_volume + cross = cross_sum - warped_mean * fixed_sum - fixed_mean * warped_sum + fixed_mean * warped_mean * win_volume + fixed_var = fixed2_sum - 2 * fixed_mean * fixed_sum + fixed_mean * fixed_mean * win_volume + warped_var = warped2_sum - 2 * warped_mean * warped_sum + warped_mean * warped_mean * win_volume + + cc = (cross * cross) / (fixed_var * warped_var + self.eps) + return -torch.mean(cc) + + + class NegativeSimilarity(nn.Module): + """把 Neurite NCC 分数转换为可最小化的 loss。""" + + def __init__(self, module: nn.Module): + super().__init__() + self.module = module + + def forward(self, fixed, warped): + return -torch.mean(self.module(fixed, warped)) + + +else: + + class LocalNCCLoss: # type: ignore[no-redef] + def __init__(self, *args, **kwargs): + require_torch() + + + class NegativeSimilarity: # type: ignore[no-redef] + def __init__(self, *args, **kwargs): + require_torch() + + +def load_nifti_tensor(path: str | Path, device: str | "torch.device" = "cpu") -> Tuple["torch.Tensor", np.ndarray]: + """读取 NIfTI 并转换为官方 VoxelMorph 输入张量。 + + NIfTI data: (X, Y, Z) + Torch tensor: (B, C, X, Y, Z) + """ + + require_torch() + nib = _require_nibabel() + img = nib.load(str(path), mmap=True) + 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) + tensor = torch.from_numpy(data.copy())[None, None].to(device=device, dtype=torch.float32) + return tensor, img.affine.copy() + + +def resolve_device(device: str): + require_torch() + if device == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(device) + + +def build_vxm_model( + nb_features: Sequence[int], + integration_steps: int, + device, +): + """构建官方 VoxelMorph Pairwise 模型。""" + + vxm, _ = require_official_voxelmorph() + model = vxm.nn.models.VxmPairwise( + ndim=3, + source_channels=1, + target_channels=1, + nb_features=list(int(v) for v in nb_features), + integration_steps=int(integration_steps), + device=str(device), + ) + return model.to(device) + + +def build_similarity_loss(loss_name: str, ncc_window: int, ncc_impl: str = "local"): + """构建相似度 loss。 + + 颈部 CT 预处理后通常有大面积零填充背景，Neurite NCC 在低方差窗口中可能 + 返回非有限值。因此默认使用带 epsilon 的本地 NCC；模型核心仍为官方 + VoxelMorph。需要对照官方依赖时，可传入 ``ncc_impl="neurite"``。 + """ + + _, ne = require_official_voxelmorph() + loss_name = loss_name.lower() + if loss_name == "ncc": + ncc_impl = ncc_impl.lower() + if ncc_impl == "neurite": + return NegativeSimilarity(ne.nn.modules.NCC(window_size=int(ncc_window))) + if ncc_impl == "local": + return LocalNCCLoss(window_size=int(ncc_window)) + raise ValueError("ncc_impl 只能是 local 或 neurite。") + if loss_name == "mse": + return ne.nn.modules.MSE() + raise ValueError("loss_name 只能是 ncc 或 mse。") + + +def build_gradient_loss(): + """官方平滑项：Neurite SpatialGradient。""" + + _, ne = require_official_voxelmorph() + return ne.nn.modules.SpatialGradient(penalty="l2") + + +def train_pair( + moving_path: str | Path, + fixed_path: str | Path, + checkpoint_path: str | Path = DEFAULT_CHECKPOINT, + epochs: int = 200, + learning_rate: float = 1e-4, + smooth_weight: float = 0.01, + image_loss: str = "ncc", + ncc_window: int = 9, + ncc_impl: str = "local", + nb_features: Sequence[int] = (16, 16, 16, 16, 16), + integration_steps: int = 0, + device: str = "auto", + use_amp: bool = True, + save_every: int = 50, +) -> List[Dict[str, float]]: + """使用一对 moving/fixed 进行无监督 VoxelMorph 训练。""" + + require_official_voxelmorph() + device_obj = resolve_device(device) + moving, _ = load_nifti_tensor(moving_path, device=device_obj) + fixed, _ = load_nifti_tensor(fixed_path, device=device_obj) + if moving.shape != fixed.shape: + raise ValueError(f"Moving 与 Fixed 形状不一致: {tuple(moving.shape)} vs {tuple(fixed.shape)}") + + model = build_vxm_model(nb_features=nb_features, integration_steps=integration_steps, device=device_obj) + similarity_loss = build_similarity_loss(image_loss, ncc_window=ncc_window, ncc_impl=ncc_impl).to(device_obj) + grad_loss_fn = build_gradient_loss().to(device_obj) + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + amp_enabled = bool(use_amp and device_obj.type == "cuda") + scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) + + checkpoint_path = Path(checkpoint_path) + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + history: List[Dict[str, float]] = [] + + for epoch in range(1, epochs + 1): + model.train() + optimizer.zero_grad(set_to_none=True) + + try: + autocast_context = torch.amp.autocast("cuda", enabled=True) if amp_enabled else nullcontext() + with autocast_context: + displacement, warped = model( + moving, + fixed, + return_warped_source=True, + return_field_type="displacement", + ) + + loss_sim = similarity_loss(fixed.float(), warped.float()) + loss_smooth = grad_loss_fn(displacement.float()) + loss = loss_sim + smooth_weight * loss_smooth + + if not torch.isfinite(loss): + raise RuntimeError( + "训练 loss 出现非有限值。可尝试使用 --ncc-impl local、" + "--image-loss mse，或检查预处理后是否有大面积常值区域。" + ) + + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + except RuntimeError as exc: + if "out of memory" in str(exc).lower(): + raise RuntimeError( + "GPU/内存不足。可尝试减小 --target-shape，或降低 --nb-features，" + "再重新预处理与训练。" + ) from exc + raise + + row = { + "epoch": float(epoch), + "loss": float(loss.detach().cpu()), + "image_loss": float(loss_sim.detach().cpu()), + "smooth_loss": float(loss_smooth.detach().cpu()), + } + history.append(row) + print( + f"Epoch {epoch:04d}/{epochs} | " + f"loss={row['loss']:.6f} image={row['image_loss']:.6f} smooth={row['smooth_loss']:.6f}" + ) + + if epoch == epochs or (save_every > 0 and epoch % save_every == 0): + save_checkpoint( + checkpoint_path, + model=model, + epoch=epoch, + input_shape_xyz=tuple(int(v) for v in moving.shape[2:]), + nb_features=nb_features, + integration_steps=integration_steps, + config={ + "moving_path": str(moving_path), + "fixed_path": str(fixed_path), + "learning_rate": learning_rate, + "smooth_weight": smooth_weight, + "image_loss": image_loss, + "ncc_window": ncc_window, + "ncc_impl": ncc_impl, + }, + history=history, + ) + + history_path = checkpoint_path.with_suffix(".history.json") + history_path.write_text(json.dumps(history, ensure_ascii=False, indent=2), encoding="utf-8") + return history + + +def save_checkpoint( + path: str | Path, + model, + epoch: int, + input_shape_xyz: Sequence[int], + nb_features: Sequence[int], + integration_steps: int, + config: Dict, + history: List[Dict[str, float]], +) -> None: + require_torch() + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "model_state_dict": model.state_dict(), + "epoch": int(epoch), + "input_shape_xyz": tuple(int(v) for v in input_shape_xyz), + "nb_features": list(int(v) for v in nb_features), + "integration_steps": int(integration_steps), + "official_core": { + "library": "voxelmorph", + "repo": OFFICIAL_VXM_REPO, + "commit": OFFICIAL_VXM_COMMIT, + "model_class": "voxelmorph.nn.models.VxmPairwise", + }, + "config": config, + "history": history, + }, + path, + ) + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="用官方 VoxelMorph 训练 3D 颈部 CT 形变配准模型。") + 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("--epochs", type=int, default=200) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--smooth-weight", type=float, default=0.01) + parser.add_argument("--image-loss", choices=["ncc", "mse"], default="ncc") + parser.add_argument("--ncc-window", type=int, default=9) + parser.add_argument("--ncc-impl", choices=["local", "neurite"], default="local") + parser.add_argument("--nb-features", type=int, nargs="+", default=[16, 16, 16, 16, 16]) + parser.add_argument("--integration-steps", type=int, default=0, help="0 为普通 dense flow；>0 使用 scaling-and-squaring。") + parser.add_argument("--device", default="auto") + parser.add_argument("--no-amp", action="store_true", help="关闭 CUDA AMP 混合精度。") + parser.add_argument("--save-every", type=int, default=50) + return parser + + +def main(argv: Iterable[str] | None = None) -> None: + args = build_arg_parser().parse_args(argv) + train_pair( + moving_path=args.moving, + fixed_path=args.fixed, + checkpoint_path=args.checkpoint, + epochs=args.epochs, + learning_rate=args.lr, + smooth_weight=args.smooth_weight, + image_loss=args.image_loss, + ncc_window=args.ncc_window, + ncc_impl=args.ncc_impl, + nb_features=args.nb_features, + integration_steps=args.integration_steps, + device=args.device, + use_amp=not args.no_amp, + save_every=args.save_every, + ) + + +if __name__ == "__main__": + main() diff --git a/outputs/checkpoints/.gitkeep b/outputs/checkpoints/.gitkeep new file mode 100644 index 0000000..5f72043 --- /dev/null +++ b/outputs/checkpoints/.gitkeep @@ -0,0 +1,2 @@ +# keep + diff --git a/outputs/inference/.gitkeep b/outputs/inference/.gitkeep new file mode 100644 index 0000000..5f72043 --- /dev/null +++ b/outputs/inference/.gitkeep @@ -0,0 +1,2 @@ +# keep + diff --git a/outputs/nifti/.gitkeep b/outputs/nifti/.gitkeep new file mode 100644 index 0000000..5f72043 --- /dev/null +++ b/outputs/nifti/.gitkeep @@ -0,0 +1,2 @@ +# keep + diff --git a/outputs/preprocessed/.gitkeep b/outputs/preprocessed/.gitkeep new file mode 100644 index 0000000..5f72043 --- /dev/null +++ b/outputs/preprocessed/.gitkeep @@ -0,0 +1,2 @@ +# keep + diff --git a/preprocess.py b/preprocess.py new file mode 100644 index 0000000..18db4c9 --- /dev/null +++ b/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() + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..65c4283 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,17 @@ +numpy>=1.24 +scipy>=1.10 +pydicom>=2.4 +nibabel>=5.2 +torch>=2.2 +streamlit>=1.33 +matplotlib>=3.8 +plotly>=5.20 +packaging>=23 +scikit-image>=0.22 +h5py>=3.10 +tqdm>=4.66 +einops>=0.7 +pystrum>=0.4 +neurite @ git+https://github.com/adalca/neurite.git@dev +voxelmorph @ git+https://github.com/voxelmorph/voxelmorph.git@db73f34b910bcefcb520f7f40a1bc4a3e0b6401d + diff --git a/scripts/setup_env.sh b/scripts/setup_env.sh new file mode 100755 index 0000000..60b61ac --- /dev/null +++ b/scripts/setup_env.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +ENV_NAME="${ENV_NAME:-voxelmorph-head-ct}" + +conda env create -f environment.yml +conda run -n "${ENV_NAME}" python -m pip install --no-cache-dir -r requirements.txt +conda run -n "${ENV_NAME}" python - <<'PY' +import os +os.environ.setdefault("NEURITE_BACKEND", "pytorch") +os.environ.setdefault("VXM_BACKEND", "pytorch") + +import torch +import voxelmorph as vxm +import neurite as ne + +print("torch:", torch.__version__) +print("cuda:", torch.cuda.is_available()) +if torch.cuda.is_available(): + print("gpu:", torch.cuda.get_device_name(0)) +print("voxelmorph VxmPairwise:", hasattr(vxm.nn.models, "VxmPairwise")) +print("neurite NCC:", hasattr(ne.nn.modules, "NCC")) +PY