torch_em.data.datasets.medical.ctspine1k

  1"""The CTSpine1K dataset contains annotations for 25 vertebrae in CT scans.
  2
  3The dataset consists of 1005 CT scans that were collected from four public collections and annotated with
  4the individual vertebrae for this release. The 'source' argument selects one of them: 'colonog' (CT
  5colonography), 'covid19', 'hnscc' (head and neck) and 'msd_liver' (the liver task of the Medical
  6Segmentation Decathlon). The label ids are 1 to 7 for the cervical vertebrae C1 to C7, 8 to 19 for the
  7thoracic vertebrae T1 to T12 and 20 to 25 for the lumbar vertebrae L1 to L6. See also `CLASS_IDS`.
  8
  9NOTE: A scan only covers a part of the spine, so each one contains a contiguous range of the label ids
 10rather than all 25.
 11
 12NOTE: The scans of the CT colonography collection are named after the id of their DICOM series, whose
 13dots hide the nifti extension from the file readers, so the data is linked under a prepared name.
 14
 15The dataset is located at https://huggingface.co/datasets/alexanderdann/CTSpine1K and is distributed
 16under the CC BY-NC-SA license.
 17This dataset is from the publication https://doi.org/10.48550/arXiv.2105.14711.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22from glob import glob
 23from natsort import natsorted
 24from typing import Union, Optional, Tuple, Literal, List
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33REPO_ID = "alexanderdann/CTSpine1K"
 34
 35SOURCES = {
 36    "colonog": "COLONOG",
 37    "covid19": "COVID-19",
 38    "hnscc": "HNSCC-3DCT-RT",
 39    "msd_liver": "MSD-T10",
 40}
 41"""Mapping from the source choice to its folder in the release."""
 42
 43CLASS_NAMES = (
 44    [f"c{i}" for i in range(1, 8)] + [f"t{i}" for i in range(1, 13)] + [f"l{i}" for i in range(1, 7)]
 45)
 46"""The vertebrae of the CTSpine1K dataset. The label id of a vertebra is its 1-based index."""
 47
 48CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 49"""Mapping from the vertebra name to its label id."""
 50
 51
 52def _prepare_name(path, prepared_dir, folder, stem):
 53    """Link a scan under a name without a dot in its stem, which the file readers need to find its extension."""
 54    out_path = os.path.join(prepared_dir, folder, f"{stem.replace('.', '_')}.nii.gz")
 55    if not os.path.exists(out_path):
 56        os.symlink(os.path.abspath(path), out_path)
 57    return out_path
 58
 59
 60def get_ctspine1k_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 61    """Download the CTSpine1K dataset.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        Filepath where the data is downloaded.
 69    """
 70    data_dir = os.path.join(path, "raw_data")
 71    if os.path.exists(data_dir) and glob(os.path.join(data_dir, "volumes", "*", "*.nii.gz")):
 72        return data_dir
 73
 74    if not download:
 75        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
 76
 77    from huggingface_hub import snapshot_download
 78
 79    os.makedirs(path, exist_ok=True)
 80    snapshot_download(repo_id=REPO_ID, repo_type="dataset", local_dir=path)
 81    return data_dir
 82
 83
 84def get_ctspine1k_paths(
 85    path: Union[os.PathLike, str],
 86    source: Optional[Literal["colonog", "covid19", "hnscc", "msd_liver"]] = None,
 87    download: bool = False,
 88) -> Tuple[List[str], List[str]]:
 89    """Get paths to the CTSpine1K data.
 90
 91    Args:
 92        path: Filepath to a folder where the data is downloaded for further processing.
 93        source: The choice of source collection. All of them are used if it is not given.
 94        download: Whether to download the data if it is not present.
 95
 96    Returns:
 97        List of filepaths for the image data.
 98        List of filepaths for the label data.
 99    """
100    if source is not None and source not in SOURCES:
101        raise ValueError(f"'{source}' is not a valid source. Choose from {list(SOURCES.keys())}.")
102
103    data_dir = get_ctspine1k_data(path, download)
104
105    folders = list(SOURCES.values()) if source is None else [SOURCES[source]]
106    prepared_dir = os.path.join(path, "prepared")
107    raw_paths, label_paths = [], []
108    for folder in folders:
109        os.makedirs(os.path.join(prepared_dir, folder), exist_ok=True)
110        for image_path in natsorted(glob(os.path.join(data_dir, "volumes", folder, "*.nii.gz"))):
111            stem = os.path.basename(image_path)[:-len(".nii.gz")]
112            label_path = os.path.join(data_dir, "labels", folder, f"{stem}_seg.nii.gz")
113            if not os.path.exists(label_path):
114                continue
115            raw_paths.append(_prepare_name(image_path, prepared_dir, folder, stem))
116            label_paths.append(_prepare_name(label_path, prepared_dir, folder, f"{stem}_seg"))
117
118    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
119
120    return raw_paths, label_paths
121
122
123def get_ctspine1k_dataset(
124    path: Union[os.PathLike, str],
125    patch_shape: Tuple[int, ...],
126    source: Optional[Literal["colonog", "covid19", "hnscc", "msd_liver"]] = None,
127    resize_inputs: bool = False,
128    download: bool = False,
129    **kwargs
130) -> Dataset:
131    """Get the CTSpine1K dataset for vertebra segmentation.
132
133    Args:
134        path: Filepath to a folder where the data is downloaded for further processing.
135        patch_shape: The patch shape to use for training.
136        source: The choice of source collection. All of them are used if it is not given.
137        resize_inputs: Whether to resize inputs to the desired patch shape.
138        download: Whether to download the data if it is not present.
139        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
140
141    Returns:
142        The segmentation dataset.
143    """
144    raw_paths, label_paths = get_ctspine1k_paths(path, source, download)
145
146    if resize_inputs:
147        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
148        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
149            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
150        )
151
152    return torch_em.default_segmentation_dataset(
153        raw_paths=raw_paths,
154        raw_key="data",
155        label_paths=label_paths,
156        label_key="data",
157        patch_shape=patch_shape,
158        is_seg_dataset=True,
159        **kwargs
160    )
161
162
163def get_ctspine1k_loader(
164    path: Union[os.PathLike, str],
165    batch_size: int,
166    patch_shape: Tuple[int, ...],
167    source: Optional[Literal["colonog", "covid19", "hnscc", "msd_liver"]] = None,
168    resize_inputs: bool = False,
169    download: bool = False,
170    **kwargs
171) -> DataLoader:
172    """Get the CTSpine1K dataloader for vertebra segmentation.
173
174    Args:
175        path: Filepath to a folder where the data is downloaded for further processing.
176        batch_size: The batch size for training.
177        patch_shape: The patch shape to use for training.
178        source: The choice of source collection. All of them are used if it is not given.
179        resize_inputs: Whether to resize inputs to the desired patch shape.
180        download: Whether to download the data if it is not present.
181        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
182
183    Returns:
184        The DataLoader.
185    """
186    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
187    dataset = get_ctspine1k_dataset(path, patch_shape, source, resize_inputs, download, **ds_kwargs)
188    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)