torch_em.data.datasets.light_microscopy.pan_multiplex

  1"""Pan-Multiplex (Pan-M) contains annotations for cell segmentation in multiplexed images of tissue.
  2
  3The dataset covers five subsets from three imaging platforms: MIBI-TOF, CODEX and Vectra.
  4Every field of view provides many protein marker channels, a nuclei channel and a cell instance mask.
  5This module builds a nuclei and a membrane composite from the marker channels, which is the common
  6input format for cell segmentation in multiplexed imaging.
  7NOTE: The cell instance masks come from the upstream studies and nobody curated them by hand.
  8NOTE: The subsets are large. See `SUBSET_PARTS` and download one subset at a time.
  9
 10The dataset is located at https://huggingface.co/datasets/JLrumberger/Pan-Multiplex.
 11This dataset is from the publication https://doi.org/10.1038/s41592-025-02826-9.
 12Please cite it if you use this dataset for your research.
 13"""
 14
 15import os
 16import gzip
 17import json
 18import shutil
 19from glob import glob
 20from warnings import warn
 21from natsort import natsorted
 22from typing import Union, Tuple, List, Literal, Optional
 23
 24import numpy as np
 25import tifffile
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30from torch_em.data import MinInstanceSampler
 31
 32from .. import util
 33
 34
 35URL = "https://huggingface.co/datasets/JLrumberger/Pan-Multiplex/resolve/main"
 36SPLIT_URL = "https://raw.githubusercontent.com/angelolab/Nimbus/main/configs"
 37
 38SUBSET_PARTS = {
 39    "mibi_decidua": ["mibi_decidua.zip"],
 40    "codex_colon": [f"codex_colon.zip.{i:03d}" for i in range(1, 5)],
 41    "vectra_colon": [f"vectra_colon.zip.{i:03d}" for i in range(1, 3)],
 42    "vectra_pancreas": [f"vectra_pancreas.zip.{i:03d}" for i in range(1, 3)],
 43    "mibi_breast": [f"mibi_breast.zip.{i:03d}" for i in range(1, 12)],
 44}
 45
 46CHECKSUMS = {
 47    "mibi_decidua.zip": "723a4c6863ca49bc063e0f6eead3a7bbf8d4d8ecb3992882739beeb7ad51eaea",
 48    "codex_colon.zip.001": "cefdc3147eb9f395130b5052ad965880fd97876af32c1fa4e2167a9cbc91c77f",
 49    "codex_colon.zip.002": "16deede6bb675c61c9c16deb967c89b316ddee8a5f7a2b0af1373cfbea470c40",
 50    "codex_colon.zip.003": "bbf6d96d3627050b1e227bc448ec3ebb773f5768b62d2e4357141e9355ba9246",
 51    "codex_colon.zip.004": "3ba68777b5cab20111118c03081b5e4aba67ae4e50e5c58fdd6732ffd49aee2c",
 52    "vectra_colon.zip.001": "fb23579672991954a0d6be32befa15bb716e14152389ec88c566eefc3c6e3dd1",
 53    "vectra_colon.zip.002": "4365420a09c6417a361382e44f85d0b8baa51001ff68fc22c46dc57b558d1667",
 54    "vectra_pancreas.zip.001": "efac85a30c1e01628ac3b0e6bf85844b724e549676ec18a7601f9ce9e1f5b1ed",
 55    "vectra_pancreas.zip.002": "16e4a5180ecaf1f184c8e1d9b2362873ff10b282d7b39599f27d9e134b63e472",
 56    "mibi_breast.zip.001": "357c146e37489bca02f52aed68b1736ac434f106771dc1345443720a72a41256",
 57    "mibi_breast.zip.002": "82ceb444dd7b3501afce3ce2078b2218047ce6bbb005ff063826be2df349720a",
 58    "mibi_breast.zip.003": "aa3361bf33a66d2358ff01c96583d6d388aaad4da06800384e8493fd1e5a0a72",
 59    "mibi_breast.zip.004": "6b61d95a37e9895f5933f55f08e65ab516a1ce663efa67928222ffb4dd53feae",
 60    "mibi_breast.zip.005": "71c250197c99e4ba640ccdbc6231a0207731d8f961b6401caea0497f69eb8474",
 61    "mibi_breast.zip.006": "77083a417465eea42dca5d60c7f1ff8bedf1fdb33d03ca147150e1a89ef8c6b0",
 62    "mibi_breast.zip.007": "bbd95daaf51ee5bcdf9c2a9e1b2517513e0fcbe544a2d8e3e702d356c7f67f53",
 63    "mibi_breast.zip.008": "d0379baef5b30f47a1dcdd4a094898a81906987bc2b9d9cc87ac95ff70b3032e",
 64    "mibi_breast.zip.009": "9d24f342ac90bd41823aa68f537128ff79be9ee1b321a0faec5fc0c1848cdef6",
 65    "mibi_breast.zip.010": "862824597edef36782ad77937725118304a77ca3e776ed682cb95facdec1b12d",
 66    "mibi_breast.zip.011": "bfcae2f5c9396558b33a0629a77105bef8ef63700e17db38968ffa741f9b1692",
 67}
 68
 69SPLIT_CONFIGS = {
 70    "mibi_decidua": "decidua_split.json",
 71    "codex_colon": "hickey_split.json",
 72    "vectra_colon": "msk_colon_split.json",
 73    "vectra_pancreas": "msk_pancreas_split.json",
 74    "mibi_breast": "tonic_split.json",
 75}
 76
 77# The channels that make up the nuclei and membrane composites, taken from the data preparation
 78# scripts of https://github.com/angelolab/Nimbus. The panel of 'codex_colon' is not documented there,
 79# so its channels come from the marker names in the archive, where DRAQ5 is the nuclear stain.
 80NUCLEI_CHANNELS = {
 81    "mibi_decidua": ["H3"],
 82    "mibi_breast": ["H3K27me3", "H3K9ac"],
 83    "vectra_colon": ["DAPI"],
 84    "vectra_pancreas": ["DAPI"],
 85    "codex_colon": ["DRAQ5"],
 86}
 87
 88MEMBRANE_CHANNELS = {
 89    "mibi_decidua": ["VIM", "HLAG", "CD3", "CD14", "CD56"],
 90    "mibi_breast": ["CD45", "ECAD", "CD14", "CD38", "CK17"],
 91    "vectra_colon": ["CD3", "CD8", "ICOS", "panCK+CK7+CAM5.2"],
 92    "vectra_pancreas": ["CD8", "CD40", "CD40-L", "panCK"],
 93    "codex_colon": ["CD45", "Podoplanin", "CD31", "Cytokeratin", "aSMA", "Vimentin"],
 94}
 95
 96# The subsets do not share a folder layout, so the image folder is named per subset.
 97IMAGE_DIRS = {
 98    "mibi_decidua": "image_data",
 99    "mibi_breast": "image_data",
100    "codex_colon": "raw_structured",
101    "vectra_colon": "raw_structured",
102    "vectra_pancreas": "raw_structured",
103}
104
105
106def _download_subset(path, subset, download):
107    """Download the parts of one subset and join them into a single archive."""
108    zip_path = os.path.join(path, f"{subset}.zip")
109    if os.path.exists(zip_path):
110        return zip_path
111
112    part_paths = []
113    for fname in SUBSET_PARTS[subset]:
114        part_path = os.path.join(path, fname)
115        util.download_source(
116            path=part_path, url=f"{URL}/{fname}", download=download, checksum=CHECKSUMS[fname]
117        )
118        part_paths.append(part_path)
119
120    if len(part_paths) == 1:
121        return part_paths[0]
122
123    # The parts are a plain byte split, so joining them yields the original archive.
124    # Each part is removed right after it is joined to limit the peak disk usage.
125    # The archive is joined under a temporary name, so that an interrupted join is not
126    # mistaken for a complete archive by the check above.
127    tmp_path = f"{zip_path}.incomplete"
128    with open(tmp_path, "wb") as dst:
129        for part_path in part_paths:
130            with open(part_path, "rb") as src:
131                shutil.copyfileobj(src, dst, length=32 * 1024 * 1024)
132            os.remove(part_path)
133    os.replace(tmp_path, zip_path)
134
135    return zip_path
136
137
138def _get_split_assignment(path, subset):
139    """Map every field of view to the split that the authors of the publication assigned to it."""
140    config = SPLIT_CONFIGS[subset]
141    config_path = os.path.join(path, config)
142    util.download_source(path=config_path, url=f"{SPLIT_URL}/{config}", download=True, checksum=None)
143
144    with open(config_path, "rb") as f:
145        payload = f.read()
146    if payload[:2] == b"\x1f\x8b":  # the server may answer with a gzip encoded body
147        payload = gzip.decompress(payload)
148    splits = json.loads(payload)
149
150    assignment = {}
151    for split, fovs in splits.items():
152        split = "val" if split == "validation" else split
153        for fov in fovs:
154            assignment[fov] = split
155    return assignment
156
157
158def _get_marker_paths(image_dir):
159    """Map every marker name in a field of view to its file.
160
161    The folder names contain brackets and commas in some subsets, so this avoids glob patterns.
162    """
163    markers = {}
164    for fname in os.listdir(image_dir):
165        if not fname.endswith((".tif", ".tiff")):
166            continue
167        name = fname[:-len(".ome.tif")] if fname.endswith(".ome.tif") else os.path.splitext(fname)[0]
168        markers[name] = os.path.join(image_dir, fname)
169    return markers
170
171
172def _get_composite(markers, channels):
173    """Sum the given marker channels after normalizing each of them by its upper quantile."""
174    stack = []
175    for name in channels:
176        if name not in markers:
177            continue
178        image = np.squeeze(tifffile.imread(markers[name])).astype("float32")
179        upper = np.quantile(image, 0.999)
180        stack.append(image / upper if upper > 0 else image)
181
182    if not stack:
183        return None
184    return np.clip(np.sum(stack, axis=0), 0, 1)
185
186
187def _find_instance_path(input_dir, subset, fov):
188    """Resolve the cell mask of a field of view, which is named differently in every subset."""
189    if subset == "mibi_decidua":
190        candidate = os.path.join(input_dir, "segmentation_data", f"{fov}_segmentation_labels.tiff")
191        return candidate if os.path.exists(candidate) else None
192
193    if subset == "mibi_breast":
194        candidate = os.path.join(input_dir, "segmentation_data", f"{fov}_feature_0.tif")
195        return candidate if os.path.exists(candidate) else None
196
197    if subset in ("vectra_colon", "vectra_pancreas"):
198        # The mask repeats the folder name of the field of view and appends the deepcell suffix.
199        seg_dir = os.path.join(input_dir, "segmentation")
200        for suffix in ("feature_0.ome.tif", "feature_0.tif"):
201            candidate = os.path.join(seg_dir, f"{fov}{suffix}")
202            if os.path.exists(candidate):
203                return candidate
204        return None
205
206    # For codex_colon the mask name differs from the field of view, so it is matched on the
207    # sample id and the region, e.g. 'B012B_reg004_X01_Y01_Z01' -> '.../B012B/B012B_..._reg004_..._labeled.ome.tif'.
208    tokens = fov.split("_")
209    sample = tokens[0]
210    region = next((t for t in tokens if t.startswith("reg")), None)
211    sample_dir = os.path.join(input_dir, "masks", sample)
212    if region is None or not os.path.isdir(sample_dir):
213        return None
214    for fname in sorted(os.listdir(sample_dir)):
215        if region in fname and fname.endswith("_labeled.ome.tif"):
216            return os.path.join(sample_dir, fname)
217    return None
218
219
220def _preprocess_data(input_dir, data_dir, subset):
221    import h5py
222
223    os.makedirs(data_dir, exist_ok=True)
224    assignment = _get_split_assignment(os.path.dirname(data_dir), subset)
225
226    image_root = os.path.join(input_dir, IMAGE_DIRS[subset])
227    if not os.path.isdir(image_root):
228        raise RuntimeError(f"Could not find the image folder '{IMAGE_DIRS[subset]}' of '{subset}' in {input_dir}.")
229
230    image_dirs = natsorted(os.path.join(image_root, name) for name in os.listdir(image_root))
231
232    n_without_mask = 0
233    for image_dir in image_dirs:
234        if not os.path.isdir(image_dir):
235            continue
236
237        fov = os.path.basename(image_dir)
238        out_path = os.path.join(data_dir, f"{fov}.h5")
239        if os.path.exists(out_path):
240            continue
241
242        instance_path = _find_instance_path(input_dir, subset, fov)
243        if instance_path is None:
244            n_without_mask += 1
245            continue
246
247        markers = _get_marker_paths(image_dir)
248        nuclei = _get_composite(markers, NUCLEI_CHANNELS[subset])
249        membrane = _get_composite(markers, MEMBRANE_CHANNELS[subset])
250        if nuclei is None:
251            raise RuntimeError(
252                f"None of the nuclei channels {NUCLEI_CHANNELS[subset]} are present for '{fov}'. "
253                f"The channels in the data are {sorted(markers)[:20]}. Please adapt 'NUCLEI_CHANNELS'."
254            )
255        if membrane is None:
256            membrane = np.zeros_like(nuclei)
257
258        instances = np.squeeze(tifffile.imread(instance_path)).astype("uint32")
259
260        # The volume is written under a temporary name and renamed once it is complete, so that an
261        # interrupted write does not leave a partial file that the check above would skip forever.
262        tmp_path = f"{out_path}.incomplete"
263        with h5py.File(tmp_path, "w") as f:
264            # A field of view that the split config does not list stays out of all three splits.
265            f.attrs["split"] = assignment.get(fov, "unassigned")
266            f.attrs["subset"] = subset
267            f.create_dataset("raw/nuclei", data=nuclei, compression="gzip")
268            f.create_dataset("raw/membrane", data=membrane, compression="gzip")
269            f.create_dataset("labels/cell", data=instances, compression="gzip")
270        os.replace(tmp_path, out_path)
271
272    if n_without_mask > 0:
273        warn(f"{n_without_mask} fields of view of '{subset}' have no cell mask and are not part of the dataset.")
274
275
276def get_pan_multiplex_data(
277    path: Union[os.PathLike, str],
278    subset: Literal["mibi_decidua", "mibi_breast", "codex_colon", "vectra_colon", "vectra_pancreas"],
279    download: bool = False,
280) -> str:
281    """Download one subset of the Pan-Multiplex dataset.
282
283    Args:
284        path: The folder where the function stores the data.
285        subset: The subset of the dataset. See `SUBSET_PARTS` for the valid choices.
286        download: Whether to download the data if it is not present.
287
288    Returns:
289        The filepath to the folder with the prepared data.
290    """
291    if subset not in SUBSET_PARTS:
292        raise ValueError(f"'{subset}' is not a valid subset. Choose one of {list(SUBSET_PARTS.keys())}.")
293
294    subset_dir = os.path.join(path, subset)
295    data_dir = os.path.join(subset_dir, "data")
296
297    # The preprocessing writes this marker once it has converted every field of view. Checking for
298    # converted volumes instead would treat an interrupted run as a complete dataset.
299    done_path = os.path.join(subset_dir, "preprocessing_done")
300    if os.path.exists(done_path):
301        return data_dir
302
303    os.makedirs(subset_dir, exist_ok=True)
304
305    input_dir = os.path.join(subset_dir, subset)
306    if not os.path.exists(input_dir):
307        # The data may have been prepared before this marker was introduced, in which case the
308        # input data can already have been removed. Then the converted volumes are all there is.
309        if glob(os.path.join(data_dir, "*.h5")):
310            open(done_path, "w").close()
311            return data_dir
312
313        zip_path = _download_subset(subset_dir, subset, download)
314        util.unzip(zip_path=zip_path, dst=subset_dir, remove=True)
315
316    _preprocess_data(input_dir, data_dir, subset)
317    open(done_path, "w").close()
318
319    return data_dir
320
321
322def get_pan_multiplex_paths(
323    path: Union[os.PathLike, str],
324    subset: Union[str, List[str]],
325    split: Optional[Literal["train", "val", "test"]] = None,
326    download: bool = False,
327) -> List[str]:
328    """Get the paths to the Pan-Multiplex data.
329
330    Args:
331        path: The folder where the function stores the data.
332        subset: One subset or a list of subsets. See `SUBSET_PARTS` for the valid choices.
333        split: The data split. The function uses all fields of view by default.
334        download: Whether to download the data if it is not present.
335
336    Returns:
337        The list of filepaths to the input data.
338    """
339    import h5py
340
341    if split is not None and split not in ("train", "val", "test"):
342        raise ValueError(f"'{split}' is not a valid split. Choose 'train', 'val' or 'test'.")
343
344    subsets = [subset] if isinstance(subset, str) else subset
345    volume_paths = []
346    for name in subsets:
347        data_dir = get_pan_multiplex_data(path, name, download)
348        volume_paths.extend(natsorted(glob(os.path.join(data_dir, "*.h5"))))
349
350    if split is not None:
351        selected, n_unassigned = [], 0
352        for volume_path in volume_paths:
353            with h5py.File(volume_path, "r") as f:
354                this_split = f.attrs.get("split")
355            if this_split == split:
356                selected.append(volume_path)
357            elif this_split == "unassigned":
358                n_unassigned += 1
359        volume_paths = selected
360
361        # The upstream split configs do not cover every field of view, e.g. 53 of the 432 fields of
362        # view of 'vectra_pancreas' are missing from theirs. Those are not part of any split.
363        if n_unassigned > 0:
364            warn(
365                f"{n_unassigned} fields of view are not listed in the split config of '{subset}' and are "
366                "left out of all splits. Load the data without a split to also use these fields of view."
367            )
368
369    assert len(volume_paths) > 0, f"Could not find data for the subset '{subset}' and split '{split}'."
370    return volume_paths
371
372
373def get_pan_multiplex_dataset(
374    path: Union[os.PathLike, str],
375    patch_shape: Tuple[int, int],
376    subset: Union[str, List[str]],
377    split: Optional[Literal["train", "val", "test"]] = None,
378    raw_channel: Literal["both", "nuclei", "membrane"] = "both",
379    download: bool = False,
380    **kwargs
381) -> Dataset:
382    """Get the Pan-Multiplex dataset for cell segmentation in multiplexed images.
383
384    Args:
385        path: The folder where the function stores the data.
386        patch_shape: The patch shape to use for training.
387        subset: One subset or a list of subsets. See `SUBSET_PARTS` for the valid choices.
388        split: The data split. The function uses all fields of view by default.
389        raw_channel: The input channels. Use 'both' for the nuclei and the membrane composite,
390            or one of 'nuclei' and 'membrane' for a single channel.
391        download: Whether to download the data if it is not present.
392        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
393            This uses a `MinInstanceSampler` by default, to avoid patches without any cell.
394
395    Returns:
396        The segmentation dataset.
397    """
398    if raw_channel not in ("both", "nuclei", "membrane"):
399        raise ValueError(f"'{raw_channel}' is not a valid raw channel. Choose 'both', 'nuclei' or 'membrane'.")
400
401    volume_paths = get_pan_multiplex_paths(path, subset, split, download)
402
403    # The two composites are stacked into channels on the fly, so that the combined input does not
404    # have to be stored a second time.
405    raw_key = ["raw/nuclei", "raw/membrane"] if raw_channel == "both" else f"raw/{raw_channel}"
406    kwargs = util.update_kwargs(kwargs, "with_channels", raw_channel == "both")
407
408    # The fields of view of some subsets are much larger than the patch shape and contain empty
409    # regions, e.g. 7.5% of the random patches of 'codex_colon' do not contain a single cell.
410    # This is only a default, so that it can be overridden with a different sampler.
411    kwargs.setdefault("sampler", MinInstanceSampler())
412
413    return torch_em.default_segmentation_dataset(
414        raw_paths=volume_paths,
415        raw_key=raw_key,
416        label_paths=volume_paths,
417        label_key="labels/cell",
418        patch_shape=patch_shape,
419        is_seg_dataset=True,
420        ndim=2,
421        **kwargs
422    )
423
424
425def get_pan_multiplex_loader(
426    path: Union[os.PathLike, str],
427    batch_size: int,
428    patch_shape: Tuple[int, int],
429    subset: Union[str, List[str]],
430    split: Optional[Literal["train", "val", "test"]] = None,
431    raw_channel: Literal["both", "nuclei", "membrane"] = "both",
432    download: bool = False,
433    **kwargs
434) -> DataLoader:
435    """Get the Pan-Multiplex dataloader for cell segmentation in multiplexed images.
436
437    Args:
438        path: The folder where the function stores the data.
439        batch_size: The batch size for training.
440        patch_shape: The patch shape to use for training.
441        subset: One subset or a list of subsets. See `SUBSET_PARTS` for the valid choices.
442        split: The data split. The function uses all fields of view by default.
443        raw_channel: The input channels. Use 'both' for the nuclei and the membrane composite,
444            or one of 'nuclei' and 'membrane' for a single channel.
445        download: Whether to download the data if it is not present.
446        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
447            This uses a `MinInstanceSampler` by default, to avoid patches without any cell.
448
449    Returns:
450        The DataLoader.
451    """
452    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
453    dataset = get_pan_multiplex_dataset(path, patch_shape, subset, split, raw_channel, download, **ds_kwargs)
454    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)