torch_em.data.datasets.electron_microscopy.nisb

  1"""NISB is a large-scale synthetic benchmark for neuron instance segmentation in connectomics.
  2
  3It comprises 9 settings with varying difficulty and imaging conditions, each providing 5 training
  4cubes, 1 validation cube, and 1 test cube. The train_100 setting is an exception with 100 training
  5cubes for scaling analysis. Cubes are 27µm side length at 9x9x20 nm voxel size (liconn: 9x9x12 nm).
  6The multichannel setting stores 8-channel embeddings instead of a single grayscale image.
  7
  8Data is streamed directly from S3 via s3fs and written to local zarr v3 stores (chunk 64^3,
  9shard 512^3, zstd compression) with (z, y, x) axis order. The source is zarr v2 with (x, y, z)
 10axis order; spatial axes are transposed and the trailing singleton channel dim on img is squeezed
 11during the write. Requires s3fs (pip install s3fs).
 12
 13The data is described in https://doi.org/10.17617/1.r2mm-1h33.
 14Please cite it if you use this dataset for a publication.
 15"""
 16
 17import os
 18import shutil
 19import warnings
 20from glob import glob
 21from typing import List, Literal, Optional, Tuple, Union
 22
 23import numpy as np
 24from tqdm import tqdm
 25from torch.utils.data import DataLoader, Dataset
 26
 27import torch_em
 28from .. import util
 29
 30
 31NISB_S3_ENDPOINT = "https://s3.nexus.mpcdf.mpg.de:443"
 32NISB_S3_BUCKET = "nisb"
 33
 34NISB_SETTINGS = [
 35    "base", "train_100", "slice_perturbed", "pos_guidance", "neg_guidance",
 36    "no_touch_thick", "touching_thin", "liconn", "multichannel",
 37]
 38
 39NISB_CHUNK_SHAPE = (64, 64, 64)
 40NISB_SHARD_SHAPE = (512, 512, 512)
 41
 42
 43# The val and test cubes are stored under these seed ids on S3.
 44NISB_SEED_IDS = {"val": (100,), "test": (101,)}
 45
 46
 47def _nisb_seed_ids(setting: str, split: str) -> Tuple[int, ...]:
 48    if split in NISB_SEED_IDS:
 49        return NISB_SEED_IDS[split]
 50    return tuple(range(100 if setting == "train_100" else 5))
 51
 52
 53def _nisb_zarr_complete(zarr_path: str) -> bool:
 54    return (
 55        os.path.isfile(os.path.join(zarr_path, "zarr.json"))
 56        and os.path.isdir(os.path.join(zarr_path, "img"))
 57        and os.path.isdir(os.path.join(zarr_path, "seg"))
 58    )
 59
 60
 61def _nisb_create_v3_array(root, name, shape, dtype, is_label):
 62    from zarr.codecs import BloscCodec
 63    shuffle = "bitshuffle" if (np.issubdtype(np.dtype(dtype), np.integer) and is_label) else "shuffle"
 64    chunks = NISB_CHUNK_SHAPE + tuple(shape[3:])
 65    shards = NISB_SHARD_SHAPE + tuple(shape[3:])
 66    return root.create_array(
 67        name, shape=shape, chunks=chunks, shards=shards, dtype=dtype,
 68        compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
 69    )
 70
 71
 72def _nisb_write_cube_v3(src, v3_path: str) -> None:
 73    """Stream a NISB cube from a zarr v2 source to a local zarr v3 store.
 74
 75    Transposes axes from (x, y, z) to (z, y, x) and squeezes the trailing singleton
 76    channel dimension on the image array.
 77    """
 78    import zarr
 79
 80    img_v2 = src["img"]
 81    seg_v2 = src["seg"]
 82
 83    squeeze_img = img_v2.ndim == 4 and img_v2.shape[-1] == 1
 84    if squeeze_img:
 85        img_shape_v3 = (img_v2.shape[2], img_v2.shape[1], img_v2.shape[0])
 86    else:
 87        img_shape_v3 = (img_v2.shape[2], img_v2.shape[1], img_v2.shape[0], img_v2.shape[3])
 88    seg_shape_v3 = (seg_v2.shape[2], seg_v2.shape[1], seg_v2.shape[0])
 89
 90    tmp_path = v3_path + ".tmp"
 91    if os.path.exists(tmp_path):
 92        shutil.rmtree(tmp_path)
 93
 94    root = zarr.open_group(tmp_path, mode="w", zarr_format=3)
 95    img_v3 = _nisb_create_v3_array(root, "img", img_shape_v3, np.dtype("uint8"), False)
 96    seg_v3 = _nisb_create_v3_array(root, "seg", seg_shape_v3, np.dtype("uint16"), True)
 97
 98    Z, Y, X = seg_shape_v3
 99    sz, sy, sx = NISB_SHARD_SHAPE
100    for z0 in range(0, Z, sz):
101        for y0 in range(0, Y, sy):
102            for x0 in range(0, X, sx):
103                z1, y1, x1 = min(z0 + sz, Z), min(y0 + sy, Y), min(x0 + sx, X)
104                block_img = np.asarray(img_v2[x0:x1, y0:y1, z0:z1])
105                if squeeze_img:
106                    block_img = block_img[..., 0]
107                img_v3[z0:z1, y0:y1, x0:x1] = np.moveaxis(block_img, [0, 2], [2, 0])
108                block_seg = np.asarray(seg_v2[x0:x1, y0:y1, z0:z1])
109                seg_v3[z0:z1, y0:y1, x0:x1] = block_seg.transpose(2, 1, 0)
110
111    shutil.move(tmp_path, v3_path)
112
113
114def _nisb_open_remote(setting: str, split: str, seed_idx: int):
115    """Open a NISB seed cube from S3 as a zarr v2 group via s3fs."""
116    try:
117        import s3fs
118    except ImportError:
119        raise ImportError("The 's3fs' package is required to download NISB data. Install it with: pip install s3fs")
120    import zarr
121
122    fs = s3fs.S3FileSystem(anon=True, endpoint_url=NISB_S3_ENDPOINT)
123    s3_path = f"{NISB_S3_BUCKET}/{setting}/{split}/seed{seed_idx}/data.zarr"
124    with warnings.catch_warnings():
125        warnings.filterwarnings("ignore", message=".*asynchronous.*")
126        store = zarr.storage.FsspecStore(fs=fs, path=s3_path)
127        return zarr.open_group(store, mode="r", zarr_format=2)
128
129
130def get_nisb_data(path: Union[os.PathLike, str], setting: str, split: str, download: bool) -> str:
131    """Stream and cache NISB data for a given setting and split from S3.
132
133    Data is read from S3 via s3fs and written to local zarr v3 stores with (z, y, x) axis
134    order, sharding (chunk 64^3, shard 512^3), and zstd compression. Already-cached seeds
135    are skipped on subsequent calls.
136
137    Args:
138        path: Filepath to a folder where the cached data will be saved.
139        setting: The NISB setting. One of NISB_SETTINGS.
140        split: The data split, one of 'train', 'val', 'test'.
141        download: Whether to stream and cache the data if it is not present.
142
143    Returns:
144        The filepath to the split directory containing seed subdirectories.
145    """
146    assert setting in NISB_SETTINGS, f"Invalid setting '{setting}'. Choose from {NISB_SETTINGS}."
147    assert split in ("train", "val", "test"), f"Invalid split '{split}'. Choose 'train', 'val', or 'test'."
148
149    split_dir = os.path.join(str(path), setting, split)
150
151    for i in tqdm(_nisb_seed_ids(setting, split), desc=f"NISB {setting}/{split}", leave=False):
152        seed_dir = os.path.join(split_dir, f"seed{i}")
153        zarr_path = os.path.join(seed_dir, "data.zarr")
154
155        if _nisb_zarr_complete(zarr_path):
156            continue
157
158        if not download:
159            raise RuntimeError(
160                f"No NISB data for setting '{setting}' split '{split}' seed {i} at '{zarr_path}'. "
161                "Set download=True to stream it from S3."
162            )
163
164        os.makedirs(seed_dir, exist_ok=True)
165        print(f"Streaming NISB {setting}/{split}/seed{i} from S3 ...")
166        src = _nisb_open_remote(setting, split, i)
167        _nisb_write_cube_v3(src, zarr_path)
168
169    return split_dir
170
171
172def get_nisb_paths(
173    path: Union[os.PathLike, str],
174    setting: str = "base",
175    split: Literal["train", "val", "test"] = "train",
176    download: bool = False,
177) -> List[str]:
178    """Get paths to NISB zarr stores for a given setting and split.
179
180    Args:
181        path: Filepath to a folder where the cached data is saved.
182        setting: The NISB setting. One of NISB_SETTINGS.
183        split: The data split, one of 'train', 'val', 'test'.
184        download: Whether to stream and cache the data if it is not present.
185
186    Returns:
187        Sorted list of filepaths to the zarr stores, one per cube/seed.
188    """
189    split_dir = get_nisb_data(path, setting, split, download)
190    paths = sorted(glob(os.path.join(split_dir, "seed*", "data.zarr")))
191    if not paths:
192        raise RuntimeError(
193            f"No zarr files found in '{split_dir}'. The download may have failed or the directory is empty."
194        )
195    return paths
196
197
198def get_nisb_dataset(
199    path: Union[os.PathLike, str],
200    patch_shape: Tuple[int, int, int],
201    setting: str = "base",
202    split: Literal["train", "val", "test"] = "train",
203    download: bool = False,
204    offsets: Optional[List[List[int]]] = None,
205    boundaries: bool = False,
206    **kwargs,
207) -> Dataset:
208    """Get the NISB dataset for neuron instance segmentation in EM.
209
210    NISB provides 9 settings of varying difficulty, each with multiple cubes at 27µm side length.
211    Image data is stored under the zarr key 'img' with shape (z, y, x) and segmentation under 'seg'.
212    The multichannel setting stores 8-channel data with shape (z, y, x, 8).
213
214    Args:
215        path: Filepath to a folder where the cached data will be saved.
216        patch_shape: The patch shape to use for training.
217        setting: The NISB setting. One of NISB_SETTINGS. Default 'base'.
218        split: The data split, one of 'train', 'val', 'test'.
219        download: Whether to stream and cache the data if it is not present.
220            Requires s3fs (pip install s3fs).
221        offsets: Offset values for affinity computation used as target.
222        boundaries: Whether to compute boundaries as the target.
223        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
224
225    Returns:
226        The segmentation dataset.
227    """
228    assert len(patch_shape) == 3
229
230    paths = get_nisb_paths(path, setting, split, download)
231
232    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
233    kwargs, _ = util.add_instance_label_transform(
234        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
235    )
236
237    return torch_em.default_segmentation_dataset(
238        raw_paths=paths,
239        raw_key="img",
240        label_paths=paths,
241        label_key="seg",
242        patch_shape=patch_shape,
243        **kwargs,
244    )
245
246
247def get_nisb_loader(
248    path: Union[os.PathLike, str],
249    patch_shape: Tuple[int, int, int],
250    batch_size: int,
251    setting: str = "base",
252    split: Literal["train", "val", "test"] = "train",
253    download: bool = False,
254    offsets: Optional[List[List[int]]] = None,
255    boundaries: bool = False,
256    **kwargs,
257) -> DataLoader:
258    """Get the DataLoader for neuron instance segmentation in the NISB dataset.
259
260    Args:
261        path: Filepath to a folder where the cached data will be saved.
262        patch_shape: The patch shape to use for training.
263        batch_size: The batch size for training.
264        setting: The NISB setting. One of NISB_SETTINGS. Default 'base'.
265        split: The data split, one of 'train', 'val', 'test'.
266        download: Whether to stream and cache the data if it is not present.
267            Requires s3fs (pip install s3fs).
268        offsets: Offset values for affinity computation used as target.
269        boundaries: Whether to compute boundaries as the target.
270        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
271
272    Returns:
273        The DataLoader.
274    """
275    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
276    ds = get_nisb_dataset(
277        path=path,
278        patch_shape=patch_shape,
279        setting=setting,
280        split=split,
281        download=download,
282        offsets=offsets,
283        boundaries=boundaries,
284        **ds_kwargs,
285    )
286    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)