Add official VoxelMorph CT registration pipeline

2026-06-03 00:30:32 +08:00
parent e8d8f2c468
commit 2dba05ae4a
16 changed files with 2130 additions and 0 deletions
--- a/data_loader.py
+++ 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()
+