torch_em.data.datasets.histopathology.tiger

  1"""The TIGER WSIROIS dataset contains semantic tissue masks for H&E breast cancer histopathology images.
  2
  3The dataset is located at https://doi.org/10.5281/zenodo.6014422. The annotations and RUMC/JB images are
  4licensed under CC BY-NC 4.0, while TCGA-derived images retain their original TCGA-BRCA rights. This dataset
  5is from the publication https://doi.org/10.1038/s41467-026-72956-x. Please cite it in your research.
  6"""
  7
  8import os
  9import stat
 10import zipfile
 11from glob import glob
 12from pathlib import PurePosixPath
 13from typing import List, Tuple, Union
 14
 15from torch.utils.data import Dataset, DataLoader
 16
 17import torch_em
 18
 19from .. import util
 20
 21
 22URL = "https://zenodo.org/api/records/6014422/files/roi-level-annotations.zip/content"
 23CHECKSUM = "94bf1a00a61b8d264a6d8d9f213000617766ce65823ab33af86498041bf866dd"
 24SUBSETS = ("tissue-bcss", "tissue-cells")
 25
 26
 27def _validate_archive(zip_path):
 28    with zipfile.ZipFile(zip_path, "r") as archive:
 29        members = archive.infolist()
 30
 31    file_members = []
 32    for member in members:
 33        member_path = PurePosixPath(member.filename)
 34        first_part = member_path.parts[0] if member_path.parts else ""
 35        if (
 36            not member_path.parts
 37            or member_path.is_absolute()
 38            or ".." in member_path.parts
 39            or "\\" in member.filename
 40            or ":" in first_part
 41            or first_part != "roi-level-annotations"
 42        ):
 43            raise RuntimeError(f"Unsafe archive member: {member.filename}")
 44
 45        file_type = stat.S_IFMT(member.external_attr >> 16)
 46        if file_type not in (0, stat.S_IFREG, stat.S_IFDIR):
 47            raise RuntimeError(f"Unsupported archive member type: {member.filename}")
 48        if not member.is_dir():
 49            file_members.append(member)
 50
 51    extracted_size = sum(member.file_size for member in file_members)
 52    archive_size = os.path.getsize(zip_path)
 53    if len(file_members) > 5000 or extracted_size > 3_000_000_000 or extracted_size > 2 * archive_size:
 54        raise RuntimeError("The TIGER archive exceeds the expected extraction limits.")
 55
 56
 57def _has_data(data_dir):
 58    return all(
 59        os.path.isdir(os.path.join(data_dir, subset, folder))
 60        for subset in SUBSETS
 61        for folder in ("images", "masks")
 62    )
 63
 64
 65def get_tiger_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 66    """Download the TIGER WSIROIS ROI-level dataset.
 67
 68    Args:
 69        path: Filepath to a folder where the downloaded data will be saved.
 70        download: Whether to download the data if it is not present.
 71
 72    Returns:
 73        The filepath to the folder with the ROI-level annotations.
 74    """
 75    data_dir = os.path.join(path, "roi-level-annotations")
 76    if _has_data(data_dir):
 77        return data_dir
 78
 79    os.makedirs(path, exist_ok=True)
 80    zip_path = os.path.join(path, "roi-level-annotations.zip")
 81    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 82    _validate_archive(zip_path)
 83    util.unzip(zip_path=zip_path, dst=path)
 84
 85    if not _has_data(data_dir):
 86        raise RuntimeError("The TIGER archive does not contain the expected image and mask folders.")
 87    return data_dir
 88
 89
 90def get_tiger_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 91    """Get paths to the TIGER WSIROIS ROI images and semantic tissue masks.
 92
 93    Args:
 94        path: Filepath to a folder where the downloaded data will be saved.
 95        download: Whether to download the data if it is not present.
 96
 97    Returns:
 98        List of filepaths for the image data.
 99        List of filepaths for the label data.
100    """
101    data_dir = get_tiger_data(path, download)
102    raw_paths, label_paths = [], []
103
104    for subset in SUBSETS:
105        subset_raw_paths = sorted(glob(os.path.join(data_dir, subset, "images", "*.png")))
106        subset_label_paths = sorted(glob(os.path.join(data_dir, subset, "masks", "*.png")))
107        if not subset_raw_paths or len(subset_raw_paths) != len(subset_label_paths):
108            raise RuntimeError(f"Invalid TIGER raw-label pairing for subset '{subset}'.")
109        if any(
110            os.path.basename(raw_path) != os.path.basename(label_path)
111            for raw_path, label_path in zip(subset_raw_paths, subset_label_paths)
112        ):
113            raise RuntimeError(f"Mismatched TIGER raw-label names for subset '{subset}'.")
114        raw_paths.extend(subset_raw_paths)
115        label_paths.extend(subset_label_paths)
116
117    return raw_paths, label_paths
118
119
120def get_tiger_dataset(
121    path: Union[os.PathLike, str],
122    patch_shape: Tuple[int, int],
123    resize_inputs: bool = False,
124    download: bool = False,
125    **kwargs,
126) -> Dataset:
127    """Get the TIGER dataset for semantic breast tissue segmentation.
128
129    The masks use label 0 for excluded pixels and labels 1 through 7 for invasive tumor,
130    tumor-associated stroma, in-situ tumor, healthy glands, necrosis, inflamed stroma, and rest.
131
132    Args:
133        path: Filepath to a folder where the downloaded data will be saved.
134        patch_shape: The patch shape to use for training.
135        resize_inputs: Whether to resize the inputs.
136        download: Whether to download the data if it is not present.
137        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
138
139    Returns:
140        The segmentation dataset.
141    """
142    raw_paths, label_paths = get_tiger_paths(path, download)
143
144    if resize_inputs:
145        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
146        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
147            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
148        )
149
150    return torch_em.default_segmentation_dataset(
151        raw_paths=raw_paths,
152        raw_key=None,
153        label_paths=label_paths,
154        label_key=None,
155        patch_shape=patch_shape,
156        is_seg_dataset=False,
157        ndim=2,
158        with_channels=True,
159        **kwargs,
160    )
161
162
163def get_tiger_loader(
164    path: Union[os.PathLike, str],
165    batch_size: int,
166    patch_shape: Tuple[int, int],
167    resize_inputs: bool = False,
168    download: bool = False,
169    **kwargs,
170) -> DataLoader:
171    """Get the TIGER dataloader for semantic breast tissue segmentation.
172
173    Args:
174        path: Filepath to a folder where the downloaded data will be saved.
175        batch_size: The batch size for training.
176        patch_shape: The patch shape to use for training.
177        resize_inputs: Whether to resize the inputs.
178        download: Whether to download the data if it is not present.
179        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
180
181    Returns:
182        The DataLoader.
183    """
184    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
185    dataset = get_tiger_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
186    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)