Add official VoxelMorph CT registration pipeline

This commit is contained in:
admin
2026-06-03 00:30:32 +08:00
parent e8d8f2c468
commit 2dba05ae4a
16 changed files with 2130 additions and 0 deletions

23
.gitignore vendored Normal file
View File

@@ -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

141
README.md Normal file
View File

@@ -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。

430
app.py Normal file
View File

@@ -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 = """
<style>
:root {
--ink: #202124;
--muted: #5f6368;
--line: #d8dee4;
--panel: #f6f8fa;
--teal: #0f766e;
--amber: #b45309;
}
.main .block-container {
padding-top: 1.1rem;
padding-bottom: 2rem;
max-width: 1480px;
}
h1, h2, h3 {
letter-spacing: 0;
}
div[data-testid="stMetric"] {
border: 1px solid var(--line);
border-radius: 6px;
padding: 0.65rem 0.75rem;
background: #ffffff;
}
section[data-testid="stSidebar"] {
border-right: 1px solid var(--line);
}
.stTabs [data-baseweb="tab-list"] {
gap: 8px;
border-bottom: 1px solid var(--line);
}
.stTabs [data-baseweb="tab"] {
border-radius: 0;
padding-left: 4px;
padding-right: 4px;
}
</style>
"""
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()

32
config.py Normal file
View File

@@ -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

354
data_loader.py Normal file
View File

@@ -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()

10
environment.yml Normal file
View File

@@ -0,0 +1,10 @@
name: voxelmorph-head-ct
channels:
- pytorch
- nvidia
- conda-forge
dependencies:
- python=3.11
- pip
- pytorch
- pytorch-cuda=12.4

242
infer.py Normal file
View File

@@ -0,0 +1,242 @@
"""独立推理模块。
本模块加载官方 ``voxelmorph.nn.models.VxmPairwise`` 训练出的权重,输出:
- warped_moving.nii.gz形变后的 Moving。
- ddf_mm.nii.gzDense Displacement Fieldshape = (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()

123
metrics.py Normal file
View File

@@ -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)),
}

400
model_and_train.py Normal file
View File

@@ -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()

View File

@@ -0,0 +1,2 @@
# keep

View File

@@ -0,0 +1,2 @@
# keep

2
outputs/nifti/.gitkeep Normal file
View File

@@ -0,0 +1,2 @@
# keep

View File

@@ -0,0 +1,2 @@
# keep

327
preprocess.py Normal file
View File

@@ -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()

17
requirements.txt Normal file
View File

@@ -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

23
scripts/setup_env.sh Executable file
View File

@@ -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