torch_em.data.datasets.medical.ctpelvic1k

  1"""The CTPelvic1K dataset contains annotations for pelvic bones in CT scans.
  2
  3This module provides the two subsets of the collection that are distributed with their own scans:
  4'clinic' with 103 annotated scans and 'clinic_metal' with 14 annotated scans of patients with metal
  5implants. The label ids are 1: sacrum, 2: hip (left), 3: hip (right), 4: lumbar vertebra.
  6See also `CLASS_IDS`.
  7
  8NOTE: 'clinic_metal' ships 75 scans but only 14 of them are annotated, so only those are returned. Its
  9scans carry a 'dataset7_' prefix that its annotations do not, unlike the consistently named 'clinic'.
 10
 11NOTE: The label ids of the two hips are assigned by the position of the annotation in the scan, which is
 12the reverse of the order that is usually quoted for this dataset. Every scan of the release is stored in
 13LPS orientation, and the annotation with id 2 lies on the left of the patient in all of them.
 14
 15NOTE: The remaining subsets of the collection hold annotations for scans of other public datasets, such
 16as the colon task of the Medical Segmentation Decathlon and KiTS, and are not provided here.
 17
 18The dataset is located at https://doi.org/10.5281/zenodo.4588403.
 19This dataset is from the publication https://doi.org/10.1007/s11548-021-02363-8.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24import re
 25from glob import glob
 26from natsort import natsorted
 27from typing import Union, Tuple, Literal, List
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36URL = "https://zenodo.org/records/4588403/files/{filename}?download=1"
 37
 38SUBSETS = {
 39    "clinic": {"images": "CTPelvic1K_dataset6_data.tar.gz", "labels": "CTPelvic1K_dataset6_Anonymized_mask.tar.gz"},
 40    "clinic_metal": {"images": "CTPelvic1K_dataset7_data.tar.gz", "labels": "CTPelvic1K_dataset7_mask.tar.gz"},
 41}
 42
 43CHECKSUMS = {
 44    "clinic": {
 45        "images": "9b8f4747d256483062aa937eb990890e59954d8be74966edd982a6dcf3f4394e",
 46        "labels": "382b0780bcbf68536631b14e9d2a42a0bccb6fc3c30c9598078ddfa0e574bf95",
 47    },
 48    "clinic_metal": {
 49        "images": "9b71fdf37c9bbb7bf99b95ea577f4bac24f9c393644e44287f59aad0b52a9cbc",
 50        "labels": "4b85e4eb7300e77ce094329933d032a88183f4f0ec65cde95e4162074bb7cca3",
 51    },
 52}
 53
 54CLASS_NAMES = ["sacrum", "hip_left", "hip_right", "lumbar_vertebra"]
 55"""The pelvic bones of the CTPelvic1K dataset. The label id of a bone is its 1-based index."""
 56
 57CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 58"""Mapping from the bone name to its label id."""
 59
 60LABEL_SUFFIX = "_mask_4label.nii.gz"
 61
 62IMAGE_SUFFIX = "_data.nii.gz"
 63
 64
 65def _case_id(filename, suffix):
 66    """Strip the suffix and the subset prefix that the scans carry but their annotations do not."""
 67    return re.sub(r"^dataset\d+_", "", filename[:-len(suffix)])
 68
 69
 70def get_ctpelvic1k_data(
 71    path: Union[os.PathLike, str], subset: Literal["clinic", "clinic_metal"] = "clinic", download: bool = False
 72) -> str:
 73    """Download the CTPelvic1K dataset.
 74
 75    Args:
 76        path: Filepath to a folder where the data is downloaded for further processing.
 77        subset: The choice of subset. Either 'clinic' with 103 scans or 'clinic_metal' with 14 scans.
 78        download: Whether to download the data if it is not present.
 79
 80    Returns:
 81        Filepath where the data is downloaded.
 82    """
 83    if subset not in SUBSETS:
 84        raise ValueError(f"'{subset}' is not a valid subset. Choose from {list(SUBSETS.keys())}.")
 85
 86    data_dir = os.path.join(path, subset)
 87    if os.path.exists(data_dir) and glob(os.path.join(data_dir, "**", f"*{LABEL_SUFFIX}"), recursive=True):
 88        return data_dir
 89
 90    os.makedirs(data_dir, exist_ok=True)
 91    for key in ("images", "labels"):
 92        filename = SUBSETS[subset][key]
 93        tar_path = os.path.join(path, filename)
 94        util.download_source(
 95            path=tar_path, url=URL.format(filename=filename), download=download, checksum=CHECKSUMS[subset][key]
 96        )
 97        util.unzip_tarfile(tar_path=tar_path, dst=data_dir, remove=False)
 98
 99    return data_dir
100
101
102def get_ctpelvic1k_paths(
103    path: Union[os.PathLike, str],
104    subset: Literal["clinic", "clinic_metal"] = "clinic",
105    download: bool = False,
106) -> Tuple[List[str], List[str]]:
107    """Get paths to the CTPelvic1K data.
108
109    Args:
110        path: Filepath to a folder where the data is downloaded for further processing.
111        subset: The choice of subset. Either 'clinic' with 103 scans or 'clinic_metal' with 14 scans.
112        download: Whether to download the data if it is not present.
113
114    Returns:
115        List of filepaths for the image data.
116        List of filepaths for the label data.
117    """
118    data_dir = get_ctpelvic1k_data(path, subset, download)
119
120    label_paths = natsorted(glob(os.path.join(data_dir, "**", f"*{LABEL_SUFFIX}"), recursive=True))
121    image_paths = {
122        _case_id(os.path.basename(p), IMAGE_SUFFIX): p
123        for p in glob(os.path.join(data_dir, "**", f"*{IMAGE_SUFFIX}"), recursive=True)
124    }
125
126    # Only a part of the scans of 'clinic_metal' is annotated, so the annotations drive the pairing.
127    raw_paths, valid_label_paths = [], []
128    for label_path in label_paths:
129        image_path = image_paths.get(_case_id(os.path.basename(label_path), LABEL_SUFFIX))
130        if image_path is not None:
131            raw_paths.append(image_path)
132            valid_label_paths.append(label_path)
133
134    assert len(raw_paths) == len(valid_label_paths) and len(raw_paths) > 0
135
136    return raw_paths, valid_label_paths
137
138
139def get_ctpelvic1k_dataset(
140    path: Union[os.PathLike, str],
141    patch_shape: Tuple[int, ...],
142    subset: Literal["clinic", "clinic_metal"] = "clinic",
143    resize_inputs: bool = False,
144    download: bool = False,
145    **kwargs
146) -> Dataset:
147    """Get the CTPelvic1K dataset for pelvic bone segmentation.
148
149    Args:
150        path: Filepath to a folder where the data is downloaded for further processing.
151        patch_shape: The patch shape to use for training.
152        subset: The choice of subset. Either 'clinic' with 103 scans or 'clinic_metal' with 14 scans.
153        resize_inputs: Whether to resize inputs to the desired patch shape.
154        download: Whether to download the data if it is not present.
155        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
156
157    Returns:
158        The segmentation dataset.
159    """
160    raw_paths, label_paths = get_ctpelvic1k_paths(path, subset, download)
161
162    if resize_inputs:
163        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
164        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
165            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
166        )
167
168    return torch_em.default_segmentation_dataset(
169        raw_paths=raw_paths,
170        raw_key="data",
171        label_paths=label_paths,
172        label_key="data",
173        patch_shape=patch_shape,
174        is_seg_dataset=True,
175        **kwargs
176    )
177
178
179def get_ctpelvic1k_loader(
180    path: Union[os.PathLike, str],
181    batch_size: int,
182    patch_shape: Tuple[int, ...],
183    subset: Literal["clinic", "clinic_metal"] = "clinic",
184    resize_inputs: bool = False,
185    download: bool = False,
186    **kwargs
187) -> DataLoader:
188    """Get the CTPelvic1K dataloader for pelvic bone segmentation.
189
190    Args:
191        path: Filepath to a folder where the data is downloaded for further processing.
192        batch_size: The batch size for training.
193        patch_shape: The patch shape to use for training.
194        subset: The choice of subset. Either 'clinic' with 103 scans or 'clinic_metal' with 14 scans.
195        resize_inputs: Whether to resize inputs to the desired patch shape.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
198
199    Returns:
200        The DataLoader.
201    """
202    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
203    dataset = get_ctpelvic1k_dataset(path, patch_shape, subset, resize_inputs, download, **ds_kwargs)
204    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)