torch_em.data.datasets.histopathology.rings

  1"""The RINGS dataset contains annotations for prostate gland and tumor region
  2segmentation in H&E stained prostate histopathological images.
  3
  4The dataset is located at https://data.mendeley.com/datasets/h8bdwrtnr5/1.
  5This dataset is from the publication https://doi.org/10.1016/j.artmed.2021.102076.
  6Please cite it if you use this dataset for your research.
  7"""
  8
  9import os
 10from glob import glob
 11from natsort import natsorted
 12from typing import List, Literal, Tuple, Union
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URLS = {
 22    "train": "https://data.mendeley.com/public-files/datasets/h8bdwrtnr5/files/8416eb6b-d1c8-4a8c-96fd-2f36a2768e46/file_downloaded",  # noqa
 23    "test": "https://data.mendeley.com/public-files/datasets/h8bdwrtnr5/files/5846b131-09fa-44c5-afa8-8fbc52adbd88/file_downloaded",  # noqa
 24}
 25CHECKSUMS = {
 26    "train": "af426249bd96d36e2c5e0110d42ceb67a8ebb79d94e2b2c15f4e727ebca38329",
 27    "test": "f8134f01ce4cbcfd703bf96a8501e6267fe87a66301e15749dac463742f8958d",
 28}
 29
 30
 31def get_rings_data(path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False) -> str:
 32    """Download the RINGS dataset.
 33
 34    Args:
 35        path: Filepath to a folder where the downloaded data will be saved.
 36        split: The choice of data split.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath to the folder where the split data is stored.
 41    """
 42    if split not in URLS:
 43        raise ValueError(f"'{split}' is not a valid split choice.")
 44
 45    data_dir = os.path.join(path, split.upper())
 46    if os.path.exists(data_dir):
 47        return data_dir
 48
 49    os.makedirs(path, exist_ok=True)
 50    zip_path = os.path.join(path, f"{split}.zip")
 51    util.download_source(path=zip_path, url=URLS[split], download=download, checksum=CHECKSUMS[split])
 52    util.unzip(zip_path=zip_path, dst=path)
 53
 54    return data_dir
 55
 56
 57def get_rings_paths(
 58    path: Union[os.PathLike, str],
 59    split: Literal["train", "test"],
 60    label_choice: Literal["glands", "tumor"] = "glands",
 61    download: bool = False,
 62) -> Tuple[List[str], List[str]]:
 63    """Get paths to the RINGS data.
 64
 65    Args:
 66        path: Filepath to a folder where the downloaded data will be saved.
 67        split: The choice of data split.
 68        label_choice: The segmentation target. Either 'glands' for gland segmentation
 69            or 'tumor' for tumor region segmentation.
 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    if label_choice not in ("glands", "tumor"):
 77        raise ValueError(f"'{label_choice}' is not a valid label choice.")
 78
 79    data_dir = get_rings_data(path, split, download)
 80    raw_paths = natsorted(glob(os.path.join(data_dir, "IMAGES", "*.png")))
 81
 82    label_folder = "MANUAL GLANDS" if label_choice == "glands" else "MANUAL TUMOR"
 83    label_paths = natsorted(glob(os.path.join(data_dir, label_folder, "*.png")))
 84
 85    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 86    assert all(
 87        os.path.basename(raw_path) == os.path.basename(label_path)
 88        for raw_path, label_path in zip(raw_paths, label_paths)
 89    )
 90
 91    return raw_paths, label_paths
 92
 93
 94def get_rings_dataset(
 95    path: Union[os.PathLike, str],
 96    patch_shape: Tuple[int, int],
 97    split: Literal["train", "test"],
 98    label_choice: Literal["glands", "tumor"] = "glands",
 99    resize_inputs: bool = False,
100    download: bool = False,
101    **kwargs,
102) -> Dataset:
103    """Get the RINGS dataset for prostate gland or tumor region segmentation.
104
105    Args:
106        path: Filepath to a folder where the downloaded data will be saved.
107        patch_shape: The patch shape to use for training.
108        split: The choice of data split.
109        label_choice: The segmentation target. Either 'glands' for gland segmentation
110            or 'tumor' for tumor region segmentation.
111        resize_inputs: Whether to resize the inputs.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
114
115    Returns:
116        The segmentation dataset.
117    """
118    raw_paths, label_paths = get_rings_paths(path, split, label_choice, download)
119
120    if resize_inputs:
121        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
122        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
123            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
124        )
125
126    return torch_em.default_segmentation_dataset(
127        raw_paths=raw_paths,
128        raw_key=None,
129        label_paths=label_paths,
130        label_key=None,
131        is_seg_dataset=False,
132        patch_shape=patch_shape,
133        ndim=2,
134        with_channels=True,
135        **kwargs,
136    )
137
138
139def get_rings_loader(
140    path: Union[os.PathLike, str],
141    batch_size: int,
142    patch_shape: Tuple[int, int],
143    split: Literal["train", "test"],
144    label_choice: Literal["glands", "tumor"] = "glands",
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs,
148) -> DataLoader:
149    """Get the RINGS dataloader for prostate gland or tumor region segmentation.
150
151    Args:
152        path: Filepath to a folder where the downloaded data will be saved.
153        batch_size: The batch size for training.
154        patch_shape: The patch shape to use for training.
155        split: The choice of data split.
156        label_choice: The segmentation target. Either 'glands' for gland segmentation
157            or 'tumor' for tumor region segmentation.
158        resize_inputs: Whether to resize the inputs.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_rings_dataset(path, patch_shape, split, label_choice, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)