torch_em.data.datasets.electron_microscopy.optic_lobe

  1"""The OpticLobe dataset contains a FIB-SEM volume of the Drosophila right optic lobe
  2with dense neuron instance segmentation, covering the right optic lobe subset of the
  3MaleCNS sample (about 50,000 neurons) at 8 nm isotropic resolution.
  4
  5The EM volume is at gs://flyem-optic-lobe/grayscale-clahe-jpeg and the segmentation is at
  6gs://flyem-optic-lobe/v1.1/segmentation.
  7
  8This dataset is from the publication https://doi.org/10.1038/s41586-025-08746-0.
  9Please cite it if you use this dataset in your research.
 10
 11The dataset is publicly available at https://www.janelia.org/project-team/flyem/optic-lobe
 12under the CC-BY license.
 13Requires cloud-volume: pip install cloud-volume.
 14
 15NOTE (on data size): the full volume is (36864, 55296, 45056) voxels at 8 nm isotropic
 16resolution. Downloading the entire volume is not feasible. Data is instead accessed by
 17specifying bounding boxes (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel
 18coordinates, streamed from GCS and cached locally as HDF5 files.
 19"""
 20
 21import hashlib
 22import os
 23from typing import List, Optional, Tuple, Union
 24
 25import numpy as np
 26from torch.utils.data import DataLoader, Dataset
 27
 28import torch_em
 29from .. import util
 30
 31
 32EM_URL = "gs://flyem-optic-lobe/grayscale-clahe-jpeg"
 33SEG_URL = "gs://flyem-optic-lobe/v1.1/segmentation"
 34
 35# A representative 1024³-voxel subvolume near the centre of the reconstructed region.
 36# Units are 8 nm voxels in (x, y, z) order, matching the CloudVolume coordinate space.
 37DEFAULT_BOUNDING_BOX = (17500, 18524, 29500, 30524, 29500, 30524)
 38
 39
 40def _bbox_to_str(bbox):
 41    return hashlib.md5("_".join(str(v) for v in bbox).encode()).hexdigest()[:12]
 42
 43
 44def get_optic_lobe_data(
 45    path: Union[os.PathLike, str],
 46    bounding_box: Tuple[int, int, int, int, int, int] = DEFAULT_BOUNDING_BOX,
 47    download: bool = False,
 48) -> str:
 49    """Stream a subvolume from the OpticLobe dataset and cache it as an HDF5 file.
 50
 51    Args:
 52        path: Filepath to a folder where the cached HDF5 file will be saved.
 53        bounding_box: The region to fetch as (x_min, x_max, y_min, y_max, z_min, z_max)
 54            in 8 nm voxel coordinates. Defaults to a central 1024³ training region.
 55        download: Whether to stream and cache the data if it is not present.
 56
 57    Returns:
 58        The filepath to the cached HDF5 file.
 59    """
 60    import h5py
 61
 62    os.makedirs(str(path), exist_ok=True)
 63    h5_path = os.path.join(str(path), f"{_bbox_to_str(bounding_box)}.h5")
 64    if os.path.exists(h5_path):
 65        return h5_path
 66
 67    if not download:
 68        raise RuntimeError(
 69            f"No cached data found at '{h5_path}'. Set download=True to stream it from GCS."
 70        )
 71
 72    try:
 73        import cloudvolume
 74    except ImportError:
 75        raise ImportError("The 'cloud-volume' package is required: pip install cloud-volume")
 76
 77    x_min, x_max, y_min, y_max, z_min, z_max = bounding_box
 78    print(f"Streaming OpticLobe EM + segmentation for bbox {bounding_box} ...")
 79
 80    em_vol = cloudvolume.CloudVolume(EM_URL, use_https=True, mip=0, progress=True)
 81    seg_vol = cloudvolume.CloudVolume(SEG_URL, use_https=True, mip=0, progress=True)
 82
 83    raw = np.array(em_vol[x_min:x_max, y_min:y_max, z_min:z_max])[..., 0].transpose(2, 1, 0)
 84    labels = np.array(seg_vol[x_min:x_max, y_min:y_max, z_min:z_max])[..., 0].transpose(2, 1, 0)
 85
 86    with h5py.File(h5_path, "w", locking=False) as f:
 87        f.attrs["bounding_box"] = bounding_box
 88        f.attrs["resolution_nm"] = em_vol.resolution.tolist()
 89        f.create_dataset("raw", data=raw.astype("uint8"), compression="gzip", chunks=True)
 90        f.create_dataset("labels", data=labels.astype("uint64"), compression="gzip", chunks=True)
 91
 92    print(f"Cached to {h5_path} (shape {raw.shape})")
 93    return h5_path
 94
 95
 96def get_optic_lobe_paths(
 97    path: Union[os.PathLike, str],
 98    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
 99    download: bool = False,
100) -> List[str]:
101    """Get paths to OpticLobe HDF5 cache files.
102
103    Args:
104        path: Filepath to a folder where the cached HDF5 files will be saved.
105        bounding_boxes: List of regions to fetch, each as
106            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
107            Defaults to [DEFAULT_BOUNDING_BOX].
108        download: Whether to stream and cache the data if it is not present.
109
110    Returns:
111        List of filepaths to the cached HDF5 files.
112    """
113    if bounding_boxes is None:
114        bounding_boxes = [DEFAULT_BOUNDING_BOX]
115    return [get_optic_lobe_data(path, bbox, download) for bbox in bounding_boxes]
116
117
118def get_optic_lobe_dataset(
119    path: Union[os.PathLike, str],
120    patch_shape: Tuple[int, int, int],
121    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
122    download: bool = False,
123    offsets: Optional[List[List[int]]] = None,
124    boundaries: bool = False,
125    **kwargs,
126) -> Dataset:
127    """Get the OpticLobe dataset for neuron instance segmentation.
128
129    Args:
130        path: Filepath to a folder where the cached HDF5 files will be saved.
131        patch_shape: The patch shape (z, y, x) to use for training.
132        bounding_boxes: List of subvolumes to use, each as
133            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
134            Defaults to [DEFAULT_BOUNDING_BOX] - a central 1024³ region.
135        download: Whether to stream and cache data if not already present.
136        offsets: Offset values for affinity computation used as target.
137        boundaries: Whether to compute boundaries as the target.
138        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
139
140    Returns:
141        The segmentation dataset.
142    """
143    assert len(patch_shape) == 3
144
145    paths = get_optic_lobe_paths(path, bounding_boxes, download)
146
147    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
148    kwargs, _ = util.add_instance_label_transform(
149        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
150    )
151
152    return torch_em.default_segmentation_dataset(
153        raw_paths=paths,
154        raw_key="raw",
155        label_paths=paths,
156        label_key="labels",
157        patch_shape=patch_shape,
158        **kwargs,
159    )
160
161
162def get_optic_lobe_loader(
163    path: Union[os.PathLike, str],
164    patch_shape: Tuple[int, int, int],
165    batch_size: int,
166    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
167    download: bool = False,
168    offsets: Optional[List[List[int]]] = None,
169    boundaries: bool = False,
170    **kwargs,
171) -> DataLoader:
172    """Get the DataLoader for neuron instance segmentation in the OpticLobe dataset.
173
174    Args:
175        path: Filepath to a folder where the cached HDF5 files will be saved.
176        patch_shape: The patch shape (z, y, x) to use for training.
177        batch_size: The batch size for training.
178        bounding_boxes: List of subvolumes to use, each as
179            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
180            Defaults to [DEFAULT_BOUNDING_BOX] - a central 1024³ region.
181        download: Whether to stream and cache data if not already present.
182        offsets: Offset values for affinity computation used as target.
183        boundaries: Whether to compute boundaries as the target.
184        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
185            or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    dataset = get_optic_lobe_dataset(
192        path, patch_shape, bounding_boxes=bounding_boxes, download=download,
193        offsets=offsets, boundaries=boundaries, **ds_kwargs
194    )
195    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)