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