torch_em.data.datasets.medical.pantrack

  1"""PanTrack is a longitudinal CT benchmark for pancreatic cancer lesion segmentation and tracking.
  2
  3It consists of 161 portal-venous phase CT scans from 45 patients with pancreatic adenocarcinoma
  4(2-11 scans per patient), collected at a single institution. The scans are annotated with instance
  5segmentation masks for 292 lesion instances (pancreatic tumors and hepatic metastases), of which 116
  6consecutive baseline-to-follow-up pairs form the longitudinal tracking benchmark (see 'tracking.json'
  7in the downloaded data for the lesion correspondences across timepoints).
  8
  9NOTE: The label legend is as follows: background: 0, non-zero integer: a lesion instance id. The
 10instance ids are non-consecutive, but consistent across all timepoints of the same patient, i.e. the
 11same lesion carries the same instance id at every scan it appears in. Not all liver lesions are
 12annotated: only hepatic metastases deemed relevant by the radiologist are included.
 13
 14The dataset is located at https://huggingface.co/datasets/mrokuss/PanTrack (CC BY-NC 4.0).
 15This dataset is from the publication https://arxiv.org/abs/2605.23118.
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20from glob import glob
 21from natsort import natsorted
 22from typing import Union, Tuple, List, Literal, Optional
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31HF_REPO = "mrokuss/PanTrack"
 32
 33
 34def get_pantrack_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the PanTrack dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the data is downloaded for further processing.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        Filepath where the data is stored.
 43    """
 44    if os.path.exists(path):
 45        return path
 46
 47    if not download:
 48        raise RuntimeError(f"Cannot find the data at '{path}', but download was set to False.")
 49
 50    try:
 51        from huggingface_hub import snapshot_download
 52    except ImportError:
 53        raise ImportError("'huggingface_hub' is required to download PanTrack. Install it via conda/pip.")
 54
 55    os.makedirs(path, exist_ok=True)
 56    snapshot_download(repo_id=HF_REPO, repo_type="dataset", local_dir=path)
 57
 58    return path
 59
 60
 61def get_pantrack_paths(
 62    path: Union[os.PathLike, str], split: Optional[Literal["train", "val"]] = None, download: bool = False,
 63) -> Tuple[List[str], List[str]]:
 64    """Get paths to the PanTrack data.
 65
 66    Args:
 67        path: Filepath to a folder where the data is downloaded for further processing.
 68        split: The choice of data split. Either 'train' or 'val', following the recommended split shipped
 69            with the dataset. If None, all scans are used.
 70        download: Whether to download the data if it is not present.
 71
 72    Returns:
 73        List of filepaths for the image data.
 74        List of filepaths for the label data.
 75    """
 76    data_dir = get_pantrack_data(path, download)
 77
 78    raw_paths = natsorted(glob(os.path.join(data_dir, "images", "*_0000.nii.gz")))
 79
 80    if split is not None:
 81        import json
 82
 83        if split not in ("train", "val"):
 84            raise ValueError(f"'{split}' is not a valid split.")
 85
 86        with open(os.path.join(data_dir, "recommended_split.json")) as f:
 87            split_ids = set(json.load(f)[0][split])
 88
 89        raw_paths = [p for p in raw_paths if os.path.basename(p).replace("_0000.nii.gz", "") in split_ids]
 90
 91    label_paths = [
 92        os.path.join(data_dir, "labels", os.path.basename(p).replace("_0000.nii.gz", ".nii.gz")) for p in raw_paths
 93    ]
 94    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 95
 96    return raw_paths, label_paths
 97
 98
 99def get_pantrack_dataset(
100    path: Union[os.PathLike, str],
101    patch_shape: Tuple[int, ...],
102    split: Optional[Literal["train", "val"]] = None,
103    resize_inputs: bool = False,
104    download: bool = False,
105    **kwargs
106) -> Dataset:
107    """Get the PanTrack dataset for pancreatic cancer lesion segmentation.
108
109    Args:
110        path: Filepath to a folder where the data is downloaded for further processing.
111        patch_shape: The patch shape to use for training.
112        split: The choice of data split. Either 'train' or 'val'. If None, all scans are used.
113        resize_inputs: Whether to resize inputs to the desired patch shape.
114        download: Whether to download the data if it is not present.
115        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
116
117    Returns:
118        The segmentation dataset.
119    """
120    raw_paths, label_paths = get_pantrack_paths(path, split, download)
121
122    if resize_inputs:
123        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
124        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
125            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
126        )
127
128    return torch_em.default_segmentation_dataset(
129        raw_paths=raw_paths,
130        raw_key="data",
131        label_paths=label_paths,
132        label_key="data",
133        patch_shape=patch_shape,
134        is_seg_dataset=True,
135        **kwargs
136    )
137
138
139def get_pantrack_loader(
140    path: Union[os.PathLike, str],
141    batch_size: int,
142    patch_shape: Tuple[int, ...],
143    split: Optional[Literal["train", "val"]] = None,
144    resize_inputs: bool = False,
145    download: bool = False,
146    **kwargs
147) -> DataLoader:
148    """Get the PanTrack dataloader for pancreatic cancer lesion segmentation.
149
150    Args:
151        path: Filepath to a folder where the data is downloaded for further processing.
152        batch_size: The batch size for training.
153        patch_shape: The patch shape to use for training.
154        split: The choice of data split. Either 'train' or 'val'. If None, all scans are used.
155        resize_inputs: Whether to resize inputs to the desired patch shape.
156        download: Whether to download the data if it is not present.
157        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
158
159    Returns:
160        The DataLoader.
161    """
162    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
163    dataset = get_pantrack_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
164    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)