torch_em.data.datasets.electron_microscopy.cerebellum_fibsem

  1"""The Cerebellum FIB-SEM dataset contains a FIB-SEM volume of mouse cerebellum (10 nm pixel size)
  2with instance segmentations of nuclei and mitochondria.
  3
  4The volume has 530 x 1010 x 1014 voxels. It contains:
  5- nuclei: 5 nucleus instances.
  6- mitochondria: ~370 mitochondria instances.
  7
  8The labels were created by students in the practical course BIO407 2026 (University of Zurich) and were not
  9necessarily proofread. The dataset is located at https://zenodo.org/records/18886823
 10and is distributed under the CC BY 4.0 license. 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
 25URL = "https://zenodo.org/api/records/18886823/files/BIO407_Sideproject_vEM_Group5.zip/content"
 26CHECKSUM = "d24a8cbfeb985c7dfa1be9d4a92884c0805c3d57c11060d06334056281c4e163"
 27
 28FILES = {
 29    "raw": "260304_FIB_SEM_aligned_nucleus_AR_BIO407.tif",
 30    "nuclei": "260304_FIB_SEM_segmentation_nucleus_AR_407.tif",
 31    "mitochondria": "260305_FIB_SEM_segmented_mitochondria_ARN_407.tif",
 32}
 33LABEL_KEYS = {"nuclei": "labels/nuclei", "mitochondria": "labels/mitochondria"}
 34
 35# The volume is split along the z-axis: the first slices are used for training, the remaining ones for validation.
 36VAL_START = 424
 37
 38
 39def _convert_to_h5(tif_dir, h5_path):
 40    import h5py
 41    import tifffile
 42
 43    with h5py.File(h5_path, "w") as f:
 44        for name, fname in FILES.items():
 45            data = tifffile.imread(os.path.join(tif_dir, fname))
 46            key = "raw" if name == "raw" else LABEL_KEYS[name]
 47            if name != "raw":
 48                data = data.astype("uint32")
 49            f.create_dataset(key, data=data, compression="gzip", chunks=(32, 256, 256))
 50
 51
 52def get_cerebellum_fibsem_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 53    """Download the Cerebellum FIB-SEM dataset.
 54
 55    Note: The data is only available as a 4.7 GB zip archive, which also contains raw volumes that are not used.
 56    Only the aligned volume and the two label volumes are extracted and converted to a single hdf5 file.
 57
 58    Args:
 59        path: Filepath to a folder where the downloaded data will be saved.
 60        download: Whether to download the data if it is not present.
 61
 62    Returns:
 63        The filepath to the hdf5 file with the data.
 64    """
 65    h5_path = os.path.join(path, "cerebellum_fibsem.h5")
 66    if os.path.exists(h5_path):
 67        return h5_path
 68
 69    os.makedirs(path, exist_ok=True)
 70    zip_path = os.path.join(path, "BIO407_Sideproject_vEM_Group5.zip")
 71    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 72
 73    import zipfile
 74    tif_dir = os.path.join(path, "tifs")
 75    with zipfile.ZipFile(zip_path, "r") as f:
 76        f.extractall(tif_dir, members=list(FILES.values()))
 77    _convert_to_h5(tif_dir, h5_path)
 78
 79    # Clean up the archive and the intermediate tif files.
 80    os.remove(zip_path)
 81    for fname in FILES.values():
 82        os.remove(os.path.join(tif_dir, fname))
 83    os.rmdir(tif_dir)
 84    return h5_path
 85
 86
 87def get_cerebellum_fibsem_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
 88    """Get the path to the Cerebellum FIB-SEM data.
 89
 90    Args:
 91        path: Filepath to a folder where the downloaded data will be saved.
 92        download: Whether to download the data if it is not present.
 93
 94    Returns:
 95        List with the filepath to the hdf5 file.
 96    """
 97    return [get_cerebellum_fibsem_data(path, download)]
 98
 99
100def get_cerebellum_fibsem_dataset(
101    path: Union[os.PathLike, str],
102    patch_shape: Tuple[int, int, int],
103    segmentation_type: Literal["nuclei", "mitochondria"] = "mitochondria",
104    split: Literal["train", "val"] = "train",
105    download: bool = False,
106    **kwargs
107) -> Dataset:
108    """Get the Cerebellum FIB-SEM dataset for nucleus or mitochondria segmentation.
109
110    Args:
111        path: Filepath to a folder where the downloaded data will be saved.
112        patch_shape: The patch shape to use for training.
113        segmentation_type: The segmentation task. Either 'nuclei' or 'mitochondria'.
114        split: The data split to use. Either 'train' (first 424 slices) or 'val' (remaining slices).
115        download: Whether to download the data if it is not present.
116        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
117
118    Returns:
119        The segmentation dataset.
120    """
121    assert segmentation_type in LABEL_KEYS, f"'{segmentation_type}' is not valid. Choose from {list(LABEL_KEYS)}."
122    assert split in ("train", "val"), f"'{split}' is not a valid data split."
123
124    paths = get_cerebellum_fibsem_paths(path, download)
125
126    roi = np.s_[:VAL_START, :, :] if split == "train" else np.s_[VAL_START:, :, :]
127    kwargs = util.update_kwargs(kwargs, "rois", [roi] * len(paths))
128    kwargs = util.update_kwargs(kwargs, "ndim", 3)
129
130    return torch_em.default_segmentation_dataset(
131        raw_paths=paths,
132        raw_key="raw",
133        label_paths=paths,
134        label_key=LABEL_KEYS[segmentation_type],
135        patch_shape=patch_shape,
136        **kwargs
137    )
138
139
140def get_cerebellum_fibsem_loader(
141    path: Union[os.PathLike, str],
142    patch_shape: Tuple[int, int, int],
143    batch_size: int,
144    segmentation_type: Literal["nuclei", "mitochondria"] = "mitochondria",
145    split: Literal["train", "val"] = "train",
146    download: bool = False,
147    **kwargs
148) -> DataLoader:
149    """Get the DataLoader for nucleus or mitochondria segmentation in the Cerebellum FIB-SEM dataset.
150
151    Args:
152        path: Filepath to a folder where the downloaded data will be saved.
153        patch_shape: The patch shape to use for training.
154        batch_size: The batch size for training.
155        segmentation_type: The segmentation task. Either 'nuclei' or 'mitochondria'.
156        split: The data split to use. Either 'train' or 'val'.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
159
160    Returns:
161        The DataLoader.
162    """
163    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
164    dataset = get_cerebellum_fibsem_dataset(path, patch_shape, segmentation_type, split, download, **ds_kwargs)
165    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)