torch_em.data.datasets.medical.hecktor

  1"""The HECKTOR dataset contains annotations for head and neck tumor segmentation in FDG-PET/CT scans.
  2
  3It comprises the training set of the HECKTOR 2022 challenge (https://hecktor.grand-challenge.org):
  4524 PET/CT studies of patients with histologically proven oropharyngeal head and neck cancer, collected
  5at 7 centers, with annotations of the primary tumor and of the involved lymph nodes. The 359 test studies
  6of the challenge are distributed without annotations and are therefore not included here.
  7
  8Every case provides a co-registered pair of a low-dose non-contrast-enhanced CT scan and an FDG-PET scan
  9(converted to standardized uptake values). The 'modality' argument selects which of the two is used as
 10the raw input, the annotation is defined on the CT grid.
 11
 12NOTE: The label legend is as follows:
 13- background: 0, primary tumor (GTVp): 1, lymph nodes (GTVn): 2
 14This is documented on https://hecktor.grand-challenge.org/Data/ and in the challenge overview paper.
 15
 16NOTE: The dataset requires registration and cannot be downloaded automatically. Please follow these steps:
 17- Visit https://hecktor.grand-challenge.org/Data/, register for the challenge on grand-challenge.org and
 18  request access to the data. The organizers grant access after approval of the request and of the signed
 19  end user agreement. NOTE: At the time of writing, the organizers state on that page that the data is
 20  temporarily unavailable while they extend the agreement with the partner hospitals.
 21- Download 'hecktor2022_training.zip' from the data download page you are given access to.
 22- Extract it into '<path>', such that
 23  '<path>/hecktor2022_training/imagesTr/<case_id>__CT.nii.gz',
 24  '<path>/hecktor2022_training/imagesTr/<case_id>__PT.nii.gz' and
 25  '<path>/hecktor2022_training/labelsTr/<case_id>.nii.gz' exist.
 26  The case ids are of the form '<center>-<number>', e.g. 'CHUM-001' or 'MDA-042'. The 524 cases are
 27  distributed over the centers as CHUM: 56, CHUP: 72, CHUS: 72, CHUV: 53, HGJ: 55, HMR: 18, MDA: 198.
 28
 29NOTE: There is no openly published copy of the 2022 release. Third-party re-uploads of HECKTOR data exist,
 30but all of the ones we are aware of are a different edition of the challenge, which must not be confused
 31with this one: HECKTOR 2020 has 201 cases, HECKTOR 2021 has 224 training cases, HECKTOR 2025 has 679
 32training cases (it adds the center USZ and many more MDA cases) and HECKTOR 2026 has 883 cases. Only the
 332022 edition has exactly 524 training cases.
 34
 35The dataset is located at https://hecktor.grand-challenge.org/Data/.
 36
 37This dataset is from the publication https://doi.org/10.1007/978-3-031-27420-6_1.
 38Please cite it if you use this dataset in your research.
 39"""
 40
 41import os
 42from glob import glob
 43from natsort import natsorted
 44from typing import Union, Tuple, Literal, List
 45
 46from torch.utils.data import Dataset, DataLoader
 47
 48import torch_em
 49
 50from .. import util
 51
 52
 53MODALITIES = {"ct": "__CT.nii.gz", "pt": "__PT.nii.gz"}
 54
 55LABEL_IDS = {"background": 0, "primary_tumor": 1, "lymph_nodes": 2}
 56
 57
 58def _find_data_dir(path):
 59    candidates = [
 60        os.path.join(path, "hecktor2022_training"), path, *glob(os.path.join(path, "*", "hecktor2022_training")),
 61    ]
 62    for candidate in candidates:
 63        if len(glob(os.path.join(candidate, "imagesTr", "*__CT.nii.gz"))) > 0:
 64            return candidate
 65    return None
 66
 67
 68def get_hecktor_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 69    """Obtain the HECKTOR dataset.
 70
 71    Args:
 72        path: Filepath to a folder where the data is stored.
 73        download: Whether to download the data if it is not present.
 74
 75    Returns:
 76        Filepath where the data is stored.
 77    """
 78    data_dir = _find_data_dir(path)
 79    if data_dir is not None:
 80        return data_dir
 81
 82    msg = f"Could not find the HECKTOR 2022 training data at '{path}'. "
 83    msg += "'torch_em' cannot download this dataset, as it requires registration for the challenge. "
 84    msg += "Please register at https://hecktor.grand-challenge.org, request access to the data as described at "
 85    msg += "https://hecktor.grand-challenge.org/Data/ and download 'hecktor2022_training.zip'. Then extract it "
 86    msg += f"into '{path}', such that "
 87    msg += f"'{os.path.join(path, 'hecktor2022_training', 'imagesTr', 'CHUM-001__CT.nii.gz')}' exists."
 88    if download:
 89        raise NotImplementedError(msg)
 90    else:
 91        raise FileNotFoundError(msg)
 92
 93
 94def get_hecktor_paths(
 95    path: Union[os.PathLike, str], modality: Literal["ct", "pt"] = "ct", download: bool = False
 96) -> Tuple[List[str], List[str]]:
 97    """Get paths to the HECKTOR data.
 98
 99    Args:
100        path: Filepath to a folder where the data is stored.
101        modality: The imaging modality to use as raw input. Either 'ct' or 'pt'.
102        download: Whether to download the data if it is not present.
103
104    Returns:
105        List of filepaths for the image data.
106        List of filepaths for the label data.
107    """
108    if modality not in MODALITIES:
109        raise ValueError(f"'{modality}' is not a valid modality. Choose one of {list(MODALITIES.keys())}.")
110
111    data_dir = get_hecktor_data(path, download)
112
113    suffix = MODALITIES[modality]
114    raw_paths = natsorted(glob(os.path.join(data_dir, "imagesTr", f"*{suffix}")))
115    label_paths = [
116        os.path.join(data_dir, "labelsTr", os.path.basename(p).replace(suffix, ".nii.gz")) for p in raw_paths
117    ]
118    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
119
120    return raw_paths, label_paths
121
122
123def get_hecktor_dataset(
124    path: Union[os.PathLike, str],
125    patch_shape: Tuple[int, ...],
126    modality: Literal["ct", "pt"] = "ct",
127    resize_inputs: bool = False,
128    download: bool = False,
129    **kwargs
130) -> Dataset:
131    """Get the HECKTOR dataset for head and neck tumor segmentation.
132
133    Args:
134        path: Filepath to a folder where the data is stored.
135        patch_shape: The patch shape to use for training.
136        modality: The imaging modality to use as raw input. Either 'ct' or 'pt'.
137        resize_inputs: Whether to resize inputs to the desired patch shape.
138        download: Whether to download the data if it is not present.
139        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
140
141    Returns:
142        The segmentation dataset.
143    """
144    raw_paths, label_paths = get_hecktor_paths(path, modality, download)
145
146    if resize_inputs:
147        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
148        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
149            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
150        )
151
152    return torch_em.default_segmentation_dataset(
153        raw_paths=raw_paths,
154        raw_key="data",
155        label_paths=label_paths,
156        label_key="data",
157        patch_shape=patch_shape,
158        is_seg_dataset=True,
159        **kwargs
160    )
161
162
163def get_hecktor_loader(
164    path: Union[os.PathLike, str],
165    batch_size: int,
166    patch_shape: Tuple[int, ...],
167    modality: Literal["ct", "pt"] = "ct",
168    resize_inputs: bool = False,
169    download: bool = False,
170    **kwargs
171) -> DataLoader:
172    """Get the HECKTOR dataloader for head and neck tumor segmentation.
173
174    Args:
175        path: Filepath to a folder where the data is stored.
176        batch_size: The batch size for training.
177        patch_shape: The patch shape to use for training.
178        modality: The imaging modality to use as raw input. Either 'ct' or 'pt'.
179        resize_inputs: Whether to resize inputs to the desired patch shape.
180        download: Whether to download the data if it is not present.
181        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
182
183    Returns:
184        The DataLoader.
185    """
186    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
187    dataset = get_hecktor_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
188    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)