torch_em.data.datasets.electron_microscopy.pancreas_fibsem

  1"""The Pancreas FIB-SEM dataset contains training data for segmenting Golgi apparatus and secretory granules
  2in FIB-SEM volumes of mouse pancreatic beta cells.
  3
  4It contains two sub-datasets with 3D annotated volumes (isotropic 16 nm resolution):
  5- golgi: Binary Golgi apparatus masks (8 train volumes of 128^3 voxels, 1 validation volume of 58 x 128 x 128 voxels).
  6- granules: Instance segmentation of secretory granules (4 train and 1 validation volume of 172^3 voxels).
  7
  8The dataset is located at https://doi.org/10.5281/zenodo.8114392 and is distributed under the CC BY 4.0 license.
  9This dataset is from the publication "A modular protocol for segmentation, spatial analysis, and visualization
 10of volume electron microscopy datasets" (Müller et al.). Please cite it if you use this dataset in your research.
 11"""
 12
 13import os
 14from glob import glob
 15from natsort import natsorted
 16from typing import Union, Tuple, List, Literal
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URLS = {
 26    "golgi": "https://zenodo.org/records/8114392/files/data_golgi.zip?download=1",
 27    "granules": "https://zenodo.org/records/8114392/files/data_granules.zip?download=1",
 28}
 29CHECKSUMS = {
 30    "golgi": "e3c377ba55384351c58b365f71af624d628aeb0a8189ef4fd5ed9a4b2da367a3",
 31    "granules": "890c903c4df7a16e9f2a84d2ae1fb6b9cd1d81ee39d497d560737933cb5c59bc",
 32}
 33
 34
 35def get_pancreas_fibsem_data(
 36    path: Union[os.PathLike, str], name: Literal["golgi", "granules"], download: bool = False
 37) -> str:
 38    """Download the Pancreas FIB-SEM dataset.
 39
 40    Args:
 41        path: Filepath to a folder where the downloaded data will be saved.
 42        name: The name of the sub-dataset. Either 'golgi' or 'granules'.
 43        download: Whether to download the data if it is not present.
 44
 45    Returns:
 46        The filepath to the folder with the data.
 47    """
 48    assert name in URLS, f"'{name}' is not a valid sub-dataset name. Choose from {list(URLS)}."
 49
 50    data_dir = os.path.join(path, f"data_{name}")
 51    if os.path.exists(data_dir):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55    zip_path = os.path.join(path, f"data_{name}.zip")
 56    util.download_source(path=zip_path, url=URLS[name], download=download, checksum=CHECKSUMS[name])
 57    util.unzip(zip_path=zip_path, dst=path)
 58    return data_dir
 59
 60
 61def get_pancreas_fibsem_paths(
 62    path: Union[os.PathLike, str],
 63    name: Literal["golgi", "granules"],
 64    split: Literal["train", "val"],
 65    download: bool = False,
 66) -> Tuple[List[str], List[str]]:
 67    """Get paths to the Pancreas FIB-SEM data.
 68
 69    Args:
 70        path: Filepath to a folder where the downloaded data will be saved.
 71        name: The name of the sub-dataset. Either 'golgi' or 'granules'.
 72        split: The data split to use. Either 'train' or 'val'.
 73        download: Whether to download the data if it is not present.
 74
 75    Returns:
 76        List of filepaths for the image data.
 77        List of filepaths for the label data.
 78    """
 79    assert split in ("train", "val"), f"'{split}' is not a valid data split."
 80
 81    data_dir = get_pancreas_fibsem_data(path, name, download)
 82    image_paths = natsorted(glob(os.path.join(data_dir, split, "images", "*.tif")))
 83    label_paths = natsorted(glob(os.path.join(data_dir, split, "masks", "*.tif")))
 84    assert len(image_paths) > 0 and len(image_paths) == len(label_paths)
 85    return image_paths, label_paths
 86
 87
 88def get_pancreas_fibsem_dataset(
 89    path: Union[os.PathLike, str],
 90    patch_shape: Tuple[int, int, int],
 91    name: Literal["golgi", "granules"],
 92    split: Literal["train", "val"],
 93    download: bool = False,
 94    **kwargs
 95) -> Dataset:
 96    """Get the Pancreas FIB-SEM dataset for Golgi apparatus or secretory granule segmentation.
 97
 98    The 'golgi' labels are binary (1 is Golgi apparatus) and the 'granules' labels are instances.
 99
100    Args:
101        path: Filepath to a folder where the downloaded data will be saved.
102        patch_shape: The patch shape to use for training.
103        name: The name of the sub-dataset. Either 'golgi' or 'granules'.
104        split: The data split to use. Either 'train' or 'val'.
105        download: Whether to download the data if it is not present.
106        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
107
108    Returns:
109        The segmentation dataset.
110    """
111    image_paths, label_paths = get_pancreas_fibsem_paths(path, name, split, download)
112
113    if name == "granules":
114        kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
115    else:
116        kwargs = util.update_kwargs(kwargs, "label_transform", torch_em.transform.label.labels_to_binary)
117    kwargs = util.update_kwargs(kwargs, "ndim", 3)
118
119    return torch_em.default_segmentation_dataset(
120        raw_paths=image_paths,
121        raw_key=None,
122        label_paths=label_paths,
123        label_key=None,
124        patch_shape=patch_shape,
125        **kwargs
126    )
127
128
129def get_pancreas_fibsem_loader(
130    path: Union[os.PathLike, str],
131    patch_shape: Tuple[int, int, int],
132    batch_size: int,
133    name: Literal["golgi", "granules"],
134    split: Literal["train", "val"],
135    download: bool = False,
136    **kwargs
137) -> DataLoader:
138    """Get the DataLoader for Golgi apparatus or secretory granule segmentation in the Pancreas FIB-SEM dataset.
139
140    Args:
141        path: Filepath to a folder where the downloaded data will be saved.
142        patch_shape: The patch shape to use for training.
143        batch_size: The batch size for training.
144        name: The name of the sub-dataset. Either 'golgi' or 'granules'.
145        split: The data split to use. Either 'train' or 'val'.
146        download: Whether to download the data if it is not present.
147        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
148
149    Returns:
150        The DataLoader.
151    """
152    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
153    dataset = get_pancreas_fibsem_dataset(path, patch_shape, name, split, download, **ds_kwargs)
154    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)