torch_em.data.datasets.electron_microscopy.cmito

  1"""Cmito dataset for dense mitochondria segmentation in whole-animal C. elegans EM volumes.
  2
  3The dataset provides dense mitochondria instance masks for five developmental stages of
  4C. elegans, built on public whole-animal EM connectome volumes:
  5
  6- l1, l2, l3, adult: from Witvliet et al. 2021 (https://doi.org/10.1038/s41586-021-03778-8),
  7  CC-BY-4.0, deposited at https://zenodo.org/records/5637219 and https://bossdb.org/project/witvliet2020.
  8- dauer: from Yim et al. 2024 (https://doi.org/10.1038/s41467-024-45943-3), whose article text
  9  is CC-BY-4.0. The mitochondria mask bucket for this stage lives in a separate GCS project
 10  (`gnd-dauer1`) than the other four stages, and no repository-level data license could be
 11  confirmed for it independently of the article license. Treat "dauer" as license-uncertain
 12  until a repository landing page or downloadable asset states its data license explicitly.
 13
 14Raw EM is streamed from BossDB (S3) at a mip level chosen per stage so that it spatially matches
 15the coarser, separately hosted mitochondria mask (on Google Cloud Storage, at native resolution).
 16Both are cached locally as zarr v3 stores in (z, y, x) axis order.
 17
 18This dataset is from the publication https://doi.org/10.1101/2024.07.19.604219.
 19Please cite it if you use this dataset in your research.
 20
 21Corresponding authors: J. Alexander Bae (jabae@snu.ac.kr), Junho Lee (elegans@snu.ac.kr).
 22Original tool repository: https://github.com/jabae/Cmito.
 23Requires cloud-volume: pip install cloud-volume.
 24"""
 25
 26import hashlib
 27import os
 28from concurrent.futures import ThreadPoolExecutor, as_completed
 29from typing import List, Literal, Optional, Sequence, Tuple, Union
 30
 31import numpy as np
 32from tqdm import tqdm
 33from torch.utils.data import Dataset, DataLoader
 34
 35import torch_em
 36from .. import util
 37
 38
 39CMITO_STAGES = {
 40    "l1": {
 41        "raw_url": "precomputed://https://bossdb-open-data.s3.amazonaws.com/witvliet2020/Dataset_2/em",
 42        "mask_url": "precomputed://https://storage.googleapis.com/gnd-neuroglancer/witvliet/dataset2/mito_seg_v3",
 43        # Raw mip3 (0.64 nm x 8 = 5.12 nm xy) spatially matches the mask's native 5.12/5.12/50 nm.
 44        "raw_mip": 3,
 45        # Full extent in nm, from the mip0 shape (26624, 22016, 368) x resolution (0.64, 0.64, 50).
 46        "bbox_nm": (0, 17039, 0, 14090, 0, 18400),
 47    },
 48    "l2": {
 49        "raw_url": "precomputed://https://bossdb-open-data.s3.amazonaws.com/witvliet2020/Dataset_5/em",
 50        "mask_url": "precomputed://https://storage.googleapis.com/gnd-neuroglancer/witvliet/dataset5/mito_seg_v4",
 51        # Raw mip2 (2 nm x 4 = 8 nm xy) spatially matches the mask's native 8/8/30 nm.
 52        "raw_mip": 2,
 53        "bbox_nm": (0, 27648, 0, 22528, 0, 25920),
 54    },
 55    "l3": {
 56        "raw_url": "precomputed://https://bossdb-open-data.s3.amazonaws.com/witvliet2020/Dataset_6/em",
 57        "mask_url": "precomputed://https://storage.googleapis.com/gnd-neuroglancer/witvliet/dataset6/mito_seg_v4",
 58        # Raw mip3 (0.768 nm x 8 = 6.144 nm xy) is the mask's voxel grid (both 4544 x 4288), although the mask
 59        # reports 8 nm. Bounding boxes are in the nm of the raw grid, see 'get_cmito_data'.
 60        "raw_mip": 3,
 61        "bbox_nm": (0, 27918, 0, 26345, 0, 21600),
 62    },
 63    "adult": {
 64        "raw_url": "precomputed://https://bossdb-open-data.s3.amazonaws.com/witvliet2020/Dataset_8/em",
 65        "mask_url": "precomputed://https://storage.googleapis.com/gnd-neuroglancer/witvliet/dataset8/mito_seg_v3",
 66        # Raw mip3 (2 nm x 8 = 16 nm xy) spatially matches the mask's native 16/16/30 nm.
 67        "raw_mip": 3,
 68        "bbox_nm": (0, 79872, 0, 44032, 0, 21120),
 69    },
 70    "dauer": {
 71        "raw_url": "precomputed://https://bossdb-open-data.s3.amazonaws.com/yim_choe_bae2023/dauer1_364/em/em",
 72        # NOTE: different GCS bucket/project than the other four stages.
 73        "mask_url": "precomputed://https://storage.googleapis.com/gnd-dauer1/dauer1_364/mito_seg_v4",
 74        # Raw mip3 (1 nm x 8 = 8 nm xy) spatially matches the mask's native 8/8/50 nm.
 75        "raw_mip": 3,
 76        "bbox_nm": (0, 9800, 0, 9792, 0, 18200),
 77    },
 78}
 79
 80CMITO_CHUNK_SHAPE = (64, 128, 128)
 81CMITO_SHARD_SHAPE = (128, 512, 512)
 82
 83
 84def _cmito_bbox_to_str(bbox):
 85    return hashlib.md5("_".join(str(v) for v in bbox).encode()).hexdigest()[:12]
 86
 87
 88def _cmito_create_array(root, name, shape, dtype, is_label):
 89    from zarr.codecs import BloscCodec
 90    shuffle = "bitshuffle" if (np.issubdtype(dtype, np.integer) and is_label) else "shuffle"
 91    return root.create_array(
 92        name,
 93        shape=shape,
 94        chunks=CMITO_CHUNK_SHAPE,
 95        shards=CMITO_SHARD_SHAPE,
 96        dtype=dtype,
 97        compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
 98    )
 99
100
101def _cmito_bbox_voxels(cv, x_min_nm, x_max_nm, y_min_nm, y_max_nm, z_min_nm, z_max_nm):
102    scale = np.array(cv.resolution)
103    x0 = int(np.floor(x_min_nm / scale[0]))
104    x1 = int(np.ceil(x_max_nm / scale[0]))
105    y0 = int(np.floor(y_min_nm / scale[1]))
106    y1 = int(np.ceil(y_max_nm / scale[1]))
107    z0 = int(np.floor(z_min_nm / scale[2]))
108    z1 = int(np.ceil(z_max_nm / scale[2]))
109    return x0, x1, y0, y1, z0, z1, (z1 - z0, y1 - y0, x1 - x0)
110
111
112def _cmito_download_to_zarr(cv, ds, x0g, y0g, z0g, name):
113    shape = ds.shape  # (z, y, x)
114    sz, sy, sx = CMITO_SHARD_SHAPE
115
116    tasks = []
117    for z0_ in range(0, shape[0], sz):
118        for y0_ in range(0, shape[1], sy):
119            for x0_ in range(0, shape[2], sx):
120                z1_ = min(z0_ + sz, shape[0])
121                y1_ = min(y0_ + sy, shape[1])
122                x1_ = min(x0_ + sx, shape[2])
123                tasks.append((
124                    (z0_, z1_), (y0_, y1_), (x0_, x1_),
125                    (x0g + x0_, x0g + x1_, y0g + y0_, y0g + y1_, z0g + z0_, z0g + z1_),
126                ))
127
128    target_dtype = np.dtype(ds.dtype)
129
130    def worker(item):
131        (z0_, z1_), (y0_, y1_), (x0_, x1_), (gx0, gx1, gy0, gy1, gz0, gz1) = item
132        block = np.asarray(cv[gx0:gx1, gy0:gy1, gz0:gz1])
133        if block.ndim == 4:
134            block = block[..., 0]
135        ds[z0_:z1_, y0_:y1_, x0_:x1_] = block.transpose(2, 1, 0).astype(target_dtype)
136
137    with ThreadPoolExecutor(max_workers=8) as ex:
138        futures = [ex.submit(worker, t) for t in tasks]
139        for fut in tqdm(as_completed(futures), total=len(futures), desc=f"Downloading '{name}'", smoothing=0.05):
140            fut.result()
141
142
143def get_cmito_data(
144    path: Union[os.PathLike, str],
145    stage: Literal["l1", "l2", "l3", "adult", "dauer"],
146    bounding_box: Optional[Tuple[float, ...]] = None,
147    download: bool = False,
148) -> str:
149    """Stream and cache one Cmito developmental-stage volume as a zarr v3 store.
150
151    The zarr store contains:
152      - raw: EM grayscale (uint8, z/y/x), at the mip level that spatially matches the mask.
153      - labels: mitochondria instance segmentation (uint16, z/y/x), at native mask resolution.
154
155    Args:
156        path: Filepath to a folder where the cached zarr store will be saved.
157        stage: The developmental stage to use. One of 'l1', 'l2', 'l3', 'adult', 'dauer'.
158            The 'dauer' stage's mask bucket has no independently confirmed data license
159            (see the module docstring); use it with that caveat in mind.
160        bounding_box: Region in nm as (x_min, x_max, y_min, y_max, z_min, z_max), in the resolution of the raw
161            data at `raw_mip`. Defaults to the full volume extent for the chosen stage.
162        download: Whether to stream and cache the data if not present.
163
164    Returns:
165        Filepath to the cached zarr store.
166    """
167    import zarr
168
169    if stage not in CMITO_STAGES:
170        raise ValueError(f"Invalid stage: '{stage}'. Choose from {list(CMITO_STAGES.keys())}.")
171
172    stage_info = CMITO_STAGES[stage]
173    os.makedirs(str(path), exist_ok=True)
174    bbox = bounding_box if bounding_box is not None else stage_info["bbox_nm"]
175    bbox_hash = _cmito_bbox_to_str(bbox)
176    zarr_path = os.path.join(str(path), f"{stage}_{bbox_hash}.zarr")
177
178    def _complete(zp):
179        return os.path.isdir(os.path.join(zp, "raw")) and os.path.isdir(os.path.join(zp, "labels"))
180
181    if _complete(zarr_path):
182        return zarr_path
183    if not download:
184        raise RuntimeError(
185            f"No cached data at '{zarr_path}'. Set download=True to stream from BossDB and GCS."
186        )
187
188    try:
189        from cloudvolume import CloudVolume
190    except ImportError:
191        raise ImportError("The 'cloud-volume' package is required: pip install cloud-volume")
192
193    x_min_nm, x_max_nm, y_min_nm, y_max_nm, z_min_nm, z_max_nm = bbox
194    raw_mip = stage_info["raw_mip"]
195    print(f"Streaming Cmito stage='{stage}' at raw_mip={raw_mip} ...")
196
197    raw_cv = CloudVolume(
198        stage_info["raw_url"], use_https=True, mip=raw_mip, progress=False, fill_missing=True,
199    )
200    mask_cv = CloudVolume(stage_info["mask_url"], mip=0, progress=False, fill_missing=True)
201
202    # For every stage the raw (at raw_mip) and the mask share one voxel grid with the same origin, but the mask does
203    # not always report the true resolution (l3 reports 8 nm on the 6.144 nm raw grid). So the box is converted to
204    # voxels once, with the raw resolution, the same voxel box is read from both, and it is clipped to both volumes.
205    assert tuple(raw_cv.voxel_offset) == tuple(mask_cv.voxel_offset), \
206        f"The raw and mask voxel grids of stage '{stage}' have different origins."
207    x0, x1, y0, y1, z0, z1, _ = _cmito_bbox_voxels(raw_cv, x_min_nm, x_max_nm, y_min_nm, y_max_nm, z_min_nm, z_max_nm)
208    x1, y1, z1 = (min(v, int(r), int(m)) for v, r, m in zip((x1, y1, z1), raw_cv.volume_size, mask_cv.volume_size))
209    shape = (z1 - z0, y1 - y0, x1 - x0)
210    if min(shape) <= 0:
211        raise ValueError(f"The bounding box {bbox} lies outside the stage '{stage}' volume.")
212
213    root = zarr.open_group(zarr_path, mode="a")
214    root.attrs["stage"] = stage
215    root.attrs["bounding_box_nm"] = list(bbox)
216    root.attrs["raw_mip"] = raw_mip
217
218    if "raw" not in root:
219        ds_raw = _cmito_create_array(root, "raw", shape, np.dtype("uint8"), is_label=False)
220        _cmito_download_to_zarr(raw_cv, ds_raw, x0, y0, z0, name="raw")
221
222    if "labels" not in root:
223        ds_lbl = _cmito_create_array(root, "labels", shape, np.dtype("uint16"), is_label=True)
224        _cmito_download_to_zarr(mask_cv, ds_lbl, x0, y0, z0, name="labels")
225
226    print(f"Cached to {zarr_path} (shape {shape})")
227    return zarr_path
228
229
230def get_cmito_paths(
231    path: Union[os.PathLike, str],
232    stages: Optional[Sequence[str]] = None,
233    bounding_box: Optional[Tuple[float, ...]] = None,
234    download: bool = False,
235) -> List[str]:
236    """Get paths to cached Cmito zarr stores.
237
238    Args:
239        path: Filepath to a folder where the cached zarr stores will be saved.
240        stages: Developmental stages to load. Defaults to all five stages.
241        bounding_box: Region in nm as (x_min, x_max, y_min, y_max, z_min, z_max), applied to
242            every requested stage. Defaults to each stage's own full volume extent.
243        download: Whether to stream and cache the data if not present.
244
245    Returns:
246        Filepaths to the cached zarr stores.
247    """
248    stages_ = list(stages) if stages is not None else list(CMITO_STAGES.keys())
249    return [get_cmito_data(path, stage, bounding_box, download) for stage in stages_]
250
251
252def get_cmito_dataset(
253    path: Union[os.PathLike, str],
254    patch_shape: Tuple[int, int, int],
255    stage: Union[str, Sequence[str]] = "l1",
256    bounding_box: Optional[Tuple[float, ...]] = None,
257    download: bool = False,
258    offsets: Optional[List[List[int]]] = None,
259    boundaries: bool = False,
260    **kwargs,
261) -> Dataset:
262    """Get the Cmito dataset for mitochondria instance segmentation in whole-animal C. elegans EM.
263
264    Args:
265        path: Filepath to a folder where the cached zarr stores will be saved.
266        patch_shape: The patch shape (z, y, x) to use for training.
267        stage: The developmental stage(s) to use, one or several of 'l1', 'l2', 'l3', 'adult',
268            'dauer'. The 'dauer' stage's mask bucket has no independently confirmed data license
269            (see the module docstring); use it with that caveat in mind.
270        bounding_box: Region in nm as (x_min, x_max, y_min, y_max, z_min, z_max), applied to
271            every requested stage. Defaults to each stage's own full volume extent.
272        download: Whether to stream and cache data if not already present.
273        offsets: Offset values for affinity computation used as target.
274        boundaries: Whether to compute boundaries as the target.
275        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
276
277    Returns:
278        The segmentation dataset.
279    """
280    assert len(patch_shape) == 3
281    stages = [stage] if isinstance(stage, str) else list(stage)
282    paths = get_cmito_paths(path, stages, bounding_box, download)
283
284    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
285    kwargs, _ = util.add_instance_label_transform(
286        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
287    )
288
289    return torch_em.default_segmentation_dataset(
290        raw_paths=paths,
291        raw_key="raw",
292        label_paths=paths,
293        label_key="labels",
294        patch_shape=patch_shape,
295        **kwargs,
296    )
297
298
299def get_cmito_loader(
300    path: Union[os.PathLike, str],
301    batch_size: int,
302    patch_shape: Tuple[int, int, int],
303    stage: Union[str, Sequence[str]] = "l1",
304    bounding_box: Optional[Tuple[float, ...]] = None,
305    download: bool = False,
306    offsets: Optional[List[List[int]]] = None,
307    boundaries: bool = False,
308    **kwargs,
309) -> DataLoader:
310    """Get the DataLoader for mitochondria instance segmentation in whole-animal C. elegans EM.
311
312    Args:
313        path: Filepath to a folder where the cached zarr stores will be saved.
314        batch_size: The batch size for training.
315        patch_shape: The patch shape (z, y, x) to use for training.
316        stage: The developmental stage(s) to use, one or several of 'l1', 'l2', 'l3', 'adult',
317            'dauer'. The 'dauer' stage's mask bucket has no independently confirmed data license
318            (see the module docstring); use it with that caveat in mind.
319        bounding_box: Region in nm as (x_min, x_max, y_min, y_max, z_min, z_max), applied to
320            every requested stage. Defaults to each stage's own full volume extent.
321        download: Whether to stream and cache data if not already present.
322        offsets: Offset values for affinity computation used as target.
323        boundaries: Whether to compute boundaries as the target.
324        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
325            or for the PyTorch DataLoader.
326
327    Returns:
328        The DataLoader.
329    """
330    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
331    ds = get_cmito_dataset(
332        path=path,
333        patch_shape=patch_shape,
334        stage=stage,
335        bounding_box=bounding_box,
336        download=download,
337        offsets=offsets,
338        boundaries=boundaries,
339        **ds_kwargs,
340    )
341    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)