torch_em.data.datasets.medical.pengwin

  1"""The PENGWIN dataset contains annotation for pelvic bone fracture and
  2fragments in CT and X-Ray images.
  3
  4The original (v1) release is from the challenge: https://pengwin.grand-challenge.org/pengwin/.
  5This release is related to the publication: https://doi.org/10.1007/978-3-031-43996-4_30.
  6
  7The updated (v2) release, PENGWIN26, is from the challenge: https://pengwin2026.grand-challenge.org/.
  8It ships 340 labeled real clinical CT cases with per-fragment instance ids, collected across seven
  9institutions: https://zenodo.org/records/19732767 (DOI: https://doi.org/10.5281/zenodo.19732767).
 10The simulated fracture cases for the separate reduction planning task are not covered by this loader.
 11
 12Please cite the respective publication(s) if you use this dataset for your research.
 13"""
 14
 15import os
 16from glob import glob
 17from natsort import natsorted
 18from typing import Union, Tuple, Literal, List
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URLS = {
 28    "CT": [
 29        "https://zenodo.org/records/10927452/files/PENGWIN_CT_train_images_part1.zip",  # inputs part 1
 30        "https://zenodo.org/records/10927452/files/PENGWIN_CT_train_images_part2.zip",  # inputs part 2
 31        "https://zenodo.org/records/10927452/files/PENGWIN_CT_train_labels.zip",  # labels
 32    ],
 33    "X-Ray": ["https://zenodo.org/records/10913196/files/train.zip"]
 34}
 35
 36CHECKSUMS = {
 37    "CT": [
 38        "e2e9f99798960607ffced1fbdeee75a626c41bf859eaf4125029a38fac6b7609",  # inputs part 1
 39        "19f3cdc5edd1daf9324c70f8ba683eed054f6ed8f2b1cc59dbd80724f8f0bbb2",  # inputs part 2
 40        "c4d3857e02d3ee5d0df6c8c918dd3cf5a7c9419135f1ec089b78215f37c6665c"  # labels
 41    ],
 42    "X-Ray": ["48d107979eb929a3c61da4e75566306a066408954cf132907bda570f2a7de725"]
 43}
 44
 45TARGET_DIRS = {
 46    "CT": ["CT/images", "CT/images", "CT/labels"],
 47    "X-Ray": ["X-Ray"]
 48}
 49
 50MODALITIES = ["CT", "X-Ray"]
 51
 52URLS_V2 = {
 53    "CT": [
 54        "https://zenodo.org/records/19732767/files/PENGWIN26_task1_2_train_part1.zip",
 55        "https://zenodo.org/records/19732767/files/PENGWIN26_task1_2_train_part2.zip",
 56        "https://zenodo.org/records/19732767/files/PENGWIN26_task1_2_train_part3.zip",
 57        "https://zenodo.org/records/19732767/files/PENGWIN26_task1_2_train_part4.zip",
 58    ]
 59}
 60
 61CHECKSUMS_V2 = {
 62    "CT": [
 63        "a9b047c7796164a4cb47f2083b67da2502070dfc7241fe42ec291a065dae0b19",
 64        "71aa69cb5620c9a5d4186de0c0e25f7bfe491b1123b4be3d18b62b4b9c267213",
 65        "26a7055993730d57f1365d3ab5e001841a048f57baa398f178a2d1298b2db636",
 66        "ff79294be9eb1827d8042c3171c460b829252e5a33c689ba856284c31518fa5d",
 67    ]
 68}
 69
 70
 71def get_pengwin_data(
 72    path: Union[os.PathLike, str],
 73    modality: Literal["CT", "X-Ray"],
 74    download: bool = False,
 75    version: Literal["v1", "v2"] = "v1",
 76) -> str:
 77    """Download the PENGWIN dataset.
 78
 79    Args:
 80        path: Filepath to a folder where the data is downloaded for further processing.
 81        modality: The choice of modality for inputs.
 82        download: Whether to download the data if it is not present.
 83        version: The version of the dataset. Either 'v1' (original PENGWIN release, CT and X-Ray) or
 84            'v2' (PENGWIN26, 340 labeled real clinical CT cases from seven institutions).
 85
 86    Returns:
 87        Filepath where the data is downloaded.
 88    """
 89    if not isinstance(modality, str) and modality in MODALITIES:
 90        raise ValueError(f"'{modality}' is not a valid modality. Please choose from {MODALITIES}.")
 91
 92    if version == "v2" and modality != "CT":
 93        raise ValueError("The 'v2' (PENGWIN26) release only provides the 'CT' modality.")
 94
 95    data_dir = os.path.join(path, "data")
 96    exists_dir = os.path.join(data_dir, modality) if version == "v1" else os.path.join(data_dir, modality, version)
 97    if os.path.exists(exists_dir):
 98        return data_dir
 99
100    os.makedirs(path, exist_ok=True)
101
102    if version == "v1":
103        for url, checksum, dst_dir in zip(URLS[modality], CHECKSUMS[modality], TARGET_DIRS[modality]):
104            zip_path = os.path.join(path, os.path.split(url)[-1])
105            util.download_source(path=zip_path, url=url, download=download, checksum=checksum)
106            util.unzip(zip_path=zip_path, dst=os.path.join(data_dir, dst_dir))
107    else:  # v2
108        dst_dir = os.path.join(data_dir, modality, version)
109        for url, checksum in zip(URLS_V2[modality], CHECKSUMS_V2[modality]):
110            zip_path = os.path.join(path, os.path.split(url)[-1])
111            util.download_source(path=zip_path, url=url, download=download, checksum=checksum)
112            util.unzip(zip_path=zip_path, dst=dst_dir)
113
114    return data_dir
115
116
117def get_pengwin_paths(
118    path: Union[os.PathLike, str],
119    modality: Literal["CT", "X-Ray"],
120    download: bool = False,
121    version: Literal["v1", "v2"] = "v1",
122) -> Tuple[List[str], List[str]]:
123    """Get paths to the PENGWIN data.
124
125    Args:
126        path: Filepath to a folder where the data is downloaded for further processing.
127        modality: The choice of modality for inputs.
128        download: Whether to download the data if it is not present.
129        version: The version of the dataset. Either 'v1' (original PENGWIN release, CT and X-Ray) or
130            'v2' (PENGWIN26, 340 labeled real clinical CT cases from seven institutions).
131
132    Returns:
133        List of filepaths for the image data.
134        List of filepaths for the label data.
135    """
136    data_dir = get_pengwin_data(path=path, modality=modality, download=download, version=version)
137
138    if version == "v2":
139        base_dir = os.path.join(data_dir, modality, version)
140        image_paths = natsorted(glob(os.path.join(base_dir, "**", "image.mha"), recursive=True))
141        gt_paths = natsorted(glob(os.path.join(base_dir, "**", "label.mha"), recursive=True))
142    elif modality == "CT":
143        image_paths = natsorted(glob(os.path.join(data_dir, modality, "images", "*.mha")))
144        gt_paths = natsorted(glob(os.path.join(data_dir, modality, "labels", "*.mha")))
145    else:  # X-Ray
146        base_dir = os.path.join(data_dir, modality, "train")
147        image_paths = natsorted(glob(os.path.join(base_dir, "input", "images", "*.tif")))
148        gt_paths = natsorted(glob(os.path.join(base_dir, "output", "images", "*.tif")))
149
150    return image_paths, gt_paths
151
152
153def get_pengwin_dataset(
154    path: Union[os.PathLike, str],
155    patch_shape: Tuple[int, ...],
156    modality: Literal["CT", "X-Ray"],
157    resize_inputs: bool = False,
158    download: bool = False,
159    version: Literal["v1", "v2"] = "v1",
160    **kwargs
161) -> Dataset:
162    """Get the PENGWIN dataset for pelvic fracture segmentation.
163
164    Args:
165        path: Filepath to a folder where the data is downloaded for further processing.
166        patch_shape: The patch shape to use for training.
167        modality: The choice of modality for inputs.
168        resize_inputs: Whether to resize inputs to the desired patch shape.
169        download: Whether to download the data if it is not present.
170        version: The version of the dataset. Either 'v1' (original PENGWIN release, CT and X-Ray) or
171            'v2' (PENGWIN26, 340 labeled real clinical CT cases from seven institutions).
172        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
173
174    Returns:
175        The segmentation dataset.
176    """
177    image_paths, gt_paths = get_pengwin_paths(path=path, modality=modality, download=download, version=version)
178
179    if resize_inputs:
180        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
181        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
182            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
183        )
184
185    return torch_em.default_segmentation_dataset(
186        raw_paths=image_paths,
187        raw_key=None,
188        label_paths=gt_paths,
189        label_key=None,
190        patch_shape=patch_shape,
191        **kwargs
192    )
193
194
195def get_pengwin_loader(
196    path: Union[os.PathLike, str],
197    batch_size: int,
198    patch_shape: Tuple[int, ...],
199    modality: Literal["CT", "X-Ray"],
200    resize_inputs: bool = False,
201    download: bool = False,
202    version: Literal["v1", "v2"] = "v1",
203    **kwargs
204) -> DataLoader:
205    """Get the PENGWIN dataloader for pelvic fracture segmentation.
206
207    Args:
208        path: Filepath to a folder where the data is downloaded for further processing.
209        batch_size: The batch size for training.
210        patch_shape: The patch shape to use for training.
211        modality: The choice of modality for inputs.
212        resize_inputs: Whether to resize inputs to the desired patch shape.
213        download: Whether to download the data if it is not present.
214        version: The version of the dataset. Either 'v1' (original PENGWIN release, CT and X-Ray) or
215            'v2' (PENGWIN26, 340 labeled real clinical CT cases from seven institutions).
216        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
217
218    Returns:
219        The DataLoader.
220    """
221    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
222    dataset = get_pengwin_dataset(path, patch_shape, modality, resize_inputs, download, version, **ds_kwargs)
223    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)