torch_em.data.datasets.medical.resect

  1"""The RESECT dataset contains pre-operative MRI and intra-operative 3D ultrasound of 23 patients
  2with low-grade gliomas, together with annotations from the RESECT-SEG extension.
  3
  4For every patient, three intra-operative ultrasound (US) volumes were acquired: before, during and after
  5tumor resection. The annotations (RESECT-SEG, an extension of the CuRIOUS 2022 challenge labels) are:
  6- 'tumor': the tumor in the US volume before resection and in the pre-operative FLAIR MRI (23 cases each),
  7- 'resection': the resection cavity in the US volumes during (21 cases) and after (22 cases) resection,
  8- 'sulci': the cerebral sulci in all three US phases (23 cases each),
  9- 'falx': the cerebral falx in the US volumes (7 to 8 cases per phase).
 10All labels are binary (0: background, 1: structure).
 11
 12The image volumes are located at https://doi.org/10.11582/2017.00004 (NIRD research data archive, CC BY 4.0)
 13and the annotations at https://osf.io/jv8bk/ (CC BY-NC-SA 4.0). Only the volumes needed for the chosen
 14source, phase and structure are downloaded.
 15
 16This dataset is from the publications https://doi.org/10.1002/mp.12268 (RESECT) and
 17https://doi.org/10.1002/mp.17317 (RESECT-SEG annotations).
 18Please cite them if you use this dataset in your research.
 19"""
 20
 21import os
 22from glob import glob
 23from natsort import natsorted
 24from typing import Union, Tuple, Literal, List, Optional
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33IMAGE_URL = "https://data.archive.sigma2.no/dataset/5686d8fa-2003-4837-8e66-8e887fabe21e/download/RESECT/NIFTI"
 34# The folder zip is generated on-the-fly by OSF, hence the checksum of the archive is not reliable.
 35LABEL_URL = "https://files.osf.io/v1/resources/jv8bk/providers/osfstorage/64cd20819cbf033b051e46c8/?zip="
 36
 37CASE_IDS = [1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 13, 14, 15, 16, 17, 18, 19, 21, 23, 24, 25, 26, 27]
 38
 39STRUCTURES = {
 40    "before": ["tumor", "sulci", "falx"],
 41    "during": ["resection", "sulci", "falx"],
 42    "after": ["resection", "sulci", "falx"],
 43    "MRI": ["tumor"],
 44}
 45
 46
 47def _get_structure(source, phase, structure):
 48    if source not in ["US", "MRI"]:
 49        raise ValueError(f"'{source}' is not a valid source. Choose either 'US' or 'MRI'.")
 50    if source == "US" and phase not in ["before", "during", "after"]:
 51        raise ValueError(f"'{phase}' is not a valid phase. Choose one of 'before', 'during' or 'after'.")
 52
 53    valid_structures = STRUCTURES["MRI" if source == "MRI" else phase]
 54    if structure is None:
 55        structure = valid_structures[0]
 56    if structure not in valid_structures:
 57        raise ValueError(
 58            f"'{structure}' is not a valid structure for source '{source}' and phase '{phase}'. "
 59            f"Choose one of {valid_structures}."
 60        )
 61    return structure
 62
 63
 64def get_resect_data(
 65    path: Union[os.PathLike, str],
 66    source: Literal["US", "MRI"] = "US",
 67    phase: Literal["before", "during", "after"] = "before",
 68    download: bool = False,
 69) -> Tuple[str, str]:
 70    """Download the RESECT dataset.
 71
 72    Args:
 73        path: Filepath to a folder where the data is downloaded for further processing.
 74        source: The imaging source. Either 'US' (intra-operative ultrasound) or 'MRI' (pre-operative FLAIR).
 75        phase: The surgical phase of the ultrasound volumes. Either 'before', 'during' or 'after' resection.
 76            Ignored for the 'MRI' source.
 77        download: Whether to download the data if it is not present.
 78
 79    Returns:
 80        Filepath to the folder with the image volumes.
 81        Filepath to the folder with the label volumes.
 82    """
 83    _get_structure(source, phase, None)
 84
 85    label_dir = os.path.join(path, "labels")
 86    if not os.path.exists(label_dir):
 87        os.makedirs(path, exist_ok=True)
 88        zip_path = os.path.join(path, "RESECT-Segmentation.zip")
 89        util.download_source(path=zip_path, url=LABEL_URL, download=download, checksum=None)
 90        util.unzip(zip_path=zip_path, dst=label_dir)
 91
 92    image_dir = os.path.join(path, "images", source)
 93    os.makedirs(image_dir, exist_ok=True)
 94    for case_id in CASE_IDS:
 95        fname = f"Case{case_id}-US-{phase}.nii.gz" if source == "US" else f"Case{case_id}-FLAIR.nii.gz"
 96        url = f"{IMAGE_URL}/Case{case_id}/{source}/{fname}"
 97        util.download_source(path=os.path.join(image_dir, fname), url=url, download=download, checksum=None)
 98
 99    return image_dir, label_dir
100
101
102def get_resect_paths(
103    path: Union[os.PathLike, str],
104    source: Literal["US", "MRI"] = "US",
105    phase: Literal["before", "during", "after"] = "before",
106    structure: Optional[Literal["tumor", "resection", "sulci", "falx"]] = None,
107    download: bool = False,
108) -> Tuple[List[str], List[str]]:
109    """Get paths to the RESECT data.
110
111    Args:
112        path: Filepath to a folder where the data is downloaded for further processing.
113        source: The imaging source. Either 'US' (intra-operative ultrasound) or 'MRI' (pre-operative FLAIR).
114        phase: The surgical phase of the ultrasound volumes. Either 'before', 'during' or 'after' resection.
115            Ignored for the 'MRI' source.
116        structure: The annotated structure. One of 'tumor' (US before resection and MRI), 'resection'
117            (resection cavity, US during and after resection), 'sulci' or 'falx' (US, all phases).
118            By default, 'tumor' for 'before' and 'MRI' and 'resection' for 'during' and 'after'.
119        download: Whether to download the data if it is not present.
120
121    Returns:
122        List of filepaths for the image data.
123        List of filepaths for the label data.
124    """
125    structure = _get_structure(source, phase, structure)
126    image_dir, label_dir = get_resect_data(path, source, phase, download)
127
128    prefix = f"US-{phase}" if source == "US" else "FLAIR"
129    label_paths = natsorted(glob(os.path.join(label_dir, "Case*", f"Case*-{prefix}-{structure}.nii.gz")))
130    raw_paths = [
131        os.path.join(image_dir, os.path.basename(p).replace(f"-{structure}.nii.gz", ".nii.gz")) for p in label_paths
132    ]
133    assert all(os.path.exists(p) for p in raw_paths), "Some image volumes are missing."
134    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
135
136    return raw_paths, label_paths
137
138
139def get_resect_dataset(
140    path: Union[os.PathLike, str],
141    patch_shape: Tuple[int, ...],
142    source: Literal["US", "MRI"] = "US",
143    phase: Literal["before", "during", "after"] = "before",
144    structure: Optional[Literal["tumor", "resection", "sulci", "falx"]] = None,
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs
148) -> Dataset:
149    """Get the RESECT dataset for brain tumor, resection cavity, sulci and falx segmentation.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        patch_shape: The patch shape to use for training.
154        source: The imaging source. Either 'US' (intra-operative ultrasound) or 'MRI' (pre-operative FLAIR).
155        phase: The surgical phase of the ultrasound volumes. Either 'before', 'during' or 'after' resection.
156            Ignored for the 'MRI' source.
157        structure: The annotated structure. One of 'tumor' (US before resection and MRI), 'resection'
158            (resection cavity, US during and after resection), 'sulci' or 'falx' (US, all phases).
159            By default, 'tumor' for 'before' and 'MRI' and 'resection' for 'during' and 'after'.
160        resize_inputs: Whether to resize inputs to the desired patch shape.
161        download: Whether to download the data if it is not present.
162        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
163
164    Returns:
165        The segmentation dataset.
166    """
167    raw_paths, label_paths = get_resect_paths(path, source, phase, structure, download)
168
169    if resize_inputs:
170        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
171        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
172            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
173        )
174
175    return torch_em.default_segmentation_dataset(
176        raw_paths=raw_paths,
177        raw_key="data",
178        label_paths=label_paths,
179        label_key="data",
180        patch_shape=patch_shape,
181        is_seg_dataset=True,
182        **kwargs
183    )
184
185
186def get_resect_loader(
187    path: Union[os.PathLike, str],
188    batch_size: int,
189    patch_shape: Tuple[int, ...],
190    source: Literal["US", "MRI"] = "US",
191    phase: Literal["before", "during", "after"] = "before",
192    structure: Optional[Literal["tumor", "resection", "sulci", "falx"]] = None,
193    resize_inputs: bool = False,
194    download: bool = False,
195    **kwargs
196) -> DataLoader:
197    """Get the RESECT dataloader for brain tumor, resection cavity, sulci and falx segmentation.
198
199    Args:
200        path: Filepath to a folder where the data is downloaded for further processing.
201        batch_size: The batch size for training.
202        patch_shape: The patch shape to use for training.
203        source: The imaging source. Either 'US' (intra-operative ultrasound) or 'MRI' (pre-operative FLAIR).
204        phase: The surgical phase of the ultrasound volumes. Either 'before', 'during' or 'after' resection.
205            Ignored for the 'MRI' source.
206        structure: The annotated structure. One of 'tumor' (US before resection and MRI), 'resection'
207            (resection cavity, US during and after resection), 'sulci' or 'falx' (US, all phases).
208            By default, 'tumor' for 'before' and 'MRI' and 'resection' for 'during' and 'after'.
209        resize_inputs: Whether to resize inputs to the desired patch shape.
210        download: Whether to download the data if it is not present.
211        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
212
213    Returns:
214        The DataLoader.
215    """
216    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
217    dataset = get_resect_dataset(path, patch_shape, source, phase, structure, resize_inputs, download, **ds_kwargs)
218    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)