torch_em.data.datasets.histopathology.beacon

  1"""The BEACON dataset contains nucleus annotations for H&E stained biopsies of Barrett's esophagus.
  2
  3The cell segmentation part of BEACON holds one tile per slide, 67 tiles in total, from two cohorts:
  4MD Anderson ('MDA', the discovery cohort) and Northwestern University ('NU', the test cohort).
  5One pathologist outlined every nucleus by hand in QuPath and assigned it to 'epithelial' or 'other'.
  6NOTE: The tiles on the server are 257 x 257 pixels, although the dataset README states 512 x 512 pixels.
  7NOTE: A few nuclei per tile carry a third, generic class 'Cell'. This module maps them to 'other'.
  8
  9The dataset is located at https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD4181 (CC BY 4.0).
 10This dataset is from the publication https://doi.org/10.1101/2025.11.11.25339952.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15import json
 16from glob import glob
 17from tqdm import tqdm
 18from natsort import natsorted
 19from typing import Union, Tuple, List, Literal, Optional
 20
 21import numpy as np
 22import pandas as pd
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32BASE_URL = "https://ftp.ebi.ac.uk/pub/databases/biostudies/S-BIAD/181/S-BIAD4181/Files"
 33FILE_LIST = "file_lists/annotation/cell_segmentation.tsv"
 34COHORTS = {"MDA (Discovery)": "mda", "NU (Test)": "nu"}
 35CLASSES = {"epithelial": 1, "other": 2, "Cell": 2}
 36
 37
 38def _download_tiles(path, download):
 39    tsv_path = os.path.join(path, "cell_segmentation.tsv")
 40    util.download_source(tsv_path, f"{BASE_URL}/{FILE_LIST}", download, checksum=None)
 41    table = pd.read_csv(tsv_path, sep="\t")
 42
 43    tiles = {}
 44    for files, source, cohort in zip(table["Files"], table["source_image"], table["Cohort"]):
 45        for remote in (files, source):
 46            local = os.path.join(path, "raw", os.path.basename(remote))
 47            if not os.path.exists(local):
 48                os.makedirs(os.path.dirname(local), exist_ok=True)
 49                util.download_source(local, f"{BASE_URL}/{remote}", download, checksum=None)
 50        tiles[os.path.splitext(os.path.basename(source))[0]] = COHORTS[cohort]
 51    return tiles
 52
 53
 54def _create_h5_files(path, tiles):
 55    import h5py
 56
 57    for tile, cohort in tqdm(tiles.items(), desc="Preprocess BEACON tiles"):
 58        out_dir = os.path.join(path, "preprocessed", cohort)
 59        os.makedirs(out_dir, exist_ok=True)
 60
 61        image = imageio.imread(os.path.join(path, "raw", f"{tile}.png"))[..., :3]
 62        instances = imageio.imread(os.path.join(path, "raw", f"{tile}_mask.tif")).astype("uint32")
 63        with open(os.path.join(path, "raw", f"{tile}_classmap.json")) as f:
 64            classmap = json.load(f)
 65
 66        semantic = np.zeros_like(instances, dtype="uint8")
 67        for name, ids in classmap.items():
 68            semantic[np.isin(instances, ids)] = CLASSES[name]
 69
 70        with h5py.File(os.path.join(out_dir, f"{tile}.h5"), "w") as f:
 71            f.create_dataset("raw", data=image.transpose(2, 0, 1), compression="gzip")
 72            f.create_dataset("labels/instances", data=instances, compression="gzip")
 73            f.create_dataset("labels/semantic", data=semantic, compression="gzip")
 74
 75
 76def get_beacon_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 77    """Download the BEACON nucleus segmentation data.
 78
 79    Args:
 80        path: Filepath to a folder where the downloaded data is saved.
 81        download: Whether to download the data if it is not present.
 82
 83    Returns:
 84        Filepath where the preprocessed data is stored.
 85    """
 86    data_dir = os.path.join(path, "preprocessed")
 87    if os.path.exists(data_dir):
 88        return data_dir
 89
 90    os.makedirs(path, exist_ok=True)
 91    tiles = _download_tiles(path, download)
 92    _create_h5_files(path, tiles)
 93    return data_dir
 94
 95
 96def get_beacon_paths(
 97    path: Union[os.PathLike, str],
 98    cohort: Optional[Literal["mda", "nu"]] = None,
 99    download: bool = False,
100) -> List[str]:
101    """Get paths to the BEACON data.
102
103    Args:
104        path: Filepath to a folder where the downloaded data is saved.
105        cohort: The choice of cohort, 'mda' or 'nu'. By default, the tiles of both cohorts are returned.
106        download: Whether to download the data if it is not present.
107
108    Returns:
109        List of filepaths for the input data.
110    """
111    data_dir = get_beacon_data(path, download)
112    if cohort is not None:
113        assert cohort in ("mda", "nu"), cohort
114    return natsorted(glob(os.path.join(data_dir, cohort or "*", "*.h5")))
115
116
117def get_beacon_dataset(
118    path: Union[os.PathLike, str],
119    patch_shape: Tuple[int, int],
120    cohort: Optional[Literal["mda", "nu"]] = None,
121    label_choice: Literal["instances", "semantic"] = "instances",
122    download: bool = False,
123    **kwargs
124) -> Dataset:
125    """Get the BEACON dataset for nucleus segmentation.
126
127    Args:
128        path: Filepath to a folder where the downloaded data is saved.
129        patch_shape: The patch shape to use for training.
130        cohort: The choice of cohort, 'mda' or 'nu'. By default, the tiles of both cohorts are used.
131        label_choice: The choice of labels, 'instances' or 'semantic' (1: epithelial, 2: other).
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
134
135    Returns:
136        The segmentation dataset.
137    """
138    assert label_choice in ("instances", "semantic"), label_choice
139    paths = get_beacon_paths(path, cohort, download)
140    return torch_em.default_segmentation_dataset(
141        raw_paths=paths,
142        raw_key="raw",
143        label_paths=paths,
144        label_key=f"labels/{label_choice}",
145        patch_shape=patch_shape,
146        ndim=2,
147        with_channels=True,
148        **kwargs
149    )
150
151
152def get_beacon_loader(
153    path: Union[os.PathLike, str],
154    batch_size: int,
155    patch_shape: Tuple[int, int],
156    cohort: Optional[Literal["mda", "nu"]] = None,
157    label_choice: Literal["instances", "semantic"] = "instances",
158    download: bool = False,
159    **kwargs
160) -> DataLoader:
161    """Get the BEACON dataloader for nucleus segmentation.
162
163    Args:
164        path: Filepath to a folder where the downloaded data is saved.
165        batch_size: The batch size for training.
166        patch_shape: The patch shape to use for training.
167        cohort: The choice of cohort, 'mda' or 'nu'. By default, the tiles of both cohorts are used.
168        label_choice: The choice of labels, 'instances' or 'semantic' (1: epithelial, 2: other).
169        download: Whether to download the data if it is not present.
170        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
171
172    Returns:
173        The DataLoader.
174    """
175    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
176    dataset = get_beacon_dataset(path, patch_shape, cohort, label_choice, download, **ds_kwargs)
177    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)