torch_em.data.datasets.electron_microscopy.platynereis

  1"""Dataset for the segmentation of different structures in EM volume of a
  2platynereis larve. Contains annotations for the segmentation of:
  3- Cuticle
  4- Cilia
  5- Cells
  6- Nuclei
  7
  8This dataset is from the publication https://doi.org/10.1016/j.cell.2021.07.017.
  9Please cite it if you use this dataset for a publication.
 10
 11The cell dataset stores corrected labels separately from the source annotations. It maps the
 12neuropil IDs in `CELL_NEUROPIL_IDS` to `ignore_label` before sampling training patches.
 13"""
 14
 15import os
 16from glob import glob
 17from tempfile import TemporaryDirectory
 18from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 19
 20import numpy as np
 21from elf.io import open_file
 22from skimage.segmentation import flood
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27from torch_em.data import ConcatDataset
 28
 29from .. import util
 30
 31
 32URLS = {
 33    "cells": "https://zenodo.org/record/3675220/files/membrane.zip",
 34    "nuclei": "https://zenodo.org/record/3675220/files/nuclei.zip",
 35    "cilia": "https://zenodo.org/record/3675220/files/cilia.zip",
 36    "cuticle": "https://zenodo.org/record/3675220/files/cuticle.zip"
 37}
 38
 39CHECKSUMS = {
 40    "cells": "30eb50c39e7e9883e1cd96e0df689fac37a56abb11e8ed088907c94a5980d6a3",
 41    "nuclei": "a05033c5fbc6a3069479ac6595b0a430070f83f5281f5b5c8913125743cf5510",
 42    "cilia": "6d2b47f63d39a671789c02d8b66cad5e4cf30eb14cdb073da1a52b7defcc5e24",
 43    "cuticle": "464f75d30133e8864958049647fe3c2216ddf2d4327569738ad72d299c991843"
 44}
 45
 46FILE_TEMPLATES = {
 47    "cells": "train_data_membrane_%02i.n5",
 48    "nuclei": "train_data_nuclei_%02i.h5",
 49    "cilia": "train_data_cilia_%02i.h5",
 50    "cuticle": "train_data_%02i.n5",
 51}
 52
 53# Default ignore label. Stays exact through a float32 round-trip, which the augmentations do,
 54# and sits far above any real instance id.
 55CELL_IGNORE_LABEL = 2 ** 24 - 1
 56# Increment this version when corrections change, so stored labels and ROI caches are regenerated.
 57CELL_LABEL_VERSION = 1
 58CELL_SOURCE_LABEL_KEY = "volumes/labels/segmentation/s1"
 59CELL_BACKGROUND_IDS = {7: (63,), 8: (5,)}
 60
 61# Ids that label neuropil rather than a single cell, found by inspecting every cell volume.
 62CELL_NEUROPIL_IDS = {
 63    1: (),
 64    2: (),
 65    3: (253,),
 66    4: (58,),
 67    5: (72,),
 68    6: (73,),
 69    7: (),
 70    8: (43,),
 71    9: (19, 1084),
 72}
 73
 74
 75def get_platynereis_cell_neuropil_ids(sample_id: int) -> Tuple[int, ...]:
 76    """Get the neuropil instance ids of a platynereis cell volume.
 77
 78    Args:
 79        sample_id: The id of the volume, between 1 and 9.
 80
 81    Returns:
 82        The instance ids that label neuropil rather than a single cell. Empty if the volume has none.
 83    """
 84    return CELL_NEUROPIL_IDS.get(sample_id, ())
 85
 86
 87def get_platynereis_cell_label_key(ignore_label: int = CELL_IGNORE_LABEL) -> str:
 88    """Get the N5 key for corrected cell labels.
 89
 90    Args:
 91        ignore_label: The value assigned to neuropil voxels.
 92
 93    Returns:
 94        The dataset key, including the correction version and ignore label.
 95    """
 96    return f"volumes/labels/segmentation_corrected/v{CELL_LABEL_VERSION}/ignore_{ignore_label}/s1"
 97
 98
 99def _split_cell_label(labels, source_id, target_id, seed, expected_size, expected_bbox):
100    if any(c >= size for c, size in zip(seed, labels.shape)) or labels[seed] != source_id:
101        raise ValueError("The muscle seed does not point to the expected source label.")
102    if np.any(labels == target_id):
103        raise ValueError(f"The muscle target label {target_id} is already assigned.")
104    muscle = flood(labels, seed, connectivity=3)
105    coordinates = np.where(muscle)
106    bbox = tuple((int(c.min()), int(c.max()) + 1) for c in coordinates)
107    if int(muscle.sum()) != expected_size or bbox != expected_bbox:
108        raise ValueError("The muscle component differs from the inspected size or bounding box.")
109    labels[muscle] = target_id
110
111
112def _correct_cell_labels(labels, sample_id, ignore_label):
113    labels = labels.astype("int64" if ignore_label < 0 else "uint64", copy=True)
114    if ignore_label != 0 and np.any(labels == ignore_label):
115        raise ValueError(f"The ignore label {ignore_label} is already assigned to an instance.")
116    if sample_id == 5:
117        if ignore_label == 256:
118            raise ValueError("The ignore label conflicts with the new muscle label 256.")
119        _split_cell_label(
120            labels, source_id=72, target_id=256, seed=(29, 70, 242), expected_size=102740,
121            expected_bbox=((15, 80), (64, 129), (185, 293)),
122        )
123    for label_id in CELL_BACKGROUND_IDS.get(sample_id, ()):
124        labels[labels == label_id] = 0
125    for label_id in get_platynereis_cell_neuropil_ids(sample_id):
126        labels[labels == label_id] = ignore_label
127    return labels
128
129
130def _prepare_cell_labels(path, sample_id, ignore_label):
131    key = get_platynereis_cell_label_key(ignore_label)
132    target = os.path.join(path, key)
133    if os.path.exists(target):
134        return
135    with open_file(path, "r") as f:
136        source = f[CELL_SOURCE_LABEL_KEY]
137        labels = _correct_cell_labels(source[:], sample_id, ignore_label)
138        chunks = source.chunks
139        spatial_attrs = {name: source.attrs[name] for name in ("offset", "global_offset") if name in source.attrs}
140
141    # Publish the complete dataset atomically. Other training ranks may prepare it at the same time.
142    with TemporaryDirectory(dir=path, prefix=".cell-labels-") as tmp:
143        tmp_path = os.path.join(tmp, "labels.n5")
144        with open_file(tmp_path, "a") as f:
145            ds = f.create_dataset("labels", data=labels, chunks=chunks, compression="gzip")
146            ds.attrs.update(spatial_attrs)
147            ds.attrs.update({
148                "correction_version": CELL_LABEL_VERSION,
149                "source_key": CELL_SOURCE_LABEL_KEY,
150                "sample_id": sample_id,
151                "ignore_label": ignore_label,
152                "neuropil_ids": list(get_platynereis_cell_neuropil_ids(sample_id)),
153                "background_ids": list(CELL_BACKGROUND_IDS.get(sample_id, ())),
154            })
155        os.makedirs(os.path.dirname(target), exist_ok=True)
156        try:
157            os.rename(os.path.join(tmp_path, "labels"), target)
158        except OSError:
159            if not os.path.isdir(target):
160                raise
161
162
163def prepare_platynereis_cell_data(
164    path: Union[os.PathLike, str],
165    sample_ids: Optional[Sequence[int]] = None,
166    download: bool = False,
167    ignore_label: int = CELL_IGNORE_LABEL,
168) -> List[str]:
169    """Prepare corrected cell labels without changing the source annotations.
170
171    Corrections split the muscle cell in volume 5, remove false foreground in volumes 7 and 8,
172    and map neuropil IDs to the ignore label. Existing labels for this version are reused.
173
174    Args:
175        path: Folder containing the membrane subfolder.
176        sample_ids: Volume IDs to prepare. By default, prepare all nine volumes.
177        download: Whether to download missing source data.
178        ignore_label: The value assigned to neuropil voxels.
179
180    Returns:
181        The N5 paths in sample ID order. Read corrected labels with `get_platynereis_cell_label_key`.
182    """
183    sample_ids = list(range(1, 10)) if sample_ids is None else sorted(sample_ids)
184    paths = get_platynereis_paths(path, sample_ids, name="cells", download=download)
185    for sample_id, data_path in zip(sample_ids, paths):
186        _prepare_cell_labels(data_path, sample_id, ignore_label)
187    return paths
188
189
190#
191# TODO data-loader for more classes:
192# - mitos
193#
194
195
196def _check_data(path, prefix, extension, n_files):
197    if not os.path.exists(path):
198        return False
199    files = glob(os.path.join(path, f"{prefix}*{extension}"))
200    return len(files) == n_files
201
202
203def get_platynereis_data(path: Union[os.PathLike, str], name: str, download: bool) -> Tuple[str, int]:
204    """Download the platynereis dataset.
205
206    Args:
207        path: Filepath to a folder where the downloaded data will be saved.
208        name: Name of the segmentation task. Available tasks: 'cuticle', 'cilia', 'cells' or 'nuclei'.
209        download: Whether to download the data if it is not present.
210
211    Returns:
212        The path to the folder where the data has been downloaded.
213        The number of files downloaded.
214    """
215    data_root = os.path.join(path, name)
216
217    if name == "cuticle":
218        ext, prefix, n_files = ".n5", "train_data_", 5
219    elif name == "cilia":
220        ext, prefix, n_files = ".h5", "train_data_cilia_", 3
221    elif name == "cells":
222        data_root = os.path.join(path, "membrane")
223        ext, prefix, n_files = ".n5", "train_data_membrane_", 9
224    elif name == "nuclei":
225        ext, prefix, n_files = ".h5", "train_data_nuclei_", 12
226    else:
227        raise ValueError(f"Invalid name {name}. Expect one of 'cuticle', 'cilia', 'cell' or 'nuclei'.")
228
229    data_is_complete = _check_data(data_root, prefix, ext, n_files)
230    if data_is_complete:
231        return data_root, n_files
232
233    os.makedirs(path, exist_ok=True)
234    url = URLS[name]
235    checksum = CHECKSUMS[name]
236
237    zip_path = os.path.join(path, f"data-{name}.zip")
238    util.download_source(zip_path, url, download=download, checksum=checksum)
239    util.unzip(zip_path, path, remove=True)
240
241    return data_root, n_files
242
243
244def get_platynereis_paths(path, sample_ids, name, rois={}, download=False, return_rois=False):
245    """Get paths to the platynereis data.
246
247    Args:
248        path: Filepath to a folder where the downloaded data will be saved.
249        sample_ids: The sample ids to use for the dataset
250        name: Name of the segmentation task. Available tasks: 'cuticle', 'cilia', 'cells' or 'nuclei'.
251        rois: The region of interests to use for the data blocks.
252        download: Whether to download the data if it is not present.
253        return_rois: Whether to return the extracted rois.
254
255    Returns:
256        The filepaths for the stored data.
257    """
258    root, n_files = get_platynereis_data(path, name, download)
259    template = os.path.join(root, FILE_TEMPLATES[name])
260
261    if sample_ids is None:
262        sample_ids = list(range(1, n_files + 1))
263    else:
264        assert min(sample_ids) >= 1 and max(sample_ids) <= n_files
265        sample_ids.sort()
266    paths = [template % sample for sample in sample_ids]
267    data_rois = [rois.get(sample, np.s_[:, :, :]) for sample in sample_ids]
268
269    if return_rois:
270        return paths, data_rois
271    else:
272        return paths
273
274
275def get_platynereis_cuticle_dataset(
276    path: Union[os.PathLike, str],
277    patch_shape: Tuple[int, int, int],
278    sample_ids: Optional[Sequence[int]] = None,
279    download: bool = False,
280    rois: Dict[int, Any] = {},
281    **kwargs
282) -> Dataset:
283    """Get the dataset for cuticle segmentation in platynereis.
284
285    Args:
286        path: Filepath to a folder where the downloaded data will be saved.
287        patch_shape: The patch shape to use for training.
288        sample_ids: The sample ids to use for the dataset
289        download: Whether to download the data if it is not present.
290        rois: The region of interests to use for the data blocks.
291        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
292
293    Returns:
294        The segmentation dataset.
295    """
296    paths, data_rois = get_platynereis_paths(
297        path=path, sample_ids=sample_ids, name="cuticle", rois=rois, download=download, return_rois=True,
298    )
299    return torch_em.default_segmentation_dataset(
300        raw_paths=paths,
301        raw_key="volumes/raw",
302        label_paths=paths,
303        label_key="volumes/labels/segmentation",
304        patch_shape=patch_shape,
305        rois=data_rois,
306        **kwargs
307    )
308
309
310def get_platynereis_cuticle_loader(
311    path: Union[os.PathLike, str],
312    patch_shape: Tuple[int, int, int],
313    batch_size: int,
314    sample_ids: Optional[Sequence[int]] = None,
315    download: bool = False,
316    rois: Dict[int, Any] = {},
317    **kwargs
318) -> DataLoader:
319    """Get the dataloader for cuticle segmentation in platynereis.
320
321    Args:
322        path: Filepath to a folder where the downloaded data will be saved.
323        patch_shape: The patch shape to use for training.
324        batch_size: The batch size for training.
325        sample_ids: The sample ids to use for the dataset
326        download: Whether to download the data if it is not present.
327        rois: The region of interests to use for the data blocks.
328        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
329
330    Returns:
331        The DataLoader.
332    """
333    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
334    ds = get_platynereis_cuticle_dataset(
335        path, patch_shape, sample_ids=sample_ids, download=download, rois=rois, **ds_kwargs,
336    )
337    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
338
339
340def get_platynereis_cilia_dataset(
341    path: Union[os.PathLike, str],
342    patch_shape: Tuple[int, int, int],
343    sample_ids: Optional[Sequence[int]] = None,
344    offsets: Optional[List[List[int]]] = None,
345    boundaries: bool = False,
346    binary: bool = False,
347    rois: Dict[int, Any] = {},
348    download: bool = False,
349    **kwargs
350) -> Dataset:
351    """Get the dataset for cilia segmentation in platynereis.
352
353    Args:
354        path: Filepath to a folder where the downloaded data will be saved.
355        patch_shape: The patch shape to use for training.
356        sample_ids: The sample ids to use for the dataset
357        offsets: Offset values for affinity computation used as target.
358        boundaries: Whether to compute boundaries as the target.
359        binary: Whether to use a binary segmentation target.
360        rois: The region of interests to use for the data blocks.
361        download: Whether to download the data if it is not present.
362        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
363
364    Returns:
365        The segmentation dataset.
366    """
367    paths, rois = get_platynereis_paths(
368        path=path, sample_ids=sample_ids, name="cilia", rois=rois, download=download, return_rois=True,
369    )
370    kwargs = util.update_kwargs(kwargs, "rois", rois)
371    kwargs, _ = util.add_instance_label_transform(
372        kwargs, add_binary_target=True, boundaries=boundaries, offsets=offsets, binary=binary,
373    )
374    return torch_em.default_segmentation_dataset(
375        raw_paths=paths,
376        raw_key="volumes/raw",
377        label_paths=paths,
378        label_key="volumes/labels/segmentation",
379        patch_shape=patch_shape,
380        **kwargs
381    )
382
383
384def get_platynereis_cilia_loader(
385    path: Union[os.PathLike, str],
386    patch_shape: Tuple[int, int, int],
387    batch_size: int,
388    sample_ids: Optional[Sequence[int]] = None,
389    offsets: Optional[List[List[int]]] = None,
390    boundaries: bool = False,
391    binary: bool = False,
392    rois: Dict[int, Any] = {},
393    download: bool = False,
394    **kwargs
395) -> DataLoader:
396    """Get the dataloader for cilia segmentation in platynereis.
397
398    Args:
399        path: Filepath to a folder where the downloaded data will be saved.
400        patch_shape: The patch shape to use for training.
401        batch_size: The batch size for training.
402        sample_ids: The sample ids to use for the dataset
403        offsets: Offset values for affinity computation used as target.
404        boundaries: Whether to compute boundaries as the target.
405        binary: Whether to return a binary segmentation target.
406        rois: The region of interests to use for the data blocks.
407        download: Whether to download the data if it is not present.
408        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
409
410    Returns:
411        The DataLoader.
412    """
413    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
414    ds = get_platynereis_cilia_dataset(
415        path, patch_shape, sample_ids=sample_ids,
416        offsets=offsets, boundaries=boundaries, binary=binary,
417        rois=rois, download=download, **ds_kwargs,
418    )
419    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
420
421
422def get_platynereis_cell_dataset(
423    path: Union[os.PathLike, str],
424    patch_shape: Tuple[int, int, int],
425    sample_ids: Optional[Sequence[int]] = None,
426    offsets: Optional[List[List[int]]] = None,
427    boundaries: bool = False,
428    rois: Dict[int, Any] = {},
429    download: bool = False,
430    ignore_label: int = CELL_IGNORE_LABEL,
431    **kwargs
432) -> Dataset:
433    """Get the dataset for cell segmentation in platynereis.
434
435    Args:
436        path: Filepath to a folder where the downloaded data will be saved.
437        patch_shape: The patch shape to use for training.
438        sample_ids: The sample ids to use for the dataset
439        offsets: Offset values for affinity computation used as target.
440        boundaries: Whether to compute boundaries as the target.
441        rois: The region of interests to use for the data blocks.
442        download: Whether to download the data if it is not present.
443        ignore_label: The value the neuropil ids of `CELL_NEUROPIL_IDS` are mapped to, so that a loss
444            can exclude them.
445        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
446
447    Returns:
448        The segmentation dataset.
449    """
450    data_paths, data_rois = get_platynereis_paths(
451        path=path, sample_ids=sample_ids, name="cells", rois=rois, download=download, return_rois=True,
452    )
453    prepare_platynereis_cell_data(path, sample_ids, download=download, ignore_label=ignore_label)
454
455    kwargs, _ = util.add_instance_label_transform(
456        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets,
457    )
458
459    ds_kwargs = dict(
460        raw_key="volumes/raw/s1", label_key=get_platynereis_cell_label_key(ignore_label), patch_shape=patch_shape,
461    )
462
463    datasets = []
464    for data_path, data_roi in zip(data_paths, data_rois):
465        datasets.append(
466            torch_em.default_segmentation_dataset(
467                raw_paths=[data_path], label_paths=[data_path], rois=[data_roi], **ds_kwargs, **kwargs
468            )
469        )
470
471    return datasets[0] if len(datasets) == 1 else ConcatDataset(*datasets)
472
473
474def get_platynereis_cell_loader(
475    path: Union[os.PathLike, str],
476    patch_shape: Tuple[int, int, int],
477    batch_size: int,
478    sample_ids: Optional[Sequence[int]] = None,
479    offsets: Optional[List[List[int]]] = None,
480    boundaries: bool = False,
481    rois: Dict[int, Any] = {},
482    download: bool = False,
483    ignore_label: int = CELL_IGNORE_LABEL,
484    **kwargs
485) -> DataLoader:
486    """Get the dataloader for cell segmentation in platynereis.
487
488    Args:
489        path: Filepath to a folder where the downloaded data will be saved.
490        patch_shape: The patch shape to use for training.
491        batch_size: The batch size for training.
492        sample_ids: The sample ids to use for the dataset
493        offsets: Offset values for affinity computation used as target.
494        boundaries: Whether to compute boundaries as the target.
495        rois: The region of interests to use for the data blocks.
496        download: Whether to download the data if it is not present.
497        ignore_label: The value the neuropil ids of `CELL_NEUROPIL_IDS` are mapped to.
498        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
499
500    Returns:
501        The DataLoader.
502    """
503    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
504    ds = get_platynereis_cell_dataset(
505        path, patch_shape, sample_ids, rois=rois,
506        offsets=offsets, boundaries=boundaries, download=download, ignore_label=ignore_label,
507        **ds_kwargs,
508    )
509    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
510
511
512def get_platynereis_nuclei_dataset(
513    path: Union[os.PathLike, str],
514    patch_shape: Tuple[int, int, int],
515    sample_ids: Optional[Sequence[int]] = None,
516    offsets: Optional[List[List[int]]] = None,
517    boundaries: bool = False,
518    binary: bool = False,
519    rois: Dict[int, Any] = {},
520    download: bool = False,
521    **kwargs
522) -> Dataset:
523    """Get the dataset for nucleus segmentation in platynereis.
524
525    Args:
526        path: Filepath to a folder where the downloaded data will be saved.
527        patch_shape: The patch shape to use for training.
528        sample_ids: The sample ids to use for the dataset
529        offsets: Offset values for affinity computation used as target.
530        boundaries: Whether to compute boundaries as the target.
531        binary: Whether to return a binary segmentation target.
532        rois: The region of interests to use for the data blocks.
533        download: Whether to download the data if it is not present.
534        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
535
536    Returns:
537        The segmentation dataset.
538    """
539    _, n_files = get_platynereis_data(path, "nuclei", download)
540
541    if sample_ids is None:
542        sample_ids = list(range(1, n_files + 1))
543    assert min(sample_ids) >= 1 and max(sample_ids) <= n_files
544    sample_ids.sort()
545
546    data_paths, data_rois = get_platynereis_paths(
547        path=path, sample_ids=sample_ids, name="nuclei", rois=rois, download=download, return_rois=True,
548    )
549
550    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
551    kwargs = util.update_kwargs(kwargs, "rois", data_rois)
552    kwargs, _ = util.add_instance_label_transform(
553        kwargs, add_binary_target=True, boundaries=boundaries, offsets=offsets, binary=binary,
554    )
555
556    return torch_em.default_segmentation_dataset(
557        raw_paths=data_paths,
558        raw_key="volumes/raw",
559        label_paths=data_paths,
560        label_key="volumes/labels/nucleus_instance_labels",
561        patch_shape=patch_shape,
562        **kwargs
563    )
564
565
566def get_platynereis_nuclei_loader(
567    path: Union[os.PathLike, str],
568    patch_shape: Tuple[int, int, int],
569    batch_size: int,
570    sample_ids: Optional[Sequence[int]] = None,
571    offsets: Optional[List[List[int]]] = None,
572    boundaries: bool = False,
573    binary: bool = False,
574    rois: Dict[int, Any] = {},
575    download: bool = False,
576    **kwargs
577) -> DataLoader:
578    """Get the dataloader for nucleus segmentation in platynereis.
579
580    Args:
581        path: Filepath to a folder where the downloaded data will be saved.
582        patch_shape: The patch shape to use for training.
583        batch_size: The batch size for training.
584        sample_ids: The sample ids to use for the dataset
585        offsets: Offset values for affinity computation used as target.
586        boundaries: Whether to compute boundaries as the target.
587        binary: Whether to return a binary segmentation target.
588        rois: The region of interests to use for the data blocks.
589        download: Whether to download the data if it is not present.
590        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
591
592    Returns:
593        The DataLoader.
594    """
595    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
596    ds = get_platynereis_nuclei_dataset(
597        path, patch_shape, sample_ids=sample_ids, rois=rois,
598        offsets=offsets, boundaries=boundaries, binary=binary, download=download,
599        **ds_kwargs,
600    )
601    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)