torch_em.data.datasets.electron_microscopy.drg_axon_mito

  1"""The DRG-Axon-Mito dataset contains mitochondrial membrane segmentation for cryo-electron tomograms of
  2axons and varicosities in cultured primary mouse dorsal root ganglion neurons.
  3
  4The data is hosted on the CryoET Data Portal at https://cryoetdataportal.cziscience.com/datasets/10512
  5(and the related datasets 10513, 10514, 10515 and 10516 from the same study).
  6
  7The dataset is part of the publication https://doi.org/10.64898/2026.07.07.737043.
  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
 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/VoxelSpacing{voxel_spacing}/"
 29RAW_URL = BASE_URL + "Tomograms/100/{run}.zarr"
 30OMM_URL = BASE_URL + "Annotations/100/source_label_omm_{omm_scale}_segmentationmask.zarr"
 31IMM_URL = BASE_URL + "Annotations/101/source_label_imm_{imm_scale}_segmentationmask.zarr"
 32
 33# The mitochondrial membrane classes and their label value in the merged mask.
 34CLASSES = {"outer_membrane": 1, "inner_membrane": 2}
 35
 36# The portal dataset, run name, voxel spacing and the annotation-file scale suffix of each class.
 37RUNS = [
 38    (10512, "Position_1", "21.600", "-1.0", "-1.0"),
 39    (10512, "Position_15", "21.600", "-1.0", "-1.0"),
 40    (10512, "Position_17", "21.600", "-1.0", "-1.0"),
 41    (10512, "Position_19", "21.600", "-1.0", "-1.0"),
 42    (10512, "Position_2", "21.600", "-1.0", "-2.0"),
 43    (10512, "Position_21", "21.600", "-1.0", "-1.0"),
 44    (10512, "Position_22", "21.600", "-1.0", "-1.0"),
 45    (10512, "Position_26", "21.600", "-1.0", "-1.0"),
 46    (10512, "Position_30", "10.800", "-1.0", "-1.0"),
 47    (10512, "Position_32", "21.600", "-1.0", "-2.0"),
 48    (10512, "Position_33", "21.600", "-1.0", "-1.0"),
 49    (10512, "Position_42", "21.600", "-1.0", "-2.0"),
 50    (10512, "Position_43", "21.600", "-1.0", "-1.0"),
 51    (10512, "Position_44", "21.600", "-1.0", "-1.0"),
 52    (10513, "Position_03", "20.579", "-1.0", "-1.0"),
 53    (10513, "Position_04", "20.579", "-1.0", "-1.0"),
 54    (10513, "Position_11", "20.579", "-1.0", "-1.0"),
 55    (10513, "Position_12", "20.579", "-1.0", "-1.0"),
 56    (10513, "Position_42", "20.579", "-1.0", "-1.0"),
 57    (10513, "Position_45", "20.579", "-1.0", "-1.0"),
 58    (10513, "Position_46", "20.579", "-1.0", "-1.0"),
 59    (10513, "Position_48", "20.579", "-1.0", "-1.0"),
 60    (10513, "Position_49", "20.579", "-1.0", "-1.0"),
 61    (10513, "Position_50", "20.579", "-1.0", "-1.0"),
 62    (10514, "Position_18", "21.600", "-1.0", "-1.0"),
 63    (10514, "Position_20", "21.600", "-1.0", "-1.0"),
 64    (10514, "Position_24", "21.600", "-1.0", "-1.0"),
 65    (10514, "Position_25", "21.600", "-1.0", "-1.0"),
 66    (10514, "Position_26", "21.600", "-1.0", "-1.0"),
 67    (10514, "Position_5", "21.600", "-1.0", "-1.0"),
 68    (10515, "Position_16B", "21.600", "-1.0", "-1.0"),
 69    (10515, "Position_17B", "21.600", "-1.0", "-1.0"),
 70    (10515, "Position_41C", "21.600", "-1.0", "-1.0"),
 71    (10515, "Position_46C", "21.600", "-1.0", "-1.0"),
 72    (10515, "Position_50C", "21.600", "-1.0", "-2.0"),
 73    (10515, "Position_51", "10.800", "-2.0", "-2.0"),
 74    (10515, "Position_52C", "21.600", "-1.0", "-1.0"),
 75    (10515, "Position_62C", "21.600", "-1.0", "-1.0"),
 76    (10515, "Position_63C", "21.600", "-1.0", "-1.0"),
 77    (10515, "Position_64", "10.800", "-2.0", "-2.0"),
 78    (10515, "Position_69C", "21.600", "-1.0", "-1.0"),
 79    (10515, "Position_84C", "21.600", "-1.0", "-1.0"),
 80    (10515, "Position_9C", "21.600", "-1.0", "-1.0"),
 81    (10516, "Position_1", "21.600", "-1.0", "-2.0"),
 82    (10516, "Position_12", "21.600", "-1.0", "-2.0"),
 83    (10516, "Position_14", "21.600", "-1.0", "-1.0"),
 84    (10516, "Position_16", "21.600", "-1.0", "-1.0"),
 85    (10516, "Position_20", "21.600", "-1.0", "-2.0"),
 86    (10516, "Position_25", "21.600", "-1.0", "-1.0"),
 87    (10516, "Position_27", "21.600", "-1.0", "-2.0"),
 88    (10516, "Position_29", "21.600", "-1.0", "-1.0"),
 89    (10516, "Position_30", "21.600", "-1.0", "-1.0"),
 90    (10516, "Position_35", "21.600", "-1.0", "-1.0"),
 91    (10516, "Position_38", "21.600", "-1.0", "-1.0"),
 92    (10516, "Position_39", "21.600", "-1.0", "-1.0"),
 93    (10516, "Position_7", "21.600", "-1.0", "-1.0"),
 94    (10516, "Position_9", "21.600", "-1.0", "-1.0"),
 95]
 96
 97
 98def _fetch(url, path, optional=False):
 99    if os.path.exists(path):
100        return True
101
102    with requests.get(url, stream=True, timeout=(20, 300)) as response:
103        # A chunk that holds only the fill value is not written by the portal.
104        if optional and response.status_code == 404:
105            return False
106        response.raise_for_status()
107        # The chunk is renamed only once it is complete, so an interrupted download is not reused.
108        tmp_path = path + ".partial"
109        with open(tmp_path, "wb") as f:
110            for block in response.iter_content(8 * 1024 ** 2):
111                f.write(block)
112
113    os.rename(tmp_path, path)
114    return True
115
116
117def _download_ome_zarr(url, out_path, download):
118    array_path = os.path.join(out_path, "0")
119    if os.path.exists(array_path):
120        return array_path
121
122    if not download:
123        raise RuntimeError(f"Cannot find the data at {out_path}, but download was set to False.")
124
125    os.makedirs(out_path, exist_ok=True)
126    for name in (".zattrs", ".zgroup"):
127        if not os.path.exists(os.path.join(out_path, name)):
128            _fetch(f"{url}/{name}", os.path.join(out_path, name))
129
130    tmp_path = os.path.join(out_path, "0.partial")
131    os.makedirs(tmp_path, exist_ok=True)
132    _fetch(f"{url}/0/.zarray", os.path.join(tmp_path, ".zarray"))
133    with open(os.path.join(tmp_path, ".zarray")) as f:
134        meta = json.load(f)
135
136    grid = [-(-size // chunk) for size, chunk in zip(meta["shape"], meta["chunks"])]
137    for z in range(grid[0]):
138        for y in range(grid[1]):
139            for x in range(grid[2]):
140                chunk_dir = os.path.join(tmp_path, str(z), str(y))
141                os.makedirs(chunk_dir, exist_ok=True)
142                _fetch(f"{url}/0/{z}/{y}/{x}", os.path.join(chunk_dir, str(x)), optional=True)
143
144    os.rename(tmp_path, array_path)
145    return array_path
146
147
148def _merge_labels(run_dir, dataset, run, voxel_spacing, omm_scale, imm_scale, download):
149    """Merge the outer- and inner-membrane masks into one multi-class volume."""
150    import zarr
151
152    label_path = os.path.join(run_dir, "labels.zarr")
153    if os.path.exists(label_path):
154        return label_path
155
156    merged = None
157    for name, scale in (("outer_membrane", omm_scale), ("inner_membrane", imm_scale)):
158        value = CLASSES[name]
159        template = OMM_URL if name == "outer_membrane" else IMM_URL
160        url = template.format(dataset=dataset, run=run, voxel_spacing=voxel_spacing, omm_scale=scale, imm_scale=scale)
161        class_dir = os.path.join(run_dir, f"class_{name}.zarr")
162        array = zarr.open_array(_download_ome_zarr(url, class_dir, download), mode="r")
163        mask = array[:] > 0
164        if merged is None:
165            merged = np.zeros(mask.shape, dtype="uint8")
166        elif mask.shape != merged.shape:
167            raise ValueError(f"The class masks of run {run} have different shapes.")
168        if np.any(merged[mask]):
169            raise ValueError(f"The class masks of run {run} overlap, so they cannot be merged.")
170        merged[mask] = value
171
172    # The store is built under a temporary name and renamed, so an interrupted merge leaves nothing behind.
173    tmp_path = label_path + ".partial"
174    if os.path.exists(tmp_path):
175        shutil.rmtree(tmp_path)
176    store = zarr.open_group(tmp_path, mode="w", zarr_format=2)
177    array = store.create_array("0", shape=merged.shape, dtype="uint8", chunks=(64, 256, 256))
178    array[:] = merged
179    os.rename(tmp_path, label_path)
180    return label_path
181
182
183def get_drg_axon_mito_data(path: Union[os.PathLike, str], download: bool = False) -> str:
184    """Download the DRG-Axon-Mito mitochondrial membrane segmentation dataset.
185
186    Args:
187        path: Filepath to a folder where the data will be downloaded.
188        download: Whether to download the data if it is not present.
189
190    Returns:
191        Filepath where the data is stored.
192    """
193    os.makedirs(path, exist_ok=True)
194
195    for dataset, run, voxel_spacing, omm_scale, imm_scale in tqdm(RUNS, desc="Downloading the DRG axon tomograms"):
196        run_dir = os.path.join(path, f"{dataset}_{run}")
197        raw_url = RAW_URL.format(dataset=dataset, run=run, voxel_spacing=voxel_spacing)
198        _download_ome_zarr(raw_url, os.path.join(run_dir, "raw.zarr"), download)
199        if not download and not os.path.exists(os.path.join(run_dir, "labels.zarr")):
200            raise RuntimeError(f"Cannot find the data at {run_dir}, but download was set to False.")
201        _merge_labels(run_dir, dataset, run, voxel_spacing, omm_scale, imm_scale, download)
202
203    return path
204
205
206def get_drg_axon_mito_paths(
207    path: Union[os.PathLike, str], download: bool = False
208) -> Tuple[List[str], List[str]]:
209    """Get paths to the DRG-Axon-Mito data.
210
211    Args:
212        path: Filepath to a folder where the data will be downloaded.
213        download: Whether to download the data if it is not present.
214
215    Returns:
216        List of filepaths to the tomograms.
217        List of filepaths to the multi-class mitochondrial membrane masks.
218    """
219    data_dir = get_drg_axon_mito_data(path, download)
220    run_dirs = [f"{dataset}_{run}" for dataset, run, _, _, _ in RUNS]
221    raw_paths = [os.path.join(data_dir, run_dir, "raw.zarr") for run_dir in run_dirs]
222    label_paths = [os.path.join(data_dir, run_dir, "labels.zarr") for run_dir in run_dirs]
223    return raw_paths, label_paths
224
225
226def get_drg_axon_mito_dataset(
227    path: Union[os.PathLike, str],
228    patch_shape: Tuple[int, int, int],
229    download: bool = False,
230    **kwargs
231) -> Dataset:
232    """Get the dataset for mitochondrial membrane segmentation in cryo-ET tomograms of DRG axons.
233
234    The labels are a multi-class mask with 1: mitochondrial outer membrane and 2: mitochondrial inner membrane.
235
236    Args:
237        path: Filepath to a folder where the data will be downloaded.
238        patch_shape: The patch shape to use for training.
239        download: Whether to download the data if it is not present.
240        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
241
242    Returns:
243        The segmentation dataset.
244    """
245    assert len(patch_shape) == 3
246
247    raw_paths, label_paths = get_drg_axon_mito_paths(path, download)
248
249    return torch_em.default_segmentation_dataset(
250        raw_paths=raw_paths,
251        raw_key="0",
252        label_paths=label_paths,
253        label_key="0",
254        patch_shape=patch_shape,
255        is_seg_dataset=True,
256        **kwargs
257    )
258
259
260def get_drg_axon_mito_loader(
261    path: Union[os.PathLike, str],
262    patch_shape: Tuple[int, int, int],
263    batch_size: int,
264    download: bool = False,
265    **kwargs
266) -> DataLoader:
267    """Get the DataLoader for mitochondrial membrane segmentation in cryo-ET tomograms of DRG axons.
268
269    Args:
270        path: Filepath to a folder where the data will be downloaded.
271        patch_shape: The patch shape to use for training.
272        batch_size: The batch size for training.
273        download: Whether to download the data if it is not present.
274        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
275
276    Returns:
277        The DataLoader.
278    """
279    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
280    dataset = get_drg_axon_mito_dataset(path, patch_shape, download=download, **ds_kwargs)
281    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)