torch_em.data.datasets.electron_microscopy.chlamydomonas_cryoet

  1"""This dataset provides organelle segmentation masks for cryo-ET tomograms of
  2Chlamydomonas reinhardtii cells, prepared by cryo-plasmaFIB milling.
  3
  4The data is hosted on the CryoET Data Portal as dataset 10301
  5(https://cryoetdataportal.czscience.com/datasets/10301, EMPIAR-11756). This loader covers three of its runs:
  6- 17072022_BrnoKrios_Arctis_p3ar_grid_Position_76 (https://cryoetdataportal.czscience.com/runs/14082)
  7- 17072022_BrnoKrios_Arctis_p3ar_grid_Position_101 (https://cryoetdataportal.czscience.com/runs/14079)
  8- 27042022_BrnoKrios_Arctis_grid9_hGIS_Position_44 (https://cryoetdataportal.czscience.com/runs/14085)
  9
 10NOTE: The expert-curated ground truth of these runs consists of point annotations (ribosomes, microtubules,
 11RubisCO) only. All dense masks used here were predicted automatically with Ais / Pom
 12(`ground_truth_status=false`, `method_type=automated`) and have not been proofread, so they are
 13pseudo-labels of varying quality. Only structures with masks in at least two of the three runs are provided.
 14
 15The data is described in https://doi.org/10.7554/eLife.98552 and https://doi.org/10.1101/2025.01.16.633326.
 16Please cite them if you use this dataset in your research.
 17"""
 18
 19import os
 20from typing import List, Tuple, Union
 21
 22from torch.utils.data import DataLoader, Dataset
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29DATASET_ID = 10301
 30VOXEL_SPACING = "VoxelSpacing7.840"
 31TOMOGRAM_FOLDER = "100"
 32
 33# The annotation folder of each mask differs between the runs.
 34RUN_LABELS = {
 35    "17072022_BrnoKrios_Arctis_p3ar_grid_Position_76": {
 36        "mitochondrion": "114", "lipid_droplet": "132", "cytoplasmic_vesicle": "125", "pyrenoid": "130",
 37        "starch_grain": "138", "nucleoplasm": "127", "endoplasmic_reticulum": "133", "golgi_apparatus": "135",
 38        "cell_wall": "136", "chloroplast_stroma": "131", "cytoplasm": "113",
 39    },
 40    "17072022_BrnoKrios_Arctis_p3ar_grid_Position_101": {
 41        "mitochondrion": "120", "lipid_droplet": "125", "cytoplasmic_vesicle": "129", "pyrenoid": "118",
 42        "starch_grain": "132", "chloroplast_stroma": "124",
 43    },
 44    "27042022_BrnoKrios_Arctis_grid9_hGIS_Position_44": {
 45        "mitochondrion": "123", "lipid_droplet": "130", "cytoplasmic_vesicle": "122", "pyrenoid": "128",
 46        "starch_grain": "136", "nucleoplasm": "125", "endoplasmic_reticulum": "131", "golgi_apparatus": "133",
 47        "cell_wall": "134", "cytoplasm": "110",
 48    },
 49}
 50RUNS = list(RUN_LABELS.keys())
 51TARGETS = (
 52    "mitochondrion", "lipid_droplet", "cytoplasmic_vesicle", "pyrenoid", "starch_grain", "nucleoplasm",
 53    "endoplasmic_reticulum", "golgi_apparatus", "cell_wall", "chloroplast_stroma", "cytoplasm",
 54)
 55
 56
 57def get_chlamydomonas_cryoet_data(
 58    path: Union[os.PathLike, str], runs: Union[str, List[str]] = RUNS, download: bool = False,
 59) -> List[str]:
 60    """Download the Chlamydomonas cryo-ET data.
 61
 62    The tomograms and their segmentation masks are streamed from the CryoET Data Portal and cached as zarr.
 63    Requires zarr (v3) and s3fs.
 64
 65    Args:
 66        path: Filepath to a folder where the downloaded data will be saved.
 67        runs: The runs (tomograms) to download. By default all runs are used.
 68        download: Whether to download the data if it is not present.
 69
 70    Returns:
 71        The filepaths to the cached data of each run.
 72    """
 73    runs = [runs] if isinstance(runs, str) else list(runs)
 74    invalid = [run for run in runs if run not in RUN_LABELS]
 75    if invalid:
 76        raise ValueError(f"Invalid runs {invalid}. Choose from {RUNS}.")
 77
 78    os.makedirs(path, exist_ok=True)
 79    return [
 80        util.stream_from_cryo_et_portal(
 81            zarr_path=os.path.join(path, f"{run}.zarr"),
 82            dataset_id=DATASET_ID,
 83            run=run,
 84            voxel_spacing=VOXEL_SPACING,
 85            tomogram_folder=TOMOGRAM_FOLDER,
 86            annotations={name: folder for name, folder in RUN_LABELS[run].items()},
 87            download=download,
 88        ) for run in runs
 89    ]
 90
 91
 92def get_chlamydomonas_cryoet_paths(
 93    path: Union[os.PathLike, str],
 94    target: str = "mitochondrion",
 95    runs: Union[str, List[str]] = RUNS,
 96    download: bool = False,
 97) -> List[str]:
 98    """Get paths to the Chlamydomonas cryo-ET data that provide the requested target.
 99
100    Args:
101        path: Filepath to a folder where the downloaded data will be saved.
102        target: The segmentation target, one of `TARGETS`.
103        runs: The runs (tomograms) to use. By default all runs providing the target are used.
104        download: Whether to download the data if it is not present.
105
106    Returns:
107        The filepaths to the cached data.
108    """
109    if target not in TARGETS:
110        raise ValueError(f"'target' must be one of {TARGETS}, got '{target}'.")
111    runs = [runs] if isinstance(runs, str) else list(runs)
112    runs = [run for run in runs if run in RUN_LABELS and target in RUN_LABELS[run]]
113    if not runs:
114        raise ValueError(f"None of the selected runs provides the target '{target}'.")
115    return get_chlamydomonas_cryoet_data(path, runs, download)
116
117
118def get_chlamydomonas_cryoet_dataset(
119    path: Union[os.PathLike, str],
120    patch_shape: Tuple[int, int, int],
121    target: str = "mitochondrion",
122    runs: Union[str, List[str]] = RUNS,
123    download: bool = False,
124    **kwargs,
125) -> Dataset:
126    """Get the dataset for organelle segmentation in Chlamydomonas cryo-ET data.
127
128    Args:
129        path: Filepath to a folder where the downloaded data will be saved.
130        patch_shape: The patch shape (z, y, x) to use for training.
131        target: The segmentation target, one of `TARGETS`.
132        runs: The runs (tomograms) to use. By default all runs providing the target are used.
133        download: Whether to download the data if it is not present.
134        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
135
136    Returns:
137        The segmentation dataset.
138    """
139    assert len(patch_shape) == 3
140
141    paths = get_chlamydomonas_cryoet_paths(path, target, runs, download)
142
143    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
144
145    return torch_em.default_segmentation_dataset(
146        raw_paths=paths,
147        raw_key="raw",
148        label_paths=paths,
149        label_key=target,
150        patch_shape=patch_shape,
151        **kwargs,
152    )
153
154
155def get_chlamydomonas_cryoet_loader(
156    path: Union[os.PathLike, str],
157    patch_shape: Tuple[int, int, int],
158    batch_size: int,
159    target: str = "mitochondrion",
160    runs: Union[str, List[str]] = RUNS,
161    download: bool = False,
162    **kwargs,
163) -> DataLoader:
164    """Get the DataLoader for organelle segmentation in Chlamydomonas cryo-ET data.
165
166    Args:
167        path: Filepath to a folder where the downloaded data will be saved.
168        patch_shape: The patch shape (z, y, x) to use for training.
169        batch_size: The batch size for training.
170        target: The segmentation target, one of `TARGETS`.
171        runs: The runs (tomograms) to use. By default all runs providing the target are used.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
174            or for the PyTorch DataLoader.
175
176    Returns:
177        The DataLoader.
178    """
179    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
180    dataset = get_chlamydomonas_cryoet_dataset(
181        path, patch_shape, target=target, runs=runs, download=download, **ds_kwargs
182    )
183    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)