torch_em.data.datasets.electron_microscopy.popsicle

  1"""The POPSICLE dataset contains annotations for bacterial compartment segmentation in cryo-ET.
  2
  3The data is hosted on the CryoET Data Portal at https://cryoetdataportal.czscience.com/depositions/10350.
  4It provides curated multi-class compartment masks for 80 tomograms from eight bacterial genera,
  5together with the matching 20 Angstrom re-binned tomograms, and follows the official train/test split.
  6
  7The dataset is part of the publication https://doi.org/10.48550/arXiv.2606.10255.
  8Please cite it if you use this dataset in your research.
  9"""
 10
 11import os
 12import json
 13import shutil
 14from typing import Union, Tuple, List, Literal
 15
 16import requests
 17from tqdm import tqdm
 18
 19import numpy as np
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28BASE_URL = "https://files.cryoetdataportal.cziscience.com/{dataset}/{run}/Reconstructions/VoxelSpacing20.000/"
 29RAW_URL = BASE_URL + "Tomograms/101/{run}.zarr"
 30LABEL_URL = BASE_URL + "Annotations/{folder}/{name}-1.0_segmentationmask.zarr"
 31
 32# The compartment classes with their label value in the merged mask and their portal annotation folder.
 33CLASSES = {
 34    "cytoplasm": (1, "100"),
 35    "membrane": (2, "102"),
 36    "periplasmic_space": (3, "104"),
 37    "bacterial_type_flagellum": (4, "101"),
 38    "dense_body": (5, "103"),
 39}
 40CORE_CLASSES = ("cytoplasm", "membrane", "periplasmic_space")
 41OPTIONAL_CLASSES = {"f": "bacterial_type_flagellum", "d": "dense_body"}
 42
 43# The portal dataset, run name, official split and the optional classes of each run.
 44RUNS = [
 45    (10053, "dga2017-01-14-21", "test", ""),
 46    (10054, "dga2016-09-09-26", "train", ""),
 47    (10054, "dga2016-09-09-32", "train", ""),
 48    (10054, "dga2016-09-09-33", "train", ""),
 49    (10054, "dga2016-09-09-35", "train", ""),
 50    (10054, "dga2016-09-09-42", "train", ""),
 51    (10054, "dga2016-09-09-51", "train", ""),
 52    (10054, "dga2016-09-09-6", "train", ""),
 53    (10054, "dga2016-09-09-78", "train", ""),
 54    (10065, "dga2016-01-13-17", "test", "fd"),
 55    (10065, "dga2016-01-13-25", "train", "fd"),
 56    (10065, "dga2016-01-13-26", "train", "fd"),
 57    (10065, "dga2016-01-13-30", "train", "fd"),
 58    (10065, "dga2016-01-13-32", "train", "d"),
 59    (10098, "dga2015-10-29-11", "test", "d"),
 60    (10098, "dga2015-10-29-31", "train", "d"),
 61    (10098, "dga2015-10-29-39", "train", "d"),
 62    (10098, "dga2015-10-29-43", "train", "d"),
 63    (10098, "dga2015-10-29-60", "train", "d"),
 64    (10155, "ycw2012-11-14-20", "test", "fd"),
 65    (10155, "ycw2012-11-14-47", "train", "fd"),
 66    (10155, "ycw2012-11-14-59", "train", "fd"),
 67    (10155, "ycw2013-05-01-34", "train", "f"),
 68    (10155, "ycw2013-08-20-24", "train", "fd"),
 69    (10155, "ycw2013-08-20-28", "train", "fd"),
 70    (10155, "ycw2013-08-20-43", "train", "f"),
 71    (10155, "ycw2013-08-20-45", "train", "fd"),
 72    (10155, "ycw2013-08-20-59", "train", "fd"),
 73    (10155, "ycw2013-09-10-28", "test", "fd"),
 74    (10155, "ycw2013-09-10-36", "train", "d"),
 75    (10155, "ycw2013-09-10-39", "train", "f"),
 76    (10155, "ycw2013-09-10-43", "train", "fd"),
 77    (10155, "ycw2013-09-10-47", "train", "fd"),
 78    (10155, "ycw2013-09-10-48", "train", "fd"),
 79    (10161, "ycw2012-09-23-21", "test", "f"),
 80    (10161, "ycw2012-09-23-31", "train", "f"),
 81    (10161, "ycw2012-09-23-39", "train", "f"),
 82    (10161, "ycw2012-09-23-46", "train", "f"),
 83    (10161, "ycw2012-09-23-55", "train", "f"),
 84    (10161, "ycw2012-09-23-66", "train", "f"),
 85    (10161, "ycw2012-09-23-67", "train", "f"),
 86    (10161, "ycw2012-09-23-70", "train", "f"),
 87    (10162, "ycw2012-03-12-3", "train", ""),
 88    (10163, "ycw2012-03-03-2", "test", "d"),
 89    (10163, "ycw2012-03-03-3", "train", ""),
 90    (10163, "ycw2012-03-12-18", "train", ""),
 91    (10163, "ycw2012-03-12-7", "train", ""),
 92    (10166, "ycw2012-09-07-23", "train", ""),
 93    (10166, "ycw2012-09-07-3", "train", ""),
 94    (10166, "ycw2012-10-08-1", "train", ""),
 95    (10166, "ycw2012-10-08-7", "train", ""),
 96    (10226, "mba2011-11-23-1", "test", "f"),
 97    (10226, "mba2011-11-23-15", "train", "f"),
 98    (10226, "mba2011-11-23-16", "train", "f"),
 99    (10226, "mba2011-11-23-20", "train", "f"),
100    (10226, "mba2011-11-23-21", "train", "f"),
101    (10226, "mba2011-11-23-22", "test", "f"),
102    (10226, "mba2011-11-23-25", "train", "f"),
103    (10226, "mba2011-11-23-26", "train", "f"),
104    (10226, "mba2011-11-23-35", "train", "f"),
105    (10226, "mba2011-11-23-7", "train", "f"),
106    (10272, "aba2015-06-04-10", "test", "f"),
107    (10272, "aba2015-06-04-16", "test", ""),
108    (10272, "aba2015-06-04-22", "train", "d"),
109    (10272, "aba2015-06-04-24", "train", ""),
110    (10272, "aba2015-06-04-26", "train", ""),
111    (10272, "aba2015-06-04-30", "train", ""),
112    (10272, "aba2015-06-04-5", "train", "d"),
113    (10272, "aba2015-06-04-9", "train", ""),
114    (10273, "aba2015-07-07-10", "train", "d"),
115    (10273, "aba2015-07-07-17", "train", "f"),
116    (10281, "aba2015-02-23-15", "train", "d"),
117    (10281, "aba2015-02-23-16", "train", "d"),
118    (10281, "aba2015-02-23-2", "test", "fd"),
119    (10281, "aba2015-02-23-20", "train", "fd"),
120    (10281, "aba2015-02-23-21", "train", "d"),
121    (10281, "aba2015-02-23-25", "train", "d"),
122    (10281, "aba2015-02-23-3", "train", "fd"),
123    (10281, "aba2015-02-23-30", "train", "d"),
124    (10281, "aba2015-02-23-9", "train", "fd"),
125]
126
127
128def _fetch(url, path, optional=False):
129    if os.path.exists(path):
130        return True
131
132    with requests.get(url, stream=True, timeout=(20, 300)) as response:
133        # A chunk that holds only the fill value is not written by the portal.
134        if optional and response.status_code == 404:
135            return False
136        response.raise_for_status()
137        # The chunk is renamed only once it is complete, so an interrupted download is not reused.
138        tmp_path = path + ".partial"
139        with open(tmp_path, "wb") as f:
140            for block in response.iter_content(8 * 1024 ** 2):
141                f.write(block)
142
143    os.rename(tmp_path, path)
144    return True
145
146
147def _download_ome_zarr(url, out_path, download):
148    array_path = os.path.join(out_path, "0")
149    if os.path.exists(array_path):
150        return array_path
151
152    if not download:
153        raise RuntimeError(f"Cannot find the data at {out_path}, but download was set to False.")
154
155    os.makedirs(out_path, exist_ok=True)
156    for name in (".zattrs", ".zgroup"):
157        if not os.path.exists(os.path.join(out_path, name)):
158            _fetch(f"{url}/{name}", os.path.join(out_path, name))
159
160    tmp_path = os.path.join(out_path, "0.partial")
161    os.makedirs(tmp_path, exist_ok=True)
162    _fetch(f"{url}/0/.zarray", os.path.join(tmp_path, ".zarray"))
163    with open(os.path.join(tmp_path, ".zarray")) as f:
164        meta = json.load(f)
165
166    grid = [-(-size // chunk) for size, chunk in zip(meta["shape"], meta["chunks"])]
167    for z in range(grid[0]):
168        for y in range(grid[1]):
169            for x in range(grid[2]):
170                chunk_dir = os.path.join(tmp_path, str(z), str(y))
171                os.makedirs(chunk_dir, exist_ok=True)
172                _fetch(f"{url}/0/{z}/{y}/{x}", os.path.join(chunk_dir, str(x)), optional=True)
173
174    os.rename(tmp_path, array_path)
175    return array_path
176
177
178def _merge_labels(run_dir, dataset, run, extras, download):
179    """Merge the per-class masks into one multi-class volume.
180
181    The classes are only mutually exclusive at the native resolution, so the coarser levels of the
182    portal label pyramid are not used.
183    """
184    import zarr
185
186    label_path = os.path.join(run_dir, "labels.zarr")
187    if os.path.exists(label_path):
188        return label_path
189
190    names = list(CORE_CLASSES) + [OPTIONAL_CLASSES[key] for key in extras]
191    merged = None
192    for name in names:
193        value, folder = CLASSES[name]
194        url = LABEL_URL.format(dataset=dataset, run=run, folder=folder, name=name)
195        class_dir = os.path.join(run_dir, f"class_{name}.zarr")
196        array = zarr.open_array(_download_ome_zarr(url, class_dir, download), mode="r")
197        mask = array[:] > 0
198        if merged is None:
199            merged = np.zeros(mask.shape, dtype="uint8")
200        elif mask.shape != merged.shape:
201            raise ValueError(f"The class masks of run {run} have different shapes.")
202        if np.any(merged[mask]):
203            raise ValueError(f"The class masks of run {run} overlap, so they cannot be merged.")
204        merged[mask] = value
205
206    # The store is built under a temporary name and renamed, so an interrupted merge leaves nothing behind.
207    # It uses the zarr format of the portal stores it sits next to.
208    tmp_path = label_path + ".partial"
209    if os.path.exists(tmp_path):
210        shutil.rmtree(tmp_path)
211    store = zarr.open_group(tmp_path, mode="w", zarr_format=2)
212    array = store.create_array("0", shape=merged.shape, dtype="uint8", chunks=(64, 256, 256))
213    array[:] = merged
214    os.rename(tmp_path, label_path)
215    return label_path
216
217
218def get_popsicle_data(
219    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
220) -> str:
221    """Download the POPSICLE bacterial segmentation dataset.
222
223    Args:
224        path: Filepath to a folder where the data will be downloaded.
225        split: The data split to download. Either 'train' or 'test'.
226        download: Whether to download the data if it is not present.
227
228    Returns:
229        Filepath where the data is stored.
230    """
231    if split not in ("train", "test"):
232        raise ValueError(f"The split must be 'train' or 'test', got '{split}'.")
233
234    data_dir = os.path.join(path, split)
235    os.makedirs(data_dir, exist_ok=True)
236
237    runs = [entry for entry in RUNS if entry[2] == split]
238    for dataset, run, _, extras in tqdm(runs, desc=f"Downloading the {split} tomograms"):
239        run_dir = os.path.join(data_dir, run)
240        _download_ome_zarr(RAW_URL.format(dataset=dataset, run=run), os.path.join(run_dir, "raw.zarr"), download)
241        if not download and not os.path.exists(os.path.join(run_dir, "labels.zarr")):
242            raise RuntimeError(f"Cannot find the data at {run_dir}, but download was set to False.")
243        _merge_labels(run_dir, dataset, run, extras, download)
244
245    return data_dir
246
247
248def get_popsicle_paths(
249    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
250) -> Tuple[List[str], List[str]]:
251    """Get paths to the POPSICLE data.
252
253    Args:
254        path: Filepath to a folder where the data will be downloaded.
255        split: The data split. Either 'train' or 'test'.
256        download: Whether to download the data if it is not present.
257
258    Returns:
259        List of filepaths to the tomograms.
260        List of filepaths to the multi-class compartment masks.
261    """
262    data_dir = get_popsicle_data(path, split, download)
263    runs = [entry[1] for entry in RUNS if entry[2] == split]
264    raw_paths = [os.path.join(data_dir, run, "raw.zarr") for run in runs]
265    label_paths = [os.path.join(data_dir, run, "labels.zarr") for run in runs]
266    return raw_paths, label_paths
267
268
269def get_popsicle_dataset(
270    path: Union[os.PathLike, str],
271    patch_shape: Tuple[int, int, int],
272    split: Literal["train", "test"],
273    download: bool = False,
274    **kwargs
275) -> Dataset:
276    """Get the dataset for bacterial compartment segmentation in cryo-ET data.
277
278    The labels are a multi-class mask with 1: cytoplasm, 2: membrane, 3: periplasmic space,
279    4: bacterial-type flagellum and 5: dense body. The last two classes are not present in every tomogram.
280
281    Args:
282        path: Filepath to a folder where the data will be downloaded.
283        patch_shape: The patch shape to use for training.
284        split: The data split. Either 'train' or 'test'.
285        download: Whether to download the data if it is not present.
286        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
287
288    Returns:
289        The segmentation dataset.
290    """
291    assert len(patch_shape) == 3
292
293    raw_paths, label_paths = get_popsicle_paths(path, split, download)
294
295    return torch_em.default_segmentation_dataset(
296        raw_paths=raw_paths,
297        raw_key="0",
298        label_paths=label_paths,
299        label_key="0",
300        patch_shape=patch_shape,
301        is_seg_dataset=True,
302        **kwargs
303    )
304
305
306def get_popsicle_loader(
307    path: Union[os.PathLike, str],
308    patch_shape: Tuple[int, int, int],
309    batch_size: int,
310    split: Literal["train", "test"],
311    download: bool = False,
312    **kwargs
313) -> DataLoader:
314    """Get the DataLoader for bacterial compartment segmentation in cryo-ET data.
315
316    Args:
317        path: Filepath to a folder where the data will be downloaded.
318        patch_shape: The patch shape to use for training.
319        batch_size: The batch size for training.
320        split: The data split. Either 'train' or 'test'.
321        download: Whether to download the data if it is not present.
322        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
323
324    Returns:
325        The DataLoader.
326    """
327    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
328    dataset = get_popsicle_dataset(path, patch_shape, split, download=download, **ds_kwargs)
329    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)