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 _squeeze_raw(raw_path):
 60    """tstvol-520-1 stores the raw as (1, z, y, x); drop the singleton so it matches the labels."""
 61    import h5py
 62
 63    # Check read-only first: opening for writing fails while the dataset holds the file open.
 64    with h5py.File(raw_path, "r") as f:
 65        if f["raw"].ndim != 4:
 66            return
 67    with h5py.File(raw_path, "a") as f:
 68        raw = f["raw"][:].reshape(f["raw"].shape[-3:])
 69        del f["raw"]
 70        f.create_dataset("raw", data=raw, compression="gzip")
 71
 72
 73def _apply_transforms(groundtruth_path):
 74    """Apply the supervoxel-to-neuron mapping from the 'transforms' dataset.
 75
 76    The groundtruth h5 files contain a 'stack' dataset and, for some samples, a 'transforms'
 77    dataset that maps supervoxels to neuron body IDs. Where it exists the mapping is applied;
 78    otherwise 'stack' already holds neuron ids. The result is saved as 'neuron_ids'.
 79    """
 80    import h5py
 81
 82    with h5py.File(groundtruth_path, "r") as f:
 83        if "neuron_ids" in f:
 84            return
 85    with h5py.File(groundtruth_path, "a") as f:
 86        stack = f["stack"][:]
 87        if stack.ndim == 4:
 88            stack = stack.reshape(stack.shape[-3:])
 89
 90        # tstvol-520-1 ships neuron ids in 'stack' and has no 'transforms'.
 91        if "transforms" not in f:
 92            neuron_ids = stack
 93        else:
 94            transforms = f["transforms"][:]
 95            # Build the mapping from supervoxel IDs to neuron body IDs.
 96            mapping = np.zeros(stack.max() + 1, dtype=stack.dtype)
 97            for src, dst in transforms:
 98                mapping[src] = dst
 99            neuron_ids = mapping[stack]
100
101        f.create_dataset("neuron_ids", data=neuron_ids, compression="gzip")
102
103
104def get_fib25_data(
105    path: Union[os.PathLike, str], samples: Tuple[str, ...], download: bool = False
106):
107    """Download the FIB-25 dataset.
108
109    Args:
110        path: Filepath to a folder where the downloaded data will be saved.
111        samples: The samples to download. Available samples are
112            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
113        download: Whether to download the data if it is not present.
114    """
115    os.makedirs(path, exist_ok=True)
116    for sample in samples:
117        assert sample in URLS, f"Invalid sample: {sample}. Choose from {SAMPLES}."
118        urls = URLS[sample]
119        checksums = CHECKSUMS[sample]
120
121        sample_dir = os.path.join(path, sample)
122        os.makedirs(sample_dir, exist_ok=True)
123
124        raw_path = os.path.join(sample_dir, "raw.h5")
125        labels_path = os.path.join(sample_dir, "groundtruth.h5")
126
127        util.download_source(raw_path, urls["raw"], download, checksum=checksums["raw"])
128        util.download_source(labels_path, urls["labels"], download, checksum=checksums["labels"])
129
130        # Apply the supervoxel-to-neuron mapping.
131        _apply_transforms(labels_path)
132        _squeeze_raw(raw_path)
133
134
135def get_fib25_paths(
136    path: Union[os.PathLike, str],
137    samples: Tuple[str, ...] = ("training_sample2",),
138    download: bool = False,
139) -> Tuple[List[str], List[str]]:
140    """Get paths to the FIB-25 data.
141
142    Args:
143        path: Filepath to a folder where the downloaded data will be saved.
144        samples: The samples to use. Available samples are
145            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
146        download: Whether to download the data if it is not present.
147
148    Returns:
149        The filepaths to the raw data and the label data.
150    """
151    get_fib25_data(path, samples, download)
152    raw_paths = [os.path.join(path, sample, "raw.h5") for sample in samples]
153    label_paths = [os.path.join(path, sample, "groundtruth.h5") for sample in samples]
154    return raw_paths, label_paths
155
156
157def get_fib25_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, int, int],
160    samples: Tuple[str, ...] = ("training_sample2",),
161    download: bool = False,
162    offsets: Optional[List[List[int]]] = None,
163    boundaries: bool = False,
164    **kwargs,
165) -> Dataset:
166    """Get the FIB-25 dataset for the segmentation of neurons in EM.
167
168    Args:
169        path: Filepath to a folder where the downloaded data will be saved.
170        patch_shape: The patch shape to use for training.
171        samples: The samples to use. Available samples are
172            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
173        download: Whether to download the data if it is not present.
174        offsets: Offset values for affinity computation used as target.
175        boundaries: Whether to compute boundaries as the target.
176        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
177
178    Returns:
179        The segmentation dataset.
180    """
181    assert len(patch_shape) == 3
182
183    raw_paths, label_paths = get_fib25_paths(path, samples, download)
184
185    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
186    kwargs, _ = util.add_instance_label_transform(
187        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
188    )
189
190    return torch_em.default_segmentation_dataset(
191        raw_paths=raw_paths,
192        raw_key="raw",
193        label_paths=label_paths,
194        label_key="neuron_ids",
195        patch_shape=patch_shape,
196        **kwargs,
197    )
198
199
200def get_fib25_loader(
201    path: Union[os.PathLike, str],
202    patch_shape: Tuple[int, int, int],
203    batch_size: int,
204    samples: Tuple[str, ...] = ("training_sample2",),
205    download: bool = False,
206    offsets: Optional[List[List[int]]] = None,
207    boundaries: bool = False,
208    **kwargs,
209) -> DataLoader:
210    """Get the DataLoader for EM neuron segmentation in the FIB-25 dataset.
211
212    Args:
213        path: Filepath to a folder where the downloaded data will be saved.
214        patch_shape: The patch shape to use for training.
215        batch_size: The batch size for training.
216        samples: The samples to use. Available samples are
217            'training_sample2', 'validation_sample', and 'tstvol-520-1'.
218        download: Whether to download the data if it is not present.
219        offsets: Offset values for affinity computation used as target.
220        boundaries: Whether to compute boundaries as the target.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
222
223    Returns:
224        The DataLoader.
225    """
226    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
227    ds = get_fib25_dataset(
228        path=path,
229        patch_shape=patch_shape,
230        samples=samples,
231        download=download,
232        offsets=offsets,
233        boundaries=boundaries,
234        **ds_kwargs,
235    )
236    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)