torch_em.data.datasets.medical.abdomenct_1k

  1"""The AbdomenCT-1K dataset contains annotations for liver, kidney, spleen and pancreas segmentation in CT scans.
  2
  3The dataset consists of 1112 abdominal CT volumes from 12 medical centers (collected from LiTS, KiTS19, MSD Spleen,
  4MSD Pancreas, NIH Pancreas and Nanjing University). The annotations of 1000 volumes are public, the annotations of
  5the remaining 112 volumes are held back for the AbdomenCT-1K benchmarks and these volumes are not used here,
  6so the dataset provides 1000 image / label pairs.
  7The label ids are: 1: liver, 2: kidney, 3: spleen, 4: pancreas (see `CLASS_IDS`).
  8
  9The dataset is located at https://zenodo.org/records/5903099 (images, part 1),
 10https://zenodo.org/records/5903846 (images, part 2) and https://zenodo.org/records/5903769 (images, part 3 and labels).
 11NOTE: The label archive is a 7z archive, so the '7z' CLI is required to extract it
 12(install it via 'conda install -c conda-forge p7zip').
 13
 14This dataset is from the publication https://doi.org/10.1109/TPAMI.2021.3100536.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, List
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URLS = {
 31    "AbdomenCT-1K-ImagePart1.zip": "https://zenodo.org/records/5903099/files/AbdomenCT-1K-ImagePart1.zip?download=1",
 32    "AbdomenCT-1K-ImagePart2.zip": "https://zenodo.org/records/5903846/files/AbdomenCT-1K-ImagePart2.zip?download=1",
 33    "AbdomenCT-1K-ImagePart3.zip": "https://zenodo.org/records/5903769/files/AbdomenCT-1K-ImagePart3.zip?download=1",
 34    "Mask.7z": "https://zenodo.org/records/5903769/files/Mask.7z?download=1",
 35}
 36
 37CHECKSUMS = {
 38    "AbdomenCT-1K-ImagePart1.zip": "3d0d8edd2a8777f8807c6715f633e2fea104f16a8514572664021ed88f97cd03",
 39    "AbdomenCT-1K-ImagePart2.zip": "3178bd3021dfed0440c4dc269b7f3125b27b60aec96e81b210d7c4c55d0fe894",
 40    "AbdomenCT-1K-ImagePart3.zip": "76bf789589a13ca1aef9b9c8b771f05c165e6d2013284bbfdb4c4656d4931c30",
 41    "Mask.7z": "cefa7cab51fb0781876550b032278ca6a1ed2a8c776c8567410ea2b2a6984e50",
 42}
 43
 44CLASS_NAMES = ["liver", "kidney", "spleen", "pancreas"]
 45"""The organs of the AbdomenCT-1K dataset. The label id of an organ is its 1-based index."""
 46
 47CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 48"""Mapping from the organ name to its label id."""
 49
 50
 51def get_abdomenct_1k_data(path: Union[os.PathLike, str], download: bool = False):
 52    """Download the AbdomenCT-1K dataset.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        download: Whether to download the data if it is not present.
 57    """
 58    os.makedirs(path, exist_ok=True)
 59
 60    for name, url in URLS.items():
 61        archive_path = os.path.join(path, name)
 62        data_dir = os.path.join(path, name.split(".")[0])
 63        if os.path.exists(data_dir):
 64            continue
 65
 66        util.download_source(path=archive_path, url=url, download=download, checksum=CHECKSUMS[name])
 67        if name.endswith(".zip"):
 68            util.unzip(zip_path=archive_path, dst=path, remove=False)
 69        else:
 70            util.unzip_7z(path_7z=archive_path, dst=data_dir, remove=False)
 71
 72
 73def get_abdomenct_1k_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 74    """Get paths to the AbdomenCT-1K data.
 75
 76    Args:
 77        path: Filepath to a folder where the data is downloaded for further processing.
 78        download: Whether to download the data if it is not present.
 79
 80    Returns:
 81        List of filepaths for the image data.
 82        List of filepaths for the label data.
 83    """
 84    get_abdomenct_1k_data(path, download)
 85
 86    # Only the cases with public annotations are used.
 87    label_paths = natsorted(glob(os.path.join(path, "Mask", "Case_*.nii.gz")))
 88    raw_paths = []
 89    for label_path in label_paths:
 90        case_id = os.path.basename(label_path)[:-len(".nii.gz")]
 91        image_paths = glob(os.path.join(path, "AbdomenCT-1K-ImagePart*", f"{case_id}_0000.nii.gz"))
 92        if len(image_paths) != 1:
 93            raise RuntimeError(f"Could not find a unique image for '{label_path}', found {image_paths}.")
 94        raw_paths.append(image_paths[0])
 95
 96    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 97    return raw_paths, label_paths
 98
 99
100def get_abdomenct_1k_dataset(
101    path: Union[os.PathLike, str],
102    patch_shape: Tuple[int, ...],
103    resize_inputs: bool = False,
104    download: bool = False,
105    **kwargs
106) -> Dataset:
107    """Get the AbdomenCT-1K dataset for abdominal organ segmentation.
108
109    Args:
110        path: Filepath to a folder where the data is downloaded for further processing.
111        patch_shape: The patch shape to use for training.
112        resize_inputs: Whether to resize inputs to the desired patch shape.
113        download: Whether to download the data if it is not present.
114        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
115
116    Returns:
117        The segmentation dataset.
118    """
119    raw_paths, label_paths = get_abdomenct_1k_paths(path, download)
120
121    if resize_inputs:
122        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
123        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
124            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
125        )
126
127    return torch_em.default_segmentation_dataset(
128        raw_paths=raw_paths,
129        raw_key="data",
130        label_paths=label_paths,
131        label_key="data",
132        patch_shape=patch_shape,
133        is_seg_dataset=True,
134        **kwargs
135    )
136
137
138def get_abdomenct_1k_loader(
139    path: Union[os.PathLike, str],
140    batch_size: int,
141    patch_shape: Tuple[int, ...],
142    resize_inputs: bool = False,
143    download: bool = False,
144    **kwargs
145) -> DataLoader:
146    """Get the AbdomenCT-1K dataloader for abdominal organ segmentation.
147
148    Args:
149        path: Filepath to a folder where the data is downloaded for further processing.
150        batch_size: The batch size for training.
151        patch_shape: The patch shape to use for training.
152        resize_inputs: Whether to resize inputs to the desired patch shape.
153        download: Whether to download the data if it is not present.
154        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
155
156    Returns:
157        The DataLoader.
158    """
159    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
160    dataset = get_abdomenct_1k_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
161    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)