torch_em.data.datasets.medical.episurg

  1"""The EPISURG dataset contains annotations for resection cavity segmentation in postoperative
  2brain MRI of epilepsy patients.
  3
  4The dataset consists of 430 postoperative T1-weighted MRI from patients who underwent resective
  5brain surgery for refractory epilepsy at the National Hospital of Neurology and Neurosurgery
  6(Queen Square, London, United Kingdom). The corresponding preoperative MRI is present for 269 of
  7these subjects. The resection cavity was manually segmented by three human raters on overlapping
  8subsets of the postoperative scans (133, 34 and 33 subjects, respectively, the second and third
  9rater re-annotating scans of the first), so 133 of the 430 subjects have a resection cavity mask.
 10
 11The dataset is located at https://doi.org/10.5522/04/9996158.v1 and is distributed under the
 12CC BY-NC-SA 4.0 license.
 13The dataset is from the publication https://doi.org/10.1007/s11548-021-02420-2.
 14Please cite it if you use this dataset for your research.
 15"""
 16
 17import os
 18from glob import glob
 19from natsort import natsorted
 20from typing import Union, Tuple, List
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URL = "https://s3-eu-west-1.amazonaws.com/pstorage-ucl-2748466690/26153588/EPISURG.zip"
 30CHECKSUM = "91c6e0698ab5a1874662e3a53ccfb509e41718ddb6ac62f0b35f06a85d1c8daf"
 31
 32
 33def get_episurg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the EPISURG dataset.
 35
 36    Args:
 37        path: Filepath to a folder where the data is downloaded for further processing.
 38        download: Whether to download the data if it is not present.
 39
 40    Returns:
 41        Filepath where the data is downloaded.
 42    """
 43    data_dir = os.path.join(path, "EPISURG")
 44    if os.path.exists(data_dir):
 45        return data_dir
 46
 47    os.makedirs(path, exist_ok=True)
 48
 49    zip_path = os.path.join(path, "EPISURG.zip")
 50    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 51    util.unzip(zip_path=zip_path, dst=path)
 52
 53    return data_dir
 54
 55
 56def get_episurg_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 57    """Get paths to the EPISURG data.
 58
 59    Only the postoperative subjects with a resection cavity mask from one of the three human
 60    raters are returned.
 61
 62    Args:
 63        path: Filepath to a folder where the data is downloaded for further processing.
 64        download: Whether to download the data if it is not present.
 65
 66    Returns:
 67        List of filepaths for the postoperative MRI data.
 68        List of filepaths for the resection cavity segmentation data.
 69    """
 70    data_dir = get_episurg_data(path, download)
 71
 72    image_paths, label_paths = [], []
 73    for label_path in natsorted(glob(os.path.join(data_dir, "subjects", "*", "postop", "*postop-seg-*.nii.gz"))):
 74        subject_dir = os.path.dirname(label_path)
 75        image_path = natsorted(glob(os.path.join(subject_dir, "*postop-t1mri-*.nii.gz")))
 76        assert len(image_path) == 1, f"Could not find a unique postop MRI for '{label_path}'."
 77        image_paths.append(image_path[0])
 78        label_paths.append(label_path)
 79
 80    return image_paths, label_paths
 81
 82
 83def get_episurg_dataset(
 84    path: Union[os.PathLike, str],
 85    patch_shape: Tuple[int, ...],
 86    resize_inputs: bool = False,
 87    download: bool = False,
 88    **kwargs
 89) -> Dataset:
 90    """Get the EPISURG dataset for resection cavity segmentation in postoperative brain MRI.
 91
 92    Args:
 93        path: Filepath to a folder where the data is downloaded for further processing.
 94        patch_shape: The patch shape to use for training.
 95        resize_inputs: Whether to resize the inputs to the expected patch shape.
 96        download: Whether to download the data if it is not present.
 97        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 98
 99    Returns:
100        The segmentation dataset.
101    """
102    image_paths, label_paths = get_episurg_paths(path, download)
103
104    if resize_inputs:
105        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
106        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
107            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
108        )
109
110    return torch_em.default_segmentation_dataset(
111        raw_paths=image_paths,
112        raw_key="data",
113        label_paths=label_paths,
114        label_key="data",
115        patch_shape=patch_shape,
116        is_seg_dataset=True,
117        **kwargs
118    )
119
120
121def get_episurg_loader(
122    path: Union[os.PathLike, str],
123    batch_size: int,
124    patch_shape: Tuple[int, ...],
125    resize_inputs: bool = False,
126    download: bool = False,
127    **kwargs
128) -> DataLoader:
129    """Get the EPISURG dataloader for resection cavity segmentation in postoperative brain MRI.
130
131    Args:
132        path: Filepath to a folder where the data is downloaded for further processing.
133        batch_size: The batch size for training.
134        patch_shape: The patch shape to use for training.
135        resize_inputs: Whether to resize the inputs to the expected patch shape.
136        download: Whether to download the data if it is not present.
137        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
138
139    Returns:
140        The DataLoader.
141    """
142    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
143    dataset = get_episurg_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
144    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)