torch_em.data.datasets.electron_microscopy.popsicle

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