torch_em.data.datasets.medical.spinal_mm

  1"""The Spinal-Multiple-Myeloma-SEG dataset contains annotations for spinal multiple myeloma lesions in CT.
  2
  3The dataset consists of 144 dual-energy CT studies with instance masks for the focal spinal lesions of
  4each patient, up to several dozen per case. The masks are stored as DICOM-SEG objects, and this module
  5pairs them with the conventional ('_konv') reconstruction of their dual-energy CT study.
  6
  7NOTE: A dual-energy study has several reconstructions on the identical voxel grid (conventional,
  8several monoenergetic keV levels, several calcium-suppression indices), so the shape and geometry of a
  9scan cannot disambiguate which one the segmentation was drawn on. The DICOM-SEG series references its
 10source CT by UID, but that series is not resolvable through TCIA anymore for some cases, so the
 11conventional reconstruction of the same study is used instead: every one of the 144 studies has exactly
 12one CT series whose description ends in '_konv', and this is the closest series in the collection to a
 13routine single-energy CT.
 14
 15NOTE: This requires the pydicom python package.
 16
 17The dataset is located at https://doi.org/10.7937/k4qv-hh78 and is distributed under the CC BY 4.0
 18license.
 19This dataset is from the publication https://doi.org/10.1038/s41597-026-08061-x.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24import json
 25from glob import glob
 26from tqdm import tqdm
 27from natsort import natsorted
 28from typing import Union, Tuple, List
 29
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .adrenal_acc import _load_dicom_volume, _load_dicom_seg, _resample_labels
 36from .. import util
 37
 38
 39COLLECTION = "Spinal-Multiple-Myeloma-SEG"
 40
 41
 42def _get_series_metadata(path, download):
 43    """Get the metadata of all series in the collection from the NBIA REST API."""
 44    import requests
 45
 46    metadata_path = os.path.join(path, "spinal_mm_series.json")
 47    if not os.path.exists(metadata_path):
 48        if not download:
 49            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
 50        response = requests.get(f"{util.NBIA_API_URL}getSeries", params={"Collection": COLLECTION})
 51        response.raise_for_status()
 52        with open(metadata_path, "w") as f:
 53            json.dump(response.json(), f, indent=2)
 54
 55    with open(metadata_path, "r") as f:
 56        return json.load(f)
 57
 58
 59def _find_konv_series(series_metadata, study_uid, patient_id):
 60    """Find the conventional reconstruction of a dual-energy study."""
 61    candidates = [
 62        series for series in series_metadata
 63        if series.get("StudyInstanceUID") == study_uid and series.get("Modality") == "CT"
 64        and (series.get("SeriesDescription") or "").lower() == f"{patient_id}_konv".lower()
 65    ]
 66    return candidates[0]["SeriesInstanceUID"] if len(candidates) == 1 else None
 67
 68
 69def _preprocess_spinal_mm(dicom_dir, series_metadata, preprocessed_dir):
 70    import h5py
 71
 72    seg_series = [series for series in series_metadata if series.get("Modality") == "SEG"]
 73
 74    os.makedirs(preprocessed_dir, exist_ok=True)
 75    for series in tqdm(seg_series, desc="Preprocess Spinal-Multiple-Myeloma-SEG"):
 76        patient_id = series["PatientID"]
 77        # A patient can have more than one study (e.g. a follow-up scan), each with its own SEG series,
 78        # so the output is keyed by the SEG series UID rather than the patient id.
 79        out_path = os.path.join(preprocessed_dir, f"{series['SeriesInstanceUID']}.h5")
 80        if os.path.exists(out_path):
 81            continue
 82
 83        seg_dir = os.path.join(dicom_dir, series["SeriesInstanceUID"])
 84        seg_paths = glob(os.path.join(seg_dir, "*.dcm"))
 85        if not seg_paths:
 86            continue
 87
 88        import pydicom
 89        study_uid = pydicom.dcmread(seg_paths[0], stop_before_pixels=True).StudyInstanceUID
 90        ct_uid = _find_konv_series(series_metadata, study_uid, patient_id)
 91        ct_dir = os.path.join(dicom_dir, ct_uid) if ct_uid else None
 92        if ct_dir is None or not glob(os.path.join(ct_dir, "*.dcm")):
 93            continue
 94
 95        volume, ct_affine = _load_dicom_volume(ct_dir)
 96        seg_labels, seg_affine = _load_dicom_seg(seg_paths[0])
 97        # The SEG object only covers the slices its lesions appear on, cropped from the full CT extent,
 98        # so it is placed on the CT's own grid rather than compared to it directly.
 99        labels = _resample_labels(seg_labels, seg_affine, volume.shape, ct_affine)
100
101        with h5py.File(out_path, "w") as f:
102            f.create_dataset("raw", data=volume, compression="gzip")
103            f.create_dataset("labels", data=labels, compression="gzip")
104
105
106def get_spinal_mm_data(path: Union[os.PathLike, str], download: bool = False) -> str:
107    """Download the Spinal-Multiple-Myeloma-SEG dataset.
108
109    Args:
110        path: Filepath to a folder where the data is downloaded for further processing.
111        download: Whether to download the data if it is not present.
112
113    Returns:
114        Filepath where the preprocessed data is stored.
115    """
116    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
117    preprocessed_dir = os.path.join(path, "preprocessed")
118
119    os.makedirs(path, exist_ok=True)
120    series_metadata = _get_series_metadata(path, download)
121
122    seg_uids = [series["SeriesInstanceUID"] for series in series_metadata if series.get("Modality") == "SEG"]
123
124    # Every SEG series belongs to one study; the conventional CT of that study is downloaded alongside it.
125    study_by_seg = {
126        series["SeriesInstanceUID"]: (series["StudyInstanceUID"], series["PatientID"])
127        for series in series_metadata if series.get("Modality") == "SEG"
128    }
129    konv_uids = [
130        _find_konv_series(series_metadata, study_uid, patient_id)
131        for study_uid, patient_id in study_by_seg.values()
132    ]
133    series_uids = sorted(set(seg_uids) | {uid for uid in konv_uids if uid})
134
135    dicom_dir = os.path.join(path, "dicom")
136    if download:  # Series that were downloaded already are skipped.
137        util.download_tcia_series(series_uids, dst=dicom_dir, csv_filename=os.path.join(path, "spinal_mm"))
138    elif not all(glob(os.path.join(dicom_dir, uid, "*.dcm")) for uid in series_uids):
139        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
140
141    _preprocess_spinal_mm(dicom_dir, series_metadata, preprocessed_dir)
142    return preprocessed_dir
143
144
145def get_spinal_mm_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
146    """Get paths to the Spinal-Multiple-Myeloma-SEG data.
147
148    Args:
149        path: Filepath to a folder where the data is downloaded for further processing.
150        download: Whether to download the data if it is not present.
151
152    Returns:
153        List of filepaths for the stored data.
154    """
155    preprocessed_dir = get_spinal_mm_data(path, download)
156    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
157    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
158    return volume_paths
159
160
161def get_spinal_mm_dataset(
162    path: Union[os.PathLike, str],
163    patch_shape: Tuple[int, ...],
164    resize_inputs: bool = False,
165    download: bool = False,
166    **kwargs
167) -> Dataset:
168    """Get the Spinal-Multiple-Myeloma-SEG dataset for spinal lesion segmentation.
169
170    Args:
171        path: Filepath to a folder where the data is downloaded for further processing.
172        patch_shape: The patch shape to use for training.
173        resize_inputs: Whether to resize inputs to the desired patch shape.
174        download: Whether to download the data if it is not present.
175        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
176
177    Returns:
178        The segmentation dataset.
179    """
180    volume_paths = get_spinal_mm_paths(path, download)
181
182    if resize_inputs:
183        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
184        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
185            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
186        )
187
188    return torch_em.default_segmentation_dataset(
189        raw_paths=volume_paths,
190        raw_key="raw",
191        label_paths=volume_paths,
192        label_key="labels",
193        patch_shape=patch_shape,
194        is_seg_dataset=True,
195        **kwargs
196    )
197
198
199def get_spinal_mm_loader(
200    path: Union[os.PathLike, str],
201    batch_size: int,
202    patch_shape: Tuple[int, ...],
203    resize_inputs: bool = False,
204    download: bool = False,
205    **kwargs
206) -> DataLoader:
207    """Get the Spinal-Multiple-Myeloma-SEG dataloader for spinal lesion segmentation.
208
209    Args:
210        path: Filepath to a folder where the data is downloaded for further processing.
211        batch_size: The batch size for training.
212        patch_shape: The patch shape to use for training.
213        resize_inputs: Whether to resize inputs to the desired patch shape.
214        download: Whether to download the data if it is not present.
215        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
216
217    Returns:
218        The DataLoader.
219    """
220    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
221    dataset = get_spinal_mm_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
222    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)