torch_em.data.datasets.histopathology.beetle

  1"""BEETLE (BrEast cancEr hisTopathoLogy sEgmentation) contains multiclass semantic
  2segmentation annotations for H&E stained breast cancer whole-slide images (WSIs),
  3collected across multiple clinical centers and scanners.
  4
  5Mask label values (see the dataset's 'label_map.json'): 0 (unannotated), 1 (other),
  62 (non-invasive epithelium), 3 (invasive epithelium), 4 (necrosis).
  7
  8NOTE: Only the 'development' set is exposed here. Its own official split uses 5
  9cross-validation folds; pass `validation_fold` to choose which fold is held out for
 10validation. The dataset also has an 'evaluation' set, but its annotations are not
 11publicly released (they are sequestered on the Grand Challenge platform), so it is
 12not included. A further subset of development slides sourced from TCGA is not bundled
 13as raw images in this Zenodo deposit (their raw whole-slide images would need to be
 14fetched separately from the NCI Genomic Data Commons) and is excluded as well.
 15
 16NOTE: The whole-slide images and masks are multi-resolution pyramidal TIFFs bundled
 17inside two large Zenodo archives (images: about 147 GB; masks: about 1.8 GB). Requested
 18slides are extracted directly from these remote archives without downloading the full
 19archives, then converted into a chunked HDF5 file at the requested pyramid level.
 20
 21This dataset is located at https://doi.org/10.5281/zenodo.16812932.
 22This dataset is from the publication https://arxiv.org/abs/2510.02037.
 23Please cite it if you use this dataset for your research.
 24"""
 25
 26import os
 27import csv
 28from pathlib import Path
 29from typing import List, Literal, Optional, Tuple, Union
 30
 31from tqdm import tqdm
 32
 33import torch
 34from torch.utils.data import Dataset, DataLoader
 35
 36import torch_em
 37
 38from .. import util
 39
 40
 41OVERVIEW_URL = "https://zenodo.org/api/records/16812932/files/data_overview.csv/content"
 42OVERVIEW_CHECKSUM = "a063a5456959cb92f3dc007844cb14208a71336d8157d006304539066d86f81b"
 43
 44IMAGES_ZIP_URL = "https://zenodo.org/api/records/16812932/files/images.zip/content"
 45ANNOTATIONS_ZIP_URL = "https://zenodo.org/api/records/16812932/files/annotations.zip/content"
 46
 47
 48def _load_manifest(path, download):
 49    csv_path = os.path.join(path, "data_overview.csv")
 50    util.download_source(path=csv_path, url=OVERVIEW_URL, download=download, checksum=OVERVIEW_CHECKSUM)
 51    with open(csv_path) as f:
 52        rows = list(csv.DictReader(f))
 53    # Only the development set ships raw whole-slide images in this deposit; the rest
 54    # (evaluation set, and a subset of development rows sourced from TCGA) are excluded.
 55    return [row for row in rows if row["split"] == "development" and row["wsi_path"]]
 56
 57
 58def _resolve_rows(path, split, validation_fold, sample_ids, download):
 59    rows = _load_manifest(path, download)
 60    fold = f"fold{validation_fold}"
 61    rows = [row for row in rows if (row["validation_fold"] == fold) == (split == "val")]
 62    if sample_ids is not None:
 63        by_name = {row["name"]: row for row in rows}
 64        missing = sorted(set(sample_ids) - set(by_name))
 65        if missing:
 66            raise ValueError(f"The following sample ids are not part of this split: {missing}")
 67        rows = [by_name[name] for name in sample_ids]
 68    return rows
 69
 70
 71def _extract_zip_member(zip_url, member, out_path):
 72    import fsspec
 73    import zipfile
 74
 75    if os.path.exists(out_path):
 76        return
 77    fs = fsspec.filesystem("http")
 78    with fs.open(zip_url) as f:
 79        data = zipfile.ZipFile(f).read(member)
 80    tmp_path = out_path + ".tmp"
 81    Path(tmp_path).write_bytes(data)
 82    os.replace(tmp_path, out_path)
 83
 84
 85def _open_level(series, level_index):
 86    import zarr
 87
 88    # The pyramidal TIFFs are natively tiled, so a zarr view reads only the requested tiles lazily.
 89    array = zarr.open(series.aszarr(), mode="r")
 90    return array if hasattr(array, "shape") else array[str(level_index)]
 91
 92
 93def _convert_slide(image_path, mask_path, output_path, resolution_level, tile=4096):
 94    import h5py
 95    import tifffile
 96
 97    image_series = tifffile.TiffFile(image_path).series[0]
 98    mask_series = tifffile.TiffFile(mask_path).series[0]
 99    height, width = mask_series.levels[resolution_level].shape[:2]
100
101    image = _open_level(image_series, resolution_level)
102    mask = _open_level(mask_series, resolution_level)
103
104    tmp_path = output_path + ".tmp"
105    with h5py.File(tmp_path, "w") as f:
106        raw = f.create_dataset(
107            "images/raw", shape=(3, height, width), dtype="uint8", compression="gzip", chunks=(1, 512, 512)
108        )
109        labels = f.create_dataset(
110            "labels/mask", shape=(height, width), dtype="uint8", compression="gzip", chunks=(512, 512)
111        )
112        for y in tqdm(range(0, height, tile), desc=f"Converting {Path(image_path).stem}"):
113            for x in range(0, width, tile):
114                th, tw = min(tile, height - y), min(tile, width - x)
115                raw[:, y:y + th, x:x + tw] = image[y:y + th, x:x + tw].transpose(2, 0, 1)
116                labels[y:y + th, x:x + tw] = mask[y:y + th, x:x + tw]
117
118    os.replace(tmp_path, output_path)
119
120
121def get_beetle_data(
122    path: Union[os.PathLike, str],
123    split: Literal["train", "val"],
124    validation_fold: int = 0,
125    sample_ids: Optional[List[str]] = None,
126    resolution_level: int = 0,
127    download: bool = False,
128) -> str:
129    """Download and preprocess the BEETLE breast cancer segmentation data.
130
131    Args:
132        path: Filepath to a folder where the data will be saved.
133        split: The split to use, either the held-out validation fold or the rest.
134        validation_fold: Which of the 5 official cross-validation folds (0-4) to use as validation.
135        sample_ids: The slide names to restrict the data to, e.g. ['patient1_wsi1'].
136            By default all slides matching `split` and `validation_fold` are used.
137        resolution_level: The pyramid level to convert. 0 is the native resolution; use a higher level to
138            reduce the size of the preprocessed data.
139        download: Whether to download the data if it is not present.
140
141    Returns:
142        Filepath to the folder where the preprocessed data is stored.
143    """
144    rows = _resolve_rows(path, split, validation_fold, sample_ids, download)
145
146    raw_dir = os.path.join(path, "raw")
147    preprocessed_dir = os.path.join(path, "preprocessed")
148    os.makedirs(raw_dir, exist_ok=True)
149    os.makedirs(preprocessed_dir, exist_ok=True)
150
151    for row in rows:
152        output_path = os.path.join(preprocessed_dir, f"{row['name']}_level{resolution_level}.h5")
153        if os.path.exists(output_path):
154            continue
155        if not download:
156            raise RuntimeError(f"Cannot find the data at {output_path}, but download was set to False")
157        image_path = os.path.join(raw_dir, f"{row['name']}.tif")
158        mask_path = os.path.join(raw_dir, f"{row['name']}_mask.tif")
159        _extract_zip_member(IMAGES_ZIP_URL, row["wsi_path"], image_path)
160        _extract_zip_member(ANNOTATIONS_ZIP_URL, row["annotation_mask_path"], mask_path)
161        _convert_slide(image_path, mask_path, output_path, resolution_level)
162
163    return preprocessed_dir
164
165
166def get_beetle_paths(
167    path: Union[os.PathLike, str],
168    split: Literal["train", "val"],
169    validation_fold: int = 0,
170    sample_ids: Optional[List[str]] = None,
171    resolution_level: int = 0,
172    download: bool = False,
173) -> List[str]:
174    """Get paths to the BEETLE breast cancer segmentation data.
175
176    Args:
177        path: Filepath to a folder where the data will be saved.
178        split: The split to use, either the held-out validation fold or the rest.
179        validation_fold: Which of the 5 official cross-validation folds (0-4) to use as validation.
180        sample_ids: The slide names to restrict the data to, e.g. ['patient1_wsi1'].
181            By default all slides matching `split` and `validation_fold` are used.
182        resolution_level: The pyramid level to convert. 0 is the native resolution; use a higher level to
183            reduce the size of the preprocessed data.
184        download: Whether to download the data if it is not present.
185
186    Returns:
187        List of filepaths to the preprocessed HDF5 files.
188    """
189    preprocessed_dir = get_beetle_data(path, split, validation_fold, sample_ids, resolution_level, download)
190    rows = _resolve_rows(path, split, validation_fold, sample_ids, download=False)
191    return [os.path.join(preprocessed_dir, f"{row['name']}_level{resolution_level}.h5") for row in rows]
192
193
194def get_beetle_dataset(
195    path: Union[os.PathLike, str],
196    patch_shape: Tuple[int, int],
197    split: Literal["train", "val"],
198    validation_fold: int = 0,
199    sample_ids: Optional[List[str]] = None,
200    resolution_level: int = 0,
201    download: bool = False,
202    label_dtype: torch.dtype = torch.int64,
203    resize_inputs: bool = False,
204    **kwargs
205) -> Dataset:
206    """Get the BEETLE dataset for multiclass breast cancer tissue segmentation in whole-slide images.
207
208    Args:
209        path: Filepath to a folder where the data will be saved.
210        patch_shape: The patch shape to use for training.
211        split: The split to use, either the held-out validation fold or the rest.
212        validation_fold: Which of the 5 official cross-validation folds (0-4) to use as validation.
213        sample_ids: The slide names to restrict the data to, e.g. ['patient1_wsi1'].
214            By default all slides matching `split` and `validation_fold` are used.
215        resolution_level: The pyramid level to convert. 0 is the native resolution; use a higher level to
216            reduce the size of the preprocessed data.
217        download: Whether to download the data if it is not present.
218        label_dtype: The datatype of the labels.
219        resize_inputs: Whether to resize the input images.
220        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
221
222    Returns:
223        The segmentation dataset.
224    """
225    volume_paths = get_beetle_paths(path, split, validation_fold, sample_ids, resolution_level, download)
226
227    if resize_inputs:
228        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
229        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
230            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
231        )
232
233    return torch_em.default_segmentation_dataset(
234        raw_paths=volume_paths,
235        raw_key="images/raw",
236        label_paths=volume_paths,
237        label_key="labels/mask",
238        patch_shape=patch_shape,
239        label_dtype=label_dtype,
240        is_seg_dataset=True,
241        with_channels=True,
242        ndim=2,
243        **kwargs
244    )
245
246
247def get_beetle_loader(
248    path: Union[os.PathLike, str],
249    batch_size: int,
250    patch_shape: Tuple[int, int],
251    split: Literal["train", "val"],
252    validation_fold: int = 0,
253    sample_ids: Optional[List[str]] = None,
254    resolution_level: int = 0,
255    download: bool = False,
256    label_dtype: torch.dtype = torch.int64,
257    resize_inputs: bool = False,
258    **kwargs
259) -> DataLoader:
260    """Get the BEETLE dataloader for multiclass breast cancer tissue segmentation in whole-slide images.
261
262    Args:
263        path: Filepath to a folder where the data will be saved.
264        batch_size: The batch size for training.
265        patch_shape: The patch shape to use for training.
266        split: The split to use, either the held-out validation fold or the rest.
267        validation_fold: Which of the 5 official cross-validation folds (0-4) to use as validation.
268        sample_ids: The slide names to restrict the data to, e.g. ['patient1_wsi1'].
269            By default all slides matching `split` and `validation_fold` are used.
270        resolution_level: The pyramid level to convert. 0 is the native resolution; use a higher level to
271            reduce the size of the preprocessed data.
272        download: Whether to download the data if it is not present.
273        label_dtype: The datatype of the labels.
274        resize_inputs: Whether to resize the input images.
275        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
276
277    Returns:
278        The DataLoader.
279    """
280    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
281    dataset = get_beetle_dataset(
282        path=path, patch_shape=patch_shape, split=split, validation_fold=validation_fold, sample_ids=sample_ids,
283        resolution_level=resolution_level, download=download, label_dtype=label_dtype, resize_inputs=resize_inputs,
284        **ds_kwargs
285    )
286    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)