torch_em.data.datasets.medical.radgenome_chestct

  1"""RadGenome-Chest CT extends CT-RATE (https://huggingface.co/datasets/ibrahimhamamci/CT-RATE, 25692 non-contrast
  2chest CT volumes from 20000 patients) with organ-level segmentation masks for 197 anatomical structures, as well
  3as grounded reports and VQA pairs (text-only; not covered by this loader).
  4
  5Unlike CT-RATE itself, which requires accepting a gated access agreement on HuggingFace before it can be
  6downloaded, the RadGenome-ChestCT repository (https://huggingface.co/datasets/RadGenome/RadGenome-ChestCT) is
  7NOT access-gated (verified via the HuggingFace API and by downloading files without any token). It bundles
  8its own already-preprocessed copy of the CT-RATE volumes (resampled to (3, 1, 1) mm spacing, cropped to the
  9foreground and stored with an identity affine, see 'processed_code/preprocess_ctrate_valid.py' in the
 10repository) under 'dataset/<split>_preprocessed/', where '<split>' is 'train' or 'valid'. This loader uses
 11these preprocessed volumes directly as the raw data, so the original (gated) CT-RATE repository is not needed.
 12
 13The per-structure binary segmentation masks are stored separately, keyed by the same volume name as the
 14preprocessed CT (eg. the CT 'dataset/valid_preprocessed/valid_1/valid_1a/valid_1_a_1.nii.gz' corresponds to the
 15mask folder 'seg_valid_1_a_1/', containing one binary '<structure_name>.nii.gz' file per structure). For the
 16validation split, all masks are bundled in a single 'dataset/valid_anatomy_mask.tar.gz' (~10.5 GB); for the
 17much larger training split, they are split into 5 parts, 'dataset/train_anatomy_mask_a{a,b,c,d,e}'
 18(~13-16 GB each, ~74 GB combined), which have to be concatenated into one tar.gz before any file inside can be
 19extracted. Because of this, downloading masks for even a single training-split case requires downloading the
 20full ~74 GB of split archive parts; there is no way around this with the way the data is packaged. The
 21validation split masks (~10.5 GB) are comparatively cheap, so this loader defaults to `split="validation"`.
 22Use `max_cases`/`case_ids` to restrict how many cases are used (and thus how many preprocessed CT files are
 23downloaded); the full mask archive of the chosen split is always downloaded once, regardless of the subset
 24size, since the masks cannot be fetched per case.
 25
 26Inspecting the actual downloaded validation-split archive shows that every one of its 1564 cases has exactly
 27210 per-structure mask files (empty structures are still stored as empty masks, as for AbdomenAtlas 3.0, see
 28`abdomen_atlas_3.py`). These 210 structure names are stored in `CLASS_NAMES` (ids 1-210 via `CLASS_IDS`, sorted
 29alphabetically); they do not exactly match the paper's "197 categories" figure, most likely because that figure
 30groups some of the numbered/lateralized variants together (eg. the 24 per-vertebra and per-rib masks, or the
 31left/right variants of paired organs). `CLASS_NAMES` was derived empirically from the downloaded validation
 32split and is used as the canonical vocabulary for both splits; if the training split contains additional
 33structure names not in this list, `merge_segmentations` skips them (a warning is printed) rather than failing.
 34
 35The dataset is licensed under CC BY 4.0. Despite the license, using it still requires agreeing to CT-RATE's
 36usage terms (non-commercial, no redistribution, no re-identification attempts), as reproduced in the dataset
 37card at https://huggingface.co/datasets/RadGenome/RadGenome-ChestCT.
 38
 39This dataset is from the publication https://arxiv.org/abs/2404.16754 (Scientific Data 2025).
 40Please also cite CT-RATE (https://arxiv.org/abs/2403.17834) and CT-CLIP if you use this dataset in your
 41research.
 42"""
 43
 44import os
 45import warnings
 46import subprocess
 47from glob import glob
 48from tqdm import tqdm
 49from natsort import natsorted
 50from typing import Union, Tuple, List, Optional
 51
 52import numpy as np
 53
 54from torch.utils.data import Dataset, DataLoader
 55
 56import torch_em
 57
 58from .. import util
 59
 60
 61REPO_ID = "RadGenome/RadGenome-ChestCT"
 62
 63SPLIT_PREFIXES = {"train": "train", "validation": "valid"}
 64
 65CLASS_NAMES = [
 66    "abdominal tissue", "adrenal gland", "aorta", "bone", "brachiocephalic trunk", "brachiocephalic vein",
 67    "breast", "bronchie", "buccal mucosa", "carotid artery", "caudate lobe", "celiac trunk",
 68    "cervical esophagus", "cervical vertebrae", "cervical vertebrae 1 (c1)", "cervical vertebrae 2 (c2)",
 69    "cervical vertebrae 3 (c3)", "cervical vertebrae 4 (c4)", "cervical vertebrae 5 (c5)",
 70    "cervical vertebrae 6 (c6)", "cervical vertebrae 7 (c7)", "clavicle", "colon", "common carotid artery",
 71    "costal cartilage", "cricopharyngeal inlet", "duodenum", "esophagus", "eustachian tube bone", "femur",
 72    "gallbladder", "head of femur", "head of left femur", "head of right femur", "heart",
 73    "heart ascending aorta", "heart atrium", "heart tissue", "heart ventricle", "humerus", "iliac artery",
 74    "iliac vena", "inferior vena cava", "internal carotid artery", "internal jugular vein", "intestine",
 75    "kidney", "kidney cyst", "kidney tumor", "larynx", "larynx glottis", "larynx supraglottis",
 76    "left adrenal gland", "left auricle of heart", "left brachiocephalic vein", "left breast",
 77    "left carotid artery", "left clavicle", "left common carotid artery", "left eustachian tube bone",
 78    "left femur", "left heart atrium", "left heart ventricle", "left humerus", "left iliac artery",
 79    "left iliac vena", "left internal carotid artery", "left internal jugular vein", "left kidney",
 80    "left kidney cyst", "left lateral inferior segment of liver", "left lateral superior segment of liver",
 81    "left lobe of liver", "left lung", "left lung lower lobe", "left lung upper lobe",
 82    "left medial segment of liver", "left rib", "left rib 1", "left rib 10", "left rib 11", "left rib 12",
 83    "left rib 2", "left rib 3", "left rib 4", "left rib 5", "left rib 6", "left rib 7", "left rib 8",
 84    "left rib 9", "left scapula", "left subclavian artery", "left thyroid", "liver", "liver tumor",
 85    "liver vessel", "lumbar vertebrae", "lumbar vertebrae 1 (l1)", "lumbar vertebrae 2 (l2)",
 86    "lumbar vertebrae 3 (l3)", "lumbar vertebrae 4 (l4)", "lumbar vertebrae 5 (l5)", "lumbar vertebrae 6 (l6)",
 87    "lung", "lung effusion", "lung lower lobe", "lung nodule", "lung tumor", "lung upper lobe", "mandible",
 88    "manubrium of sternum", "mediastinal tissue", "muscle", "myocardium", "pancreas", "pancreas tumor",
 89    "portal vein and splenic vein", "prostate", "pulmonary artery", "pulmonary embolism", "pulmonary vein",
 90    "rectum", "renal artery", "renal vein", "rib", "rib 1", "rib 10", "rib 11", "rib 12", "rib 2", "rib 3",
 91    "rib 4", "rib 5", "rib 6", "rib 7", "rib 8", "rib 9", "rib cartilage", "right adrenal gland",
 92    "right anterior inferior segment of liver", "right anterior superior segment of liver",
 93    "right brachiocephalic vein", "right breast", "right carotid artery", "right clavicle",
 94    "right common carotid artery", "right eustachian tube bone", "right femur", "right heart atrium",
 95    "right heart ventricle", "right humerus", "right iliac artery", "right iliac vena",
 96    "right internal carotid artery", "right internal jugular vein", "right kidney", "right kidney cyst",
 97    "right lobe of liver", "right lung", "right lung lower lobe", "right lung middle lobe",
 98    "right lung upper lobe", "right posterior inferior segment of liver",
 99    "right posterior superior segment of liver", "right rib", "right rib 1", "right rib 10", "right rib 11",
100    "right rib 12", "right rib 2", "right rib 3", "right rib 4", "right rib 5", "right rib 6", "right rib 7",
101    "right rib 8", "right rib 9", "right scapula", "right subclavian artery", "right thyroid",
102    "sacral vertebrae 1 (s1)", "scapula", "skin", "small bowel", "spinal canal", "spinal cord", "spleen",
103    "sternum", "stomach", "subclavian artery", "superior vena cava", "thoracic cavity", "thoracic vertebrae",
104    "thoracic vertebrae 1 (t1)", "thoracic vertebrae 10 (t10)", "thoracic vertebrae 11 (t11)",
105    "thoracic vertebrae 12 (t12)", "thoracic vertebrae 2 (t2)", "thoracic vertebrae 3 (t3)",
106    "thoracic vertebrae 4 (t4)", "thoracic vertebrae 5 (t5)", "thoracic vertebrae 6 (t6)",
107    "thoracic vertebrae 7 (t7)", "thoracic vertebrae 8 (t8)", "thoracic vertebrae 9 (t9)", "thymus", "thyroid",
108    "thyroid gland", "trachea", "vertebrae",
109]
110"""The 210 anatomical structure names found in the RadGenome-ChestCT masks (validation split), sorted
111alphabetically. See the module docstring for how this relates to the paper's "197 categories" figure."""
112
113CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
114"""Mapping from structure name to its label id in the combined label volumes created by `merge_segmentations`."""
115
116
117def _volume_name(relpath: str) -> str:
118    return os.path.splitext(os.path.splitext(os.path.basename(relpath))[0])[0]
119
120
121def _find_mask_dir(path: str, volume_name: str) -> Optional[str]:
122    matches = glob(os.path.join(path, "**", f"seg_{volume_name}"), recursive=True)
123    return matches[0] if matches else None
124
125
126def _find_image_path(path: str, volume_name: str) -> Optional[str]:
127    matches = glob(os.path.join(path, "**", f"{volume_name}.nii.gz"), recursive=True)
128    matches = [m for m in matches if os.sep + "preprocessed" + os.sep in m or "preprocessed" in m]
129    return matches[0] if matches else None
130
131
132def merge_segmentations(mask_dir: str) -> str:
133    """Merge the per-structure binary masks of one RadGenome-ChestCT case into a single semantic label volume.
134
135    The merged volume is stored as 'combined_labels.nii.gz' in the mask folder. If it already exists,
136    it is not recomputed.
137
138    Args:
139        mask_dir: The 'seg_<volume_name>' folder holding the per-structure masks.
140
141    Returns:
142        The filepath to the merged label volume.
143    """
144    import nibabel as nib
145
146    label_path = os.path.join(mask_dir, "combined_labels.nii.gz")
147    if os.path.exists(label_path):
148        return label_path
149
150    labels, affine = None, None
151    for mask_path in natsorted(glob(os.path.join(mask_dir, "*.nii.gz"))):
152        class_name = os.path.splitext(os.path.splitext(os.path.basename(mask_path))[0])[0]
153        if class_name not in CLASS_IDS:
154            warnings.warn(f"Skipping unknown structure '{class_name}' in '{mask_dir}'.")
155            continue
156        nifti = nib.load(mask_path)
157        mask = np.asarray(nifti.dataobj) > 0
158        if labels is None:
159            labels, affine = np.zeros(mask.shape, dtype="uint8"), nifti.affine
160        labels[mask] = CLASS_IDS[class_name]
161
162    if labels is None:
163        raise RuntimeError(f"Could not find any segmentation masks in '{mask_dir}'.")
164
165    nib.save(nib.Nifti1Image(labels, affine), label_path)
166    return label_path
167
168
169def _concatenate_train_mask_shards(path: str) -> str:
170    combined_path = os.path.join(path, "dataset", "train_anatomy_mask.tar.gz")
171    if os.path.exists(combined_path):
172        return combined_path
173
174    part_paths = natsorted(glob(os.path.join(path, "dataset", "train_anatomy_mask_a?")))
175    assert len(part_paths) == 5, f"Expected 5 'train_anatomy_mask_a*' parts, found {len(part_paths)} in '{path}'."
176
177    print("Concatenating the 5 train anatomy mask shards into one archive (~74 GB). This can take a while.")
178    with open(combined_path, "wb") as dst:
179        subprocess.run(["cat"] + part_paths, stdout=dst, check=True)
180    return combined_path
181
182
183def get_radgenome_chestct_data(
184    path: Union[os.PathLike, str],
185    split: str = "validation",
186    max_cases: Optional[int] = None,
187    case_ids: Optional[List[str]] = None,
188    download: bool = False,
189) -> List[str]:
190    """Download the RadGenome-ChestCT dataset.
191
192    The dataset is not gated on HuggingFace and does not require an access token. It is, however, very
193    large: even the validation split requires downloading the full ~10.5 GB mask archive regardless of how
194    many cases are requested (masks cannot be fetched per case), and the training split requires downloading
195    ~74 GB of mask archive shards for the same reason. Use `max_cases`/`case_ids` to at least limit how many
196    preprocessed CT volumes are downloaded; leave both at their default (None) to use the full split.
197
198    Args:
199        path: Filepath to a folder where the data is downloaded for further processing.
200        split: The data split, either 'train' or 'validation'.
201        max_cases: The maximum number of cases to use, taken in natural sort order. Mutually exclusive
202            with `case_ids`.
203        case_ids: Explicit list of case / volume ids (eg. ['valid_1_a_1']) to use. Mutually exclusive
204            with `max_cases`.
205        download: Whether to download the data if it is not present.
206
207    Returns:
208        The volume names (eg. 'valid_1_a_1') of the available cases.
209    """
210    assert split in SPLIT_PREFIXES, f"'split' must be one of {list(SPLIT_PREFIXES)}, got '{split}'."
211    assert max_cases is None or case_ids is None, "'max_cases' and 'case_ids' are mutually exclusive."
212    prefix = SPLIT_PREFIXES[split]
213
214    def _available_volumes():
215        image_dir = os.path.join(path, "dataset", f"{prefix}_preprocessed")
216        volumes = [
217            _volume_name(p) for p in glob(os.path.join(image_dir, "*", "*", "*.nii.gz"))
218            if _find_mask_dir(path, _volume_name(p)) is not None
219        ]
220        return natsorted(volumes)
221
222    volumes = _available_volumes()
223    if volumes:
224        if max_cases is not None:
225            volumes = volumes[:max_cases]
226        elif case_ids is not None:
227            volumes = [v for v in volumes if v in case_ids]
228        if volumes:
229            return volumes
230
231    if not download:
232        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
233
234    from huggingface_hub import HfApi, snapshot_download
235
236    os.makedirs(path, exist_ok=True)
237
238    api = HfApi()
239    files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
240    image_files = [f for f in files if f.startswith(f"dataset/{prefix}_preprocessed/") and f.endswith(".nii.gz")]
241    assert image_files, f"Could not find any preprocessed volumes for split '{split}' in the repository."
242    volname_to_relpath = {_volume_name(f): f for f in image_files}
243    all_volumes = natsorted(volname_to_relpath.keys())
244
245    if case_ids is not None:
246        target_volumes = [v for v in all_volumes if v in case_ids]
247    elif max_cases is not None:
248        target_volumes = all_volumes[:max_cases]
249    else:
250        target_volumes = all_volumes
251        print(f"Downloading the full RadGenome-ChestCT '{split}' split ({len(all_volumes)} cases).")
252
253    image_patterns = [volname_to_relpath[v] for v in target_volumes]
254    if split == "train":
255        mask_patterns = [f"dataset/train_anatomy_mask_a{c}" for c in "abcde"]
256        print(
257            "Downloading the training split's segmentation masks requires all 5 mask archive shards "
258            "(~74 GB combined), regardless of how many cases are requested."
259        )
260    else:
261        mask_patterns = ["dataset/valid_anatomy_mask.tar.gz"]
262        print("Downloading the validation split's segmentation mask archive (~10.5 GB).")
263
264    print(f"Downloading {len(target_volumes)} preprocessed CT volume(s) for split '{split}'.")
265    snapshot_download(
266        repo_id=REPO_ID, repo_type="dataset", local_dir=path, allow_patterns=image_patterns + mask_patterns
267    )
268
269    if split == "train":
270        tar_path = _concatenate_train_mask_shards(path)
271    else:
272        tar_path = os.path.join(path, "dataset", "valid_anatomy_mask.tar.gz")
273
274    import tarfile
275
276    mask_extract_dir = os.path.join(path, "dataset")
277    wanted = [f"seg_{v}/" for v in target_volumes]
278    with tarfile.open(tar_path) as tf:
279        members = [m for m in tf.getmembers() if any(w in m.name for w in wanted)]
280        tf.extractall(path=mask_extract_dir, members=members)
281
282    volumes = [v for v in target_volumes if _find_mask_dir(path, v) is not None]
283    if not volumes:
284        raise RuntimeError(f"Could not find any RadGenome-ChestCT masks for the requested cases in '{path}'.")
285
286    return volumes
287
288
289def get_radgenome_chestct_paths(
290    path: Union[os.PathLike, str],
291    split: str = "validation",
292    max_cases: Optional[int] = None,
293    case_ids: Optional[List[str]] = None,
294    download: bool = False,
295) -> Tuple[List[str], List[str]]:
296    """Get paths to the RadGenome-ChestCT data.
297
298    Args:
299        path: Filepath to a folder where the data is downloaded for further processing.
300        split: The data split, either 'train' or 'validation'.
301        max_cases: The maximum number of cases to use. See `get_radgenome_chestct_data` for details.
302        case_ids: Explicit list of case ids to use. See `get_radgenome_chestct_data` for details.
303        download: Whether to download the data if it is not present.
304
305    Returns:
306        List of filepaths for the image data.
307        List of filepaths for the label data.
308    """
309    volumes = get_radgenome_chestct_data(path, split, max_cases, case_ids, download)
310
311    raw_paths, label_paths = [], []
312    for volume in tqdm(volumes, desc="Preparing RadGenome-ChestCT labels"):
313        image_path = _find_image_path(path, volume)
314        mask_dir = _find_mask_dir(path, volume)
315        if image_path is None or mask_dir is None:
316            continue
317        raw_paths.append(image_path)
318        label_paths.append(merge_segmentations(mask_dir))
319
320    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
321    return raw_paths, label_paths
322
323
324def get_radgenome_chestct_dataset(
325    path: Union[os.PathLike, str],
326    patch_shape: Tuple[int, ...],
327    split: str = "validation",
328    max_cases: Optional[int] = None,
329    case_ids: Optional[List[str]] = None,
330    resize_inputs: bool = False,
331    download: bool = False,
332    **kwargs
333) -> Dataset:
334    """Get the RadGenome-ChestCT dataset for chest CT organ-level segmentation.
335
336    Args:
337        path: Filepath to a folder where the data is downloaded for further processing.
338        patch_shape: The patch shape to use for training.
339        split: The data split, either 'train' or 'validation'.
340        max_cases: The maximum number of cases to use. See `get_radgenome_chestct_data` for details.
341        case_ids: Explicit list of case ids to use. See `get_radgenome_chestct_data` for details.
342        resize_inputs: Whether to resize inputs to the desired patch shape.
343        download: Whether to download the data if it is not present.
344        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
345
346    Returns:
347        The segmentation dataset.
348    """
349    raw_paths, label_paths = get_radgenome_chestct_paths(path, split, max_cases, case_ids, download)
350
351    if resize_inputs:
352        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
353        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
354            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
355        )
356
357    return torch_em.default_segmentation_dataset(
358        raw_paths=raw_paths,
359        raw_key="data",
360        label_paths=label_paths,
361        label_key="data",
362        patch_shape=patch_shape,
363        is_seg_dataset=True,
364        **kwargs
365    )
366
367
368def get_radgenome_chestct_loader(
369    path: Union[os.PathLike, str],
370    batch_size: int,
371    patch_shape: Tuple[int, ...],
372    split: str = "validation",
373    max_cases: Optional[int] = None,
374    case_ids: Optional[List[str]] = None,
375    resize_inputs: bool = False,
376    download: bool = False,
377    **kwargs
378) -> DataLoader:
379    """Get the RadGenome-ChestCT dataloader for chest CT organ-level segmentation.
380
381    Args:
382        path: Filepath to a folder where the data is downloaded for further processing.
383        batch_size: The batch size for training.
384        patch_shape: The patch shape to use for training.
385        split: The data split, either 'train' or 'validation'.
386        max_cases: The maximum number of cases to use. See `get_radgenome_chestct_data` for details.
387        case_ids: Explicit list of case ids to use. See `get_radgenome_chestct_data` for details.
388        resize_inputs: Whether to resize inputs to the desired patch shape.
389        download: Whether to download the data if it is not present.
390        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch
391            DataLoader.
392
393    Returns:
394        The DataLoader.
395    """
396    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
397    dataset = get_radgenome_chestct_dataset(
398        path, patch_shape, split, max_cases, case_ids, resize_inputs, download, **ds_kwargs
399    )
400    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)