torch_em.data.datasets.electron_microscopy.manc

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