torch_em.data.datasets.medical.enhance_pet

  1"""The ENHANCE.PET 1.6k dataset contains annotations for 130 anatomical structures in whole-/total-body CT
  2volumes, co-acquired with [18F]FDG-PET as part of routine oncological imaging (~1,597 studies).
  3
  4The per-structure binary masks are grouped into 7 category volumes ('Body-Composition', 'Cardiac', 'Muscles',
  5'Organs', 'Peripheral-Bones', 'Ribs', 'Vertebrae'), each registered to the CT grid, with structure ids
  6restarting at 1 within each category. `get_enhance_pet_data` merges the 7 category volumes of a case into a
  7single semantic label volume, where the label id of a structure is its (1-based) position in `CLASS_NAMES`
  8(see `CLASS_NAMES_BY_CATEGORY` for the per-category structure lists this is built from). The PET volumes have
  9a different grid than the CT / label volumes (no shared voxel grid without resampling) and are not exposed by
 10this loader, only the CT volumes are used as raw data.
 11
 12The data is located at the public, anonymously readable S3 bucket 'enhance-pet-1-6k'
 13(https://registry.opendata.aws/enhance-pet-1-6k/, region 'us-west-2'), released under a CC-BY-4.0 license.
 14
 15This dataset is from the publication https://doi.org/10.1038/s41597-026-07218-y.
 16Please cite it if you use this dataset for your research.
 17"""
 18
 19import os
 20from glob import glob
 21from concurrent import futures
 22from typing import Union, Tuple, List, Optional
 23
 24import numpy as np
 25from tqdm import tqdm
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34BUCKET_URL = "https://enhance-pet-1-6k.s3.us-west-2.amazonaws.com"
 35
 36CLASS_NAMES_BY_CATEGORY = {
 37    "Body-Composition": ["skeletal_muscle", "subcutaneous_fat", "visceral_fat"],
 38    "Cardiac": [
 39        "heart_myocardium", "heart_atrium_left", "heart_atrium_right", "heart_ventricle_left",
 40        "heart_ventricle_right", "aorta", "iliac_artery_left", "iliac_artery_right", "iliac_vena_left",
 41        "iliac_vena_right", "inferior_vena_cava", "portal_splenic_vein", "pulmonary_artery",
 42    ],
 43    "Muscles": [
 44        "autochthon_left", "autochthon_right", "gluteus_maximus_left", "gluteus_maximus_right",
 45        "gluteus_medius_left", "gluteus_medius_right", "gluteus_minimus_left", "gluteus_minimus_right",
 46        "iliopsoas_left", "iliopsoas_right",
 47    ],
 48    "Organs": [
 49        "adrenal_gland_left", "adrenal_gland_right", "bladder", "brain", "gallbladder", "kidney_left",
 50        "kidney_right", "liver", "lung_lower_lobe_left", "lung_lower_lobe_right", "lung_middle_lobe_right",
 51        "lung_upper_lobe_left", "lung_upper_lobe_right", "pancreas", "spleen", "stomach", "thyroid_left",
 52        "thyroid_right",
 53    ],
 54    "Peripheral-Bones": [
 55        "carpal_left", "carpal_right", "clavicle_left", "clavicle_right", "femur_left", "femur_right",
 56        "fibula_left", "fibula_right", "fingers_left", "fingers_right", "humerus_left", "humerus_right",
 57        "metacarpal_left", "metacarpal_right", "metatarsal_left", "metatarsal_right", "patella_left",
 58        "patella_right", "radius_left", "radius_right", "scapula_left", "scapula_right", "skull", "tarsal_left",
 59        "tarsal_right", "tibia_left", "tibia_right", "toes_left", "toes_right", "ulna_left", "ulna_right",
 60    ],
 61    "Ribs": [f"rib_left_{i}" for i in range(1, 14)] + [f"rib_right_{i}" for i in range(1, 14)] + ["sternum"],
 62    "Vertebrae": (
 63        [f"vertebra_C{i}" for i in range(1, 8)] + [f"vertebra_T{i}" for i in range(1, 13)]
 64        + [f"vertebra_L{i}" for i in range(1, 7)] + ["hip_left", "hip_right", "sacrum"]
 65    ),
 66}
 67"""The anatomical structures of the ENHANCE.PET 1.6k dataset, grouped by their ground-truth category volume."""
 68
 69CATEGORIES = sorted(CLASS_NAMES_BY_CATEGORY.keys())
 70
 71CLASS_NAMES = [name for category in CATEGORIES for name in CLASS_NAMES_BY_CATEGORY[category]]
 72"""The anatomical structures of the ENHANCE.PET 1.6k dataset. The label id of a structure is its 1-based index."""
 73
 74CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 75"""Mapping from the name of an anatomical structure to its label id in the merged label volumes."""
 76
 77
 78def _list_case_ids(n_workers=None):
 79    import requests
 80    from xml.etree import ElementTree
 81
 82    case_ids, token = [], None
 83    while True:
 84        params = {"list-type": "2", "prefix": "imaging-data/images/CT/"}
 85        if token:
 86            params["continuation-token"] = token
 87        response = requests.get(BUCKET_URL, params=params)
 88        response.raise_for_status()
 89        root = ElementTree.fromstring(response.content)
 90        ns = {"s3": "http://s3.amazonaws.com/doc/2006-03-01/"}
 91        for content in root.findall("s3:Contents", ns):
 92            key = content.find("s3:Key", ns).text
 93            case_ids.append(os.path.basename(key).replace(".nii.gz", ""))
 94        is_truncated = root.find("s3:IsTruncated", ns).text == "true"
 95        if not is_truncated:
 96            break
 97        token = root.find("s3:NextContinuationToken", ns).text
 98
 99    return sorted(case_ids)
100
101
102def _merge_segmentations(path: str, case_id: str) -> str:
103    """Merge the 7 per-category label volumes of one case into a single semantic label volume."""
104    import nibabel as nib
105
106    label_path = os.path.join(path, "imaging-data", "labels", f"{case_id}.nii.gz")
107    if os.path.exists(label_path):
108        return label_path
109
110    labels, affine, header, offset = None, None, None, 0
111    for category in CATEGORIES:
112        category_path = os.path.join(path, "imaging-data", "ground-truth", category, f"{case_id}.nii.gz")
113        category_nii = nib.load(category_path)
114        category_labels = np.round(np.asarray(category_nii.dataobj)).astype("uint8")
115        if labels is None:
116            labels = np.zeros(category_labels.shape, dtype="uint8")
117            affine, header = category_nii.affine, category_nii.header
118        foreground = category_labels > 0
119        labels[foreground] = category_labels[foreground] + offset
120        offset += len(CLASS_NAMES_BY_CATEGORY[category])
121
122    os.makedirs(os.path.dirname(label_path), exist_ok=True)
123    tmp_path = label_path.replace(".nii.gz", ".incomplete.nii.gz")
124    nib.save(nib.Nifti1Image(labels, affine, header), tmp_path)
125    os.replace(tmp_path, label_path)
126    return label_path
127
128
129def _download_case(case_id: str, path: str) -> None:
130    image_path = os.path.join(path, "imaging-data", "images", "CT", f"{case_id}.nii.gz")
131    util.download_source(
132        path=image_path, url=f"{BUCKET_URL}/imaging-data/images/CT/{case_id}.nii.gz", download=True,
133    )
134    for category in CATEGORIES:
135        category_path = os.path.join(path, "imaging-data", "ground-truth", category, f"{case_id}.nii.gz")
136        util.download_source(
137            path=category_path,
138            url=f"{BUCKET_URL}/imaging-data/ground-truth/{category}/{case_id}.nii.gz", download=True,
139        )
140    _merge_segmentations(path, case_id)
141
142
143def get_enhance_pet_data(
144    path: Union[os.PathLike, str], n_cases: Optional[int] = None, n_workers: Optional[int] = None,
145    download: bool = False,
146) -> str:
147    """Download the ENHANCE.PET 1.6k dataset and merge the per-category masks into semantic label volumes.
148
149    NOTE: The full collection is about 250 GB. Use `n_cases` to only download a subset for a quick start.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        n_cases: The number of cases to download, sorted by case id. By default all ~1597 cases are downloaded.
154        n_workers: The number of parallel download / merging workers. By default the number of CPUs (at most
155            16) is used.
156        download: Whether to download the data if it is not present.
157
158    Returns:
159        Filepath where the data is downloaded.
160    """
161    os.makedirs(path, exist_ok=True)
162    case_ids = _list_case_ids()
163    if n_cases is not None:
164        case_ids = case_ids[:n_cases]
165
166    missing = [
167        case_id for case_id in case_ids
168        if not os.path.exists(os.path.join(path, "imaging-data", "labels", f"{case_id}.nii.gz"))
169    ]
170    if missing and not download:
171        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
172
173    if n_workers is None:
174        n_cpus = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)
175        n_workers = min(16, n_cpus)
176
177    with futures.ThreadPoolExecutor(n_workers) as pool:
178        tasks = [pool.submit(_download_case, case_id, path) for case_id in missing]
179        for task in tqdm(futures.as_completed(tasks), total=len(tasks), desc="Download ENHANCE.PET cases"):
180            task.result()
181
182    return path
183
184
185def get_enhance_pet_paths(
186    path: Union[os.PathLike, str], n_cases: Optional[int] = None, download: bool = False,
187) -> Tuple[List[str], List[str]]:
188    """Get paths to the ENHANCE.PET 1.6k data.
189
190    Args:
191        path: Filepath to a folder where the data is downloaded for further processing.
192        n_cases: The number of cases to use, sorted by case id. By default all downloaded cases are used.
193        download: Whether to download the data if it is not present.
194
195    Returns:
196        List of filepaths for the CT image data.
197        List of filepaths for the label data.
198    """
199    get_enhance_pet_data(path, n_cases, download=download)
200
201    raw_paths = sorted(glob(os.path.join(path, "imaging-data", "images", "CT", "*.nii.gz")))
202    label_paths = [
203        os.path.join(path, "imaging-data", "labels", os.path.basename(p)) for p in raw_paths
204    ]
205
206    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
207    assert all(os.path.exists(p) for p in label_paths)
208    return raw_paths, label_paths
209
210
211def get_enhance_pet_dataset(
212    path: Union[os.PathLike, str],
213    patch_shape: Tuple[int, int, int],
214    n_cases: Optional[int] = None,
215    resize_inputs: bool = False,
216    download: bool = False,
217    **kwargs
218) -> Dataset:
219    """Get the ENHANCE.PET 1.6k dataset for segmentation of anatomical structures in whole-body CT.
220
221    Args:
222        path: Filepath to a folder where the data is downloaded for further processing.
223        patch_shape: The patch shape to use for training.
224        n_cases: The number of cases to use, sorted by case id. By default all downloaded cases are used.
225        resize_inputs: Whether to resize the inputs to the patch shape.
226        download: Whether to download the data if it is not present.
227        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
228
229    Returns:
230        The segmentation dataset.
231    """
232    raw_paths, label_paths = get_enhance_pet_paths(path, n_cases, download)
233
234    if resize_inputs:
235        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
236        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
237            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
238        )
239
240    return torch_em.default_segmentation_dataset(
241        raw_paths=raw_paths,
242        raw_key="data",
243        label_paths=label_paths,
244        label_key="data",
245        is_seg_dataset=True,
246        patch_shape=patch_shape,
247        ndim=3,
248        **kwargs
249    )
250
251
252def get_enhance_pet_loader(
253    path: Union[os.PathLike, str],
254    batch_size: int,
255    patch_shape: Tuple[int, int, int],
256    n_cases: Optional[int] = None,
257    resize_inputs: bool = False,
258    download: bool = False,
259    **kwargs
260) -> DataLoader:
261    """Get the ENHANCE.PET 1.6k dataloader for segmentation of anatomical structures in whole-body CT.
262
263    Args:
264        path: Filepath to a folder where the data is downloaded for further processing.
265        batch_size: The batch size for training.
266        patch_shape: The patch shape to use for training.
267        n_cases: The number of cases to use, sorted by case id. By default all downloaded cases are used.
268        resize_inputs: Whether to resize the inputs to the patch shape.
269        download: Whether to download the data if it is not present.
270        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
271
272    Returns:
273        The DataLoader.
274    """
275    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
276    dataset = get_enhance_pet_dataset(path, patch_shape, n_cases, resize_inputs, download, **ds_kwargs)
277    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)