torch_em.data.datasets.electron_microscopy.saber

  1"""The SABER dataset contains annotations for bacterial cell segmentation in cryo-ET.
  2
  3The data is hosted on the CryoET Data Portal at https://cryoetdataportal.czscience.com/depositions/10331.
  4This loader provides the expert-reviewed subset of the deposition: the semi-manual instance masks of
  5Legionella pneumophila cell interiors. The automated SABER predictions are not part of it.
  6
  7The dataset is part of the publication https://doi.org/10.2139/ssrn.6754411.
  8Please cite it if you use this dataset in your research.
  9"""
 10
 11import os
 12import json
 13from typing import Union, Tuple, List
 14
 15import requests
 16from tqdm import tqdm
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25BASE_URL = "https://files.cryoetdataportal.cziscience.com/{dataset}/{run}/Reconstructions/{spacing}/"
 26RAW_URL = BASE_URL + "Tomograms/100/{run}.zarr"
 27LABEL_URL = BASE_URL + "Annotations/{annotation}/intracellular_anatomical_structure-1.0_instancesegmentationmask.zarr"
 28
 29# The portal dataset, run name, voxel spacing folder and annotation folder of each ground-truth run.
 30RUNS = [
 31    (10062, "dga2017-02-02-12", "VoxelSpacing15.600", "104"),
 32    (10062, "dga2017-02-02-13", "VoxelSpacing15.600", "104"),
 33    (10062, "dga2017-02-02-23", "VoxelSpacing15.600", "105"),
 34    (10062, "dga2017-02-02-24", "VoxelSpacing15.600", "105"),
 35    (10062, "dga2017-02-02-25", "VoxelSpacing15.600", "105"),
 36    (10062, "dga2017-02-02-26", "VoxelSpacing15.600", "105"),
 37    (10062, "dga2017-02-02-27", "VoxelSpacing15.600", "104"),
 38    (10062, "dga2017-02-02-29", "VoxelSpacing15.600", "105"),
 39    (10062, "dga2017-02-02-31", "VoxelSpacing15.600", "105"),
 40    (10064, "dga2017-02-08-101", "VoxelSpacing15.600", "106"),
 41    (10064, "dga2017-02-08-111", "VoxelSpacing15.600", "106"),
 42    (10064, "dga2017-02-08-14", "VoxelSpacing15.600", "105"),
 43    (10064, "dga2017-02-08-16", "VoxelSpacing15.600", "105"),
 44    (10064, "dga2017-02-08-17", "VoxelSpacing15.600", "108"),
 45    (10064, "dga2017-02-08-18", "VoxelSpacing15.600", "106"),
 46    (10064, "dga2017-02-08-27", "VoxelSpacing15.600", "107"),
 47    (10064, "dga2017-02-08-28", "VoxelSpacing15.600", "105"),
 48    (10064, "dga2017-02-08-49", "VoxelSpacing15.600", "106"),
 49    (10077, "dga2016-01-12-12", "VoxelSpacing16.800", "105"),
 50    (10077, "dga2016-03-31-4", "VoxelSpacing16.800", "106"),
 51    (10077, "dga2016-03-31-41", "VoxelSpacing16.800", "105"),
 52    (10077, "dga2016-03-31-45", "VoxelSpacing16.800", "106"),
 53]
 54
 55SCALES = (0, 1, 2)
 56
 57
 58def _fetch(url, path, optional=False):
 59    if os.path.exists(path):
 60        return True
 61
 62    with requests.get(url, stream=True, timeout=(20, 300)) as response:
 63        # A chunk that holds only the fill value is not written by the portal.
 64        if optional and response.status_code == 404:
 65            return False
 66        response.raise_for_status()
 67        # The chunk is renamed only once it is complete, so an interrupted download is not reused.
 68        tmp_path = path + ".partial"
 69        with open(tmp_path, "wb") as f:
 70            for block in response.iter_content(8 * 1024 ** 2):
 71                f.write(block)
 72
 73    os.rename(tmp_path, path)
 74    return True
 75
 76
 77def _download_ome_zarr(url, out_path, scale, download):
 78    array_path = os.path.join(out_path, str(scale))
 79    if os.path.exists(array_path):
 80        return array_path
 81
 82    if not download:
 83        raise RuntimeError(f"Cannot find the data at {out_path}, but download was set to False.")
 84
 85    os.makedirs(out_path, exist_ok=True)
 86    for name in (".zattrs", ".zgroup"):
 87        if not os.path.exists(os.path.join(out_path, name)):
 88            _fetch(f"{url}/{name}", os.path.join(out_path, name))
 89
 90    # The scale is downloaded under a temporary name so that an interrupted download is not treated as complete.
 91    tmp_path = os.path.join(out_path, f"{scale}.partial")
 92    os.makedirs(tmp_path, exist_ok=True)
 93    _fetch(f"{url}/{scale}/.zarray", os.path.join(tmp_path, ".zarray"))
 94    with open(os.path.join(tmp_path, ".zarray")) as f:
 95        meta = json.load(f)
 96
 97    grid = [-(-size // chunk) for size, chunk in zip(meta["shape"], meta["chunks"])]
 98    for z in range(grid[0]):
 99        for y in range(grid[1]):
100            for x in range(grid[2]):
101                chunk_dir = os.path.join(tmp_path, str(z), str(y))
102                os.makedirs(chunk_dir, exist_ok=True)
103                _fetch(f"{url}/{scale}/{z}/{y}/{x}", os.path.join(chunk_dir, str(x)), optional=True)
104
105    os.rename(tmp_path, array_path)
106    return array_path
107
108
109def get_saber_data(path: Union[os.PathLike, str], scale: int = 0, download: bool = False) -> str:
110    """Download the SABER cryo-ET dataset.
111
112    Args:
113        path: Filepath to a folder where the data will be downloaded.
114        scale: The resolution level of the multiscale data. 0 is the native resolution.
115        download: Whether to download the data if it is not present.
116
117    Returns:
118        Filepath where the data is stored.
119    """
120    if scale not in SCALES:
121        raise ValueError(f"The scale must be one of {SCALES}, got {scale}.")
122
123    data_dir = os.path.join(path, "data")
124    os.makedirs(data_dir, exist_ok=True)
125
126    for dataset, run, spacing, annotation in tqdm(RUNS, desc="Downloading tomograms"):
127        urls = {
128            "raw": RAW_URL.format(dataset=dataset, run=run, spacing=spacing),
129            "labels": LABEL_URL.format(dataset=dataset, run=run, spacing=spacing, annotation=annotation),
130        }
131        for name, url in urls.items():
132            _download_ome_zarr(url, os.path.join(data_dir, run, f"{name}.zarr"), scale, download)
133
134    return data_dir
135
136
137def get_saber_paths(
138    path: Union[os.PathLike, str], scale: int = 0, download: bool = False
139) -> Tuple[List[str], List[str]]:
140    """Get paths to the SABER data.
141
142    Args:
143        path: Filepath to a folder where the data will be downloaded.
144        scale: The resolution level of the multiscale data. 0 is the native resolution.
145        download: Whether to download the data if it is not present.
146
147    Returns:
148        List of filepaths to the tomograms.
149        List of filepaths to the instance masks.
150    """
151    data_dir = get_saber_data(path, scale, download)
152    raw_paths = [os.path.join(data_dir, run, "raw.zarr") for _, run, _, _ in RUNS]
153    label_paths = [os.path.join(data_dir, run, "labels.zarr") for _, run, _, _ in RUNS]
154    return raw_paths, label_paths
155
156
157def get_saber_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, int, int],
160    scale: int = 0,
161    download: bool = False,
162    **kwargs
163) -> Dataset:
164    """Get the dataset for bacterial cell segmentation in cryo-ET data.
165
166    Args:
167        path: Filepath to a folder where the data will be downloaded.
168        patch_shape: The patch shape to use for training.
169        scale: The resolution level of the multiscale data. 0 is the native resolution.
170        download: Whether to download the data if it is not present.
171        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
172
173    Returns:
174        The segmentation dataset.
175    """
176    assert len(patch_shape) == 3
177
178    raw_paths, label_paths = get_saber_paths(path, scale, download)
179
180    return torch_em.default_segmentation_dataset(
181        raw_paths=raw_paths,
182        raw_key=str(scale),
183        label_paths=label_paths,
184        label_key=str(scale),
185        patch_shape=patch_shape,
186        is_seg_dataset=True,
187        **kwargs
188    )
189
190
191def get_saber_loader(
192    path: Union[os.PathLike, str],
193    patch_shape: Tuple[int, int, int],
194    batch_size: int,
195    scale: int = 0,
196    download: bool = False,
197    **kwargs
198) -> DataLoader:
199    """Get the DataLoader for bacterial cell segmentation in cryo-ET data.
200
201    Args:
202        path: Filepath to a folder where the data will be downloaded.
203        patch_shape: The patch shape to use for training.
204        batch_size: The batch size for training.
205        scale: The resolution level of the multiscale data. 0 is the native resolution.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
208
209    Returns:
210        The DataLoader.
211    """
212    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
213    dataset = get_saber_dataset(path, patch_shape, scale=scale, download=download, **ds_kwargs)
214    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)