torch_em.data.datasets.medical.crossmoda

  1"""The crossMoDA dataset contains annotations for vestibular schwannoma and cochlea segmentation
  2in contrast-enhanced T1-weighted (ceT1) MRI.
  3
  4The data comes from the crossMoDA 2022 challenge (https://crossmoda2022.grand-challenge.org/) for unsupervised
  5cross-modality domain adaptation. The training set contains 210 annotated ceT1 scans (the source domain,
  6105 from London and 105 from Tilburg) and 210 unpaired, unannotated high-resolution T2 scans (the target domain).
  7Only the annotated source domain is used for the segmentation dataset and loader, the target domain images
  8can be accessed via `get_crossmoda_paths` with `domain='target'`.
  9
 10The label ids are: background: 0, vestibular schwannoma (tumor): 1 and cochlea: 2.
 11
 12The dataset is located at https://zenodo.org/records/6504722.
 13
 14This dataset is from the publications https://doi.org/10.7937/TCIA.9YTJ-5Q73 and
 15https://doi.org/10.1016/j.media.2022.102628.
 16Please cite them if you use this dataset in your research.
 17"""
 18
 19import os
 20from glob import glob
 21from natsort import natsorted
 22from typing import Union, Tuple, Literal, List, Optional
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URL = "https://zenodo.org/records/6504722/files/crossmoda2022_training.zip?download=1"
 32CHECKSUM = "d3db17e04fd7b4c7bfc8cd569f63e01953cc3a218d5bc754f916729483ef0cdb"
 33
 34# The center is encoded in the filenames: 'ldn' for London and 'etz' for Tilburg.
 35CENTERS = {"london": "ldn", "tilburg": "etz"}
 36
 37
 38def get_crossmoda_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 39    """Download the crossMoDA 2022 training data.
 40
 41    Args:
 42        path: Filepath to a folder where the data is downloaded for further processing.
 43        download: Whether to download the data if it is not present.
 44
 45    Returns:
 46        Filepath to the folder with the downloaded data.
 47    """
 48    source_dir = os.path.join(path, "training_source")
 49    target_dir = os.path.join(path, "training_target")
 50    if os.path.exists(source_dir) and os.path.exists(target_dir):
 51        return path
 52
 53    os.makedirs(path, exist_ok=True)
 54
 55    zip_path = os.path.join(path, "crossmoda2022_training.zip")
 56    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 57    util.unzip(zip_path=zip_path, dst=path)
 58
 59    return path
 60
 61
 62def get_crossmoda_paths(
 63    path: Union[os.PathLike, str],
 64    center: Optional[Literal["london", "tilburg"]] = None,
 65    domain: Literal["source", "target"] = "source",
 66    download: bool = False,
 67) -> Tuple[List[str], Optional[List[str]]]:
 68    """Get paths to the crossMoDA data.
 69
 70    Args:
 71        path: Filepath to a folder where the data is downloaded for further processing.
 72        center: The center the scans come from. Either 'london' or 'tilburg'. By default, both centers are used.
 73        domain: The domain of the scans. Either 'source' (annotated ceT1) or 'target' (unannotated hrT2).
 74        download: Whether to download the data if it is not present.
 75
 76    Returns:
 77        List of filepaths for the image data.
 78        List of filepaths for the label data. None for the 'target' domain, which has no annotations.
 79    """
 80    if center is not None and center not in CENTERS:
 81        raise ValueError(f"'{center}' is not a valid center.")
 82    if domain not in ("source", "target"):
 83        raise ValueError(f"'{domain}' is not a valid domain.")
 84
 85    data_dir = get_crossmoda_data(path, download)
 86
 87    center_pattern = "*" if center is None else CENTERS[center]
 88    sequence = "ceT1" if domain == "source" else "hrT2"
 89    pattern = os.path.join(data_dir, f"training_{domain}", f"crossmoda*_{center_pattern}_*_{sequence}.nii.gz")
 90    raw_paths = natsorted(glob(pattern))
 91    assert len(raw_paths) > 0, f"No crossMoDA volumes found at '{pattern}'."
 92
 93    if domain == "target":
 94        return raw_paths, None
 95
 96    label_paths = [p.replace(f"_{sequence}.nii.gz", "_Label.nii.gz") for p in raw_paths]
 97    assert all(os.path.exists(p) for p in label_paths)
 98
 99    return raw_paths, label_paths
100
101
102def get_crossmoda_dataset(
103    path: Union[os.PathLike, str],
104    patch_shape: Tuple[int, ...],
105    center: Optional[Literal["london", "tilburg"]] = None,
106    resize_inputs: bool = False,
107    download: bool = False,
108    **kwargs
109) -> Dataset:
110    """Get the crossMoDA dataset for vestibular schwannoma and cochlea segmentation in ceT1 MRI.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded for further processing.
114        patch_shape: The patch shape to use for training.
115        center: The center the scans come from. Either 'london' or 'tilburg'. By default, both centers are used.
116        resize_inputs: Whether to resize inputs to the desired patch shape.
117        download: Whether to download the data if it is not present.
118        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
119
120    Returns:
121        The segmentation dataset.
122    """
123    raw_paths, label_paths = get_crossmoda_paths(path, center, "source", download)
124
125    if resize_inputs:
126        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
127        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
128            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
129        )
130
131    return torch_em.default_segmentation_dataset(
132        raw_paths=raw_paths,
133        raw_key="data",
134        label_paths=label_paths,
135        label_key="data",
136        patch_shape=patch_shape,
137        is_seg_dataset=True,
138        **kwargs
139    )
140
141
142def get_crossmoda_loader(
143    path: Union[os.PathLike, str],
144    batch_size: int,
145    patch_shape: Tuple[int, ...],
146    center: Optional[Literal["london", "tilburg"]] = None,
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> DataLoader:
151    """Get the crossMoDA dataloader for vestibular schwannoma and cochlea segmentation in ceT1 MRI.
152
153    Args:
154        path: Filepath to a folder where the data is downloaded for further processing.
155        batch_size: The batch size for training.
156        patch_shape: The patch shape to use for training.
157        center: The center the scans come from. Either 'london' or 'tilburg'. By default, both centers are used.
158        resize_inputs: Whether to resize inputs to the desired patch shape.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_crossmoda_dataset(path, patch_shape, center, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)