torch_em.data.datasets.medical.c3ro

  1"""The C3RO dataset contains crowdsourced radiotherapy contours on CT scans of five cancer sites: breast,
  2gastrointestinal (GI), gynecologic (GYN), head and neck (H&N) and sarcoma.
  3
  4NOTE: This is NOT a large training set. The dataset has only five CT scans, one per cancer site. Its value is the
  5number of annotators: every scan was contoured by many independent raters, split into experts (4 to 15 per site)
  6and non-experts (26 to 124 per site), for 3 to 11 target structures per site (e.g. heart, parotid glands, GTV, CTV).
  7For every structure and annotator group a consensus mask, computed with STAPLE, is provided as well. This makes
  8the dataset useful for research on inter-rater variability and label uncertainty.
  9
 10Each item of this loader pairs the CT scan of a site with the mask of a single structure. By default, this is the
 11STAPLE consensus of the expert raters. Use `annotator` to select the expert or non-expert group, and `rater` to
 12select the masks of a single rater of that group instead of the consensus. Individual raters additionally contoured
 13some auxiliary structures that have no consensus mask (e.g. 'body' or 'ptv'), which are then included as well.
 14Masks that do not match the shape of the CT are skipped (none of the 2477 NIfTI masks is affected).
 15
 16The data is located at https://doi.org/10.6084/m9.figshare.21074182, released under a CC-BY-4.0 license.
 17
 18Please cite the associated publication and the dataset if you use it for 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
 33URL = "https://ndownloader.figshare.com/files/42025569"
 34CHECKSUM = "74c94faa18ef2e78c512b1ec2a71755f9b93cba37d0aff19afeeb47a77f18887"
 35
 36SITES = ["Breast", "GI", "GYN", "H&N", "Sarcoma"]
 37ANNOTATORS = {"expert": "Expert", "non_expert": "Non-Expert"}
 38
 39
 40def get_c3ro_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 41    """Download the C3RO dataset.
 42
 43    Args:
 44        path: Filepath to a folder where the data is downloaded for further processing.
 45        download: Whether to download the data if it is not present.
 46
 47    Returns:
 48        Filepath where the data is downloaded.
 49    """
 50    data_dir = os.path.join(path, "Organized_files_v4")
 51    if os.path.exists(data_dir):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55
 56    zip_path = os.path.join(path, "Organized_files_v4.zip")
 57    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 58    util.unzip(zip_path=zip_path, dst=path, remove=False)
 59
 60    assert os.path.exists(data_dir), f"The extraction of the C3RO archive did not create '{data_dir}'."
 61
 62    return data_dir
 63
 64
 65def get_c3ro_paths(
 66    path: Union[os.PathLike, str],
 67    site: Optional[Literal["Breast", "GI", "GYN", "H&N", "Sarcoma"]] = None,
 68    annotator: Literal["expert", "non_expert"] = "expert",
 69    rater: Optional[str] = None,
 70    download: bool = False,
 71) -> Tuple[List[str], List[str]]:
 72    """Get paths to the C3RO data.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        site: The choice of cancer site. By default, all five sites are used.
 77        annotator: The choice of annotator group. Either 'expert' or 'non_expert'.
 78        rater: The id of a single rater of the annotator group. By default, the STAPLE consensus masks are used.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        List of filepaths for the image data, which contains the CT scan of a site once per structure.
 83        List of filepaths for the label data.
 84    """
 85    import nibabel as nib
 86
 87    if site is not None and site not in SITES:
 88        raise ValueError(f"'{site}' is not a valid site. Choose one of {SITES}.")
 89    if annotator not in ANNOTATORS:
 90        raise ValueError(f"'{annotator}' is not a valid annotator group. Choose one of {list(ANNOTATORS)}.")
 91
 92    data_dir = get_c3ro_data(path, download)
 93
 94    raw_paths, label_paths = [], []
 95    for site_name in ([site] if site is not None else SITES):
 96        ct_path = os.path.join(data_dir, site_name, "CT", "NIFTI", f"Image_CT_{site_name}.nii.gz")
 97        seg_dir = os.path.join(data_dir, site_name, "Segmentations", ANNOTATORS[annotator])
 98        if rater is None:
 99            mask_paths = natsorted(glob(os.path.join(seg_dir, "Consensus", "*.nii.gz")))
100        else:
101            mask_paths = natsorted(glob(os.path.join(seg_dir, str(rater), "NIFTI", "*.nii.gz")))
102
103        ct_shape = nib.load(ct_path).shape
104        for mask_path in mask_paths:
105            if nib.load(mask_path).shape == ct_shape:
106                raw_paths.append(ct_path)
107                label_paths.append(mask_path)
108
109    if len(raw_paths) == 0:
110        raise ValueError(f"No masks were found for site '{site}', annotator '{annotator}' and rater '{rater}'.")
111
112    return raw_paths, label_paths
113
114
115def get_c3ro_dataset(
116    path: Union[os.PathLike, str],
117    patch_shape: Tuple[int, int, int],
118    site: Optional[Literal["Breast", "GI", "GYN", "H&N", "Sarcoma"]] = None,
119    annotator: Literal["expert", "non_expert"] = "expert",
120    rater: Optional[str] = None,
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> Dataset:
125    """Get the C3RO dataset for organ and target volume segmentation in radiotherapy planning CT.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        patch_shape: The patch shape to use for training.
130        site: The choice of cancer site. By default, all five sites are used.
131        annotator: The choice of annotator group. Either 'expert' or 'non_expert'.
132        rater: The id of a single rater of the annotator group. By default, the STAPLE consensus masks are used.
133        resize_inputs: Whether to resize the inputs to the patch shape.
134        download: Whether to download the data if it is not present.
135        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
136
137    Returns:
138        The segmentation dataset.
139    """
140    raw_paths, label_paths = get_c3ro_paths(path, site, annotator, rater, download)
141
142    if resize_inputs:
143        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
144        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
145            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
146        )
147
148    return torch_em.default_segmentation_dataset(
149        raw_paths=raw_paths,
150        raw_key="data",
151        label_paths=label_paths,
152        label_key="data",
153        is_seg_dataset=True,
154        patch_shape=patch_shape,
155        ndim=3,
156        **kwargs
157    )
158
159
160def get_c3ro_loader(
161    path: Union[os.PathLike, str],
162    batch_size: int,
163    patch_shape: Tuple[int, int, int],
164    site: Optional[Literal["Breast", "GI", "GYN", "H&N", "Sarcoma"]] = None,
165    annotator: Literal["expert", "non_expert"] = "expert",
166    rater: Optional[str] = None,
167    resize_inputs: bool = False,
168    download: bool = False,
169    **kwargs
170) -> DataLoader:
171    """Get the C3RO dataloader for organ and target volume segmentation in radiotherapy planning CT.
172
173    Args:
174        path: Filepath to a folder where the data is downloaded for further processing.
175        batch_size: The batch size for training.
176        patch_shape: The patch shape to use for training.
177        site: The choice of cancer site. By default, all five sites are used.
178        annotator: The choice of annotator group. Either 'expert' or 'non_expert'.
179        rater: The id of a single rater of the annotator group. By default, the STAPLE consensus masks are used.
180        resize_inputs: Whether to resize the inputs to the patch shape.
181        download: Whether to download the data if it is not present.
182        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
183
184    Returns:
185        The DataLoader.
186    """
187    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
188    dataset = get_c3ro_dataset(path, patch_shape, site, annotator, rater, resize_inputs, download, **ds_kwargs)
189    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)