torch_em.data.datasets.electron_microscopy.sponge_em

  1"""This dataset contains volume EM data of a sponge chamber with
  2segmentation annotations for cells, cilia and microvilli.
  3
  4It contains three annotated volumes. The dataset is part of the publication
  5https://doi.org/10.1126/science.abj2949. Please cite this publication of you use the
  6dataset in your research.
  7"""
  8
  9import os
 10from glob import glob
 11from typing import Optional, Sequence, Tuple, Union, List
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20URL = "https://zenodo.org/record/8150818/files/sponge_em_train_data.zip?download=1"
 21CHECKSUM = "f1df616cd60f81b91d7642933e9edd74dc6c486b2e546186a7c1e54c67dd32a5"
 22
 23
 24def get_sponge_em_data(path: Union[os.PathLike, str], download: bool) -> Tuple[str, int]:
 25    """Download the SpongeEM training data.
 26
 27    Args:
 28        path: Filepath to a folder where the downloaded data will be saved.
 29        download: Whether to download the data if it is not present.
 30
 31    Returns:
 32        The path to the downloaded data.
 33        The number of downloaded volumes.
 34    """
 35    n_files = len(glob(os.path.join(path, "*.h5")))
 36    if n_files == 3:
 37        return path, n_files
 38    elif n_files == 0:
 39        pass
 40    else:
 41        raise RuntimeError(
 42            f"Invalid number of downloaded files in {path}. Please remove this folder and rerun this function."
 43        )
 44
 45    os.makedirs(path, exist_ok=True)
 46    zip_path = os.path.join(path, "data.zip")
 47    util.download_source(zip_path, URL, download, CHECKSUM)
 48    util.unzip(zip_path, path)
 49
 50    n_files = len(glob(os.path.join(path, "*.h5")))
 51    assert n_files == 3
 52    return path, n_files
 53
 54
 55def get_sponge_em_paths(
 56    path: Union[os.PathLike, str], sample_ids: Optional[Sequence[int]], download: bool = False
 57) -> List[str]:
 58    """Get paths to the SpongeEM data.
 59
 60    Args:
 61        path: Filepath to a folder where the downloaded data will saved.
 62        sample_ids: The sample to download, valid ids are 1, 2 and 3.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        The filepaths to the stored data.
 67    """
 68    data_folder, n_files = get_sponge_em_data(path, download)
 69
 70    if sample_ids is None:
 71        sample_ids = range(1, n_files + 1)
 72
 73    paths = [os.path.join(data_folder, f"train_data_0{i}.h5") for i in sample_ids]
 74    return paths
 75
 76
 77def get_sponge_em_dataset(
 78    path: Union[os.PathLike, str],
 79    mode: str,
 80    patch_shape: Tuple[int, int, int],
 81    sample_ids: Optional[Sequence[int]] = None,
 82    download: bool = False,
 83    **kwargs
 84) -> Dataset:
 85    """Get the SpongeEM dataset for the segmentation of structures in EM.
 86
 87    Args:
 88        path: Filepath to a folder where the downloaded data will be saved.
 89        mode: Choose the segmentation task, either 'semantic' or 'instances'.
 90        patch_shape: The patch shape to use for training.
 91        sample_ids: The sample to download, valid ids are 1, 2 and 3.
 92        download: Whether to download the data if it is not present.
 93        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 94
 95    Returns:
 96       The segmentation dataset.
 97    """
 98    assert mode in ("semantic", "instances")
 99
100    paths = get_sponge_em_paths(path, sample_ids, download)
101
102    return torch_em.default_segmentation_dataset(
103        raw_paths=paths,
104        raw_key="volumes/raw",
105        label_paths=paths,
106        label_key=f"volumes/labels/{mode}",
107        patch_shape=patch_shape,
108        **kwargs
109    )
110
111
112def get_sponge_em_loader(
113    path: Union[os.PathLike, str],
114    mode: str,
115    patch_shape: Tuple[int, int, int],
116    batch_size: int,
117    sample_ids: Optional[Sequence[int]] = None,
118    download: bool = False,
119    **kwargs
120) -> DataLoader:
121    """Get the SpongeEM dataloader for the segmentation of structures in EM.
122
123    Args:
124        path: Filepath to a folder where the downloaded data will be saved.
125        mode: Choose the segmentation task, either 'semantic' or 'instances'.
126        patch_shape: The patch shape to use for training.
127        batch_size: The batch size for training.
128        sample_ids: The sample to download, valid ids are 1, 2 and 3.
129        download: Whether to download the data if it is not present.
130        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
131
132    Returns:
133       The DataLoader.
134    """
135    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
136    ds = get_sponge_em_dataset(path, mode, patch_shape, sample_ids=sample_ids, download=download, **ds_kwargs)
137    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)