torch_em.data.datasets.medical.abdomen_atlas

  1"""The AbdomenAtlas 1.1 Mini dataset contains annotations for 25 anatomical structures in abdominal CT scans.
  2
  3The dataset consists of 9262 cases with per-structure binary masks and a combined semantic label volume.
  4Only the 5195 cases BDMAP_00000001 to BDMAP_00005195 contain the CT scan; for the remaining cases the CT scans have
  5to be obtained from the RSNA 2023 Abdominal Trauma Detection challenge (see the dataset page), so these cases are
  6skipped here. The label ids of the combined label volume ('combined_labels.nii.gz') are given in `CLASS_IDS`:
  71: aorta, 2: gall bladder, 3: kidney (left), 4: kidney (right), 5: liver, 6: pancreas, 7: postcava, 8: spleen,
  89: stomach, 10: adrenal gland (left), 11: adrenal gland (right), 12: bladder, 13: celiac trunk, 14: colon,
  915: duodenum, 16: esophagus, 17: femur (left), 18: femur (right), 19: hepatic vessel, 20: intestine, 21: lung (left),
 1022: lung (right), 23: portal vein and splenic vein, 24: prostate, 25: rectum.
 11If the combined label volume is missing for a case, it is created by merging the per-structure masks in
 12'segmentations/<structure>.nii.gz' in the order of `CLASS_NAMES` (a structure with a higher id takes precedence).
 13
 14The dataset is located at https://huggingface.co/datasets/AbdomenAtlas/_AbdomenAtlas1.1Mini. It is gated:
 15to download it, create a HuggingFace account, accept the terms and conditions on the dataset page and create an
 16access token (https://huggingface.co/settings/tokens). Pass the token via the `token` argument or the `HF_TOKEN`
 17environment variable. Alternatively, download and extract the dataset manually (see the dataset page) and pass the
 18folder that contains the 'BDMAP_XXXXXXXX' case folders (or their parent folder) as `path`.
 19The dataset is licensed under CC BY-NC-SA 4.0.
 20
 21This dataset is from the publication https://doi.org/10.1016/j.media.2024.103285.
 22Please cite it if you use this dataset in your research.
 23"""
 24
 25import os
 26from glob import glob
 27from tqdm import tqdm
 28from natsort import natsorted
 29from typing import Union, Tuple, List, Optional
 30
 31import numpy as np
 32
 33from torch.utils.data import Dataset, DataLoader
 34
 35import torch_em
 36
 37from .. import util
 38
 39
 40REPO_ID = "AbdomenAtlas/_AbdomenAtlas1.1Mini"
 41
 42CLASS_NAMES = [
 43    "aorta", "gall_bladder", "kidney_left", "kidney_right", "liver", "pancreas", "postcava", "spleen", "stomach",
 44    "adrenal_gland_left", "adrenal_gland_right", "bladder", "celiac_trunk", "colon", "duodenum", "esophagus",
 45    "femur_left", "femur_right", "hepatic_vessel", "intestine", "lung_left", "lung_right",
 46    "portal_vein_and_splenic_vein", "prostate", "rectum",
 47]
 48"""The anatomical structures of the AbdomenAtlas 1.1 dataset. The label id of a structure is its 1-based index."""
 49
 50CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 51"""Mapping from the name of an anatomical structure to its label id in the combined label volumes."""
 52
 53
 54def merge_segmentations(case_dir: str) -> str:
 55    """Merge the per-structure binary masks of one AbdomenAtlas case into a single semantic label volume.
 56
 57    The merged volume is stored as 'combined_labels.nii.gz' in the case folder. If it already exists,
 58    it is not recomputed.
 59
 60    Args:
 61        case_dir: The folder of the case, which contains the 'segmentations' sub-folder.
 62
 63    Returns:
 64        The filepath to the merged label volume.
 65    """
 66    import nibabel as nib
 67
 68    label_path = os.path.join(case_dir, "combined_labels.nii.gz")
 69    if os.path.exists(label_path):
 70        return label_path
 71
 72    labels, affine = None, None
 73    for class_name in CLASS_NAMES:
 74        mask_path = os.path.join(case_dir, "segmentations", f"{class_name}.nii.gz")
 75        if not os.path.exists(mask_path):
 76            continue
 77        nifti = nib.load(mask_path)
 78        mask = np.asarray(nifti.dataobj) > 0
 79        if labels is None:
 80            labels, affine = np.zeros(mask.shape, dtype="uint8"), nifti.affine
 81        labels[mask] = CLASS_IDS[class_name]
 82
 83    if labels is None:
 84        raise RuntimeError(f"Could not find any segmentation masks in '{case_dir}'.")
 85
 86    nib.save(nib.Nifti1Image(labels, affine), label_path)
 87    return label_path
 88
 89
 90def _find_case_dirs(path):
 91    # NOTE: The archives do not all extract to the same depth. Most of them place the 'BDMAP_XXXXXXXX' folders
 92    # directly in the extraction folder, while the last two keep them inside a folder named after the archive.
 93    # So all depths have to be searched, otherwise the cases of the nested archives are silently missed.
 94    case_dirs = {}
 95    for pattern in ["BDMAP_*", os.path.join("*", "BDMAP_*"), os.path.join("*", "*", "BDMAP_*")]:
 96        for case_dir in glob(os.path.join(path, pattern)):
 97            if not os.path.isdir(case_dir):
 98                continue
 99
100            # If a case is found at multiple depths, then the folder that contains the image is preferred.
101            case_name = os.path.basename(case_dir)
102            if case_name not in case_dirs or os.path.exists(os.path.join(case_dir, "ct.nii.gz")):
103                case_dirs[case_name] = case_dir
104
105    # The folders are sorted by the case name, so that the order does not depend on the extraction depth.
106    return [case_dirs[case_name] for case_name in natsorted(case_dirs)]
107
108
109def get_abdomen_atlas_data(
110    path: Union[os.PathLike, str], token: Optional[str] = None, download: bool = False
111) -> List[str]:
112    """Download the AbdomenAtlas 1.1 Mini dataset.
113
114    Args:
115        path: Filepath to a folder where the data is downloaded for further processing.
116        token: The HuggingFace access token. By default, the 'HF_TOKEN' environment variable is used.
117        download: Whether to download the data if it is not present.
118
119    Returns:
120        The filepaths to the case folders.
121    """
122    case_dirs = _find_case_dirs(path)
123    if case_dirs:
124        return case_dirs
125
126    if not download:
127        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
128
129    token = os.environ.get("HF_TOKEN") if token is None else token
130    if token is None:
131        raise RuntimeError(
132            "The AbdomenAtlas 1.1 Mini dataset is gated on HuggingFace. To download it: create a HuggingFace account, "
133            f"accept the terms and conditions at https://huggingface.co/datasets/{REPO_ID}, create an access token at "
134            "https://huggingface.co/settings/tokens and pass it via the 'token' argument or the 'HF_TOKEN' environment "
135            "variable."
136        )
137
138    from huggingface_hub import snapshot_download
139
140    os.makedirs(path, exist_ok=True)
141    print("The AbdomenAtlas 1.1 Mini data is not available yet and will be downloaded.")
142    print("Note that this dataset is very large (~300 GB), so this step can take several hours.")
143    try:
144        snapshot_download(
145            repo_id=REPO_ID, repo_type="dataset", token=token, local_dir=path, allow_patterns=["*.tar.gz", "*.csv"]
146        )
147    except Exception as e:
148        raise RuntimeError(
149            f"The download of the AbdomenAtlas 1.1 Mini dataset failed ({e}). Please make sure that you have accepted "
150            f"the terms and conditions at https://huggingface.co/datasets/{REPO_ID} with the account of the token."
151        )
152
153    for tar_path in natsorted(glob(os.path.join(path, "*.tar.gz"))):
154        util.unzip_tarfile(tar_path=tar_path, dst=os.path.join(path, "uncompressed"), remove=False)
155
156    case_dirs = _find_case_dirs(path)
157    if not case_dirs:
158        raise RuntimeError(f"Could not find the 'BDMAP_XXXXXXXX' case folders of the AbdomenAtlas dataset in '{path}'.")
159    return case_dirs
160
161
162def get_abdomen_atlas_paths(
163    path: Union[os.PathLike, str], token: Optional[str] = None, download: bool = False
164) -> Tuple[List[str], List[str]]:
165    """Get paths to the AbdomenAtlas 1.1 Mini data.
166
167    Args:
168        path: Filepath to a folder where the data is downloaded for further processing.
169        token: The HuggingFace access token. By default, the 'HF_TOKEN' environment variable is used.
170        download: Whether to download the data if it is not present.
171
172    Returns:
173        List of filepaths for the image data.
174        List of filepaths for the label data.
175    """
176    case_dirs = get_abdomen_atlas_data(path, token, download)
177
178    raw_paths, label_paths = [], []
179    for case_dir in tqdm(case_dirs, desc="Preparing AbdomenAtlas labels"):
180        raw_path = os.path.join(case_dir, "ct.nii.gz")
181        if not os.path.exists(raw_path):  # The cases without CT are skipped.
182            continue
183        raw_paths.append(raw_path)
184        label_paths.append(merge_segmentations(case_dir))
185
186    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
187    return raw_paths, label_paths
188
189
190def get_abdomen_atlas_dataset(
191    path: Union[os.PathLike, str],
192    patch_shape: Tuple[int, ...],
193    token: Optional[str] = None,
194    resize_inputs: bool = False,
195    download: bool = False,
196    **kwargs
197) -> Dataset:
198    """Get the AbdomenAtlas 1.1 Mini dataset for abdominal organ segmentation.
199
200    Args:
201        path: Filepath to a folder where the data is downloaded for further processing.
202        patch_shape: The patch shape to use for training.
203        token: The HuggingFace access token. By default, the 'HF_TOKEN' environment variable is used.
204        resize_inputs: Whether to resize inputs to the desired patch shape.
205        download: Whether to download the data if it is not present.
206        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
207
208    Returns:
209        The segmentation dataset.
210    """
211    raw_paths, label_paths = get_abdomen_atlas_paths(path, token, download)
212
213    if resize_inputs:
214        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
215        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
216            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
217        )
218
219    return torch_em.default_segmentation_dataset(
220        raw_paths=raw_paths,
221        raw_key="data",
222        label_paths=label_paths,
223        label_key="data",
224        patch_shape=patch_shape,
225        is_seg_dataset=True,
226        **kwargs
227    )
228
229
230def get_abdomen_atlas_loader(
231    path: Union[os.PathLike, str],
232    batch_size: int,
233    patch_shape: Tuple[int, ...],
234    token: Optional[str] = None,
235    resize_inputs: bool = False,
236    download: bool = False,
237    **kwargs
238) -> DataLoader:
239    """Get the AbdomenAtlas 1.1 Mini dataloader for abdominal organ segmentation.
240
241    Args:
242        path: Filepath to a folder where the data is downloaded for further processing.
243        batch_size: The batch size for training.
244        patch_shape: The patch shape to use for training.
245        token: The HuggingFace access token. By default, the 'HF_TOKEN' environment variable is used.
246        resize_inputs: Whether to resize inputs to the desired patch shape.
247        download: Whether to download the data if it is not present.
248        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
249
250    Returns:
251        The DataLoader.
252    """
253    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
254    dataset = get_abdomen_atlas_dataset(path, patch_shape, token, resize_inputs, download, **ds_kwargs)
255    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)