354 lines
13 KiB
Python
354 lines
13 KiB
Python
"""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] = col_cos * col_spacing
|
||
affine[:3, 1] = row_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()
|