torch_em.data.datasets.electron_microscopy.vnc

  1"""The VNC dataset contains segmentation annotations for mitochondria in EM.
  2It contains two volumes from TEM of the drosophila brain.
  3
  4Please cite https://doi.org/10.6084/m9.figshare.856713.v1 if you use this dataset in your publication.
  5"""
  6
  7import os
  8from glob import glob
  9from shutil import rmtree
 10from typing import List, Optional, Union, Tuple
 11
 12import imageio
 13import numpy as np
 14from bioimage_cpp.segmentation import label
 15
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23URL = "https://github.com/unidesigner/groundtruth-drosophila-vnc/archive/refs/heads/master.zip"
 24CHECKSUM = "f7bd0db03c86b64440a16b60360ad60c0a4411f89e2c021c7ee2c8d6af3d7e86"
 25
 26
 27def _create_volume(f, key, pattern, process=None):
 28    images = glob(pattern)
 29    images.sort()
 30    data = np.concatenate([imageio.imread(im)[None] for im in images], axis=0)
 31    if process is not None:
 32        data = process(data)
 33    f.create_dataset(key, data=data, compression="gzip")
 34
 35
 36def get_vnc_data(path: Union[os.PathLike, str], download: bool) -> str:
 37    """Download the VNC training data.
 38
 39    Args:
 40        path: Filepath to a folder where the downloaded data will be saved.
 41        download: Whether to download the data if it is not present.
 42
 43    Returns:
 44        The path to the downloaded data.
 45    """
 46    import h5py
 47
 48    train_path = os.path.join(path, "vnc_train.h5")
 49    test_path = os.path.join(path, "vnc_test.h5")
 50    if os.path.exists(train_path) and os.path.exists(test_path):
 51        return path
 52
 53    os.makedirs(path, exist_ok=True)
 54    zip_path = os.path.join(path, "vnc.zip")
 55    util.download_source(zip_path, URL, download, CHECKSUM)
 56    util.unzip(zip_path, path, remove=True)
 57
 58    root = os.path.join(path, "groundtruth-drosophila-vnc-master")
 59    assert os.path.exists(root)
 60
 61    with h5py.File(train_path, "w") as f:
 62        _create_volume(f, "raw", os.path.join(root, "stack1", "raw", "*.tif"))
 63        _create_volume(f, "labels/mitochondria", os.path.join(root, "stack1", "mitochondria", "*.png"), process=label)
 64        _create_volume(f, "labels/synapses", os.path.join(root, "stack1", "synapses", "*.png"), process=label)
 65        # TODO find the post-processing to go from neuron labels to membrane labels
 66        # _create_volume(f, "labels/neurons", os.path.join(root, "stack1", "membranes", "*.png"))
 67
 68    with h5py.File(test_path, "w") as f:
 69        _create_volume(f, "raw", os.path.join(root, "stack2", "raw", "*.tif"))
 70
 71    rmtree(root)
 72    return path
 73
 74
 75def get_vnc_mito_paths(path: Union[os.PathLike, str], download: bool = False) -> str:
 76    """Get path to the VNC data.
 77
 78    Args:
 79        path: Filepath to a folder where the downloaded data is saved.
 80        download: Whether to download the data if it is not present.
 81
 82    Returns:
 83        The filepath to the stored data.
 84    """
 85    get_vnc_data(path, download)
 86    data_path = os.path.join(path, "vnc_train.h5")
 87    return data_path
 88
 89
 90def get_vnc_mito_dataset(
 91    path: Union[os.PathLike, str],
 92    patch_shape: Tuple[int, int, int],
 93    offsets: Optional[List[List[int]]] = None,
 94    boundaries: bool = False,
 95    binary: bool = False,
 96    download: bool = False,
 97    **kwargs
 98) -> Dataset:
 99    """Get the VNC dataset for segmentating mitochondria in EM.
100
101    Args:
102        path: Filepath to a folder where the downloaded data will be saved.
103        patch_shape: The patch shape to use for training.
104        offsets: Offset values for affinity computation used as target.
105        boundaries: Whether to compute boundaries as the target.
106        binary: Whether to return a binary segmentation target.
107        download: Whether to download the data if it is not present.
108        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
109
110    Returns:
111       The segmentation dataset.
112    """
113    data_path = get_vnc_mito_paths(path, download)
114
115    kwargs, _ = util.add_instance_label_transform(
116        kwargs, add_binary_target=True, boundaries=boundaries, offsets=offsets, binary=binary,
117    )
118
119    return torch_em.default_segmentation_dataset(
120        raw_paths=data_path,
121        raw_key="raw",
122        label_paths=data_path,
123        label_key="labels/mitochondria",
124        patch_shape=patch_shape,
125        **kwargs
126    )
127
128
129def get_vnc_mito_loader(
130    path: Union[os.PathLike, str],
131    patch_shape: Tuple[int, int, int],
132    batch_size: int,
133    offsets: Optional[List[List[int]]] = None,
134    boundaries: bool = False,
135    binary: bool = False,
136    download: bool = False,
137    **kwargs
138) -> DataLoader:
139    """Get the VNC dataloader for segmentating mitochondria in EM.
140
141    Args:
142        path: Filepath to a folder where the downloaded data will be saved.
143        patch_shape: The patch shape to use for training.
144        batch_size: The batch size for training.
145        offsets: Offset values for affinity computation used as target.
146        boundaries: Whether to compute boundaries as the target.
147        binary: Whether to return a binary segmentation target.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
150
151    Returns:
152       The DataLoader.
153    """
154    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
155    ds = get_vnc_mito_dataset(
156        path, patch_shape, download=download, offsets=offsets, boundaries=boundaries, binary=binary, **ds_kwargs
157    )
158    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
159
160
161# TODO implement
162def get_vnc_neuron_loader(path, patch_shape, download=False, **kwargs):
163    raise NotImplementedError