torch_em.data.datasets.medical.ct_org

  1"""The CT-ORG dataset contains annotations for multiple organs in CT.
  2
  3It consists of 140 CT volumes (abdominal and full-body, with and without contrast, including 9 PET-CT scans)
  4with semantic labels for: 1: liver, 2: bladder, 3: lungs, 4: kidneys, 5: bone, 6: brain (only in the minority
  5of scans which show the brain). Cases 0-20 form the official test split (all organs annotated manually), cases 21-139
  6form the training split (lungs and bones were segmented with morphological algorithms).
  7
  8The dataset is located at https://www.cancerimagingarchive.net/collection/ct-org/ and is only offered via
  9IBM Aspera there. We download the data from a mirror of the original nifti files at
 10https://huggingface.co/datasets/MedOtter/ct-org instead.
 11
 12This dataset is from the publication https://doi.org/10.1038/s41597-020-00715-8.
 13The data was released at https://doi.org/10.7937/tcia.2019.tt7f4v7o.
 14Please cite it if you use this dataset in your research.
 15"""
 16
 17import os
 18from tqdm import tqdm
 19from typing import Union, Tuple, Literal, List, Optional
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://huggingface.co/datasets/MedOtter/ct-org/resolve/main/"
 29
 30# The data is downloaded as 280 individual files, which are not checksummed.
 31CHECKSUM = None
 32
 33NUM_VOLUMES = 140
 34TEST_IDS = list(range(21))
 35TRAIN_IDS = list(range(21, NUM_VOLUMES))
 36
 37LABEL_IDS = {"liver": 1, "bladder": 2, "lungs": 3, "kidneys": 4, "bone": 5, "brain": 6}
 38
 39
 40def get_ct_org_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 41    """Download the CT-ORG 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, "data")
 51    os.makedirs(os.path.join(data_dir, "volumes"), exist_ok=True)
 52    os.makedirs(os.path.join(data_dir, "labels"), exist_ok=True)
 53
 54    for i in tqdm(range(NUM_VOLUMES), desc="Download CT-ORG"):
 55        for folder, name in [("volumes", "volume"), ("labels", "labels")]:
 56            fname = f"{name}-{i}.nii.gz"
 57            fpath = os.path.join(data_dir, folder, fname)
 58            util.download_source(path=fpath, url=f"{URL}{folder}/{fname}", download=download, checksum=CHECKSUM)
 59
 60    return data_dir
 61
 62
 63def get_ct_org_paths(
 64    path: Union[os.PathLike, str], split: Optional[Literal["train", "test"]] = None, download: bool = False
 65) -> Tuple[List[str], List[str]]:
 66    """Get paths to the CT-ORG data.
 67
 68    Args:
 69        path: Filepath to a folder where the data is downloaded for further processing.
 70        split: The choice of data split. Either 'train' or 'test'. If None, all volumes are returned.
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        List of filepaths for the image data.
 75        List of filepaths for the label data.
 76    """
 77    data_dir = get_ct_org_data(path, download)
 78
 79    if split is None:
 80        ids = TRAIN_IDS + TEST_IDS
 81    elif split == "train":
 82        ids = TRAIN_IDS
 83    elif split == "test":
 84        ids = TEST_IDS
 85    else:
 86        raise ValueError(f"'{split}' is not a valid split.")
 87
 88    raw_paths = [os.path.join(data_dir, "volumes", f"volume-{i}.nii.gz") for i in sorted(ids)]
 89    label_paths = [os.path.join(data_dir, "labels", f"labels-{i}.nii.gz") for i in sorted(ids)]
 90    return raw_paths, label_paths
 91
 92
 93def get_ct_org_dataset(
 94    path: Union[os.PathLike, str],
 95    patch_shape: Tuple[int, ...],
 96    split: Optional[Literal["train", "test"]] = None,
 97    resize_inputs: bool = False,
 98    download: bool = False,
 99    **kwargs
100) -> Dataset:
101    """Get the CT-ORG dataset for organ segmentation.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        patch_shape: The patch shape to use for training.
106        split: The choice of data split. Either 'train' or 'test'. If None, all volumes are returned.
107        resize_inputs: Whether to resize inputs to the desired patch shape.
108        download: Whether to download the data if it is not present.
109        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
110
111    Returns:
112        The segmentation dataset.
113    """
114    raw_paths, label_paths = get_ct_org_paths(path, split, download)
115
116    if resize_inputs:
117        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
118        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
119            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
120        )
121
122    return torch_em.default_segmentation_dataset(
123        raw_paths=raw_paths,
124        raw_key="data",
125        label_paths=label_paths,
126        label_key="data",
127        patch_shape=patch_shape,
128        is_seg_dataset=True,
129        **kwargs
130    )
131
132
133def get_ct_org_loader(
134    path: Union[os.PathLike, str],
135    batch_size: int,
136    patch_shape: Tuple[int, ...],
137    split: Optional[Literal["train", "test"]] = None,
138    resize_inputs: bool = False,
139    download: bool = False,
140    **kwargs
141) -> DataLoader:
142    """Get the CT-ORG dataloader for organ segmentation.
143
144    Args:
145        path: Filepath to a folder where the data is downloaded for further processing.
146        batch_size: The batch size for training.
147        patch_shape: The patch shape to use for training.
148        split: The choice of data split. Either 'train' or 'test'. If None, all volumes are returned.
149        resize_inputs: Whether to resize inputs to the desired patch shape.
150        download: Whether to download the data if it is not present.
151        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
152
153    Returns:
154        The DataLoader.
155    """
156    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
157    dataset = get_ct_org_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
158    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)