torch_em.data.datasets.electron_microscopy.deepict_pombe

  1"""The DeePiCt S. pombe dataset contains expert-verified segmentations of cellular structures in cryo-ET.
  2It consists of 10 tomograms of S. pombe cryo-FIB lamellae (collected with defocus) at 13.48 Angstrom voxel spacing,
  3which are annotated with up to 12 different structures: cytoplasm, vesicle, mitochondrion, microtubule,
  4endoplasmic reticulum, nuclear envelope, nucleus, vacuole, lipid droplet, Golgi apparatus, multivesicular body
  5and membrane. Not all structures are annotated in every tomogram, see `RUN_CLASSES`.
  6
  7The data is hosted on the CryoET Data Portal at https://cryoetdataportal.czscience.com/datasets/10000.
  8The dataset is part of the publication https://doi.org/10.1038/s41592-022-01746-2.
  9Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13from typing import List, Literal, Tuple, Union
 14
 15import numpy as np
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24DATASET_ID = 10000
 25VOXEL_SPACING = "VoxelSpacing13.480"
 26TOMOGRAM_FOLDER = "100"
 27
 28# The structures with their portal annotation folder. The same folder is used for all runs.
 29# The order defines the label ids of the 'all' target (starting at 1); later structures overwrite earlier ones.
 30CLASSES = {
 31    "cytoplasm": "102",
 32    "mitochondrion": "103",
 33    "vesicle": "104",
 34    "microtubule": "105",
 35    "endoplasmic_reticulum": "106",
 36    "nuclear_envelope": "107",
 37    "nucleus": "108",
 38    "vacuole": "109",
 39    "lipid_droplet": "110",
 40    "golgi_apparatus": "111",
 41    "multivesicular_body": "112",
 42    "membrane": "114",
 43}
 44
 45# The membrane overlaps with the other structures and is only annotated in some runs.
 46# It is therefore not part of the 'all' target, but only available as an individual target.
 47ALL_CLASSES = tuple(name for name in CLASSES if name != "membrane")
 48
 49# The ground-truth structures that are annotated in each run.
 50RUN_CLASSES = {
 51    "TS_026": ("cytoplasm", "vesicle", "endoplasmic_reticulum", "vacuole", "multivesicular_body", "membrane"),
 52    "TS_027": (
 53        "cytoplasm", "vesicle", "endoplasmic_reticulum", "nuclear_envelope", "nucleus", "vacuole", "golgi_apparatus"
 54    ),
 55    "TS_028": ("cytoplasm", "vesicle", "vacuole", "lipid_droplet", "golgi_apparatus", "multivesicular_body"),
 56    "TS_029": ("cytoplasm", "mitochondrion", "vesicle", "endoplasmic_reticulum", "vacuole"),
 57    "TS_030": ("cytoplasm", "vesicle", "endoplasmic_reticulum", "golgi_apparatus", "membrane"),
 58    "TS_034": (
 59        "cytoplasm", "mitochondrion", "vesicle", "endoplasmic_reticulum", "vacuole", "golgi_apparatus",
 60        "multivesicular_body", "membrane"
 61    ),
 62    "TS_037": (
 63        "cytoplasm", "vesicle", "nuclear_envelope", "nucleus", "vacuole", "multivesicular_body", "membrane"
 64    ),
 65    "TS_041": (
 66        "cytoplasm", "vesicle", "microtubule", "endoplasmic_reticulum", "vacuole", "lipid_droplet", "membrane"
 67    ),
 68    "TS_043": ("cytoplasm", "vesicle", "nuclear_envelope", "nucleus", "golgi_apparatus"),
 69    "TS_045": (
 70        "cytoplasm", "mitochondrion", "vesicle", "endoplasmic_reticulum", "nuclear_envelope", "nucleus",
 71        "golgi_apparatus"
 72    ),
 73}
 74RUNS = tuple(RUN_CLASSES.keys())
 75TARGETS = ("all",) + tuple(CLASSES.keys())
 76
 77
 78def _check_runs(runs):
 79    invalid = [run for run in runs if run not in RUN_CLASSES]
 80    if invalid:
 81        raise ValueError(f"Invalid runs {invalid}. Choose from {RUNS}.")
 82
 83
 84def _add_all_labels(zarr_path, run):
 85    """Merge the structures of a run into a single semantic segmentation (0 is background)."""
 86    import zarr
 87    from zarr.codecs import BloscCodec
 88
 89    root = zarr.open_group(zarr_path, mode="a")
 90    if "all" in root:
 91        return
 92
 93    names = [name for name in ALL_CLASSES if name in RUN_CLASSES[run]]
 94    shape = root["raw"].shape
 95    merged = np.zeros(shape, dtype="uint8")
 96    for name in names:
 97        merged[root[name][:] > 0] = ALL_CLASSES.index(name) + 1
 98
 99    arr = root.create_array(
100        "all", shape=shape, chunks=tuple(min(128, s) for s in shape), dtype=merged.dtype,
101        compressors=BloscCodec(cname="zstd", clevel=6, shuffle="bitshuffle"),
102    )
103    arr[:] = merged
104
105
106def get_deepict_pombe_data(
107    path: Union[os.PathLike, str], runs: Union[str, List[str]] = RUNS, download: bool = False,
108) -> List[str]:
109    """Download the DeePiCt S. pombe data.
110
111    The tomograms and their segmentation masks are streamed from the CryoET Data Portal and cached as zarr.
112
113    Args:
114        path: Filepath to a folder where the downloaded data will be saved.
115        runs: The runs (tomograms) to download. By default all runs are used.
116        download: Whether to download the data if it is not present.
117
118    Returns:
119        The filepaths to the cached data of each run.
120    """
121    runs = [runs] if isinstance(runs, str) else list(runs)
122    _check_runs(runs)
123
124    os.makedirs(path, exist_ok=True)
125    zarr_paths = []
126    for run in runs:
127        zarr_path = util.stream_from_cryo_et_portal(
128            zarr_path=os.path.join(path, f"{run}.zarr"),
129            dataset_id=DATASET_ID,
130            run=run,
131            voxel_spacing=VOXEL_SPACING,
132            tomogram_folder=TOMOGRAM_FOLDER,
133            annotations={name: CLASSES[name] for name in RUN_CLASSES[run]},
134            download=download,
135        )
136        _add_all_labels(zarr_path, run)
137        zarr_paths.append(zarr_path)
138
139    return zarr_paths
140
141
142def get_deepict_pombe_paths(
143    path: Union[os.PathLike, str],
144    target: str = "all",
145    runs: Union[str, List[str]] = RUNS,
146    download: bool = False,
147) -> List[str]:
148    """Get paths to the DeePiCt S. pombe data.
149
150    Args:
151        path: Filepath to a folder where the downloaded data will be saved.
152        target: The segmentation target. Either 'all' for the semantic segmentation of all structures except
153            the membrane, or the name of a single structure for its binary segmentation (see `CLASSES`).
154        runs: The runs (tomograms) to use. By default all runs that provide the target are used.
155        download: Whether to download the data if it is not present.
156
157    Returns:
158        The filepaths to the stored data.
159    """
160    if target not in TARGETS:
161        raise ValueError(f"'target' must be one of {TARGETS}, got '{target}'.")
162    runs = [runs] if isinstance(runs, str) else list(runs)
163    _check_runs(runs)
164
165    if target != "all":
166        runs = [run for run in runs if target in RUN_CLASSES[run]]
167        if not runs:
168            raise ValueError(f"None of the selected runs is annotated with '{target}'.")
169
170    return get_deepict_pombe_data(path, runs, download)
171
172
173def get_deepict_pombe_dataset(
174    path: Union[os.PathLike, str],
175    patch_shape: Tuple[int, int, int],
176    target: Literal["all", "cytoplasm", "mitochondrion", "vesicle", "microtubule", "endoplasmic_reticulum",
177                    "nuclear_envelope", "nucleus", "vacuole", "lipid_droplet", "golgi_apparatus",
178                    "multivesicular_body", "membrane"] = "all",
179    runs: Union[str, List[str]] = RUNS,
180    download: bool = False,
181    **kwargs
182) -> Dataset:
183    """Get the dataset for structure segmentation in the DeePiCt S. pombe cryo-ET data.
184
185    Note: Not all structures are annotated in every tomogram. For the 'all' target, structures that were not
186    annotated in a tomogram are part of its background.
187
188    Args:
189        path: Filepath to a folder where the downloaded data will be saved.
190        patch_shape: The patch shape (z, y, x) to use for training.
191        target: The segmentation target. Either 'all' for the semantic segmentation of all structures except
192            the membrane, or the name of a single structure for its binary segmentation (see `CLASSES`).
193        runs: The runs (tomograms) to use. By default all runs that provide the target are used.
194        download: Whether to download the data if it is not present.
195        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
196
197    Returns:
198        The segmentation dataset.
199    """
200    assert len(patch_shape) == 3
201
202    paths = get_deepict_pombe_paths(path, target, runs, download)
203
204    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
205
206    return torch_em.default_segmentation_dataset(
207        raw_paths=paths,
208        raw_key="raw",
209        label_paths=paths,
210        label_key=target,
211        patch_shape=patch_shape,
212        **kwargs
213    )
214
215
216def get_deepict_pombe_loader(
217    path: Union[os.PathLike, str],
218    patch_shape: Tuple[int, int, int],
219    batch_size: int,
220    target: Literal["all", "cytoplasm", "mitochondrion", "vesicle", "microtubule", "endoplasmic_reticulum",
221                    "nuclear_envelope", "nucleus", "vacuole", "lipid_droplet", "golgi_apparatus",
222                    "multivesicular_body", "membrane"] = "all",
223    runs: Union[str, List[str]] = RUNS,
224    download: bool = False,
225    **kwargs
226) -> DataLoader:
227    """Get the DataLoader for structure segmentation in the DeePiCt S. pombe cryo-ET data.
228
229    Args:
230        path: Filepath to a folder where the downloaded data will be saved.
231        patch_shape: The patch shape (z, y, x) to use for training.
232        batch_size: The batch size for training.
233        target: The segmentation target. Either 'all' for the semantic segmentation of all structures except
234            the membrane, or the name of a single structure for its binary segmentation (see `CLASSES`).
235        runs: The runs (tomograms) to use. By default all runs that provide the target are used.
236        download: Whether to download the data if it is not present.
237        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
238
239    Returns:
240        The DataLoader.
241    """
242    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
243    dataset = get_deepict_pombe_dataset(
244        path, patch_shape, target=target, runs=runs, download=download, **ds_kwargs
245    )
246    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)