torch_em.data.datasets.electron_microscopy.fafb

  1"""The FAFB (Full Adult Fly Brain) dataset contains a serial-section TEM volume of the
  2full adult female Drosophila brain with dense neuron instance segmentation from FlyWire.
  3
  4The EM (FAFB v14) is a ssTEM dataset. The native 4 x 4 x 40 nm mip level is a
  5placeholder with no data and mip=1 (8 x 8 x 40 nm) holds EM only. mip=2 (16 x 16 x 40 nm)
  6is the finest level of the FlyWire neuron segmentation (materialization v783, Nature 2024
  7paper), so both are used at 16 x 16 x 40 nm.
  8
  9Bounding boxes are specified in 16 x 16 x 40 nm voxel coordinates
 10(x_min, x_max, y_min, y_max, z_min, z_max).
 11Valid coordinate overlap between EM (mip=2) and seg: x=[5100,59200], y=[1440,29600], z=[16,7062].
 12
 13The EM is at gs://microns-seunglab/drosophila_v0/alignment/image_rechunked (mip=2) and
 14the neuron segmentation (v783) is at gs://flywire_v141_m783.
 15
 16This dataset is from the publication https://doi.org/10.1038/s41586-024-07558-y.
 17Please cite it if you use this dataset in your research.
 18
 19The dataset is publicly available at https://flywire.ai.
 20Requires cloud-volume: pip install cloud-volume.
 21
 22NOTE (on data size): the full seg volume is (54100, 28160, 7046) voxels at 16 x 16 x 40 nm.
 23Downloading the entire volume is not feasible. Data is streamed from GCS and cached
 24locally as zarr v3 stores by specifying bounding boxes.
 25
 26NOTE (AA): The data annotations are amazing, I personally think that the segmentation
 27resolution is too low. If we wanna use it, we should go one resolution higher
 28(we are at s2 atm).
 29"""
 30
 31import hashlib
 32import os
 33from typing import List, Optional, Tuple, Union
 34
 35import numpy as np
 36from torch.utils.data import DataLoader, Dataset
 37
 38import torch_em
 39from .. import util
 40
 41
 42EM_URL = "gs://microns-seunglab/drosophila_v0/alignment/image_rechunked"
 43SEG_URL = "gs://flywire_v141_m783"
 44# mip=2 gives 16x16x40nm, matching the seg resolution; mip=0 is a placeholder with no data.
 45EM_MIP = 2
 46
 47# 1024x1024x410-voxel crops (16 x 16 x 16 um) inside brain tissue at three depths; the brain fills only part of
 48# the coordinate range, so a box must be checked against the data.
 49DEFAULT_BOUNDING_BOXES = [
 50    (35840, 36864, 6656, 7680, 1500, 1910),  # right, dorsal, anterior
 51    (32768, 33792, 18944, 19968, 1500, 1910),  # midline, ventral, anterior
 52    (15360, 16384, 17920, 18944, 3500, 3910),  # left optic lobe, ventral, mid-depth
 53    (48128, 49152, 18944, 19968, 3500, 3910),  # right optic lobe, ventral, mid-depth
 54    (24576, 25600, 11776, 12800, 3500, 3910),  # left, dorsal, mid-depth
 55    (41984, 43008, 12800, 13824, 3500, 3910),  # right, mid-height, mid-depth
 56    (18432, 19456, 11776, 12800, 5500, 5910),  # left, dorsal, posterior
 57    (21504, 22528, 17920, 18944, 5500, 5910),  # left, ventral, posterior
 58    (23552, 24576, 14848, 15872, 5500, 5910),  # left, mid-height, posterior
 59]
 60DEFAULT_BOUNDING_BOX = DEFAULT_BOUNDING_BOXES[4]
 61
 62FAFB_CHUNK_SHAPE = (64, 256, 256)
 63
 64
 65def _bbox_to_str(bbox):
 66    return hashlib.md5("_".join(str(v) for v in bbox).encode()).hexdigest()[:12]
 67
 68
 69def _create_array(root, name, shape, dtype, is_label):
 70    from zarr.codecs import BloscCodec
 71    shuffle = "bitshuffle" if (np.issubdtype(dtype, np.integer) and is_label) else "shuffle"
 72    return root.create_array(
 73        name,
 74        shape=shape,
 75        chunks=FAFB_CHUNK_SHAPE,
 76        dtype=dtype,
 77        compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
 78    )
 79
 80
 81def get_fafb_data(
 82    path: Union[os.PathLike, str],
 83    bounding_box: Tuple[int, int, int, int, int, int] = DEFAULT_BOUNDING_BOX,
 84    download: bool = False,
 85) -> str:
 86    """Stream a subvolume from the FAFB dataset and cache it as a zarr v3 store.
 87
 88    Args:
 89        path: Filepath to a folder where the cached zarr store will be saved.
 90        bounding_box: The region to fetch as (x_min, x_max, y_min, y_max, z_min, z_max)
 91            in 16 nm voxel coordinates. Defaults to DEFAULT_BOUNDING_BOXES, 1024x1024x410 crops inside brain tissue.
 92        download: Whether to stream and cache the data if it is not present.
 93
 94    Returns:
 95        The filepath to the cached zarr store.
 96    """
 97    import zarr
 98
 99    os.makedirs(str(path), exist_ok=True)
100    zarr_path = os.path.join(str(path), f"{_bbox_to_str(bounding_box)}.zarr")
101
102    root = zarr.open_group(zarr_path, mode="a")
103    if "raw" in root and "labels" in root:
104        return zarr_path
105
106    if not download:
107        raise RuntimeError(
108            f"No cached data found at '{zarr_path}'. Set download=True to stream it from GCS."
109        )
110
111    try:
112        import cloudvolume
113    except ImportError:
114        raise ImportError("The 'cloud-volume' package is required: pip install cloud-volume")
115
116    x_min, x_max, y_min, y_max, z_min, z_max = bounding_box
117    print(f"Streaming FAFB EM + FlyWire segmentation for bbox {bounding_box} ...")
118
119    em_vol = cloudvolume.CloudVolume(EM_URL, use_https=True, mip=EM_MIP, progress=True)
120    seg_vol = cloudvolume.CloudVolume(SEG_URL, use_https=True, mip=0, progress=True)
121
122    raw = np.array(em_vol[x_min:x_max, y_min:y_max, z_min:z_max])[..., 0].transpose(2, 1, 0)
123    labels = np.array(seg_vol[x_min:x_max, y_min:y_max, z_min:z_max])[..., 0].transpose(2, 1, 0)
124
125    # The brain fills only part of the coordinate range and the servers return zeros outside it.
126    if not raw.any() or len(np.unique(labels)) < 2:
127        raise RuntimeError(
128            f"The bounding box {bounding_box} holds no tissue or no segmentation. "
129            "Pick a box inside the brain, e.g. one of DEFAULT_BOUNDING_BOXES."
130        )
131
132    # FlyWire IDs are large uint64 values - relabel to consecutive integers.
133    _, labels = np.unique(labels, return_inverse=True)
134    labels = labels.reshape(raw.shape).astype("uint64")
135
136    shape = tuple(min(r, l) for r, l in zip(raw.shape, labels.shape))
137    raw = raw[:shape[0], :shape[1], :shape[2]]
138    labels = labels[:shape[0], :shape[1], :shape[2]]
139
140    root.attrs["bounding_box"] = list(bounding_box)
141    root.attrs["resolution_nm"] = [16, 16, 40]
142
143    if "raw" not in root:
144        ds_raw = _create_array(root, "raw", shape, np.dtype("uint8"), is_label=False)
145        ds_raw[:] = raw
146    if "labels" not in root:
147        ds_lbl = _create_array(root, "labels", shape, np.dtype("uint64"), is_label=True)
148        ds_lbl[:] = labels
149
150    print(f"Cached to {zarr_path} (shape {shape})")
151    return zarr_path
152
153
154def get_fafb_paths(
155    path: Union[os.PathLike, str],
156    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
157    download: bool = False,
158) -> List[str]:
159    """Get paths to FAFB zarr stores.
160
161    Args:
162        path: Filepath to a folder where the cached zarr stores will be saved.
163        bounding_boxes: List of regions to fetch, each as
164            (x_min, x_max, y_min, y_max, z_min, z_max) in 16 nm voxel coordinates.
165            Defaults to DEFAULT_BOUNDING_BOXES, 1024x1024x410 crops inside brain tissue.
166        download: Whether to stream and cache the data if it is not present.
167
168    Returns:
169        List of filepaths to the cached zarr stores.
170    """
171    if bounding_boxes is None:
172        bounding_boxes = DEFAULT_BOUNDING_BOXES
173    return [get_fafb_data(path, bbox, download) for bbox in bounding_boxes]
174
175
176def get_fafb_dataset(
177    path: Union[os.PathLike, str],
178    patch_shape: Tuple[int, int, int],
179    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
180    download: bool = False,
181    offsets: Optional[List[List[int]]] = None,
182    boundaries: bool = False,
183    **kwargs,
184) -> Dataset:
185    """Get the FAFB dataset for neuron instance segmentation.
186
187    Args:
188        path: Filepath to a folder where the cached zarr stores will be saved.
189        patch_shape: The patch shape (z, y, x) to use for training.
190        bounding_boxes: List of subvolumes to use, each as
191            (x_min, x_max, y_min, y_max, z_min, z_max) in 16 nm voxel coordinates.
192            Defaults to DEFAULT_BOUNDING_BOXES, 1024x1024x410 crops inside brain tissue.
193        download: Whether to stream and cache data if not already present.
194        offsets: Offset values for affinity computation used as target.
195        boundaries: Whether to compute boundaries as the target.
196        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
197
198    Returns:
199        The segmentation dataset.
200    """
201    assert len(patch_shape) == 3
202
203    paths = get_fafb_paths(path, bounding_boxes, download)
204
205    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
206    kwargs, _ = util.add_instance_label_transform(
207        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
208    )
209
210    return torch_em.default_segmentation_dataset(
211        raw_paths=paths,
212        raw_key="raw",
213        label_paths=paths,
214        label_key="labels",
215        patch_shape=patch_shape,
216        **kwargs,
217    )
218
219
220def get_fafb_loader(
221    path: Union[os.PathLike, str],
222    patch_shape: Tuple[int, int, int],
223    batch_size: int,
224    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
225    download: bool = False,
226    offsets: Optional[List[List[int]]] = None,
227    boundaries: bool = False,
228    **kwargs,
229) -> DataLoader:
230    """Get the DataLoader for neuron instance segmentation in the FAFB dataset.
231
232    Args:
233        path: Filepath to a folder where the cached zarr stores will be saved.
234        patch_shape: The patch shape (z, y, x) to use for training.
235        batch_size: The batch size for training.
236        bounding_boxes: List of subvolumes to use, each as
237            (x_min, x_max, y_min, y_max, z_min, z_max) in 16 nm voxel coordinates.
238            Defaults to DEFAULT_BOUNDING_BOXES, 1024x1024x410 crops inside brain tissue.
239        download: Whether to stream and cache data if not already present.
240        offsets: Offset values for affinity computation used as target.
241        boundaries: Whether to compute boundaries as the target.
242        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
243            or for the PyTorch DataLoader.
244
245    Returns:
246        The DataLoader.
247    """
248    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
249    dataset = get_fafb_dataset(
250        path, patch_shape, bounding_boxes=bounding_boxes,
251        download=download, offsets=offsets, boundaries=boundaries, **ds_kwargs
252    )
253    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)