torch_em.data.datasets.light_microscopy.liconn

  1"""The LICONN dataset contains a dense connectomic reconstruction of mouse hippocampal
  2CA1 neuropil acquired by spinning-disk confocal microscopy of expansion-microscopy-processed
  3tissue (~16x expansion), yielding a native voxel resolution of 9x9x12 nm (XYZ) at mip=0.
  4All neuronal structures are densely annotated as instance segmentations: 18,268 axons
  5(342 mm total length), 1,643 dendrites (119 mm total), and 71,269 spines.
  6
  7Two volumes are available, selected via the 'volume' argument:
  8- 'fig2' (default, ExPID82_1): the densely reconstructed CA1 volume from Fig. 2 described above,
  9  at 18x18x24 nm. Two segmentation variants are provided:
 10    - 'proofread': manually proofread segmentation (higher accuracy).
 11    - 'agglomerated': automatically agglomerated segmentation.
 12  Stored in liconn.zarr with arrays 'raw', 'seg_proofread', and 'seg_agglomerated'.
 13- 'fig1' (ExPID124): the large hippocampal volume from Fig. 1a, at 20x20x25 nm. Its labeled region
 14  is 20156x5724x997 voxels (XYZ), ~8x larger than 'fig2'. NOTE: neither of its segmentations
 15  is proofread dense ground truth:
 16    - 'sparse': reconstructions of ~80 selected neurons (somata and neurites), covering ~6% of
 17      the volume. Everything else is 0, so unlabeled neurites are NOT background.
 18    - 'dense_auto': the dense automatic segmentation, which is unproofread and over-split
 19      (e.g. nuclei are broken into many fragments).
 20  Stored in liconn_fig1.zarr with arrays 'raw', 'seg_sparse', and 'seg_dense_auto'.
 21
 22The data is served as Neuroglancer precomputed volumes from Google Cloud Storage
 23(gs://liconn-public) and requires the cloudvolume package to download.
 24For 'fig1' only the bounding box that contains labels is downloaded, which drops the unannotated
 25margins. 'fig2' is downloaded in full to keep its voxel coordinates stable; its labeled bounding box
 26is available as FIG2_LABELED_ROI and can be passed as 'roi'. All volumes are stored in zarr v3 stores
 27with sharding.
 28
 29This dataset is from the following publication:
 30- Velicky et al. (2025): https://doi.org/10.1038/s41586-025-08985-1
 31Please cite it if you use this dataset in your research.
 32"""
 33
 34import os
 35import json
 36from concurrent.futures import ThreadPoolExecutor, as_completed
 37from typing import List, Optional, Tuple, Union
 38
 39import numpy as np
 40from tqdm import tqdm
 41
 42from torch.utils.data import Dataset, DataLoader
 43
 44import torch_em
 45
 46from .. import util
 47
 48
 49IMG_URL = "precomputed://https://storage.googleapis.com/liconn-public/ExPID82_1/image_230130b"
 50SEG_PR_URL = "precomputed://https://storage.googleapis.com/liconn-public/ExPID82_1/segmentation/231030_agg_240123"
 51SEG_AGG_URL = "precomputed://https://storage.googleapis.com/liconn-public/ExPID82_1/segmentation/231030_agg_230921_cmpl"
 52
 53FIG1_URL = "precomputed://https://storage.googleapis.com/liconn-public/ExPID124"
 54FIG1_IMG_URL = f"{FIG1_URL}/image"
 55FIG1_SEG_SPARSE_URL = f"{FIG1_URL}/seg_230907_agglo231128_combined2"
 56FIG1_SEG_DENSE_URL = f"{FIG1_URL}/seg_230907"
 57
 58ZARR_FNAME = "liconn.zarr"
 59
 60# The (Z, Y, X) region of 'fig2' that contains labels, computed from the full resolution segmentations.
 61FIG2_LABELED_ROI = (slice(72, 627), slice(185, 4535), slice(140, 3690))
 62
 63# 'bbox' is the downloaded (half-open, ZYX) voxel range; 'fig2' stays at full extent to keep its coordinates.
 64VOLUMES = {
 65    "fig2": {
 66        "zarr": ZARR_FNAME,
 67        "image": IMG_URL,
 68        "segmentations": {"proofread": SEG_PR_URL, "agglomerated": SEG_AGG_URL},
 69        "bbox": ((0, 795), (0, 4870), (0, 3825)),
 70        "accept_unflagged": True,
 71    },
 72    "fig1": {
 73        "zarr": "liconn_fig1.zarr",
 74        "image": FIG1_IMG_URL,
 75        "segmentations": {"sparse": FIG1_SEG_SPARSE_URL, "dense_auto": FIG1_SEG_DENSE_URL},
 76        "bbox": ((0, 997), (204, 5928), (1180, 21336)),
 77    },
 78}
 79SHARD_SHAPE = (64, 256, 256)
 80CHUNK_SHAPE = (32, 64, 64)
 81
 82# Marks arrays whose download finished, so that interrupted downloads are detected.
 83COMPLETE_ATTR = "download_complete"
 84
 85
 86def _to_zyx(a: np.ndarray) -> np.ndarray:
 87    # CloudVolume returns (X, Y, Z[, C]); squeeze trailing channel dim if present.
 88    if a.ndim == 4:
 89        a = a.squeeze(axis=-1)
 90    if a.ndim != 3:
 91        raise ValueError(f"Expected 3D block, got shape {a.shape}")
 92    return a.transpose(2, 1, 0)
 93
 94
 95def _create_array(root, name: str, shape, dtype, is_label: bool):
 96    from zarr.codecs import BloscCodec
 97
 98    shuffle = "bitshuffle" if (np.issubdtype(dtype, np.integer) and is_label) else "shuffle"
 99    ds = root.create_array(
100        name,
101        shape=shape,
102        chunks=CHUNK_SHAPE,
103        shards=SHARD_SHAPE,
104        dtype=dtype,
105        compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
106        overwrite=True,  # Replaces leftovers from an interrupted download.
107    )
108    ds.attrs[COMPLETE_ATTR] = False
109    return ds
110
111
112def _array_complete(zarr_path: str, name: str, shape: Tuple[int, ...], accept_unflagged: bool = False) -> bool:
113    meta_path = os.path.join(zarr_path, name, "zarr.json")
114    if not os.path.isfile(meta_path):
115        return False
116    with open(meta_path) as f:
117        meta = json.load(f)
118    # A different shape means the array was not cropped to the current bounding box.
119    if tuple(meta["shape"]) != tuple(shape):
120        return False
121    complete = meta.get("attributes", {}).get(COMPLETE_ATTR)
122    if complete is None and accept_unflagged:
123        # Downloaded before the completion flag existed.
124        return len(os.listdir(os.path.dirname(meta_path))) > 1
125    return complete is True
126
127
128def _download_ng_volume(vol, ds, name: str, bbox) -> None:
129    # The bounding box is given relative to the start of the volume.
130    vx0, vy0, vz0 = map(int, vol.bounds.minpt)
131    (bz0, bz1), (by0, by1), (bx0, bx1) = bbox
132    x0, y0, z0 = vx0 + bx0, vy0 + by0, vz0 + bz0
133    shape = (bz1 - bz0, by1 - by0, bx1 - bx0)
134
135    tasks = []
136    for z0_ in range(0, shape[0], SHARD_SHAPE[0]):
137        for y0_ in range(0, shape[1], SHARD_SHAPE[1]):
138            for x0_ in range(0, shape[2], SHARD_SHAPE[2]):
139                z1_ = min(z0_ + SHARD_SHAPE[0], shape[0])
140                y1_ = min(y0_ + SHARD_SHAPE[1], shape[1])
141                x1_ = min(x0_ + SHARD_SHAPE[2], shape[2])
142                tasks.append((
143                    (z0_, z1_), (y0_, y1_), (x0_, x1_),
144                    (x0 + x0_, x0 + x1_, y0 + y0_, y0 + y1_, z0 + z0_, z0 + z1_)
145                ))
146
147    # os.cpu_count() counts all cores of the node; more than ~32 threads only costs memory.
148    n_cpus = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 4)
149    max_workers = min(32, 4 * n_cpus)
150
151    def worker(item):
152        (z0_, z1_), (y0_, y1_), (x0_, x1_), (gx0, gx1, gy0, gy1, gz0, gz1) = item
153        block = np.asarray(vol[gx0:gx1, gy0:gy1, gz0:gz1])
154        ds[z0_:z1_, y0_:y1_, x0_:x1_] = _to_zyx(block)
155
156    with ThreadPoolExecutor(max_workers=max_workers) as ex:
157        futures = [ex.submit(worker, t) for t in tasks]
158        for fut in tqdm(as_completed(futures), total=len(futures), desc=f"Downloading '{name}'", smoothing=0.05):
159            fut.result()
160
161    ds.attrs[COMPLETE_ATTR] = True
162
163
164def _get_segmentation(volume: str, segmentation: Optional[str]) -> str:
165    if volume not in VOLUMES:
166        raise ValueError(f"'{volume}' is not a valid volume. Choose from {tuple(VOLUMES)}.")
167    choices = tuple(VOLUMES[volume]["segmentations"])
168    if segmentation is None:
169        return choices[0]
170    if segmentation not in choices:
171        raise ValueError(f"'{segmentation}' is not a valid segmentation for '{volume}'. Choose from {choices}.")
172    return segmentation
173
174
175def get_liconn_data(
176    path: Union[os.PathLike, str],
177    segmentation: Optional[str] = None,
178    download: bool = False,
179    volume: str = "fig2",
180) -> None:
181    """Download the LICONN image and segmentation into a single zarr v3 store with sharding.
182
183    The full 'fig2' volume and the bounding box containing all labels of 'fig1' are downloaded (image at
184    mip=1, i.e. 18x18x24 nm for 'fig2' and 20x20x25 nm for 'fig1'; segmentation at mip=0, same voxel grid).
185    Further ROI-based sub-region selection is not supported at download time - use the roi parameter in
186    get_liconn_dataset to restrict patch sampling to a sub-region after the data is on disk.
187
188    The 'fig2' arrays are stored in liconn.zarr ('raw', 'seg_proofread', 'seg_agglomerated'), the 'fig1'
189    arrays in liconn_fig1.zarr ('raw', 'seg_sparse', 'seg_dense_auto'). Each array uses zarr v3 sharding
190    (shard shape SHARD_SHAPE, inner chunk shape CHUNK_SHAPE) with zstd+blosc compression.
191    An array only counts as present once its download has finished; an interrupted download
192    is restarted from scratch.
193
194    Args:
195        path: Filepath to a folder where the data will be saved.
196        segmentation: Which segmentation variant to download. Either 'proofread' or 'agglomerated' for 'fig2',
197            or 'sparse' or 'dense_auto' for 'fig1'. By default the first of these is used.
198        download: Whether to download the data if it is not present.
199        volume: Which volume to download. Either 'fig2' or 'fig1'.
200    """
201    segmentation = _get_segmentation(volume, segmentation)
202    config = VOLUMES[volume]
203
204    os.makedirs(path, exist_ok=True)
205    zarr_path = os.path.join(str(path), config["zarr"])
206    label_key = f"seg_{segmentation}"
207    bbox = config["bbox"]
208    shape = tuple(stop - start for start, stop in bbox)
209
210    accept_unflagged = config.get("accept_unflagged", False)
211    raw_missing = not _array_complete(zarr_path, "raw", shape, accept_unflagged)
212    label_missing = not _array_complete(zarr_path, label_key, shape, accept_unflagged)
213
214    if not raw_missing and not label_missing:
215        return
216
217    if not download:
218        missing = [k for k, m in [("raw", raw_missing), (label_key, label_missing)] if m]
219        raise RuntimeError(
220            f"LICONN arrays {missing} not found or incomplete in {zarr_path}. Pass download=True to download them."
221        )
222
223    # Only the download needs cloudvolume, so data already on disk loads without it.
224    try:
225        from cloudvolume import CloudVolume
226    except ImportError:
227        raise ImportError(
228            "cloudvolume is required to download the LICONN data. Install it with: pip install cloud-volume"
229        )
230
231    import zarr
232
233    root = zarr.open_group(zarr_path, mode="a")
234
235    if raw_missing:
236        img_cv = CloudVolume(config["image"], mip=1, progress=False, cache=False, fill_missing=True)
237        ds = _create_array(root, "raw", shape, np.dtype(img_cv.dtype), is_label=False)
238        _download_ng_volume(img_cv, ds, name="raw", bbox=bbox)
239
240    if label_missing:
241        seg_url = config["segmentations"][segmentation]
242        seg_cv = CloudVolume(seg_url, mip=0, progress=False, cache=False, fill_missing=True)
243        ds = _create_array(root, label_key, shape, np.dtype(seg_cv.dtype), is_label=True)
244        _download_ng_volume(seg_cv, ds, name=label_key, bbox=bbox)
245
246
247def get_liconn_paths(
248    path: Union[os.PathLike, str],
249    segmentation: Optional[str] = None,
250    download: bool = False,
251    volume: str = "fig2",
252) -> str:
253    """Get the filepath to the LICONN zarr store.
254
255    The 'fig2' store (liconn.zarr) contains arrays 'raw', 'seg_proofread', and 'seg_agglomerated',
256    the 'fig1' store (liconn_fig1.zarr) contains arrays 'raw', 'seg_sparse', and 'seg_dense_auto'.
257
258    Args:
259        path: Filepath to a folder where the data will be saved.
260        segmentation: Which segmentation variant to ensure is present. Either 'proofread' or 'agglomerated'
261            for 'fig2', or 'sparse' or 'dense_auto' for 'fig1'. By default the first of these is used.
262        download: Whether to download the data if it is not present.
263        volume: Which volume to use. Either 'fig2' or 'fig1'.
264
265    Returns:
266        Filepath to the zarr store.
267    """
268    get_liconn_data(path, segmentation, download, volume)
269    return os.path.join(str(path), VOLUMES[volume]["zarr"])
270
271
272def get_liconn_dataset(
273    path: Union[os.PathLike, str],
274    patch_shape: Tuple[int, int, int],
275    segmentation: Optional[str] = None,
276    roi: Optional[Tuple[slice, ...]] = None,
277    download: bool = False,
278    offsets: Optional[List[List[int]]] = None,
279    boundaries: bool = False,
280    volume: str = "fig2",
281    **kwargs,
282) -> Dataset:
283    """Get the LICONN dataset for neuron instance segmentation in expansion microscopy.
284
285    Args:
286        path: Filepath to a folder where the data will be saved.
287        patch_shape: The patch shape to use for training.
288        segmentation: Which segmentation variant to use. Either 'proofread' or 'agglomerated' for 'fig2',
289            or 'sparse' or 'dense_auto' for 'fig1'. By default the first of these is used.
290        roi: Optional region-of-interest as a tuple of slices (Z, Y, X) restricting which part
291            of the already-downloaded data is used for patch sampling. The data is always downloaded
292            regardless of this parameter. For 'fig2', FIG2_LABELED_ROI restricts sampling to the labeled region.
293        download: Whether to download the data if it is not present.
294        offsets: Offset values for affinity computation used as target.
295        boundaries: Whether to compute boundaries as the target.
296        volume: Which volume to use. Either 'fig2' (dense proofread CA1 volume) or 'fig1'
297            (large volume with sparse or unproofread segmentations, see the module docstring).
298        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
299
300    Returns:
301        The segmentation dataset.
302    """
303    assert len(patch_shape) == 3
304
305    segmentation = _get_segmentation(volume, segmentation)
306    zarr_path = get_liconn_paths(path, segmentation, download, volume)
307    label_key = f"seg_{segmentation}"
308
309    kwargs, _ = util.add_instance_label_transform(
310        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
311    )
312    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
313
314    return torch_em.default_segmentation_dataset(
315        raw_paths=zarr_path,
316        raw_key="raw",
317        label_paths=zarr_path,
318        label_key=label_key,
319        patch_shape=patch_shape,
320        rois=roi,
321        **kwargs,
322    )
323
324
325def get_liconn_loader(
326    path: Union[os.PathLike, str],
327    batch_size: int,
328    patch_shape: Tuple[int, int, int],
329    segmentation: Optional[str] = None,
330    roi: Optional[Tuple[slice, ...]] = None,
331    download: bool = False,
332    offsets: Optional[List[List[int]]] = None,
333    boundaries: bool = False,
334    volume: str = "fig2",
335    **kwargs,
336) -> DataLoader:
337    """Get the DataLoader for the LICONN dataset for neuron instance segmentation.
338
339    Args:
340        path: Filepath to a folder where the data will be saved.
341        batch_size: The batch size for training.
342        patch_shape: The patch shape to use for training.
343        segmentation: Which segmentation variant to use. Either 'proofread' or 'agglomerated' for 'fig2',
344            or 'sparse' or 'dense_auto' for 'fig1'. By default the first of these is used.
345        roi: Optional region-of-interest as a tuple of slices (Z, Y, X) restricting which part
346            of the already-downloaded data is used for patch sampling. The data is always downloaded
347            regardless of this parameter. For 'fig2', FIG2_LABELED_ROI restricts sampling to the labeled region.
348        download: Whether to download the data if it is not present.
349        offsets: Offset values for affinity computation used as target.
350        boundaries: Whether to compute boundaries as the target.
351        volume: Which volume to use. Either 'fig2' (dense proofread CA1 volume) or 'fig1'
352            (large volume with sparse or unproofread segmentations, see the module docstring).
353        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
354
355    Returns:
356        The DataLoader.
357    """
358    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
359    dataset = get_liconn_dataset(
360        path, patch_shape, segmentation, roi, download, offsets, boundaries, volume, **ds_kwargs
361    )
362    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)