torch_em.data.datasets.medical.qin_prostate

  1"""The QIN-PROSTATE-Repeatability dataset contains annotations for the prostate and its peripheral
  2zone in MRI.
  3
  4The dataset consists of a test-retest repeatability study: 15 patients were each scanned twice, with
  5expert segmentations of the whole prostate, its peripheral zone, a lesion and normal tissue (see
  6`SEGMENT_IDS`) drawn on three different MRI series per scan (T2-weighted, an ADC map, and a DCE
  7subtraction image), each as a separate DICOM-SEG object paired here with the exact series it
  8references. Not every segment is present in every file (e.g. a case without a lesion has no 'Lesion'
  9segment), and a DICOM-SEG's own segment numbering only reflects the segments it actually contains, so
 10this module remaps every segment to `SEGMENT_IDS` by its name rather than trusting that numbering.
 11
 12NOTE: This requires the pydicom python package.
 13
 14The dataset is located at https://doi.org/10.7937/K9/TCIA.2018.MR1CKGND and is distributed under the
 15CC BY 4.0 license.
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20import json
 21from glob import glob
 22from tqdm import tqdm
 23from natsort import natsorted
 24from typing import Union, Tuple, List
 25
 26import numpy as np
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .adrenal_acc import _load_dicom_volume, _load_dicom_seg, _resample_labels
 33from .. import util
 34
 35
 36COLLECTION = "QIN-PROSTATE-Repeatability"
 37
 38SEGMENT_IDS = {"normal": 1, "peripheral_zone": 2, "lesion": 3, "prostate": 4}
 39"""Mapping from the anatomical structure to its label id."""
 40
 41_SEGMENT_NAME_TO_ID = {
 42    "normal": SEGMENT_IDS["normal"],
 43    "peripheral zone of the prostate": SEGMENT_IDS["peripheral_zone"],
 44    "lesion": SEGMENT_IDS["lesion"],
 45    "prostate": SEGMENT_IDS["prostate"],
 46}
 47
 48
 49def _get_series_metadata(path, download):
 50    """Get the metadata of all series in the collection from the NBIA REST API."""
 51    import requests
 52
 53    metadata_path = os.path.join(path, "qin_prostate_series.json")
 54    if not os.path.exists(metadata_path):
 55        if not download:
 56            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
 57        response = requests.get(f"{util.NBIA_API_URL}getSeries", params={"Collection": COLLECTION})
 58        response.raise_for_status()
 59        with open(metadata_path, "w") as f:
 60            json.dump(response.json(), f, indent=2)
 61
 62    with open(metadata_path, "r") as f:
 63        return json.load(f)
 64
 65
 66def _referenced_series_uid(seg):
 67    return str(seg.ReferencedSeriesSequence[0].SeriesInstanceUID)
 68
 69
 70def _remap_segments_by_name(seg_labels, seg):
 71    """Remap a DICOM-SEG label volume from its own (file-specific) segment numbers to `SEGMENT_IDS`,
 72    since a file that omits a segment does not preserve the numbering of the others.
 73    """
 74    remapped = np.zeros_like(seg_labels)
 75    for segment in seg.SegmentSequence:
 76        canonical_id = _SEGMENT_NAME_TO_ID.get(str(segment.SegmentLabel).lower())
 77        if canonical_id is not None:
 78            remapped[seg_labels == int(segment.SegmentNumber)] = canonical_id
 79    return remapped
 80
 81
 82def _preprocess_qin_prostate(dicom_dir, series_metadata, preprocessed_dir):
 83    import h5py
 84    import pydicom
 85
 86    seg_series = [series for series in series_metadata if series.get("Modality") == "SEG"]
 87
 88    os.makedirs(preprocessed_dir, exist_ok=True)
 89    for series in tqdm(seg_series, desc="Preprocess QIN-PROSTATE-Repeatability"):
 90        seg_paths = glob(os.path.join(dicom_dir, series["SeriesInstanceUID"], "*.dcm"))
 91        if not seg_paths:
 92            continue
 93
 94        out_path = os.path.join(preprocessed_dir, f"{series['SeriesInstanceUID']}.h5")
 95        if os.path.exists(out_path):
 96            continue
 97
 98        seg = pydicom.dcmread(seg_paths[0])
 99        image_dir = os.path.join(dicom_dir, _referenced_series_uid(seg))
100        if not glob(os.path.join(image_dir, "*.dcm")):
101            continue
102
103        volume, image_affine = _load_dicom_volume(image_dir)
104        seg_labels, seg_affine = _load_dicom_seg(seg_paths[0])
105        seg_labels = _remap_segments_by_name(seg_labels, seg)
106        labels = _resample_labels(seg_labels, seg_affine, volume.shape, image_affine)
107        if labels.max() == 0:
108            continue
109
110        with h5py.File(out_path, "w") as f:
111            f.create_dataset("raw", data=volume, compression="gzip")
112            f.create_dataset("labels", data=labels, compression="gzip")
113
114
115def get_qin_prostate_data(path: Union[os.PathLike, str], download: bool = False) -> str:
116    """Download the QIN-PROSTATE-Repeatability dataset.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        download: Whether to download the data if it is not present.
121
122    Returns:
123        Filepath where the preprocessed data is stored.
124    """
125    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
126    preprocessed_dir = os.path.join(path, "preprocessed")
127
128    os.makedirs(path, exist_ok=True)
129    series_metadata = _get_series_metadata(path, download)
130
131    seg_uids = [series["SeriesInstanceUID"] for series in series_metadata if series.get("Modality") == "SEG"]
132
133    dicom_dir = os.path.join(path, "dicom")
134    if download:  # The SEG series are downloaded first, so the image series they reference can be found.
135        util.download_tcia_series(seg_uids, dst=dicom_dir, csv_filename=os.path.join(path, "qin_prostate_seg"))
136
137        import pydicom
138        image_uids = set()
139        for uid in seg_uids:
140            seg_paths = glob(os.path.join(dicom_dir, uid, "*.dcm"))
141            if seg_paths:
142                seg = pydicom.dcmread(seg_paths[0])
143                image_uids.add(_referenced_series_uid(seg))
144        util.download_tcia_series(
145            sorted(image_uids), dst=dicom_dir, csv_filename=os.path.join(path, "qin_prostate_image")
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_qin_prostate(dicom_dir, series_metadata, preprocessed_dir)
151    return preprocessed_dir
152
153
154def get_qin_prostate_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
155    """Get paths to the QIN-PROSTATE-Repeatability 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_qin_prostate_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_qin_prostate_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 QIN-PROSTATE-Repeatability dataset for prostate and peripheral zone 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_qin_prostate_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_qin_prostate_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 QIN-PROSTATE-Repeatability dataloader for prostate and peripheral zone 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_qin_prostate_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
231    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)