torch_em.data.datasets.electron_microscopy.aimseg

  1"""AimSeg contanins annotations for semantic segmentation of myelins, axons and inner cytoplasmic tongue,
  2and for instance segmentation of myelinated axons in TEM images of corpus callosum (CC) in adult mice.
  3
  4This dataset is available at https://zenodo.org/records/8351731.
  5The dataset is from the publication https://doi.org/10.1371/journal.pcbi.1010845.
  6Please cite them if you use this dataset in your research.
  7"""
  8
  9import os
 10from glob import glob
 11from pathlib import Path
 12from natsort import natsorted
 13from typing import Tuple, List, Union, Optional, Literal
 14
 15from torch.utils.data import Dataset, DataLoader
 16
 17import torch_em
 18
 19from .. import util
 20
 21
 22URLS = [
 23    "https://zenodo.org/records/8351731/files/Control_Dataset.rar",
 24    "https://zenodo.org/records/8351731/files/Validation_Dataset_v1.rar",
 25]
 26
 27CHECKSUMS = [
 28    "ecd569a5f91166a09d93d29a10e2ddd2eaa3e82df531785b7aa243e426467673",
 29    "647216eb09a644be8980224a52d8168fa2fa5a1fd0537fb1e5d6102ec30e396d"
 30]
 31
 32
 33def get_aimseg_data(path: Union[os.PathLike, str], download: bool = False):
 34    """Get the AimSeg data.
 35
 36    Args:
 37        path: Path to a folder where the data is downloaded.
 38        download: Whether to download the data if it is not present.
 39    """
 40    for url, checksum in zip(URLS, CHECKSUMS):
 41        rarfname = url.rsplit("/")[-1]
 42        dirname = Path(rarfname).stem
 43
 44        if os.path.exists(os.path.join(path, dirname)):
 45            continue
 46
 47        os.makedirs(path, exist_ok=True)
 48
 49        util.download_source(path=os.path.join(path, rarfname), url=url, download=download, checksum=checksum)
 50        util.unzip_rarfile(rar_path=os.path.join(path, rarfname), dst=path)
 51
 52
 53def get_aimseg_paths(
 54    path: Union[os.PathLike, str],
 55    split: Optional[Literal["control", "validation"]] = None,
 56    targets: Literal["instances", "semantic"] = "instances",
 57    download: bool = False,
 58) -> Tuple[List[str], List[str]]:
 59    """Get paths to the AimSeg data.
 60
 61    Args:
 62        path: Filepath to a folder where the data is downloaded.
 63        split: The split of the data to be used for training.
 64            Either `control` focused on healthy control specimen,
 65            or `validation` focused on mice undergoing remyelination.
 66        targets: The choice of support labels for the task.
 67            Either `instances` for annotated myelinated axons or `semantic` for axons, inner tongue and myelins.
 68        download: Whether to download the data if it is not present.
 69
 70    Returns:
 71        List of filepaths for the image data.
 72        List of filepaths for the label data.
 73    """
 74    # Download the AimSeg data.
 75    get_aimseg_data(path, download)
 76
 77    # Get the directory name for desired targets.
 78    if targets == "instances":
 79        dirname = "GroundTruth_Instance"
 80    elif targets == "semantic":
 81        dirname = "GroundTruth_Semantic"
 82    else:
 83        raise ValueError(f"'{targets}' is not a valid target choice. Please choose from 'instances' / 'semantic'.")
 84
 85    # Get the paths to image and corresponding labels
 86    raw_paths, label_paths = [], []
 87    if split and split not in ["control", "validation"]:
 88        raise ValueError(f"'{split}' is not a valid split choice. Please choose from 'control' / 'validation'.")
 89
 90    if split != "validation":
 91        raw_paths.extend(natsorted(glob(os.path.join(path, "Control_Dataset", "Images", "*.tif"))))
 92        label_paths.extend(natsorted(glob(os.path.join(path, "Control_Dataset", dirname, "*.tif"))))
 93
 94    if split != "control":
 95        raw_paths.extend(natsorted(glob(os.path.join(path, "Validation_Dataset_v1", "Images", "*.tif"))))
 96        label_paths.extend(natsorted(glob(os.path.join(path, "Validation_Dataset_v1", dirname, "*.tif"))))
 97
 98    assert raw_paths and len(raw_paths) == len(label_paths)
 99
100    return raw_paths, label_paths
101
102
103def get_aimseg_dataset(
104    path: Union[os.PathLike, str],
105    patch_shape: Tuple[int, int],
106    targets: Literal["instances", "semantic"] = "instances",
107    download: bool = False,
108    **kwargs
109) -> Dataset:
110    """Get the AimSeg dataset for axon and myelin segmentation.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded.
114        patch_shape: The patch shape to use for training.
115        targets: The choice of support labels for the task.
116            Either `instances` for annotated myelinated axons or `semantic` for axons, inner tongue and myelins.
117        download: Whether to download the data if it is not present.
118        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
119
120    Returns:
121        The segmentation dataset.
122    """
123    raw_paths, label_paths = get_aimseg_paths(path, None, targets, download)
124
125    return torch_em.default_segmentation_dataset(
126        raw_paths=raw_paths,
127        raw_key=None,
128        label_paths=label_paths,
129        label_key=None,
130        patch_shape=patch_shape,
131        **kwargs
132    )
133
134
135def get_aimseg_loader(
136    path: Union[os.PathLike, str],
137    batch_size: int,
138    patch_shape: Tuple[int, int],
139    targets: Literal["instances", "semantic"] = "instances",
140    download: bool = False,
141    **kwargs
142) -> DataLoader:
143    """Get the AimSeg dataset for axon and myelin segmentation.
144
145    Args:
146        path: Filepath to a folder where the data is downloaded.
147        batch_size: The batch size for training.
148        patch_shape: The patch shape to use for training.
149        targets: The choice of support labels for the task.
150            Either 'instances' for annotated myelinated axons or 'semantic' for axons, inner tongue and myelins.
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
153
154    Returns:
155        The DataLoader
156    """
157    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
158    dataset = get_aimseg_dataset(path, patch_shape, targets, download, **ds_kwargs)
159    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)