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 _recover_saturated_ids(instances, rgb_path):
221    """Recover the cell ids that a uint16 mask clips at 65535 from the RGB rendering of the mask.
222
223    Four CODEX colon fields hold more than 65534 cells. Their '_labeled.ome.tif' is uint16, so every cell
224    from id 65535 on shares that value. The '_masks.png' next to it draws each cell in its own colour, with a
225    one pixel edge ring in the colour of a neighbour. The cells are the colour components with an interior
226    pixel inside the saturated region. The rings go to the nearest such component.
227    """
228    from PIL import Image
229    from skimage.measure import label
230    from skimage.segmentation import expand_labels
231
232    saturated_id = np.iinfo(instances.dtype).max
233    saturated = instances == saturated_id
234
235    Image.MAX_IMAGE_PIXELS = None
236    rgb = np.asarray(Image.open(rgb_path)).astype("uint32")
237    colour = rgb[..., 0] + (rgb[..., 1] << 8) + (rgb[..., 2] << 16)
238    components = label(np.where(saturated, colour + 1, 0), connectivity=1)
239
240    interior = components > 0
241    for axis in (0, 1):
242        for shift in (1, -1):
243            interior &= np.roll(components, shift, axis=axis) == components
244    keep = np.zeros(components.max() + 1, dtype=bool)
245    keep[np.unique(components[interior])] = True
246    keep[0] = False
247    cores = np.unique(np.where(keep[components], components, 0), return_inverse=True)[1].reshape(components.shape)
248
249    cells = expand_labels(cores, distance=3)
250    cells[~saturated] = 0
251    recovered = instances.astype("uint32")
252    recovered[saturated] = 0
253    recovered[cells > 0] = cells[cells > 0] + saturated_id - 1
254    return recovered
255
256
257def _load_instances(instance_path):
258    """Load a cell mask and recover the ids that its uint16 encoding clipped."""
259    instances = np.squeeze(tifffile.imread(instance_path))
260    rgb_path = instance_path.replace("_labeled.ome.tif", "_masks.png")
261    if instances.dtype == np.uint16 and instances.max() == np.iinfo(np.uint16).max and os.path.exists(rgb_path):
262        instances = _recover_saturated_ids(instances, rgb_path)
263    return instances.astype("uint32")
264
265
266def _preprocess_data(input_dir, data_dir, subset):
267    import h5py
268
269    os.makedirs(data_dir, exist_ok=True)
270    assignment = _get_split_assignment(os.path.dirname(data_dir), subset)
271
272    image_root = os.path.join(input_dir, IMAGE_DIRS[subset])
273    if not os.path.isdir(image_root):
274        raise RuntimeError(f"Could not find the image folder '{IMAGE_DIRS[subset]}' of '{subset}' in {input_dir}.")
275
276    image_dirs = natsorted(os.path.join(image_root, name) for name in os.listdir(image_root))
277
278    n_without_mask = 0
279    for image_dir in image_dirs:
280        if not os.path.isdir(image_dir):
281            continue
282
283        fov = os.path.basename(image_dir)
284        out_path = os.path.join(data_dir, f"{fov}.h5")
285        if os.path.exists(out_path):
286            continue
287
288        instance_path = _find_instance_path(input_dir, subset, fov)
289        if instance_path is None:
290            n_without_mask += 1
291            continue
292
293        markers = _get_marker_paths(image_dir)
294        nuclei = _get_composite(markers, NUCLEI_CHANNELS[subset])
295        membrane = _get_composite(markers, MEMBRANE_CHANNELS[subset])
296        if nuclei is None:
297            raise RuntimeError(
298                f"None of the nuclei channels {NUCLEI_CHANNELS[subset]} are present for '{fov}'. "
299                f"The channels in the data are {sorted(markers)[:20]}. Please adapt 'NUCLEI_CHANNELS'."
300            )
301        if membrane is None:
302            membrane = np.zeros_like(nuclei)
303
304        instances = _load_instances(instance_path)
305
306        # The volume is written under a temporary name and renamed once it is complete, so that an
307        # interrupted write does not leave a partial file that the check above would skip forever.
308        tmp_path = f"{out_path}.incomplete"
309        with h5py.File(tmp_path, "w") as f:
310            # A field of view that the split config does not list stays out of all three splits.
311            f.attrs["split"] = assignment.get(fov, "unassigned")
312            f.attrs["subset"] = subset
313            f.create_dataset("raw/nuclei", data=nuclei, compression="gzip")
314            f.create_dataset("raw/membrane", data=membrane, compression="gzip")
315            f.create_dataset("labels/cell", data=instances, compression="gzip")
316        os.replace(tmp_path, out_path)
317
318    if n_without_mask > 0:
319        warn(f"{n_without_mask} fields of view of '{subset}' have no cell mask and are not part of the dataset.")
320
321
322def get_pan_multiplex_data(
323    path: Union[os.PathLike, str],
324    subset: Literal["mibi_decidua", "mibi_breast", "codex_colon", "vectra_colon", "vectra_pancreas"],
325    download: bool = False,
326) -> str:
327    """Download one subset of the Pan-Multiplex dataset.
328
329    Args:
330        path: The folder where the function stores the data.
331        subset: The subset of the dataset. See `SUBSET_PARTS` for the valid choices.
332        download: Whether to download the data if it is not present.
333
334    Returns:
335        The filepath to the folder with the prepared data.
336    """
337    if subset not in SUBSET_PARTS:
338        raise ValueError(f"'{subset}' is not a valid subset. Choose one of {list(SUBSET_PARTS.keys())}.")
339
340    subset_dir = os.path.join(path, subset)
341    data_dir = os.path.join(subset_dir, "data")
342
343    # The preprocessing writes this marker once it has converted every field of view. Checking for
344    # converted volumes instead would treat an interrupted run as a complete dataset.
345    done_path = os.path.join(subset_dir, "preprocessing_done")
346    if os.path.exists(done_path):
347        return data_dir
348
349    os.makedirs(subset_dir, exist_ok=True)
350
351    input_dir = os.path.join(subset_dir, subset)
352    if not os.path.exists(input_dir):
353        # The data may have been prepared before this marker was introduced, in which case the
354        # input data can already have been removed. Then the converted volumes are all there is.
355        if glob(os.path.join(data_dir, "*.h5")):
356            open(done_path, "w").close()
357            return data_dir
358
359        zip_path = _download_subset(subset_dir, subset, download)
360        util.unzip(zip_path=zip_path, dst=subset_dir, remove=True)
361
362    _preprocess_data(input_dir, data_dir, subset)
363    open(done_path, "w").close()
364
365    return data_dir
366
367
368def get_pan_multiplex_paths(
369    path: Union[os.PathLike, str],
370    subset: Union[str, List[str]],
371    split: Optional[Literal["train", "val", "test"]] = None,
372    download: bool = False,
373) -> List[str]:
374    """Get the paths to the Pan-Multiplex data.
375
376    Args:
377        path: The folder where the function stores the data.
378        subset: One subset or a list of subsets. See `SUBSET_PARTS` for the valid choices.
379        split: The data split. The function uses all fields of view by default.
380        download: Whether to download the data if it is not present.
381
382    Returns:
383        The list of filepaths to the input data.
384    """
385    import h5py
386
387    if split is not None and split not in ("train", "val", "test"):
388        raise ValueError(f"'{split}' is not a valid split. Choose 'train', 'val' or 'test'.")
389
390    subsets = [subset] if isinstance(subset, str) else subset
391    volume_paths = []
392    for name in subsets:
393        data_dir = get_pan_multiplex_data(path, name, download)
394        volume_paths.extend(natsorted(glob(os.path.join(data_dir, "*.h5"))))
395
396    if split is not None:
397        selected, n_unassigned = [], 0
398        for volume_path in volume_paths:
399            with h5py.File(volume_path, "r") as f:
400                this_split = f.attrs.get("split")
401            if this_split == split:
402                selected.append(volume_path)
403            elif this_split == "unassigned":
404                n_unassigned += 1
405        volume_paths = selected
406
407        # The upstream split configs do not cover every field of view, e.g. 53 of the 432 fields of
408        # view of 'vectra_pancreas' are missing from theirs. Those are not part of any split.
409        if n_unassigned > 0:
410            warn(
411                f"{n_unassigned} fields of view are not listed in the split config of '{subset}' and are "
412                "left out of all splits. Load the data without a split to also use these fields of view."
413            )
414
415    assert len(volume_paths) > 0, f"Could not find data for the subset '{subset}' and split '{split}'."
416    return volume_paths
417
418
419def get_pan_multiplex_dataset(
420    path: Union[os.PathLike, str],
421    patch_shape: Tuple[int, int],
422    subset: Union[str, List[str]],
423    split: Optional[Literal["train", "val", "test"]] = None,
424    raw_channel: Literal["both", "nuclei", "membrane"] = "both",
425    download: bool = False,
426    **kwargs
427) -> Dataset:
428    """Get the Pan-Multiplex dataset for cell segmentation in multiplexed images.
429
430    Args:
431        path: The folder where the function stores the data.
432        patch_shape: The patch shape to use for training.
433        subset: One subset or a list of subsets. See `SUBSET_PARTS` for the valid choices.
434        split: The data split. The function uses all fields of view by default.
435        raw_channel: The input channels. Use 'both' for the nuclei and the membrane composite,
436            or one of 'nuclei' and 'membrane' for a single channel.
437        download: Whether to download the data if it is not present.
438        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
439            This uses a `MinInstanceSampler` by default, to avoid patches without any cell.
440
441    Returns:
442        The segmentation dataset.
443    """
444    if raw_channel not in ("both", "nuclei", "membrane"):
445        raise ValueError(f"'{raw_channel}' is not a valid raw channel. Choose 'both', 'nuclei' or 'membrane'.")
446
447    volume_paths = get_pan_multiplex_paths(path, subset, split, download)
448
449    # The two composites are stacked into channels on the fly, so that the combined input does not
450    # have to be stored a second time.
451    raw_key = ["raw/nuclei", "raw/membrane"] if raw_channel == "both" else f"raw/{raw_channel}"
452    kwargs = util.update_kwargs(kwargs, "with_channels", raw_channel == "both")
453
454    # The fields of view of some subsets are much larger than the patch shape and contain empty
455    # regions, e.g. 7.5% of the random patches of 'codex_colon' do not contain a single cell.
456    # This is only a default, so that it can be overridden with a different sampler.
457    kwargs.setdefault("sampler", MinInstanceSampler())
458
459    return torch_em.default_segmentation_dataset(
460        raw_paths=volume_paths,
461        raw_key=raw_key,
462        label_paths=volume_paths,
463        label_key="labels/cell",
464        patch_shape=patch_shape,
465        is_seg_dataset=True,
466        ndim=2,
467        **kwargs
468    )
469
470
471def get_pan_multiplex_loader(
472    path: Union[os.PathLike, str],
473    batch_size: int,
474    patch_shape: Tuple[int, int],
475    subset: Union[str, List[str]],
476    split: Optional[Literal["train", "val", "test"]] = None,
477    raw_channel: Literal["both", "nuclei", "membrane"] = "both",
478    download: bool = False,
479    **kwargs
480) -> DataLoader:
481    """Get the Pan-Multiplex dataloader for cell segmentation in multiplexed images.
482
483    Args:
484        path: The folder where the function stores the data.
485        batch_size: The batch size for training.
486        patch_shape: The patch shape to use for training.
487        subset: One subset or a list of subsets. See `SUBSET_PARTS` for the valid choices.
488        split: The data split. The function uses all fields of view by default.
489        raw_channel: The input channels. Use 'both' for the nuclei and the membrane composite,
490            or one of 'nuclei' and 'membrane' for a single channel.
491        download: Whether to download the data if it is not present.
492        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
493            This uses a `MinInstanceSampler` by default, to avoid patches without any cell.
494
495    Returns:
496        The DataLoader.
497    """
498    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
499    dataset = get_pan_multiplex_dataset(path, patch_shape, subset, split, raw_channel, download, **ds_kwargs)
500    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)