torch_em.data.datasets.medical.acrin_6698_dwi

  1"""The ACRIN-6698 dataset contains annotations for the whole breast tumor in diffusion-weighted MRI.
  2
  3This module covers the 1103 'DWI SEG' objects of the ACRIN-6698 / I-SPY2 collection: manual whole-tumor
  4segmentations on the apparent diffusion coefficient (ADC) map of breast cancer patients undergoing
  5neoadjuvant chemotherapy. Each DICOM-SEG object is named '...DWI SEG: from S<n>: Whole Tumor Manual',
  6where '<n>' is the DICOM series number of the ADC map it was drawn on; this module resolves that
  7reference by series number rather than a DICOM reference, since the SEG objects in this collection
  8carry none (neither a `ReferencedSeriesSequence` nor a per-frame derivation link).
  9
 10NOTE: The collection's other, unrelated family of 1110 'ISPY2: VOLSER' functional-tumor-volume masks
 11(thresholded from a signal-enhancement-ratio map, not a manual DWI segmentation) is covered separately
 12by `medical.acrin_6698`.
 13
 14NOTE: This requires the pydicom python package.
 15
 16The dataset is located at https://doi.org/10.7937/TCIA.kk02-6d95 and is distributed under the
 17CC BY 4.0 license.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import re
 23import json
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from typing import Union, Tuple, List
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .adrenal_acc import _load_dicom_volume, _load_dicom_seg, _resample_labels
 34from .. import util
 35
 36
 37COLLECTION = "ACRIN-6698"
 38
 39
 40def _get_series_metadata(path, download):
 41    """Get the metadata of all series in the collection from the NBIA REST API."""
 42    import requests
 43
 44    metadata_path = os.path.join(path, "acrin_6698_dwi_series.json")
 45    if not os.path.exists(metadata_path):
 46        if not download:
 47            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
 48        response = requests.get(f"{util.NBIA_API_URL}getSeries", params={"Collection": COLLECTION})
 49        response.raise_for_status()
 50        with open(metadata_path, "w") as f:
 51            json.dump(response.json(), f, indent=2)
 52
 53    with open(metadata_path, "r") as f:
 54        return json.load(f)
 55
 56
 57def _is_dwi_seg(series):
 58    return series.get("Modality") == "SEG" and "DWI SEG" in (series.get("SeriesDescription") or "")
 59
 60
 61def _referenced_series_number(series):
 62    match = re.search(r"from S(\d+)", series["SeriesDescription"])
 63    return match.group(1) if match else None
 64
 65
 66def _find_adc_series(series_metadata, study_uid, series_number):
 67    """Find the ADC map a DWI SEG was drawn on, matched by the series number in its own name."""
 68    candidates = [
 69        series for series in series_metadata
 70        if series.get("StudyInstanceUID") == study_uid and series.get("Modality") == "MR"
 71        and f"ADC: from S{series_number}:" in (series.get("SeriesDescription") or "")
 72    ]
 73    return candidates[0]["SeriesInstanceUID"] if len(candidates) == 1 else None
 74
 75
 76def _preprocess_acrin_6698_dwi(dicom_dir, series_metadata, preprocessed_dir):
 77    import h5py
 78
 79    seg_series = [series for series in series_metadata if _is_dwi_seg(series)]
 80
 81    os.makedirs(preprocessed_dir, exist_ok=True)
 82    for series in tqdm(seg_series, desc="Preprocess ACRIN-6698 (DWI)"):
 83        seg_paths = glob(os.path.join(dicom_dir, series["SeriesInstanceUID"], "*.dcm"))
 84        if not seg_paths:
 85            continue
 86
 87        out_path = os.path.join(preprocessed_dir, f"{series['SeriesInstanceUID']}.h5")
 88        if os.path.exists(out_path):
 89            continue
 90
 91        series_number = _referenced_series_number(series)
 92        if series_number is None:
 93            continue
 94        adc_uid = _find_adc_series(series_metadata, series["StudyInstanceUID"], series_number)
 95        adc_dir = os.path.join(dicom_dir, adc_uid) if adc_uid else None
 96        if adc_dir is None or not glob(os.path.join(adc_dir, "*.dcm")):
 97            continue
 98
 99        volume, adc_affine = _load_dicom_volume(adc_dir)
100        seg_labels, seg_affine = _load_dicom_seg(seg_paths[0])
101        # The SEG only covers the slices where the tumor is present, cropped from the full ADC extent,
102        # so it is placed on the ADC's own grid rather than compared to it directly.
103        labels = _resample_labels(seg_labels, seg_affine, volume.shape, adc_affine)
104        if labels.max() == 0:
105            continue
106
107        with h5py.File(out_path, "w") as f:
108            f.create_dataset("raw", data=volume, compression="gzip")
109            f.create_dataset("labels", data=labels, compression="gzip")
110
111
112def get_acrin_6698_dwi_data(path: Union[os.PathLike, str], download: bool = False) -> str:
113    """Download the ACRIN-6698 dataset.
114
115    Args:
116        path: Filepath to a folder where the data is downloaded for further processing.
117        download: Whether to download the data if it is not present.
118
119    Returns:
120        Filepath where the preprocessed data is stored.
121    """
122    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
123    preprocessed_dir = os.path.join(path, "preprocessed")
124
125    os.makedirs(path, exist_ok=True)
126    series_metadata = _get_series_metadata(path, download)
127
128    seg_uids = [series["SeriesInstanceUID"] for series in series_metadata if _is_dwi_seg(series)]
129
130    dicom_dir = os.path.join(path, "dicom")
131    if download:  # The SEG series are downloaded first, so the ADC series they reference can be found.
132        util.download_tcia_series(seg_uids, dst=dicom_dir, csv_filename=os.path.join(path, "acrin_6698_dwi_seg"))
133
134        seg_by_uid = {series["SeriesInstanceUID"]: series for series in series_metadata if _is_dwi_seg(series)}
135        adc_uids = set()
136        for uid in seg_uids:
137            series = seg_by_uid[uid]
138            series_number = _referenced_series_number(series)
139            if series_number is None:
140                continue
141            adc_uid = _find_adc_series(series_metadata, series["StudyInstanceUID"], series_number)
142            if adc_uid:
143                adc_uids.add(adc_uid)
144        util.download_tcia_series(
145            sorted(adc_uids), dst=dicom_dir, csv_filename=os.path.join(path, "acrin_6698_dwi_adc")
146        )
147    elif not glob(os.path.join(dicom_dir, "*", "*.dcm")):
148        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
149
150    _preprocess_acrin_6698_dwi(dicom_dir, series_metadata, preprocessed_dir)
151    return preprocessed_dir
152
153
154def get_acrin_6698_dwi_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
155    """Get paths to the ACRIN-6698 data.
156
157    Args:
158        path: Filepath to a folder where the data is downloaded for further processing.
159        download: Whether to download the data if it is not present.
160
161    Returns:
162        List of filepaths for the stored data.
163    """
164    preprocessed_dir = get_acrin_6698_dwi_data(path, download)
165    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
166    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
167    return volume_paths
168
169
170def get_acrin_6698_dwi_dataset(
171    path: Union[os.PathLike, str],
172    patch_shape: Tuple[int, ...],
173    resize_inputs: bool = False,
174    download: bool = False,
175    **kwargs
176) -> Dataset:
177    """Get the ACRIN-6698 dataset for breast tumor segmentation.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        patch_shape: The patch shape to use for training.
182        resize_inputs: Whether to resize inputs to the desired patch shape.
183        download: Whether to download the data if it is not present.
184        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
185
186    Returns:
187        The segmentation dataset.
188    """
189    volume_paths = get_acrin_6698_dwi_paths(path, download)
190
191    if resize_inputs:
192        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
193        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
194            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
195        )
196
197    return torch_em.default_segmentation_dataset(
198        raw_paths=volume_paths,
199        raw_key="raw",
200        label_paths=volume_paths,
201        label_key="labels",
202        patch_shape=patch_shape,
203        is_seg_dataset=True,
204        **kwargs
205    )
206
207
208def get_acrin_6698_dwi_loader(
209    path: Union[os.PathLike, str],
210    batch_size: int,
211    patch_shape: Tuple[int, ...],
212    resize_inputs: bool = False,
213    download: bool = False,
214    **kwargs
215) -> DataLoader:
216    """Get the ACRIN-6698 dataloader for breast tumor segmentation.
217
218    Args:
219        path: Filepath to a folder where the data is downloaded for further processing.
220        batch_size: The batch size for training.
221        patch_shape: The patch shape to use for training.
222        resize_inputs: Whether to resize inputs to the desired patch shape.
223        download: Whether to download the data if it is not present.
224        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
225
226    Returns:
227        The DataLoader.
228    """
229    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
230    dataset = get_acrin_6698_dwi_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
231    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)