torch_em.data.datasets.electron_microscopy.fib25

  1"""FIB-25 is a dataset for neuron segmentation in EM.
  2
  3It contains FIB-SEM data and segmentation ground truth from the Drosophila medulla,
  4as part of the FlyEM project at Janelia Research Campus.
  5
  6The dataset is from the publication https://doi.org/10.1073/pnas.1509820112.
  7Please cite this publication if you use the dataset in your research.
  8
  9The data is hosted at https://github.com/google/ffn via Google Cloud Storage.
 10"""
 11
 12import os
 13from typing import List, Optional, Tuple, Union
 14
 15import numpy as np
 16
 17import torch_em
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21from .. import util
 22
 23
 24GCS_BUCKET = "https://storage.googleapis.com/ffn-flyem-fib25"
 25
 26URLS = {
 27    "training_sample2": {
 28        "raw": f"{GCS_BUCKET}/training_sample2/grayscale_maps.h5",
 29        "labels": f"{GCS_BUCKET}/training_sample2/groundtruth.h5",
 30    },
 31    "validation_sample": {
 32        "raw": f"{GCS_BUCKET}/validation_sample/grayscale_maps.h5",
 33        "labels": f"{GCS_BUCKET}/validation_sample/groundtruth.h5",
 34    },
 35    "tstvol-520-1": {
 36        "raw": f"{GCS_BUCKET}/tstvol-520-1/raw.h5",
 37        "labels": f"{GCS_BUCKET}/tstvol-520-1/groundtruth.h5",
 38    },
 39}
 40
 41CHECKSUMS = {
 42    "training_sample2": {
 43        "raw": "ea031c98ee2de778a9a3a1e6d410df5de73e4ac28022df8e7255d84e3394cafa",
 44        "labels": "fd508e7aee1fe51ac9ae0460db4a841d275236f013c1f2552314b4f21b1010ea",
 45    },
 46    "validation_sample": {
 47        "raw": "400ccb2a7268a3880c63656e0d794f8e6252e62031869455cc8caeef245b2a83",
 48        "labels": "2c5e31af0af5476bc9669b88d01a4570a26eb020799eaf6131aa75f2f7d92e98",
 49    },
 50    "tstvol-520-1": {
 51        "raw": "0667e701c8b4464003d8a6cb0cf9deb2aa79fb415ec51deeac92e5f9c67a5a66",
 52        "labels": "ae61ae78a9874eb35ae8e5ed29b4cbfe7bbd07a61789ddb70aef4deb2532eb4e",
 53    },
 54}
 55
 56SAMPLES = list(URLS.keys())
 57
 58
 59def _apply_transforms(groundtruth_path):
 60    """Apply the supervoxel-to-neuron mapping from the 'transforms' dataset.
 61
 62    The groundtruth h5 files contain a 'stack' dataset with supervoxel IDs
 63    and a 'transforms' dataset that maps supervoxels to neuron body IDs.
 64    This function applies the mapping and saves the result as 'neuron_ids'.
 65    """
 66    import h5py
 67
 68    with h5py.File(groundtruth_path, "a") as f:
 69        if "neuron_ids" in f:
 70            return
 71
 72        stack = f["stack"][:]
 73        transforms = f["transforms"][:]
 74
 75        # Build the mapping from supervoxel IDs to neuron body IDs.
 76        mapping = np.zeros(stack.max() + 1, dtype=stack.dtype)
 77        for src, dst in transforms:
 78            mapping[src] = dst
 79        neuron_ids = mapping[stack]
 80
 81        f.create_dataset("neuron_ids", data=neuron_ids, compression="gzip")
 82
 83
 84def get_fib25_data(
 85    path: Union[os.PathLike, str], samples: Tuple[str, ...], download: bool = False
 86):
 87    """Download the FIB-25 dataset.
 88
 89    Args:
 90        path: Filepath to a folder where the downloaded data will be saved.
 91        samples: The samples to download. Available samples are
 92            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
 93        download: Whether to download the data if it is not present.
 94    """
 95    os.makedirs(path, exist_ok=True)
 96    for sample in samples:
 97        assert sample in URLS, f"Invalid sample: {sample}. Choose from {SAMPLES}."
 98        urls = URLS[sample]
 99        checksums = CHECKSUMS[sample]
100
101        sample_dir = os.path.join(path, sample)
102        os.makedirs(sample_dir, exist_ok=True)
103
104        raw_path = os.path.join(sample_dir, "raw.h5")
105        labels_path = os.path.join(sample_dir, "groundtruth.h5")
106
107        util.download_source(raw_path, urls["raw"], download, checksum=checksums["raw"])
108        util.download_source(labels_path, urls["labels"], download, checksum=checksums["labels"])
109
110        # Apply the supervoxel-to-neuron mapping.
111        _apply_transforms(labels_path)
112
113
114def get_fib25_paths(
115    path: Union[os.PathLike, str],
116    samples: Tuple[str, ...] = ("training_sample2",),
117    download: bool = False,
118) -> Tuple[List[str], List[str]]:
119    """Get paths to the FIB-25 data.
120
121    Args:
122        path: Filepath to a folder where the downloaded data will be saved.
123        samples: The samples to use. Available samples are
124            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
125        download: Whether to download the data if it is not present.
126
127    Returns:
128        The filepaths to the raw data and the label data.
129    """
130    get_fib25_data(path, samples, download)
131    raw_paths = [os.path.join(path, sample, "raw.h5") for sample in samples]
132    label_paths = [os.path.join(path, sample, "groundtruth.h5") for sample in samples]
133    return raw_paths, label_paths
134
135
136def get_fib25_dataset(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int, int],
139    samples: Tuple[str, ...] = ("training_sample2",),
140    download: bool = False,
141    offsets: Optional[List[List[int]]] = None,
142    boundaries: bool = False,
143    **kwargs,
144) -> Dataset:
145    """Get the FIB-25 dataset for the segmentation of neurons in EM.
146
147    Args:
148        path: Filepath to a folder where the downloaded data will be saved.
149        patch_shape: The patch shape to use for training.
150        samples: The samples to use. Available samples are
151            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
152        download: Whether to download the data if it is not present.
153        offsets: Offset values for affinity computation used as target.
154        boundaries: Whether to compute boundaries as the target.
155        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
156
157    Returns:
158        The segmentation dataset.
159    """
160    assert len(patch_shape) == 3
161
162    raw_paths, label_paths = get_fib25_paths(path, samples, download)
163
164    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
165    kwargs, _ = util.add_instance_label_transform(
166        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
167    )
168
169    return torch_em.default_segmentation_dataset(
170        raw_paths=raw_paths,
171        raw_key="raw",
172        label_paths=label_paths,
173        label_key="neuron_ids",
174        patch_shape=patch_shape,
175        **kwargs,
176    )
177
178
179def get_fib25_loader(
180    path: Union[os.PathLike, str],
181    patch_shape: Tuple[int, int, int],
182    batch_size: int,
183    samples: Tuple[str, ...] = ("training_sample2",),
184    download: bool = False,
185    offsets: Optional[List[List[int]]] = None,
186    boundaries: bool = False,
187    **kwargs,
188) -> DataLoader:
189    """Get the DataLoader for EM neuron segmentation in the FIB-25 dataset.
190
191    Args:
192        path: Filepath to a folder where the downloaded data will be saved.
193        patch_shape: The patch shape to use for training.
194        batch_size: The batch size for training.
195        samples: The samples to use. Available samples are
196            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
197        download: Whether to download the data if it is not present.
198        offsets: Offset values for affinity computation used as target.
199        boundaries: Whether to compute boundaries as the target.
200        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
201
202    Returns:
203        The DataLoader.
204    """
205    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
206    ds = get_fib25_dataset(
207        path=path,
208        patch_shape=patch_shape,
209        samples=samples,
210        download=download,
211        offsets=offsets,
212        boundaries=boundaries,
213        **ds_kwargs,
214    )
215    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)