torch_em.data.datasets.electron_microscopy.asem

  1"""ASEM is a dataset for segmentation of cellular structures in FIB-SEM.
  2
  3The data is hosted in the public `asem-project` S3 bucket (https://open.quiltdata.com/b/asem-project).
  4Each cell is stored as `datasets/<id>/<id>.zarr`, with the raw data under `volumes/raw` and the ground-truth
  5annotations under `volumes/labels/<organelle>`. The labels are semantic (a single foreground value).
  6Some label arrays span the whole cell, others are crops with an offset into the raw volume;
  7only the raw region covered by the labels of the chosen organelles is downloaded.
  8A bounding box (z0, z1, y0, y1, x0, x1) in raw voxels can be given per volume. Then only this region is
  9downloaded, into a separate store `<id>/<id>_z<z0>-<z1>_y<y0>-<y1>_x<x0>-<x1>.zarr` with the box in its attributes.
 10
 11The dataset was publised in https://doi.org/10.1083/jcb.202208005.
 12Please cite this publication if you use the dataset in your research.
 13"""
 14
 15import os
 16from concurrent import futures
 17from typing import Dict, Union, Tuple, Optional, List
 18
 19import numpy as np
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26from ... import ConcatDataset
 27
 28
 29BUCKET = "s3://asem-project/datasets"
 30
 31# The volume ids are the ids in the bucket. They correspond to the cells of the publication as follows:
 32# 46 and 58 are cells 1 and 2 (HEK293A, the bucket does not say which is which), 61 is cell 3 (BSC-1),
 33# 64 is cell 6 (SUM159), and 1E, 2E and 3A are cells 12, 13 and 13a (high-pressure frozen).
 34VOLUMES = ["46", "58", "61", "64", "1E", "2E", "3A"]
 35
 36ORGANELLES = {
 37    "mito": ["46", "58", "61", "64", "2E", "3A"],
 38    "golgi": ["46", "58", "61", "64"],
 39    "er": ["46", "58", "61", "64"],
 40    "ccp": ["1E", "2E"],
 41    "np": ["3A"],
 42    "np_bottom": ["3A"],
 43}
 44
 45# The mito labels of 61 span the whole cell, but only this block (in voxels of the label array) is annotated.
 46# The block is the hold-out region the publication evaluates on (the extent of `volumes/predictions/mito`).
 47ANNOTATED_REGIONS = {("61", "mito"): np.s_[0:700, 190:730, 2630:3380]}
 48
 49BLOCK_SHAPE = (128, 512, 512)
 50N_THREADS = 8
 51
 52
 53def _check_volume_ids(volume_ids, organelle):
 54    for volume_id in volume_ids:
 55        if volume_id.startswith("cell_"):
 56            raise ValueError(
 57                f"The ASEM volume ids are the ids of the S3 bucket, for example '64' instead of 'cell_6', "
 58                f"got '{volume_id}'. See the comment above 'VOLUMES' for the cell numbers of the publication."
 59            )
 60        if volume_id not in ORGANELLES[organelle]:
 61            raise ValueError(
 62                f"The volume '{volume_id}' has no '{organelle}' labels. Choose from {ORGANELLES[organelle]}."
 63            )
 64
 65
 66def _open_remote(volume_id, key):
 67    import fsspec
 68    import zarr
 69
 70    store = fsspec.get_mapper(f"{BUCKET}/{volume_id}/{volume_id}.zarr/volumes/{key}", anon=True)
 71    return zarr.open_array(store, mode="r", zarr_format=2)
 72
 73
 74def _copy_region(src, dst, region):
 75    """Copy a region of the remote array into the local array, block by block."""
 76    starts = [r.start for r in region]
 77    grid = [int(np.ceil(s / b)) for s, b in zip(dst.shape, BLOCK_SHAPE)]
 78    blocks = [
 79        tuple(slice(i * b, min((i + 1) * b, s)) for i, b, s in zip(index, BLOCK_SHAPE, dst.shape))
 80        for index in np.ndindex(*grid)
 81    ]
 82
 83    def copy_block(block):
 84        dst[block] = src[tuple(slice(start + b.start, start + b.stop) for start, b in zip(starts, block))]
 85
 86    with futures.ThreadPoolExecutor(N_THREADS) as tp:
 87        list(tp.map(copy_block, blocks))
 88
 89
 90def _get_raw_key(volume_path, organelle):
 91    return f"raw_{organelle}" if os.path.exists(os.path.join(volume_path, "volumes", f"raw_{organelle}")) else "raw"
 92
 93
 94def _have_organelle(volume_path, organelle):
 95    volumes = os.path.join(volume_path, "volumes")
 96    labels = os.path.join(volumes, "labels", organelle)
 97    return os.path.exists(labels) and os.path.exists(os.path.join(volumes, _get_raw_key(volume_path, organelle)))
 98
 99
100def _get_volume_path(path, volume_id, bounding_box):
101    if bounding_box is None:
102        return os.path.join(path, volume_id, f"{volume_id}.zarr")
103    z0, z1, y0, y1, x0, x1 = bounding_box
104    return os.path.join(path, volume_id, f"{volume_id}_z{z0}-{z1}_y{y0}-{y1}_x{x0}-{x1}.zarr")
105
106
107def _download_organelle(volume_path, volume_id, organelle, bounding_box=None):
108    """Download the labels of one organelle and the raw region they cover.
109
110    The raw data is stored as 'volumes/raw' if the labels cover the whole raw volume (shared by all organelles),
111    and as 'volumes/raw_<organelle>' otherwise. If a bounding box is given, the raw data and the labels of this
112    box are stored, and the raw data is always 'volumes/raw'. Each array is written to a temporary folder and
113    moved in place once it is complete, so that an interrupted download is not mistaken for a complete one.
114    """
115    import zarr
116    from zarr.codecs import BloscCodec
117
118    raw, labels = _open_remote(volume_id, "raw"), _open_remote(volume_id, f"labels/{organelle}")
119    resolution = np.array(raw.attrs["resolution"])
120    offset = (np.array(labels.attrs["offset"]) - np.array(raw.attrs["offset"])) // resolution
121    if bounding_box is None:
122        label_region = ANNOTATED_REGIONS.get((volume_id, organelle), tuple(slice(0, s) for s in labels.shape))
123        raw_region = tuple(slice(int(o) + r.start, int(o) + r.stop) for o, r in zip(offset, label_region))
124    else:
125        raw_region = tuple(slice(start, stop) for start, stop in zip(bounding_box[::2], bounding_box[1::2]))
126        label_region = tuple(slice(r.start - int(o), r.stop - int(o)) for o, r in zip(offset, raw_region))
127        if any(r.start < 0 or r.stop > s or r.start >= r.stop for r, s in zip(label_region, labels.shape)):
128            label_box = [(int(o), int(o) + s) for o, s in zip(offset, labels.shape)]
129            raise ValueError(
130                f"The bounding box {bounding_box} is not inside the '{organelle}' labels of the ASEM volume "
131                f"'{volume_id}', which cover the raw region {label_box}."
132            )
133    full_raw = all(r.start == 0 and r.stop == s for r, s in zip(raw_region, raw.shape))
134    raw_key = "raw" if full_raw or bounding_box is not None else f"raw_{organelle}"
135
136    root = zarr.open_group(volume_path, mode="a")
137    if bounding_box is not None:
138        root.attrs["bounding_box"] = list(bounding_box)
139    group = root.require_group("volumes")
140    for key, src, region in [(raw_key, raw, raw_region), (f"labels/{organelle}", labels, label_region)]:
141        out_path = os.path.join(volume_path, "volumes", key)
142        if os.path.exists(out_path):
143            continue
144        print(f"Downloading 'volumes/{key}' of the ASEM volume '{volume_id}' ...")
145        tmp_key = f"{key}_incomplete"
146        dst = group.create_array(
147            tmp_key, shape=tuple(r.stop - r.start for r in region), chunks=(64, 256, 256), dtype=src.dtype,
148            compressors=BloscCodec(cname="zstd", clevel=5), overwrite=True,
149        )
150        _copy_region(src, dst, region)
151        dst.attrs.update({
152            "resolution": resolution.tolist(),
153            "offset": (np.array(raw.attrs["offset"]) + np.array([r.start for r in raw_region]) * resolution).tolist(),
154        })
155        os.replace(os.path.join(volume_path, "volumes", tmp_key), out_path)
156
157
158def get_asem_data(
159    path: Union[os.PathLike, str],
160    volume_ids: List[str],
161    download: bool = False,
162    organelles: Optional[Union[List[str], str]] = None,
163    bounding_boxes: Optional[Dict[str, Tuple[int, int, int, int, int, int]]] = None,
164):
165    """Download the ASEM dataset.
166
167    The data is streamed from the public S3 bucket at s3://asem-project, which requires `s3fs`.
168
169    Args:
170        path: Filepath to a folder where the downloaded data will be saved.
171        volume_ids: List of volumes to download.
172        download: Whether to download the data if it is not present.
173        organelles: The organelles to download for each volume. By default all organelles annotated in the volume.
174        bounding_boxes: The bounding box per volume id, as (z0, z1, y0, y1, x0, x1) in raw voxels.
175            Only this region of the volume is downloaded. Volumes without a bounding box are downloaded in full.
176    """
177    if organelles is None:
178        organelles = list(ORGANELLES)
179    elif isinstance(organelles, str):
180        organelles = [organelles]
181
182    for volume_id in volume_ids:
183        volume_organelles = [organelle for organelle in organelles if volume_id in ORGANELLES[organelle]]
184        for organelle in volume_organelles:
185            _check_volume_ids([volume_id], organelle)
186        if not volume_organelles:
187            raise ValueError(f"The volume '{volume_id}' has none of the organelles {organelles}.")
188
189        bounding_box = None if bounding_boxes is None else bounding_boxes.get(volume_id)
190        volume_path = _get_volume_path(path, volume_id, bounding_box)
191        for organelle in volume_organelles:
192            if _have_organelle(volume_path, organelle):
193                continue
194            if not download:
195                raise FileNotFoundError(
196                    f"The '{organelle}' data of the ASEM volume '{volume_id}' is not found at '{volume_path}', "
197                    "and 'download' is set to False."
198                )
199            _download_organelle(volume_path, volume_id, organelle, bounding_box)
200
201
202def get_asem_paths(
203    path: Union[os.PathLike, str],
204    volume_ids: List[str],
205    download: bool = False,
206    organelles: Optional[Union[List[str], str]] = None,
207    bounding_boxes: Optional[Dict[str, Tuple[int, int, int, int, int, int]]] = None,
208) -> List[str]:
209    """Get paths to the ASEM data.
210
211    Args:
212        path: Filepath to a folder where the downloaded data will be saved.
213        volume_ids: List of volumes to download.
214        download: Whether to download the data if it is not present.
215        organelles: The organelles to download for each volume. By default all organelles annotated in the volume.
216        bounding_boxes: The bounding box per volume id, as (z0, z1, y0, y1, x0, x1) in raw voxels.
217            Only this region of the volume is downloaded. Volumes without a bounding box are downloaded in full.
218
219    Returns:
220        List of paths for all volume ids.
221    """
222    get_asem_data(path, volume_ids, download, organelles, bounding_boxes)
223    bounding_boxes = {} if bounding_boxes is None else bounding_boxes
224    return [_get_volume_path(path, volume_id, bounding_boxes.get(volume_id)) for volume_id in volume_ids]
225
226
227def _check_input_args(input_arg, default_values):
228    if input_arg is None:
229        input_arg = default_values
230    else:
231        if isinstance(input_arg, str):
232            assert input_arg in default_values
233            input_arg = [input_arg]
234
235    return input_arg
236
237
238def get_asem_dataset(
239    path: Union[os.PathLike, str],
240    patch_shape: Tuple[int, ...],
241    download: bool = False,
242    organelles: Optional[Union[List[str], str]] = None,
243    volume_ids: Optional[Union[List[str], str]] = None,
244    bounding_boxes: Optional[Dict[str, Tuple[int, int, int, int, int, int]]] = None,
245    **kwargs
246) -> Dataset:
247    """Get dataset for segmentation of organelles in FIB-SEM cells.
248
249    The labels are semantic, use a label transform (e.g. connected components) to obtain instances.
250    For the mito labels of volume 61 only the annotated block is used (see `ANNOTATED_REGIONS`),
251    so any `rois` passed for it are relative to that block. A bounding box replaces this block.
252
253    Args:
254        path: Filepath to a folder where the downloaded data will be saved.
255        patch_shape: The patch shape to use for training.
256        download: Whether to download the data if it is not present.
257        organelles: The choice of organelles.
258        volume_ids: The choice of volumes.
259        bounding_boxes: The bounding box per volume id, as (z0, z1, y0, y1, x0, x1) in raw voxels.
260            Only this region of the volume is used. Volumes without a bounding box are used in full.
261        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
262
263    Returns:
264        The segmentation dataset.
265    """
266    organelles = _check_input_args(organelles, ORGANELLES)
267
268    all_datasets = []
269    for organelle in organelles:
270        if volume_ids is None:
271            organelle_volume_ids = ORGANELLES[organelle]
272        else:
273            organelle_volume_ids = [volume_ids] if isinstance(volume_ids, str) else list(volume_ids)
274            _check_volume_ids(organelle_volume_ids, organelle)
275
276        volume_paths = get_asem_paths(path, organelle_volume_ids, download, [organelle], bounding_boxes)
277        for volume_path in volume_paths:
278            dataset = torch_em.default_segmentation_dataset(
279                raw_paths=volume_path,
280                raw_key=f"volumes/{_get_raw_key(volume_path, organelle)}",
281                label_paths=volume_path,
282                label_key=f"volumes/labels/{organelle}",
283                patch_shape=patch_shape,
284                is_seg_dataset=True,
285                **kwargs
286            )
287            dataset.max_sampling_attempts = 5000
288            all_datasets.append(dataset)
289
290    return ConcatDataset(*all_datasets)
291
292
293def get_asem_loader(
294    path: Union[os.PathLike, str],
295    patch_shape: Tuple[int, ...],
296    batch_size: int,
297    download: bool = False,
298    organelles: Optional[Union[List[str], str]] = None,
299    volume_ids: Optional[Union[List[str], str]] = None,
300    bounding_boxes: Optional[Dict[str, Tuple[int, int, int, int, int, int]]] = None,
301    **kwargs
302) -> DataLoader:
303    """Get dataloader for the segmentation of organelles in FIB-SEM cells.
304
305    Args:
306        path: Filepath to a folder where the downloaded data will be saved.
307        patch_shape: The patch shape to use for training.
308        batch_size: The batch size for training.
309        download: Whether to download the data if it is not present.
310        organelles: The choice of organelles.
311        volume_ids: The choice of volumes.
312        bounding_boxes: The bounding box per volume id, as (z0, z1, y0, y1, x0, x1) in raw voxels.
313            Only this region of the volume is used. Volumes without a bounding box are used in full.
314        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
315
316    Returns:
317        The DataLoader.
318    """
319    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
320    ds = get_asem_dataset(path, patch_shape, download, organelles, volume_ids, bounding_boxes, **ds_kwargs)
321    loader = torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
322    return loader