torch_em.data.datasets.medical.abdomen_atlas_3

  1"""The AbdomenAtlas 3.0 Mini dataset contains annotations for 44 anatomical structures
  2(organs, blood vessels, liver/pancreas sub-segments and tumor lesions) in abdominal CT scans.
  3
  4The dataset consists of 9262 cases (BDMAP_00000001 to BDMAP_00009262), each with a CT scan
  5('ct.nii.gz') and per-structure binary masks in a 'segmentations' sub-folder. Unlike AbdomenAtlas 1.1 Mini
  6(see `abdomen_atlas.py`), it is a public, ungated release: no HuggingFace account or access token is needed.
  7The data is hosted as 40 shards of (mostly) 232 cases each, separately for the CT scans ('image_only/*.tar.gz',
  8~14 GB per shard) and the segmentation masks ('mask_only/*.tar.gz', ~320 MB per shard), for a combined size of
  9about 586 GB. Downloading a case therefore always downloads the whole shard it belongs to; use `max_cases` or
 10`case_ids` to only fetch the shards required for a small subset of cases (see below).
 11
 12The label ids of the combined semantic label volume created by `merge_segmentations` are given in `CLASS_IDS`:
 131: adrenal_gland_left, 2: adrenal_gland_right, 3: bladder, 4: colon, 5: duodenum, 6: esophagus, 7: femur_left,
 148: femur_right, 9: gall_bladder, 10: intestine, 11: kidney_left, 12: kidney_right, 13: liver, 14: lung_left,
 1515: lung_right, 16: pancreas, 17: prostate, 18: rectum, 19: spleen, 20: stomach, 21: aorta, 22: celiac_aa,
 1623: celiac_trunk, 24: common_bile_duct, 25: hepatic_vessel, 26: portal_vein_and_splenic_vein, 27: postcava,
 1728: superior_mesenteric_artery, 29: veins, 30: liver_segment_1, 31: liver_segment_2, 32: liver_segment_3,
 1833: liver_segment_4, 34: liver_segment_5, 35: liver_segment_6, 36: liver_segment_7, 37: liver_segment_8,
 1938: pancreas_head, 39: pancreas_body, 40: pancreas_tail, 41: liver_lesion, 42: kidney_lesion,
 2043: pancreatic_lesion, 44: colon_lesion.
 21The tumor lesion classes (41-44) are merged last, so they take precedence over the organ they lie in wherever
 22the lesion and organ masks overlap. 'colon_lesion' is only present for a small subset of cases (the mask files
 23for the other 43 structures are provided for every case, even if empty).
 24If the combined label volume is missing for a case, it is created by merging the per-structure masks in
 25'segmentations/<structure>.nii.gz' in the order of `CLASS_NAMES` (a structure with a higher id takes precedence).
 26
 27The dataset is located at https://huggingface.co/datasets/AbdomenAtlas/AbdomenAtlas3.0Mini. It is licensed
 28under CC BY-NC-SA 4.0.
 29
 30This dataset is from the publication https://arxiv.org/abs/2501.04678.
 31Please cite it if you use this dataset in your research.
 32"""
 33
 34import os
 35import re
 36from glob import glob
 37from tqdm import tqdm
 38from natsort import natsorted
 39from typing import Union, Tuple, List, Optional
 40
 41import numpy as np
 42
 43from torch.utils.data import Dataset, DataLoader
 44
 45import torch_em
 46
 47from .. import util
 48from .abdomen_atlas import _find_case_dirs
 49
 50
 51REPO_ID = "AbdomenAtlas/AbdomenAtlas3.0Mini"
 52
 53ORGAN_NAMES = [
 54    "adrenal_gland_left", "adrenal_gland_right", "bladder", "colon", "duodenum", "esophagus", "femur_left",
 55    "femur_right", "gall_bladder", "intestine", "kidney_left", "kidney_right", "liver", "lung_left", "lung_right",
 56    "pancreas", "prostate", "rectum", "spleen", "stomach",
 57]
 58"""The organs of the AbdomenAtlas 3.0 dataset."""
 59
 60VESSEL_NAMES = [
 61    "aorta", "celiac_aa", "celiac_trunk", "common_bile_duct", "hepatic_vessel", "portal_vein_and_splenic_vein",
 62    "postcava", "superior_mesenteric_artery", "veins",
 63]
 64"""The blood vessels and bile duct of the AbdomenAtlas 3.0 dataset."""
 65
 66SUBSEGMENT_NAMES = [
 67    "liver_segment_1", "liver_segment_2", "liver_segment_3", "liver_segment_4", "liver_segment_5",
 68    "liver_segment_6", "liver_segment_7", "liver_segment_8", "pancreas_head", "pancreas_body", "pancreas_tail",
 69]
 70"""The liver and pancreas sub-segments of the AbdomenAtlas 3.0 dataset."""
 71
 72LESION_NAMES = ["liver_lesion", "kidney_lesion", "pancreatic_lesion", "colon_lesion"]
 73"""The tumor lesion classes of the AbdomenAtlas 3.0 dataset. These are merged last, ie. they take precedence
 74over the organ they lie in wherever the lesion and organ masks overlap."""
 75
 76CLASS_NAMES = ORGAN_NAMES + VESSEL_NAMES + SUBSEGMENT_NAMES + LESION_NAMES
 77"""All structures of the AbdomenAtlas 3.0 dataset, in the order used to build the combined label volume."""
 78
 79CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 80"""Mapping from the name of a structure to its label id in the combined label volumes."""
 81
 82
 83def merge_segmentations(case_dir: str) -> str:
 84    """Merge the per-structure binary masks of one AbdomenAtlas 3.0 case into a single semantic label volume.
 85
 86    The merged volume is stored as 'combined_labels.nii.gz' in the case folder. If it already exists,
 87    it is not recomputed.
 88
 89    Args:
 90        case_dir: The folder of the case, which contains the 'segmentations' sub-folder.
 91
 92    Returns:
 93        The filepath to the merged label volume.
 94    """
 95    import nibabel as nib
 96
 97    label_path = os.path.join(case_dir, "combined_labels.nii.gz")
 98    if os.path.exists(label_path):
 99        return label_path
100
101    labels, affine = None, None
102    for class_name in CLASS_NAMES:
103        mask_path = os.path.join(case_dir, "segmentations", f"{class_name}.nii.gz")
104        if not os.path.exists(mask_path):
105            continue
106        nifti = nib.load(mask_path)
107        mask = np.asarray(nifti.dataobj) > 0
108        if labels is None:
109            labels, affine = np.zeros(mask.shape, dtype="uint8"), nifti.affine
110        labels[mask] = CLASS_IDS[class_name]
111
112    if labels is None:
113        raise RuntimeError(f"Could not find any segmentation masks in '{case_dir}'.")
114
115    nib.save(nib.Nifti1Image(labels, affine), label_path)
116    return label_path
117
118
119def _parse_shards(files, subdir):
120    # Parse the ('AbdomenAtlas3_images_BDMAP_BDMAP_00000001_BDMAP_00000232.tar.gz'-style) shard filenames
121    # into (start_case, end_case, filename) tuples, so that the shard(s) covering a given case can be found.
122    shards = []
123    for fpath in files:
124        if not (fpath.startswith(f"{subdir}/") and fpath.endswith(".tar.gz")):
125            continue
126        numbers = re.findall(r"BDMAP_(\d{8})", os.path.basename(fpath))
127        assert len(numbers) == 2, f"Could not parse the case range from '{fpath}'."
128        shards.append((int(numbers[0]), int(numbers[1]), fpath))
129    assert shards, f"Could not find any shards under '{subdir}' in the AbdomenAtlas 3.0 repository."
130    return sorted(shards)
131
132
133def _shards_for_cases(shards, case_numbers):
134    patterns = set()
135    for start, end, fpath in shards:
136        if any(start <= number <= end for number in case_numbers):
137            patterns.add(fpath)
138    return patterns
139
140
141def get_abdomen_atlas_3_data(
142    path: Union[os.PathLike, str],
143    max_cases: Optional[int] = None,
144    case_ids: Optional[List[str]] = None,
145    download: bool = False,
146) -> List[str]:
147    """Download the AbdomenAtlas 3.0 Mini dataset.
148
149    The dataset is ungated and does not require a HuggingFace account or access token. It is, however, very
150    large (~586 GB for all 9262 cases, split into 40 shards of ~14 GB (images) + ~320 MB (masks) each), and a
151    case can only be downloaded together with the full shard (of up to 232 cases) it belongs to. Use `max_cases`
152    or `case_ids` to only download the shards required for a small subset of cases; leave both at their default
153    (None) to download the full dataset.
154
155    Args:
156        path: Filepath to a folder where the data is downloaded for further processing.
157        max_cases: The maximum number of cases to download, taken in order starting from 'BDMAP_00000001'.
158            Only the shards covering these cases are downloaded. Mutually exclusive with `case_ids`.
159        case_ids: Explicit list of case ids (eg. ['BDMAP_00000001', 'BDMAP_00000002']) to download. Only the
160            shards covering these cases are downloaded. Mutually exclusive with `max_cases`.
161        download: Whether to download the data if it is not present.
162
163    Returns:
164        The filepaths to the case folders.
165    """
166    assert max_cases is None or case_ids is None, "'max_cases' and 'case_ids' are mutually exclusive."
167
168    case_dirs = _find_case_dirs(path)
169    if case_dirs:
170        if max_cases is not None:
171            case_dirs = case_dirs[:max_cases]
172        elif case_ids is not None:
173            case_dirs = [c for c in case_dirs if os.path.basename(c) in case_ids]
174        if case_dirs:
175            return case_dirs
176
177    if not download:
178        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
179
180    from huggingface_hub import HfApi, snapshot_download
181
182    os.makedirs(path, exist_ok=True)
183
184    api = HfApi()
185    files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
186    image_shards = _parse_shards(files, "image_only")
187    mask_shards = _parse_shards(files, "mask_only")
188
189    if case_ids is not None:
190        case_numbers = [int(cid.split("_")[-1]) for cid in case_ids]
191    elif max_cases is not None:
192        last_case = image_shards[-1][1]
193        case_numbers = list(range(1, min(max_cases, last_case) + 1))
194    else:
195        case_numbers = None  # Download everything.
196
197    if case_numbers is None:
198        allow_patterns = ["*.tar.gz", "*.csv"]
199        print("The AbdomenAtlas 3.0 Mini data is not available yet and will be downloaded.")
200        print("Note that this dataset is very large (~586 GB), so this step can take several hours.")
201    else:
202        allow_patterns = sorted(
203            _shards_for_cases(image_shards, case_numbers) | _shards_for_cases(mask_shards, case_numbers)
204        ) + ["*.csv"]
205        print(f"Downloading the AbdomenAtlas 3.0 Mini shards required for {len(case_numbers)} case(s).")
206
207    snapshot_download(repo_id=REPO_ID, repo_type="dataset", local_dir=path, allow_patterns=allow_patterns)
208
209    for tar_path in natsorted(glob(os.path.join(path, "*", "*.tar.gz"))):
210        util.unzip_tarfile(tar_path=tar_path, dst=os.path.join(path, "uncompressed"), remove=False)
211
212    case_dirs = _find_case_dirs(path)
213    if not case_dirs:
214        raise RuntimeError(
215            f"Could not find the 'BDMAP_XXXXXXXX' case folders of the AbdomenAtlas 3.0 dataset in '{path}'."
216        )
217
218    if max_cases is not None:
219        case_dirs = case_dirs[:max_cases]
220    elif case_ids is not None:
221        case_dirs = [c for c in case_dirs if os.path.basename(c) in case_ids]
222
223    return case_dirs
224
225
226def get_abdomen_atlas_3_paths(
227    path: Union[os.PathLike, str],
228    max_cases: Optional[int] = None,
229    case_ids: Optional[List[str]] = None,
230    download: bool = False,
231) -> Tuple[List[str], List[str]]:
232    """Get paths to the AbdomenAtlas 3.0 Mini data.
233
234    Args:
235        path: Filepath to a folder where the data is downloaded for further processing.
236        max_cases: The maximum number of cases to use, taken in order starting from 'BDMAP_00000001'.
237            See `get_abdomen_atlas_3_data` for details on how this restricts the download volume.
238        case_ids: Explicit list of case ids to use. See `get_abdomen_atlas_3_data` for details.
239        download: Whether to download the data if it is not present.
240
241    Returns:
242        List of filepaths for the image data.
243        List of filepaths for the label data.
244    """
245    case_dirs = get_abdomen_atlas_3_data(path, max_cases, case_ids, download)
246
247    raw_paths, label_paths = [], []
248    for case_dir in tqdm(case_dirs, desc="Preparing AbdomenAtlas 3.0 labels"):
249        raw_path = os.path.join(case_dir, "ct.nii.gz")
250        if not os.path.exists(raw_path):
251            continue
252        raw_paths.append(raw_path)
253        label_paths.append(merge_segmentations(case_dir))
254
255    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
256    return raw_paths, label_paths
257
258
259def get_abdomen_atlas_3_dataset(
260    path: Union[os.PathLike, str],
261    patch_shape: Tuple[int, ...],
262    max_cases: Optional[int] = None,
263    case_ids: Optional[List[str]] = None,
264    resize_inputs: bool = False,
265    download: bool = False,
266    **kwargs
267) -> Dataset:
268    """Get the AbdomenAtlas 3.0 Mini dataset for abdominal organ, vessel and tumor segmentation.
269
270    Args:
271        path: Filepath to a folder where the data is downloaded for further processing.
272        patch_shape: The patch shape to use for training.
273        max_cases: The maximum number of cases to use. See `get_abdomen_atlas_3_data` for details.
274        case_ids: Explicit list of case ids to use. See `get_abdomen_atlas_3_data` for details.
275        resize_inputs: Whether to resize inputs to the desired patch shape.
276        download: Whether to download the data if it is not present.
277        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
278
279    Returns:
280        The segmentation dataset.
281    """
282    raw_paths, label_paths = get_abdomen_atlas_3_paths(path, max_cases, case_ids, download)
283
284    if resize_inputs:
285        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
286        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
287            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
288        )
289
290    return torch_em.default_segmentation_dataset(
291        raw_paths=raw_paths,
292        raw_key="data",
293        label_paths=label_paths,
294        label_key="data",
295        patch_shape=patch_shape,
296        is_seg_dataset=True,
297        **kwargs
298    )
299
300
301def get_abdomen_atlas_3_loader(
302    path: Union[os.PathLike, str],
303    batch_size: int,
304    patch_shape: Tuple[int, ...],
305    max_cases: Optional[int] = None,
306    case_ids: Optional[List[str]] = None,
307    resize_inputs: bool = False,
308    download: bool = False,
309    **kwargs
310) -> DataLoader:
311    """Get the AbdomenAtlas 3.0 Mini dataloader for abdominal organ, vessel and tumor segmentation.
312
313    Args:
314        path: Filepath to a folder where the data is downloaded for further processing.
315        batch_size: The batch size for training.
316        patch_shape: The patch shape to use for training.
317        max_cases: The maximum number of cases to use. See `get_abdomen_atlas_3_data` for details.
318        case_ids: Explicit list of case ids to use. See `get_abdomen_atlas_3_data` for details.
319        resize_inputs: Whether to resize inputs to the desired patch shape.
320        download: Whether to download the data if it is not present.
321        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
322
323    Returns:
324        The DataLoader.
325    """
326    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
327    dataset = get_abdomen_atlas_3_dataset(path, patch_shape, max_cases, case_ids, resize_inputs, download, **ds_kwargs)
328    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)