torch_em.data.datasets.medical.autopet_organ

  1"""The AutoPET-Organ dataset contains annotations for 11 organs in whole-body FDG-PET scans.
  2
  3The dataset consists of 100 studies of the AutoPET collection with organ annotations that were added and
  4expert-examined for the SegAnyPET publication, so the annotations are downloaded from there and the
  5scans are taken from the existing autopet loader. The label ids are 1: liver, 2: kidney (left),
  63: kidney (right), 4: heart, 5: spleen, 6: aorta, 7: lung lower lobe (left), 8: lung lower lobe (right),
  79: lung upper lobe (left), 10: lung upper lobe (right), 11: lung middle lobe (right).
  8See also `CLASS_IDS`.
  9
 10NOTE: The release stores the right middle lobe under two ids, 11 and 12, which are perfectly
 11complementary: of the 100 studies, 64 use id 11, 36 use id 12, and no study uses both or neither. The
 12two carry the same structure, matching in side, in position relative to the other organs of the same
 13study, and in size. 'harmonize' maps id 12 onto id 11 so that one organ has one id, which is the
 14default because the raw ids split one class in two and make a broken training target.
 15
 16NOTE: The source describes a prostate annotation, but no study of this release contains a twelfth
 17structure once ids 11 and 12 are merged, so the annotations cover 11 organs rather than 12.
 18
 19NOTE: The label ids were assigned by measuring each annotation in the data, since the source lists the
 20organs without their ids. Every study is stored in LPS orientation, and the ids follow from the side and
 21the position of each annotation relative to the other organs of the same study.
 22
 23The annotations are located at https://github.com/YichiZhang98/SegAnyPET and the scans at
 24https://autopet.grand-challenge.org.
 25This dataset is from the publication https://doi.org/10.48550/arXiv.2502.14351.
 26Please cite it, and the AutoPET publication https://doi.org/10.1038/s41597-022-01718-3, if you use this
 27dataset in your research.
 28"""
 29
 30import os
 31import re
 32from glob import glob
 33from natsort import natsorted
 34from typing import Union, Optional, Tuple, List
 35
 36import numpy as np
 37
 38from torch.utils.data import Dataset, DataLoader
 39
 40import torch_em
 41
 42from .autopet import get_autopet_data
 43from .. import util
 44
 45
 46URL = "https://github.com/YichiZhang98/SegAnyPET/raw/main/AutoPET-OrganlabelsTr.zip"
 47
 48CHECKSUM = "8794137a0f8df0bc192aa8bb98009b3bc706e4591cbe8f289f8d38e1c5574142"
 49
 50CLASS_NAMES = [
 51    "liver", "kidney_left", "kidney_right", "heart", "spleen", "aorta",
 52    "lung_lower_lobe_left", "lung_lower_lobe_right", "lung_upper_lobe_left", "lung_upper_lobe_right",
 53    "lung_middle_lobe_right",
 54]
 55"""The organs of the AutoPET-Organ dataset. The label id of an organ is its 1-based index."""
 56
 57CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 58"""Mapping from the organ name to its label id."""
 59
 60MIDDLE_LOBE_ALIAS = 12
 61"""The second id that the release uses for the right middle lobe, which `harmonize` maps onto id 11."""
 62
 63
 64def _prepare_labels(label_path, out_path, harmonize):
 65    """Store the annotation under a filename without a dot in its stem, optionally harmonizing its ids.
 66
 67    Two of the study names abbreviate 'nativ und mit' as 'nativ u. mit', and the dot of that abbreviation
 68    makes the '.nii.gz' extension unrecognizable to the file readers.
 69    """
 70    import nibabel as nib
 71
 72    if os.path.exists(out_path):
 73        return out_path
 74
 75    image = nib.load(label_path)
 76    labels = np.asarray(image.dataobj)
 77    if harmonize:
 78        labels = np.where(labels == MIDDLE_LOBE_ALIAS, CLASS_IDS["lung_middle_lobe_right"], labels)
 79
 80    nib.save(nib.Nifti1Image(labels.astype("uint8"), image.affine, image.header), out_path)
 81    return out_path
 82
 83
 84def get_autopet_organ_data(
 85    path: Union[os.PathLike, str], autopet_path: Optional[Union[os.PathLike, str]] = None, download: bool = False
 86) -> Tuple[str, str]:
 87    """Download the AutoPET-Organ dataset.
 88
 89    Args:
 90        path: Filepath to a folder where the data is downloaded for further processing.
 91        autopet_path: Filepath to an existing AutoPET download. The scans are downloaded to `path` if it
 92            is not given, which takes several hundred gigabytes.
 93        download: Whether to download the data if it is not present.
 94
 95    Returns:
 96        Filepath where the annotations are downloaded.
 97        Filepath where the scans are downloaded.
 98    """
 99    label_dir = os.path.join(path, "AutoPET-OrganlabelsTr")
100    if not os.path.exists(label_dir):
101        os.makedirs(path, exist_ok=True)
102        zip_path = os.path.join(path, "AutoPET-OrganlabelsTr.zip")
103        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
104        util.unzip(zip_path=zip_path, dst=path, remove=False)
105
106    # The scans are the whole-body PET studies of the AutoPET collection, which are several hundred
107    # gigabytes, so an existing download of them is used when it is given.
108    autopet_path = path if autopet_path is None else autopet_path
109    get_autopet_data(path=autopet_path, download=download)
110    image_dir = os.path.join(autopet_path, "AutoPET-II", "FDG-PET-CT-Lesions")
111
112    return label_dir, image_dir
113
114
115def get_autopet_organ_paths(
116    path: Union[os.PathLike, str],
117    harmonize: bool = True,
118    autopet_path: Optional[Union[os.PathLike, str]] = None,
119    download: bool = False,
120) -> Tuple[List[str], List[str]]:
121    """Get paths to the AutoPET-Organ data.
122
123    Args:
124        path: Filepath to a folder where the data is downloaded for further processing.
125        harmonize: Whether to map the second id of the right middle lobe onto its first one.
126        autopet_path: Filepath to an existing AutoPET download.
127        download: Whether to download the data if it is not present.
128
129    Returns:
130        List of filepaths for the image data.
131        List of filepaths for the label data.
132    """
133    label_dir, image_dir = get_autopet_organ_data(path, autopet_path, download)
134
135    prepared_dir = os.path.join(path, "harmonized" if harmonize else "prepared")
136    os.makedirs(prepared_dir, exist_ok=True)
137
138    raw_paths, label_paths = [], []
139    for label_path in natsorted(glob(os.path.join(label_dir, "*.nii.gz"))):
140        # The annotations are named '<patient id>_<study id>_.nii.gz' for the study folder they belong to.
141        stem = os.path.basename(label_path)[:-len(".nii.gz")].rstrip("_")
142        match = re.match(r"^(PETCT_[0-9a-f]+)_(.+)$", stem)
143        if match is None:
144            continue
145
146        image_path = os.path.join(image_dir, match.group(1), match.group(2), "SUV.nii.gz")
147        if not os.path.exists(image_path):
148            continue
149
150        out_path = os.path.join(prepared_dir, f"{stem.replace('.', '')}.nii.gz")
151        raw_paths.append(image_path)
152        label_paths.append(_prepare_labels(label_path, out_path, harmonize))
153
154    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
155
156    return raw_paths, label_paths
157
158
159def get_autopet_organ_dataset(
160    path: Union[os.PathLike, str],
161    patch_shape: Tuple[int, ...],
162    harmonize: bool = True,
163    autopet_path: Optional[Union[os.PathLike, str]] = None,
164    resize_inputs: bool = False,
165    download: bool = False,
166    **kwargs
167) -> Dataset:
168    """Get the AutoPET-Organ dataset for organ segmentation in whole-body PET.
169
170    Args:
171        path: Filepath to a folder where the data is downloaded for further processing.
172        patch_shape: The patch shape to use for training.
173        harmonize: Whether to map the second id of the right middle lobe onto its first one.
174        autopet_path: Filepath to an existing AutoPET download.
175        resize_inputs: Whether to resize inputs to the desired patch shape.
176        download: Whether to download the data if it is not present.
177        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
178
179    Returns:
180        The segmentation dataset.
181    """
182    raw_paths, label_paths = get_autopet_organ_paths(path, harmonize, autopet_path, download)
183
184    if resize_inputs:
185        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
186        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
187            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
188        )
189
190    return torch_em.default_segmentation_dataset(
191        raw_paths=raw_paths,
192        raw_key="data",
193        label_paths=label_paths,
194        label_key="data",
195        patch_shape=patch_shape,
196        is_seg_dataset=True,
197        **kwargs
198    )
199
200
201def get_autopet_organ_loader(
202    path: Union[os.PathLike, str],
203    batch_size: int,
204    patch_shape: Tuple[int, ...],
205    harmonize: bool = True,
206    autopet_path: Optional[Union[os.PathLike, str]] = None,
207    resize_inputs: bool = False,
208    download: bool = False,
209    **kwargs
210) -> DataLoader:
211    """Get the AutoPET-Organ dataloader for organ segmentation in whole-body PET.
212
213    Args:
214        path: Filepath to a folder where the data is downloaded for further processing.
215        batch_size: The batch size for training.
216        patch_shape: The patch shape to use for training.
217        harmonize: Whether to map the second id of the right middle lobe onto its first one.
218        autopet_path: Filepath to an existing AutoPET download.
219        resize_inputs: Whether to resize inputs to the desired patch shape.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
222
223    Returns:
224        The DataLoader.
225    """
226    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
227    dataset = get_autopet_organ_dataset(
228        path, patch_shape, harmonize, autopet_path, resize_inputs, download, **ds_kwargs
229    )
230    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)