torch_em.data.datasets.histopathology.ocelot

  1"""The OCELOT dataset contains tissue segmentation masks for H&E histopathology images
  2sourced from TCGA, covering bladder, endometrium, head-and-neck, kidney, prostate, and
  3stomach cancer.
  4
  5The dataset is located at https://zenodo.org/records/8417503. The data is licensed under
  6CC BY-NC 4.0. This dataset is from the publication https://doi.org/10.1109/CVPR52729.2023.02289.
  7Please cite it in your research.
  8
  9NOTE: OCELOT also ships point annotations for cell detection, which this module does not expose,
 10since torch-em only integrates the tissue segmentation masks.
 11"""
 12
 13import os
 14import stat
 15import zipfile
 16from glob import glob
 17from pathlib import PurePosixPath
 18from typing import List, Literal, Optional, Tuple, Union
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://zenodo.org/records/8417503/files/ocelot2023_v1.0.1.zip?download=1"
 28CHECKSUM = "74f46b79e3c4076ca0012d403629ab1e1e412591faf79a910d4d8bdd92c47920"
 29SPLITS = ("train", "val", "test")
 30
 31
 32def _validate_archive(zip_path):
 33    with zipfile.ZipFile(zip_path, "r") as archive:
 34        members = archive.infolist()
 35
 36    file_members = []
 37    for member in members:
 38        member_path = PurePosixPath(member.filename)
 39        first_part = member_path.parts[0] if member_path.parts else ""
 40        if (
 41            not member_path.parts
 42            or member_path.is_absolute()
 43            or ".." in member_path.parts
 44            or "\\" in member.filename
 45            or ":" in first_part
 46            or first_part != "ocelot2023_v1.0.1"
 47        ):
 48            raise RuntimeError(f"Unsafe archive member: {member.filename}")
 49
 50        file_type = stat.S_IFMT(member.external_attr >> 16)
 51        if file_type not in (0, stat.S_IFREG, stat.S_IFDIR):
 52            raise RuntimeError(f"Unsupported archive member type: {member.filename}")
 53        if not member.is_dir():
 54            file_members.append(member)
 55
 56    extracted_size = sum(member.file_size for member in file_members)
 57    archive_size = os.path.getsize(zip_path)
 58    if len(file_members) > 20000 or extracted_size > 3_000_000_000 or extracted_size > 5 * archive_size:
 59        raise RuntimeError("The OCELOT archive exceeds the expected extraction limits.")
 60
 61
 62def _has_data(data_dir):
 63    return all(
 64        os.path.isdir(os.path.join(data_dir, "images", split, "tissue"))
 65        and os.path.isdir(os.path.join(data_dir, "annotations", split, "tissue"))
 66        for split in SPLITS
 67    )
 68
 69
 70def get_ocelot_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 71    """Download the OCELOT dataset.
 72
 73    Args:
 74        path: Filepath to a folder where the downloaded data will be saved.
 75        download: Whether to download the data if it is not present.
 76
 77    Returns:
 78        Filepath to the folder where the data is stored.
 79    """
 80    data_dir = os.path.join(path, "ocelot2023_v1.0.1")
 81    if _has_data(data_dir):
 82        return data_dir
 83
 84    os.makedirs(path, exist_ok=True)
 85    zip_path = os.path.join(path, "ocelot2023_v1.0.1.zip")
 86    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 87    _validate_archive(zip_path)
 88    util.unzip(zip_path=zip_path, dst=path)
 89
 90    if not _has_data(data_dir):
 91        raise RuntimeError("The OCELOT archive does not contain the expected image and mask folders.")
 92    return data_dir
 93
 94
 95def get_ocelot_paths(
 96    path: Union[os.PathLike, str],
 97    split: Optional[Literal["train", "val", "test"]] = None,
 98    download: bool = False,
 99) -> Tuple[List[str], List[str]]:
100    """Get paths to the OCELOT tissue images and semantic segmentation masks.
101
102    Args:
103        path: Filepath to a folder where the downloaded data will be saved.
104        split: The split to use. By default all splits ('train', 'val', 'test') are used.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        List of filepaths for the image data.
109        List of filepaths for the label data.
110    """
111    data_dir = get_ocelot_data(path, download)
112    splits = SPLITS if split is None else (split,)
113
114    raw_paths, label_paths = [], []
115    for this_split in splits:
116        split_raw_paths = sorted(glob(os.path.join(data_dir, "images", this_split, "tissue", "*.jpg")))
117        split_label_paths = sorted(glob(os.path.join(data_dir, "annotations", this_split, "tissue", "*.png")))
118        if not split_raw_paths or len(split_raw_paths) != len(split_label_paths):
119            raise RuntimeError(f"Invalid OCELOT raw-label pairing for split '{this_split}'.")
120        if any(
121            os.path.splitext(os.path.basename(raw_path))[0] != os.path.splitext(os.path.basename(label_path))[0]
122            for raw_path, label_path in zip(split_raw_paths, split_label_paths)
123        ):
124            raise RuntimeError(f"Mismatched OCELOT raw-label names for split '{this_split}'.")
125        raw_paths.extend(split_raw_paths)
126        label_paths.extend(split_label_paths)
127
128    return raw_paths, label_paths
129
130
131def get_ocelot_dataset(
132    path: Union[os.PathLike, str],
133    patch_shape: Tuple[int, int],
134    split: Optional[Literal["train", "val", "test"]] = None,
135    resize_inputs: bool = False,
136    download: bool = False,
137    **kwargs,
138) -> Dataset:
139    """Get the OCELOT dataset for tissue segmentation.
140
141    The masks use label 1 for background, 2 for cancer area, and 255 for unlabeled pixels.
142
143    Args:
144        path: Filepath to a folder where the downloaded data will be saved.
145        patch_shape: The patch shape to use for training.
146        split: The split to use. By default all splits ('train', 'val', 'test') are used.
147        resize_inputs: Whether to resize the inputs.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
150
151    Returns:
152        The segmentation dataset.
153    """
154    raw_paths, label_paths = get_ocelot_paths(path, split, download)
155
156    if resize_inputs:
157        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
158        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
159            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
160        )
161
162    return torch_em.default_segmentation_dataset(
163        raw_paths=raw_paths,
164        raw_key=None,
165        label_paths=label_paths,
166        label_key=None,
167        patch_shape=patch_shape,
168        is_seg_dataset=False,
169        ndim=2,
170        with_channels=True,
171        **kwargs,
172    )
173
174
175def get_ocelot_loader(
176    path: Union[os.PathLike, str],
177    batch_size: int,
178    patch_shape: Tuple[int, int],
179    split: Optional[Literal["train", "val", "test"]] = None,
180    resize_inputs: bool = False,
181    download: bool = False,
182    **kwargs,
183) -> DataLoader:
184    """Get the OCELOT dataloader for tissue segmentation.
185
186    Args:
187        path: Filepath to a folder where the downloaded data will be saved.
188        batch_size: The batch size for training.
189        patch_shape: The patch shape to use for training.
190        split: The split to use. By default all splits ('train', 'val', 'test') are used.
191        resize_inputs: Whether to resize the inputs.
192        download: Whether to download the data if it is not present.
193        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
194
195    Returns:
196        The DataLoader.
197    """
198    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
199    dataset = get_ocelot_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
200    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)