torch_em.data.datasets.electron_microscopy.axonem

  1"""AxomEM is a datast for segmenting axons in electron microscopy.
  2It contains two large annotated volumes, one from mouse cortex, the other from human cortex.
  3This dataset was used for the AxonEM Challenge: https://axonem.grand-challenge.org/.
  4
  5Please cite the publication https://arxiv.org/abs/2107.05451 if you use this dataset for your research.
  6"""
  7
  8import os
  9from glob import glob
 10from typing import Union, Sequence, List, Tuple
 11
 12from torch.utils.data import Dataset, DataLoader
 13
 14import torch_em
 15
 16from .. import util
 17
 18
 19URLS = {
 20    "human": "https://huggingface.co/datasets/pytc/AxonEM/resolve/main/EM30-H-train-9vol-pad-20-512-512.zip",
 21    "mouse": "https://huggingface.co/datasets/pytc/AxonEM/resolve/main/EM30-M-train-9vol-pad-20-512-512.zip",
 22}
 23
 24CHECKSUMS = {
 25    "human": "0b53d155ff62f5e24c552bf90adce329fcf9a8fefd5c697f8bcd0312a61fda60",
 26    "mouse": "dae06b5dabe388ab7a0ff4e51548174f041a338d0d06bd665586aa7fdd43bac2",
 27}
 28
 29
 30def get_axonem_data(path: Union[os.PathLike, str], samples: Sequence[str], download: bool = False):
 31    """Download the AxonEM training data.
 32
 33    Args:
 34        path: Filepath to a folder where the downloaded data will be saved.
 35        samples: The samples to download. The available samples are 'human' and 'mouse'.
 36        download: Whether to download the data if it is not present.
 37    """
 38    if isinstance(samples, str):
 39        samples = [samples]
 40
 41    assert len(set(samples) - {"human", "mouse"}) == 0, f"{samples}"
 42    os.makedirs(path, exist_ok=True)
 43
 44    for sample in samples:
 45        dst = os.path.join(path, sample)
 46        if os.path.exists(dst):
 47            continue
 48
 49        os.makedirs(dst, exist_ok=True)
 50
 51        # Download the zipfile.
 52        zip_path = os.path.join(path, f"{sample}.zip")
 53        util.download_source(path=zip_path, url=URLS[sample], download=download, checksum=CHECKSUMS[sample])
 54
 55        # Extract the h5 crops from the zipfile.
 56        util.unzip(zip_path=zip_path, dst=dst, remove=True)
 57
 58        if sample == "mouse":
 59            # NOTE: We need to make a hotfix by removing a crop which does not have masks.
 60            label_path = os.path.join(path, "mouse", "valid_mask.h5")
 61            os.remove(label_path)
 62
 63            # And the additional volume with no corresponding mask.
 64            image_path = os.path.join(path, "mouse", "im_675-800-800_pad.h5")
 65            os.remove(image_path)
 66
 67
 68def get_axonem_paths(
 69    path: Union[os.PathLike, str], samples: Sequence[str], download: bool = False,
 70) -> Tuple[List[str], List[str]]:
 71    """Get paths for the AxonEM training data.
 72
 73    Args:
 74        path: Filepath to a folder where the downloaded data will be saved.
 75        samples: The samples to download. The available samples are 'human' and 'mouse'.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image volumes.
 80        List of filepaths for the label volumes.
 81    """
 82    get_axonem_data(path, samples, download)
 83
 84    if isinstance(samples, str):
 85        samples = [samples]
 86
 87    image_paths, label_paths = [], []
 88    for sample in samples:
 89        curr_image_paths = glob(os.path.join(path, sample, "im_*.h5"))
 90        image_paths.extend(curr_image_paths)
 91        label_paths.extend([p.replace("im_", "seg_") for p in curr_image_paths])
 92
 93    return image_paths, label_paths
 94
 95
 96def get_axonem_dataset(
 97    path: Union[os.PathLike, str],
 98    patch_shape: Tuple[int, ...],
 99    samples: Sequence[str] = ("human", "mouse"),
100    download: bool = False,
101    **kwargs
102) -> Dataset:
103    """Get the AxonEM dataset for the segmentation of axons in EM.
104
105    Args:
106        path: Filepath to a folder where the downloaded data will be saved.
107        patch_shape: The patch shape to use for training.
108        samples: The samples to download. The available samples are 'human' and 'mouse'.
109        download: Whether to download the data if it is not present.
110        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
111
112    Returns:
113       The segmentation dataset.
114    """
115    image_paths, label_paths = get_axonem_paths(path, samples, download)
116
117    return torch_em.default_segmentation_dataset(
118        raw_paths=image_paths,
119        raw_key="main",
120        label_paths=label_paths,
121        label_key="main",
122        patch_shape=patch_shape,
123        **kwargs
124    )
125
126
127def get_axonem_loader(
128    path: Union[os.PathLike, str],
129    batch_size: int,
130    patch_shape: Tuple[int, ...],
131    samples: Sequence[str] = ("human", "mouse"),
132    download: bool = False,
133    **kwargs
134) -> DataLoader:
135    """Get the AxonEM dataloader for the segmentation of axons in EM.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        batch_size: The batch size for training.
140        patch_shape: The patch shape to use for training.
141        samples: The samples to download. The available samples are 'human' and 'mouse'.
142        download: Whether to download the data if it is not present.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
144
145    Returns:
146        The DataLoader.
147    """
148    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
149    dataset = get_axonem_dataset(path, patch_shape, samples, download, **ds_kwargs)
150    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)