torch_em.data.datasets.histopathology.gleason2019

  1"""The Gleason2019 dataset contains pixel-level Gleason grade annotations for H&E
  2stained prostate cancer tissue microarray (TMA) cores.
  3
  4Mask label values: 0 (benign), 1 (Gleason pattern 3), 2 (Gleason pattern 4),
  53 (Gleason pattern 5), 4 (unlabelled).
  6
  7The test split (TMA 80) carries two independent pathologist annotations.
  8
  9This dataset is located at https://doi.org/10.7910/DVN/OCYCMP.
 10This dataset is from the publication https://doi.org/10.1038/s41598-018-30535-1.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 17from typing import Union, Literal, Tuple, List
 18
 19import numpy as np
 20import imageio.v3 as imageio
 21from PIL import Image
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30RAW_URLS = {
 31    "ZT111_4_A": "https://dataverse.harvard.edu/api/access/datafile/3201629",
 32    "ZT111_4_B": "https://dataverse.harvard.edu/api/access/datafile/3201630",
 33    "ZT111_4_C": "https://dataverse.harvard.edu/api/access/datafile/3201631",
 34    "ZT199_1_A": "https://dataverse.harvard.edu/api/access/datafile/3201632",
 35    "ZT199_1_B": "https://dataverse.harvard.edu/api/access/datafile/3201633",
 36    "ZT204_6_A": "https://dataverse.harvard.edu/api/access/datafile/3201634",
 37    "ZT204_6_B": "https://dataverse.harvard.edu/api/access/datafile/3201635",
 38    "ZT76_39_A": "https://dataverse.harvard.edu/api/access/datafile/3201623",
 39    "ZT76_39_B": "https://dataverse.harvard.edu/api/access/datafile/3201625",
 40    "ZT80_38_A": "https://dataverse.harvard.edu/api/access/datafile/3201626",
 41    "ZT80_38_B": "https://dataverse.harvard.edu/api/access/datafile/3201627",
 42    "ZT80_38_C": "https://dataverse.harvard.edu/api/access/datafile/3201628",
 43}
 44
 45RAW_CHECKSUMS = {
 46    "ZT111_4_A": "4e87448fe2db959a757c792069df5d49aaf08f305484a829832cb2d44e60fba0",
 47    "ZT111_4_B": "6e898136490c5fc4d46bb66541519a4bdfe77bcb7fcc7d2f39440195e8fe01ce",
 48    "ZT111_4_C": "85eeeeefd89f55d4aa12ce498f573c55f2ed83bf957b73a736a446cdb7540b45",
 49    "ZT199_1_A": "2e1f4cc38097d36be6d33c76688ec85df87c7f53812b205e34c1d5c47a720b45",
 50    "ZT199_1_B": "fa2d79e4891c5043f4dbb3db92e8d75de8d2f7c9afdf51d48c80dfb4f1d0b507",
 51    "ZT204_6_A": "ad02fe27bbdae2d2af0bf4e9725d8dabe1c15bdb49125e9e0131705401cdb610",
 52    "ZT204_6_B": "b89d3b8422cb1c2a9f3034d4f146ba0d3ce82260ab512c5ea5fc983506205c14",
 53    "ZT76_39_A": "054be24c8522bdacef2916a48e2463da238e39efbb96c4aae6e2f802168cc14f",
 54    "ZT76_39_B": "9d18c74ea72e6936a7ec6481373191b0c3da4f1a915280439b25efa7bcb645dc",
 55    "ZT80_38_A": "3ae19e2a3b00c4158171487da7cab10f05cbaddbb256a41ad6be73f4be091972",
 56    "ZT80_38_B": "406211b0822d165555346ae459db6ddb1daa32c9257f1578aeb7c2da03f3c910",
 57    "ZT80_38_C": "16bda74826ac8c907c85460b72f713f4f63d765167c18b2729743eed7e806048",
 58}
 59
 60MASK_URLS = {
 61    "train": "https://dataverse.harvard.edu/api/access/datafile/3201636",
 62    "test_pathologist1": "https://dataverse.harvard.edu/api/access/datafile/3201654",
 63    "test_pathologist2": "https://dataverse.harvard.edu/api/access/datafile/3201655",
 64}
 65
 66MASK_CHECKSUMS = {
 67    "train": "f4a3ce4a599b210d60cddc6ad3cd8d28d58546fbdb82ee59ba52fc66e90698a6",
 68    "test_pathologist1": "265c31f163072c7cdab9bcd850748b9196e79ac68497449dad542a3095f388ff",
 69    "test_pathologist2": "2c996b8479aad8ff442ec0e23c80fec4aa9cc39961751252145058266e1267a2",
 70}
 71
 72SPLIT_TMAS = {
 73    "train": ["ZT111_4_A", "ZT111_4_B", "ZT111_4_C", "ZT199_1_A", "ZT199_1_B", "ZT204_6_A", "ZT204_6_B"],
 74    "val": ["ZT76_39_A", "ZT76_39_B"],
 75    "test": ["ZT80_38_A", "ZT80_38_B", "ZT80_38_C"],
 76}
 77
 78
 79def get_gleason2019_data(
 80    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False,
 81) -> str:
 82    """Download the Gleason2019 data.
 83
 84    Args:
 85        path: Filepath to a folder where the downloaded data will be saved.
 86        split: The choice of data split.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        Filepath where the dataset is downloaded and stored for further preprocessing.
 91    """
 92    os.makedirs(path, exist_ok=True)
 93
 94    for stem in SPLIT_TMAS[split]:
 95        if os.path.exists(os.path.join(path, stem)):
 96            continue
 97        tar_path = os.path.join(path, f"{stem}.tar.gz")
 98        util.download_source(path=tar_path, url=RAW_URLS[stem], download=download, checksum=RAW_CHECKSUMS[stem])
 99        util.unzip_tarfile(tar_path=tar_path, dst=path)
100
101    mask_key = "train" if split in ("train", "val") else "test_pathologist1"
102    if not os.path.exists(os.path.join(path, f"Gleason_masks_{mask_key}")):
103        tar_path = os.path.join(path, f"Gleason_masks_{mask_key}.tar.gz")
104        util.download_source(
105            path=tar_path, url=MASK_URLS[mask_key], download=download, checksum=MASK_CHECKSUMS[mask_key]
106        )
107        util.unzip_tarfile(tar_path=tar_path, dst=path)
108
109    if split == "test" and not os.path.exists(os.path.join(path, "Gleason_masks_test_pathologist2")):
110        tar_path = os.path.join(path, "Gleason_masks_test_pathologist2.tar.gz")
111        util.download_source(
112            path=tar_path, url=MASK_URLS["test_pathologist2"], download=download,
113            checksum=MASK_CHECKSUMS["test_pathologist2"],
114        )
115        util.unzip_tarfile(tar_path=tar_path, dst=path)
116
117    return path
118
119
120def get_gleason2019_paths(
121    path: Union[os.PathLike, str],
122    split: Literal["train", "val", "test"],
123    test_pathologist: Literal[1, 2] = 1,
124    download: bool = False,
125) -> Tuple[List[str], List[str]]:
126    """Get paths to the Gleason2019 data.
127
128    Args:
129        path: Filepath to a folder where the downloaded data will be saved.
130        split: The choice of data split.
131        test_pathologist: The choice of pathologist annotation to use for the 'test' split.
132        download: Whether to download the data if it is not present.
133
134    Returns:
135        List of filepaths to the image data.
136        List of filepaths to the label data.
137    """
138    get_gleason2019_data(path, split, download)
139
140    if split == "test":
141        mask_dir = os.path.join(path, f"Gleason_masks_test_pathologist{test_pathologist}")
142        mask_prefix = f"mask{test_pathologist}_"
143    else:
144        mask_dir = os.path.join(path, "Gleason_masks_train")
145        mask_prefix = "mask_"
146
147    converted_dir = f"{mask_dir}_indexed"
148    os.makedirs(converted_dir, exist_ok=True)
149
150    raw_paths, label_paths = [], []
151    for stem in SPLIT_TMAS[split]:
152        for raw_path in natsorted(glob(os.path.join(path, stem, "*.jpg"))):
153            mask_name = f"{mask_prefix}{os.path.basename(raw_path)[:-len('.jpg')]}.png"
154            mask_path = os.path.join(mask_dir, mask_name)
155            if not os.path.exists(mask_path):
156                continue
157            label_path = os.path.join(converted_dir, mask_name)
158            if not os.path.exists(label_path):
159                # The released masks are palette-indexed PNGs (0-4); keep the raw
160                # palette index instead of the RGB colors most readers expand them to.
161                mask = Image.open(mask_path)
162                assert mask.mode == "P", f"Expected a palette-indexed mask, got mode '{mask.mode}': {mask_path}"
163                imageio.imwrite(label_path, np.array(mask).astype("uint8"))
164            raw_paths.append(raw_path)
165            label_paths.append(label_path)
166
167    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
168    return raw_paths, label_paths
169
170
171def get_gleason2019_dataset(
172    path: Union[os.PathLike, str],
173    patch_shape: Tuple[int, int],
174    split: Literal["train", "val", "test"],
175    test_pathologist: Literal[1, 2] = 1,
176    resize_inputs: bool = False,
177    download: bool = False,
178    **kwargs
179) -> Dataset:
180    """Get the Gleason2019 dataset for Gleason pattern segmentation.
181
182    Args:
183        path: Filepath to a folder where the downloaded data will be saved.
184        patch_shape: The patch shape to use for training.
185        split: The choice of data split.
186        test_pathologist: The choice of pathologist annotation to use for the 'test' split.
187        resize_inputs: Whether to resize the inputs.
188        download: Whether to download the data if it is not present.
189        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
190
191    Returns:
192        The segmentation dataset.
193    """
194    raw_paths, label_paths = get_gleason2019_paths(path, split, test_pathologist, download)
195
196    if resize_inputs:
197        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
198        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
199            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
200        )
201
202    return torch_em.default_segmentation_dataset(
203        raw_paths=raw_paths,
204        raw_key=None,
205        label_paths=label_paths,
206        label_key=None,
207        is_seg_dataset=False,
208        patch_shape=patch_shape,
209        with_channels=True,
210        ndim=2,
211        **kwargs
212    )
213
214
215def get_gleason2019_loader(
216    path: Union[os.PathLike, str],
217    batch_size: int,
218    patch_shape: Tuple[int, int],
219    split: Literal["train", "val", "test"],
220    test_pathologist: Literal[1, 2] = 1,
221    resize_inputs: bool = False,
222    download: bool = False,
223    **kwargs
224) -> DataLoader:
225    """Get the Gleason2019 dataloader for Gleason pattern segmentation.
226
227    Args:
228        path: Filepath to a folder where the downloaded data will be saved.
229        batch_size: The batch size for training.
230        patch_shape: The patch shape to use for training.
231        split: The choice of data split.
232        test_pathologist: The choice of pathologist annotation to use for the 'test' split.
233        resize_inputs: Whether to resize the inputs.
234        download: Whether to download the data if it is not present.
235        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
236
237    Returns:
238        The DataLoader.
239    """
240    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
241    dataset = get_gleason2019_dataset(path, patch_shape, split, test_pathologist, resize_inputs, download, **ds_kwargs)
242    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)