torch_em.data.datasets.histopathology.spatialproteomics_bnhl

  1"""This dataset contains cell instance segmentation annotations for highly multiplexed
  2immunofluorescence imaging of reactive lymph nodes and various B cell Non-Hodgkin lymphomas (BNHL),
  3acquired with the Akoya Phenocycler.
  4
  5The data is from the publication https://doi.org/10.1038/s41592-026-03155-1 and hosted on the
  6BioImage Archive at https://www.ebi.ac.uk/biostudies/studies/S-BIAD2100. Please cite it if you use
  7this dataset in your research.
  8
  9This loader covers the 250 tissue microarray (TMA) cores. Each core is a 3000x3000 pixel crop with
 1056 immunofluorescence channels and a matching per-cell instance segmentation mask, generated with
 11cellpose on the DAPI channel. The BioImage Archive record also hosts 5 whole-slide images
 12(24-63 GB each), which are out of scope for this loader.
 13
 14On first use, each TMA core is converted into a single HDF5 file with the following layout:
 15    - 'raw/all': the (56, H, W) stack of all channels.
 16    - 'raw/channels/<channel>': each individual channel (H, W), see `CHANNELS` for the full list.
 17    - 'labels/instances': the instance segmentation.
 18"""
 19
 20import os
 21from glob import glob
 22from typing import List, Optional, Sequence, Tuple, Union
 23
 24import pandas as pd
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33BASE_URL = "https://ftp.ebi.ac.uk/biostudies/fire/S-BIAD/100/S-BIAD2100/Files/data_for_publication"
 34MANIFEST_URL = f"{BASE_URL}/derived_data/tma_sample_overview.csv"
 35MANIFEST_CHECKSUM = "c2bb69d29af94c09cb7cd5b3b10e3acfdffb6beea1d7e8dd46b93cc30f55b0f8"
 36
 37CHANNELS = (
 38    "DAPI", "Helios", "CD10", "TCF7/TCF1", "PD-L1", "BCL-6", "FOXP3", "CD69", "Perforin", "CD19",
 39    "LAG3", "CD21", "CD62L", "c-myc", "CD138", "CD15", "BCL-2", "CD56", "IKZF3", "CD25", "NOXA",
 40    "Tim3", "Serpin B9", "Podoplanin", "CD38", "SPARC", "ICOS", "CXCR5", "CD163", "FADD", "p53",
 41    "Collagen IV", "CD4", "CD7", "Kappa", "CD20", "CD34", "PAX5", "PD-1", "CD45RA", "CD11b",
 42    "Lambda", "CD57", "CD11c", "CD90", "HLA DR", "CD68", "CD31", "CD45", "CD3", "Cytokeratin",
 43    "CD45RO", "CD8", "Granzyme B", "CD79a", "Ki-67",
 44)
 45
 46
 47def _sanitize_channel(name):
 48    return name.replace("/", "-")
 49
 50
 51def _get_manifest(path, download):
 52    manifest_path = os.path.join(path, "tma_sample_overview.csv")
 53    util.download_source(manifest_path, MANIFEST_URL, download, MANIFEST_CHECKSUM)
 54    return pd.read_csv(manifest_path, index_col=0)
 55
 56
 57def _convert_sample(zip_path, output_path):
 58    import h5py
 59    import zarr
 60
 61    store = zarr.storage.ZipStore(zip_path, mode="r")
 62    group = zarr.open_group(store=store, mode="r")
 63    channels = [str(c) for c in group["channels"][:]]
 64    if channels != list(CHANNELS):
 65        raise RuntimeError(f"Unexpected channel order in {zip_path}.")
 66
 67    raw = group["_image_raw"][:]
 68    instances = group["_segmentation"][:]
 69    store.close()
 70
 71    tmp_path = output_path + ".tmp"
 72    with h5py.File(tmp_path, "w") as f:
 73        f.create_dataset("raw/all", data=raw, compression="gzip", chunks=(len(CHANNELS), 512, 512))
 74        for i, name in enumerate(CHANNELS):
 75            f.create_dataset(f"raw/channels/{_sanitize_channel(name)}", data=raw[i], compression="gzip")
 76        f.create_dataset("labels/instances", data=instances, compression="gzip")
 77
 78    os.replace(tmp_path, output_path)
 79
 80
 81def get_spatialproteomics_bnhl_data(
 82    path: Union[os.PathLike, str],
 83    samples: Optional[Sequence[str]] = None,
 84    download: bool = False,
 85) -> str:
 86    """Download and preprocess the spatialproteomics BNHL tissue microarray (TMA) data.
 87
 88    Args:
 89        path: Filepath to a folder where the downloaded data will be saved.
 90        samples: The TMA sample ids to prepare. By default all 250 TMA cores are prepared, which
 91            requires downloading about 60 GB of data.
 92        download: Whether to download the data if it is not present.
 93
 94    Returns:
 95        Filepath to the folder where the preprocessed data is stored.
 96    """
 97    os.makedirs(path, exist_ok=True)
 98    manifest = _get_manifest(path, download)
 99    valid_samples = set(manifest["sample_id"])
100
101    if samples is None:
102        samples = sorted(valid_samples)
103    else:
104        invalid = sorted(set(samples) - valid_samples)
105        if invalid:
106            raise ValueError(f"Invalid sample id(s) {invalid}.")
107
108    preprocessed_dir = os.path.join(path, "preprocessed")
109    os.makedirs(preprocessed_dir, exist_ok=True)
110
111    zip_dir = os.path.join(path, "tmas")
112    os.makedirs(zip_dir, exist_ok=True)
113
114    for sample_id in samples:
115        output_path = os.path.join(preprocessed_dir, f"{sample_id}.h5")
116        if os.path.exists(output_path):
117            continue
118
119        zip_path = os.path.join(zip_dir, f"{sample_id}.zarr.zip")
120        util.download_source(zip_path, f"{BASE_URL}/tmas/{sample_id}.zarr.zip", download, checksum=None)
121        _convert_sample(zip_path, output_path)
122
123    return preprocessed_dir
124
125
126def get_spatialproteomics_bnhl_paths(
127    path: Union[os.PathLike, str],
128    samples: Optional[Sequence[str]] = None,
129    download: bool = False,
130) -> List[str]:
131    """Get paths to the preprocessed spatialproteomics BNHL TMA data.
132
133    Args:
134        path: Filepath to a folder where the downloaded data will be saved.
135        samples: The TMA sample ids to load. By default all 250 TMA cores are loaded.
136        download: Whether to download the data if it is not present.
137
138    Returns:
139        List of filepaths to the preprocessed HDF5 files.
140    """
141    preprocessed_dir = get_spatialproteomics_bnhl_data(path, samples, download)
142    if samples is None:
143        paths = sorted(glob(os.path.join(preprocessed_dir, "*.h5")))
144    else:
145        paths = [os.path.join(preprocessed_dir, f"{sample_id}.h5") for sample_id in samples]
146
147    missing = [p for p in paths if not os.path.exists(p)]
148    if missing:
149        raise RuntimeError(f"Could not find the data at {missing}.")
150
151    return paths
152
153
154def get_spatialproteomics_bnhl_dataset(
155    path: Union[os.PathLike, str],
156    patch_shape: Tuple[int, int],
157    samples: Optional[Sequence[str]] = None,
158    channel: str = "all",
159    download: bool = False,
160    resize_inputs: bool = False,
161    **kwargs
162) -> Dataset:
163    """Get the spatialproteomics BNHL dataset for cell instance segmentation in multiplexed
164    immunofluorescence images of B cell Non-Hodgkin lymphomas.
165
166    Args:
167        path: Filepath to a folder where the downloaded data will be saved.
168        patch_shape: The patch shape to use for training.
169        samples: The TMA sample ids to load. By default all 250 TMA cores are loaded.
170        channel: The raw input. Either 'all' for the full (56, H, W) channel stack, or the name of
171            a single channel, see `CHANNELS` for the full list, e.g. 'CD20'.
172        download: Whether to download the data if it is not present.
173        resize_inputs: Whether to resize the input images.
174        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
175
176    Returns:
177        The segmentation dataset.
178    """
179    if channel == "all":
180        raw_key, with_channels = "raw/all", True
181    elif channel in CHANNELS:
182        raw_key, with_channels = f"raw/channels/{_sanitize_channel(channel)}", False
183    else:
184        raise ValueError(f"'{channel}' is not a valid channel. Choose 'all' or one of {CHANNELS}.")
185
186    paths = get_spatialproteomics_bnhl_paths(path, samples, download)
187
188    if resize_inputs:
189        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
190        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
191            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
192        )
193
194    return torch_em.default_segmentation_dataset(
195        raw_paths=paths,
196        raw_key=raw_key,
197        label_paths=paths,
198        label_key="labels/instances",
199        patch_shape=patch_shape,
200        is_seg_dataset=True,
201        with_channels=with_channels,
202        ndim=2,
203        **kwargs
204    )
205
206
207def get_spatialproteomics_bnhl_loader(
208    path: Union[os.PathLike, str],
209    patch_shape: Tuple[int, int],
210    batch_size: int,
211    samples: Optional[Sequence[str]] = None,
212    channel: str = "all",
213    download: bool = False,
214    resize_inputs: bool = False,
215    **kwargs
216) -> DataLoader:
217    """Get the spatialproteomics BNHL dataloader for cell instance segmentation in multiplexed
218    immunofluorescence images of B cell Non-Hodgkin lymphomas.
219
220    Args:
221        path: Filepath to a folder where the downloaded data will be saved.
222        patch_shape: The patch shape to use for training.
223        batch_size: The batch size for training.
224        samples: The TMA sample ids to load. By default all 250 TMA cores are loaded.
225        channel: The raw input. Either 'all' for the full (56, H, W) channel stack, or the name of
226            a single channel, see `CHANNELS` for the full list, e.g. 'CD20'.
227        download: Whether to download the data if it is not present.
228        resize_inputs: Whether to resize the input images.
229        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the
230            PyTorch DataLoader.
231
232    Returns:
233        The DataLoader.
234    """
235    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
236    dataset = get_spatialproteomics_bnhl_dataset(
237        path, patch_shape, samples=samples, channel=channel, download=download,
238        resize_inputs=resize_inputs, **ds_kwargs
239    )
240    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)