torch_em.data.datasets.medical.pants

  1"""PanTS (The Pancreatic Tumor Segmentation Dataset) contains annotations for pancreatic tumor and
  2sub-segment segmentation, together with 24 surrounding anatomical structures, in abdominal CT scans.
  3
  4The dataset consists of 36,390 CT volumes from 145 medical centers in total, of which 9,901 (case ids
  5'PanTS_00000001' to 'PanTS_00009901') form the public release used by this module: 9,000 training cases
  6('PanTS_00000001'-'PanTS_00009000') and 901 public in-distribution test cases
  7('PanTS_00009001'-'PanTS_00009901'). The remaining cases are proprietary external test sets and are not
  8released. Each case has a CT scan ('ct.nii.gz') and a combined semantic label volume ('combined_labels.nii.gz')
  9with the label ids described in `CLASS_IDS`, covering the pancreas and its sub-segments (head, body, tail),
 10the pancreatic duct, the pancreatic lesion (tumor), and 24 further abdominal/thoracic organs and vascular /
 11skeletal structures.
 12
 13The dataset is hosted in two parts:
 14- The CT scans are hosted, ungated, on HuggingFace at https://huggingface.co/datasets/BodyMaps/PanTSMini, split
 15  into 9 training shards of 1000 cases each (~30-36 GB per shard) and 1 test shard of 901 cases (~28 GB).
 16- The labels are hosted, ungated, as a single combined archive (~15.5 GB) at
 17  https://www.cs.jhu.edu/~zongwei/dataset/PanTSMini_Label.tar.gz, which covers all 9901 cases at once (unlike
 18  the per-range image shards).
 19Downloading a case therefore always downloads the whole image shard it belongs to; use `max_cases` or
 20`case_ids` to only fetch the shard(s) required for a small subset of cases. The label archive, however, is not
 21sharded, so fetching even a single case's label requires streaming through the (compressed) archive until that
 22case is found; the module only ever holds the matching labels in memory / on disk, but still has to receive
 23all bytes up to the last requested case, which in the worst case is the whole ~15.5 GB archive.
 24The data is licensed under CC BY-NC-SA 4.0.
 25
 26This dataset is from the publication https://doi.org/10.48550/arXiv.2507.01291.
 27Please cite it if you use this dataset in your research.
 28"""
 29
 30import os
 31import re
 32import tarfile
 33from glob import glob
 34from tqdm import tqdm
 35from natsort import natsorted
 36from typing import Union, Tuple, List, Optional, Literal
 37
 38import requests
 39
 40from torch.utils.data import Dataset, DataLoader
 41
 42import torch_em
 43
 44from .. import util
 45
 46
 47REPO_ID = "BodyMaps/PanTSMini"
 48LABEL_URL = "https://www.cs.jhu.edu/~zongwei/dataset/PanTSMini_Label.tar.gz"
 49
 50CLASS_NAMES = [
 51    "adrenal_gland_left", "adrenal_gland_right", "aorta", "bladder", "celiac_artery", "colon",
 52    "common_bile_duct", "duodenum", "femur_left", "femur_right", "gall_bladder", "kidney_left", "kidney_right",
 53    "liver", "lung_left", "lung_right", "pancreas", "pancreas_body", "pancreas_head", "pancreas_tail",
 54    "pancreatic_duct", "postcava", "prostate", "spleen", "stomach", "superior_mesenteric_artery", "veins",
 55    "pancreatic_lesion",
 56]
 57"""The anatomical structures of the PanTS dataset, in the order of their label id."""
 58
 59CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 60"""Mapping from the name of a structure to its label id in the combined label volumes."""
 61
 62LAST_TRAIN_CASE = 9000
 63LAST_TEST_CASE = 9901
 64
 65
 66def _case_name(number: int) -> str:
 67    return f"PanTS_{number:08d}"
 68
 69
 70def _case_number(case_name: str) -> int:
 71    return int(case_name.split("_")[-1])
 72
 73
 74def _find_case_dirs(image_dir):
 75    return natsorted(
 76        d for d in glob(os.path.join(image_dir, "PanTS_*")) if os.path.isfile(os.path.join(d, "ct.nii.gz"))
 77    )
 78
 79
 80def _shards_for_cases(files, case_numbers, split):
 81    tag = "ImageTr" if split == "train" else "ImageTe"
 82    shards = set()
 83    for fpath in files:
 84        if not fpath.endswith(".tar.gz"):
 85            continue
 86        match = re.search(rf"{tag}_(\d{{8}})_(\d{{8}})\.tar\.gz$", os.path.basename(fpath))
 87        if match is None:
 88            continue
 89        start, end = int(match.group(1)), int(match.group(2))
 90        if case_numbers is None or any(start <= number <= end for number in case_numbers):
 91            shards.add(fpath)
 92    return shards
 93
 94
 95def _download_images(path, split, max_cases, case_ids, download):
 96    image_dir = os.path.join(path, "ImageTr" if split == "train" else "ImageTe")
 97    last_case = LAST_TRAIN_CASE if split == "train" else LAST_TEST_CASE
 98    first_case = 1 if split == "train" else (LAST_TRAIN_CASE + 1)
 99
100    if case_ids is not None:
101        case_numbers = [_case_number(cid) for cid in case_ids]
102    elif max_cases is not None:
103        case_numbers = list(range(first_case, min(first_case + max_cases - 1, last_case) + 1))
104    else:
105        case_numbers = None
106
107    case_dirs = _find_case_dirs(image_dir) if os.path.exists(image_dir) else []
108    if case_dirs and case_numbers is not None:
109        have = {_case_number(os.path.basename(d)) for d in case_dirs}
110        if set(case_numbers).issubset(have):
111            return image_dir
112    elif case_dirs and case_numbers is None:
113        return image_dir
114
115    if not download:
116        raise RuntimeError(f"Cannot find the data at {image_dir}, but download was set to False")
117
118    from huggingface_hub import HfApi, snapshot_download
119
120    os.makedirs(path, exist_ok=True)
121    api = HfApi()
122    files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
123    shards = _shards_for_cases(files, case_numbers, split)
124    if not shards:
125        raise RuntimeError(f"Could not find any PanTSMini image shards for the '{split}' split.")
126
127    print(f"Downloading {len(shards)} PanTSMini image shard(s) for the '{split}' split.")
128    snapshot_download(repo_id=REPO_ID, repo_type="dataset", local_dir=path, allow_patterns=sorted(shards))
129
130    for tar_path in natsorted(glob(os.path.join(path, "*.tar.gz"))):
131        util.unzip_tarfile(tar_path=tar_path, dst=image_dir, remove=False)
132
133    case_dirs = _find_case_dirs(image_dir)
134    if not case_dirs:
135        raise RuntimeError(f"Could not find any 'PanTS_XXXXXXXX' case folders under '{image_dir}'.")
136
137    return image_dir
138
139
140def _extract_labels_for_cases(path, case_numbers):
141    train_dir = os.path.join(path, "LabelTr")
142    test_dir = os.path.join(path, "LabelTe")
143
144    def _missing(numbers):
145        missing = []
146        for number in numbers:
147            dst = os.path.join(train_dir if number <= LAST_TRAIN_CASE else test_dir, _case_name(number))
148            if not os.path.exists(os.path.join(dst, "combined_labels.nii.gz")):
149                missing.append(number)
150        return missing
151
152    if case_numbers is not None:
153        missing = _missing(case_numbers)
154        if not missing:
155            return train_dir, test_dir
156        wanted = set(missing)
157    else:
158        wanted = None
159
160    os.makedirs(train_dir, exist_ok=True)
161    os.makedirs(test_dir, exist_ok=True)
162
163    print("Streaming the PanTS label archive (~15.5 GB, covers all cases, cannot be sharded)...")
164    with requests.get(LABEL_URL, stream=True, allow_redirects=True) as r:
165        r.raise_for_status()
166        r.raw.decode_content = True
167        with tarfile.open(fileobj=r.raw, mode="r|gz") as tf:
168            found = set()
169            for member in tqdm(tf, desc="Extracting PanTS labels"):
170                top = member.name.split("/")[0]
171                if not top.startswith("PanTS_"):
172                    continue
173                number = _case_number(top)
174                if wanted is not None and number not in wanted:
175                    continue
176                dst = train_dir if number <= LAST_TRAIN_CASE else test_dir
177                tf.extract(member, dst)
178                if wanted is not None:
179                    found.add(number)
180                    if found == wanted:
181                        break
182
183    return train_dir, test_dir
184
185
186def get_pants_data(
187    path: Union[os.PathLike, str],
188    split: Literal["train", "test"],
189    max_cases: Optional[int] = None,
190    case_ids: Optional[List[str]] = None,
191    download: bool = False,
192) -> Tuple[str, str]:
193    """Download the PanTS dataset.
194
195    The dataset is ungated and does not require a HuggingFace account or access token. It is, however, very
196    large (~300 GB for the images alone), and a case can only be downloaded together with the full image shard
197    (of up to 1000 cases) it belongs to. Use `max_cases` or `case_ids` to only download the shard(s) required
198    for a small subset of cases; leave both at their default (None) to download the full split.
199
200    Args:
201        path: Filepath to a folder where the data is downloaded for further processing.
202        split: The choice of data split. Either 'train' (9000 cases) or 'test' (901 cases).
203        max_cases: The maximum number of cases to download, taken in order starting from the first case of
204            `split`. Only the shard(s) covering these cases are downloaded. Mutually exclusive with `case_ids`.
205        case_ids: Explicit list of case ids (eg. ['PanTS_00000001']) to download. Only the shard(s) covering
206            these cases are downloaded. Mutually exclusive with `max_cases`.
207        download: Whether to download the data if it is not present.
208
209    Returns:
210        Filepath to the folder with the image data.
211        Filepath to the folder with the label data.
212    """
213    if split not in ("train", "test"):
214        raise ValueError(f"'{split}' is not a valid split. Please choose one of 'train' or 'test'.")
215    assert max_cases is None or case_ids is None, "'max_cases' and 'case_ids' are mutually exclusive."
216
217    image_dir = _download_images(path, split, max_cases, case_ids, download)
218    case_dirs = _find_case_dirs(image_dir)
219    if case_ids is not None:
220        case_dirs = [d for d in case_dirs if os.path.basename(d) in case_ids]
221    elif max_cases is not None:
222        case_dirs = case_dirs[:max_cases]
223    case_numbers = [_case_number(os.path.basename(d)) for d in case_dirs]
224
225    if not download and not os.path.exists(os.path.join(path, "LabelTr" if split == "train" else "LabelTe")):
226        raise RuntimeError(f"Cannot find the label data at {path}, but download was set to False")
227
228    train_dir, test_dir = _extract_labels_for_cases(path, case_numbers if case_ids or max_cases else None)
229    label_dir = train_dir if split == "train" else test_dir
230
231    return image_dir, label_dir
232
233
234def get_pants_paths(
235    path: Union[os.PathLike, str],
236    split: Literal["train", "test"],
237    max_cases: Optional[int] = None,
238    case_ids: Optional[List[str]] = None,
239    download: bool = False,
240) -> Tuple[List[str], List[str]]:
241    """Get paths to the PanTS data.
242
243    Args:
244        path: Filepath to a folder where the data is downloaded for further processing.
245        split: The choice of data split. Either 'train' (9000 cases) or 'test' (901 cases).
246        max_cases: The maximum number of cases to use. See `get_pants_data` for details.
247        case_ids: Explicit list of case ids to use. See `get_pants_data` for details.
248        download: Whether to download the data if it is not present.
249
250    Returns:
251        List of filepaths for the image data.
252        List of filepaths for the label data.
253    """
254    image_dir, label_dir = get_pants_data(path, split, max_cases, case_ids, download)
255
256    case_dirs = _find_case_dirs(image_dir)
257    if case_ids is not None:
258        case_dirs = [d for d in case_dirs if os.path.basename(d) in case_ids]
259    elif max_cases is not None:
260        case_dirs = case_dirs[:max_cases]
261
262    raw_paths, label_paths = [], []
263    for case_dir in case_dirs:
264        case_name = os.path.basename(case_dir)
265        label_path = os.path.join(label_dir, case_name, "combined_labels.nii.gz")
266        if not os.path.exists(label_path):
267            continue
268        raw_paths.append(os.path.join(case_dir, "ct.nii.gz"))
269        label_paths.append(label_path)
270
271    if len(raw_paths) == 0 or len(raw_paths) != len(label_paths):
272        raise RuntimeError("Something went wrong with fetching the image and label paths.")
273
274    return raw_paths, label_paths
275
276
277def get_pants_dataset(
278    path: Union[os.PathLike, str],
279    patch_shape: Tuple[int, ...],
280    split: Literal["train", "test"],
281    max_cases: Optional[int] = None,
282    case_ids: Optional[List[str]] = None,
283    resize_inputs: bool = False,
284    download: bool = False,
285    **kwargs
286) -> Dataset:
287    """Get the PanTS dataset for pancreatic tumor and abdominal anatomy segmentation.
288
289    Args:
290        path: Filepath to a folder where the data is downloaded for further processing.
291        patch_shape: The patch shape to use for training.
292        split: The choice of data split. Either 'train' (9000 cases) or 'test' (901 cases).
293        max_cases: The maximum number of cases to use. See `get_pants_data` for details.
294        case_ids: Explicit list of case ids to use. See `get_pants_data` for details.
295        resize_inputs: Whether to resize inputs to the desired patch shape.
296        download: Whether to download the data if it is not present.
297        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
298
299    Returns:
300        The segmentation dataset.
301    """
302    raw_paths, label_paths = get_pants_paths(path, split, max_cases, case_ids, download)
303
304    if resize_inputs:
305        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
306        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
307            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
308        )
309
310    return torch_em.default_segmentation_dataset(
311        raw_paths=raw_paths,
312        raw_key="data",
313        label_paths=label_paths,
314        label_key="data",
315        patch_shape=patch_shape,
316        is_seg_dataset=True,
317        **kwargs
318    )
319
320
321def get_pants_loader(
322    path: Union[os.PathLike, str],
323    batch_size: int,
324    patch_shape: Tuple[int, ...],
325    split: Literal["train", "test"],
326    max_cases: Optional[int] = None,
327    case_ids: Optional[List[str]] = None,
328    resize_inputs: bool = False,
329    download: bool = False,
330    **kwargs
331) -> DataLoader:
332    """Get the PanTS dataloader for pancreatic tumor and abdominal anatomy segmentation.
333
334    Args:
335        path: Filepath to a folder where the data is downloaded for further processing.
336        batch_size: The batch size for training.
337        patch_shape: The patch shape to use for training.
338        split: The choice of data split. Either 'train' (9000 cases) or 'test' (901 cases).
339        max_cases: The maximum number of cases to use. See `get_pants_data` for details.
340        case_ids: Explicit list of case ids to use. See `get_pants_data` for details.
341        resize_inputs: Whether to resize inputs to the desired patch shape.
342        download: Whether to download the data if it is not present.
343        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
344
345    Returns:
346        The DataLoader.
347    """
348    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
349    dataset = get_pants_dataset(
350        path, patch_shape, split, max_cases, case_ids, resize_inputs, download, **ds_kwargs
351    )
352    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)