torch_em.data.datasets.electron_microscopy.brain_pfibsem

  1"""The Brain PFIB-SEM dataset contains a PFIB-SEM volume of mouse brain (5 x 5 x 5 nm voxel size)
  2with expert annotations for mitochondria and synapses.
  3
  4The volume has 181 x 1300 x 1300 voxels and was acquired at two dwell times: 0.5 us (noisy) and 2 us (reference).
  5The annotations (binary semantic masks) were made on the 2 us reference stack.
  6
  7The dataset also contains further volumes (SBF-SEM and PFIB-SEM at different dwell times) for denoising,
  8which are not used here. The dataset is located at https://zenodo.org/records/20139642
  9and is distributed under the CC BY 4.0 license. It is part of the publication
 10https://doi.org/10.1101/2025.08.26.672334. Please cite it if you use this dataset in your research.
 11"""
 12
 13import os
 14from typing import Union, Tuple, List, Literal
 15
 16import numpy as np
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25BASE_URL = "https://zenodo.org/api/records/20139642/files"
 26FILES = {
 27    "noisy": "PFIB_seg_0d5us_Raw_data_5x5x5nm.tif",
 28    "reference": "PFIB_seg_2us_Reference_5x5x5nm.tif",
 29    "mitochondria": "mitos_human_label.tif",
 30    "synapses": "synapses_human_label.tif",
 31}
 32CHECKSUMS = {
 33    "noisy": "c9f799e9a3de701c4bbce0c98c3ba3ea7b26d68f70361b589804e30ec99f7dbd",
 34    "reference": "eb4517d99010f04d79acc0ad4cd0ae4ff7880df833548f1cc34ed5078a75bacb",
 35    "mitochondria": "8ecb184e465b94459fdf3fa0fd36313010d51b7ba56faa077dcd99fa13f4e2f6",
 36    "synapses": "11e851d8b074b03f0c79320457b22b3a2cb38759779820c0bcc610839e1b3e5d",
 37}
 38LABEL_KEYS = {"mitochondria": "labels/mitochondria", "synapses": "labels/synapses"}
 39
 40# The volume is split along the z-axis: the first slices are used for training, the remaining ones for validation.
 41VAL_START = 144
 42
 43
 44def _convert_to_h5(path, h5_path):
 45    import h5py
 46    import tifffile
 47
 48    with h5py.File(h5_path, "w") as f:
 49        for name, fname in FILES.items():
 50            data = tifffile.imread(os.path.join(path, fname))
 51            if name in ("noisy", "reference"):
 52                key = f"raw/{name}"
 53            else:
 54                key = LABEL_KEYS[name]
 55                data = (data > 0).astype("uint8")
 56            f.create_dataset(key, data=data, compression="gzip", chunks=(32, 256, 256))
 57
 58
 59def get_brain_pfibsem_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 60    """Download the Brain PFIB-SEM dataset.
 61
 62    Note: Only the PFIB-SEM segmentation volumes (4 x ~300 MB) are downloaded and converted to a single hdf5 file.
 63
 64    Args:
 65        path: Filepath to a folder where the downloaded data will be saved.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        The filepath to the hdf5 file with the data.
 70    """
 71    h5_path = os.path.join(path, "brain_pfibsem.h5")
 72    if os.path.exists(h5_path):
 73        return h5_path
 74
 75    os.makedirs(path, exist_ok=True)
 76    for name, fname in FILES.items():
 77        util.download_source(
 78            path=os.path.join(path, fname), url=f"{BASE_URL}/{fname}/content",
 79            download=download, checksum=CHECKSUMS[name],
 80        )
 81
 82    _convert_to_h5(path, h5_path)
 83
 84    # Clean up the intermediate tif files.
 85    for fname in FILES.values():
 86        os.remove(os.path.join(path, fname))
 87    return h5_path
 88
 89
 90def get_brain_pfibsem_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
 91    """Get the path to the Brain PFIB-SEM data.
 92
 93    Args:
 94        path: Filepath to a folder where the downloaded data will be saved.
 95        download: Whether to download the data if it is not present.
 96
 97    Returns:
 98        List with the filepath to the hdf5 file.
 99    """
100    return [get_brain_pfibsem_data(path, download)]
101
102
103def get_brain_pfibsem_dataset(
104    path: Union[os.PathLike, str],
105    patch_shape: Tuple[int, int, int],
106    segmentation_type: Literal["mitochondria", "synapses"] = "mitochondria",
107    input_type: Literal["noisy", "reference"] = "reference",
108    split: Literal["train", "val"] = "train",
109    download: bool = False,
110    **kwargs
111) -> Dataset:
112    """Get the Brain PFIB-SEM dataset for mitochondria or synapse segmentation.
113
114    Args:
115        path: Filepath to a folder where the downloaded data will be saved.
116        patch_shape: The patch shape to use for training.
117        segmentation_type: The segmentation task. Either 'mitochondria' or 'synapses'.
118        input_type: The input volume. Either 'reference' (2 us dwell time, on which the labels were made)
119            or 'noisy' (0.5 us dwell time).
120        split: The data split to use. Either 'train' (first 144 slices) or 'val' (remaining slices).
121        download: Whether to download the data if it is not present.
122        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
123
124    Returns:
125        The segmentation dataset.
126    """
127    assert segmentation_type in LABEL_KEYS, f"'{segmentation_type}' is not valid. Choose from {list(LABEL_KEYS)}."
128    assert input_type in ("noisy", "reference"), f"'{input_type}' is not a valid input type."
129    assert split in ("train", "val"), f"'{split}' is not a valid data split."
130
131    paths = get_brain_pfibsem_paths(path, download)
132
133    roi = np.s_[:VAL_START, :, :] if split == "train" else np.s_[VAL_START:, :, :]
134    kwargs = util.update_kwargs(kwargs, "rois", [roi] * len(paths))
135    kwargs = util.update_kwargs(kwargs, "ndim", 3)
136
137    return torch_em.default_segmentation_dataset(
138        raw_paths=paths,
139        raw_key=f"raw/{input_type}",
140        label_paths=paths,
141        label_key=LABEL_KEYS[segmentation_type],
142        patch_shape=patch_shape,
143        **kwargs
144    )
145
146
147def get_brain_pfibsem_loader(
148    path: Union[os.PathLike, str],
149    patch_shape: Tuple[int, int, int],
150    batch_size: int,
151    segmentation_type: Literal["mitochondria", "synapses"] = "mitochondria",
152    input_type: Literal["noisy", "reference"] = "reference",
153    split: Literal["train", "val"] = "train",
154    download: bool = False,
155    **kwargs
156) -> DataLoader:
157    """Get the DataLoader for mitochondria or synapse segmentation in the Brain PFIB-SEM dataset.
158
159    Args:
160        path: Filepath to a folder where the downloaded data will be saved.
161        patch_shape: The patch shape to use for training.
162        batch_size: The batch size for training.
163        segmentation_type: The segmentation task. Either 'mitochondria' or 'synapses'.
164        input_type: The input volume. Either 'reference' or 'noisy'.
165        split: The data split to use. Either 'train' or 'val'.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
168
169    Returns:
170        The DataLoader.
171    """
172    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
173    dataset = get_brain_pfibsem_dataset(
174        path, patch_shape, segmentation_type, input_type, split, download, **ds_kwargs
175    )
176    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)