torch_em.data.datasets.light_microscopy.neurons_3d_multispecies

  1"""This dataset contains 3D label-free microscopy volumes of neuronal cell
  2bodies (somata) from human, mouse and rat brain tissue, imaged with oblique
  3illumination or Dodt gradient contrast (DGC) microscopy, with manually
  4annotated instance segmentation masks and an official train / val / test split.
  5
  6The dataset is located at https://zenodo.org/records/20797635.
  7Please cite it if you use this dataset in your research.
  8"""
  9
 10import os
 11import json
 12from glob import glob
 13from natsort import natsorted
 14from typing import List, Literal, Optional, Tuple, Union
 15
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23URLS = {
 24    "human_oblique": "https://zenodo.org/records/20797635/files/human_neurons_oblique.zip",
 25    "human_dodt": "https://zenodo.org/records/20797635/files/human_neurons_dodt.zip",
 26    "mouse_oblique": "https://zenodo.org/records/20797635/files/mice_neurons_oblique.zip",
 27    "mouse_dodt": "https://zenodo.org/records/20797635/files/mice_neurons_dodt.zip",
 28    "rat_oblique": "https://zenodo.org/records/20797635/files/rat_neurons_oblique.zip",
 29    "rat_dodt": "https://zenodo.org/records/20797635/files/rat_neurons_dodt.zip",
 30}
 31
 32CHECKSUMS = {
 33    "human_oblique": "8b2c37b74fe2b890a3d941097d619651d56fa457a37c010c1c31d24604acdebc",
 34    "human_dodt": "1b4a72302e0e85dfb3169fd753dcfca6c4d200cea135ce7b2d8401af009d3621",
 35    "mouse_oblique": "31dd601cf114359a66e39266b688e51448f70bf6a516474b8905f3071ce372fa",
 36    "mouse_dodt": "3debfdd5534509b8161bbf2f952cba1bd7629aa22c78d5f32bf585d10edb3b1c",
 37    "rat_oblique": "0b7be7b0a9436d74ed1f2f2676256c9ad0072996f89fd95ce2b01a4a2cc64f86",
 38    "rat_dodt": "431269c513bb558f7434018fc087c22e31c1e7cbdccc4da2d399d0bf3b85715a",
 39}
 40
 41# Zenodo archive names differ from the species / modality keys used here.
 42ARCHIVE_NAMES = {
 43    "human_oblique": "human_neurons_oblique",
 44    "human_dodt": "human_neurons_dodt",
 45    "mouse_oblique": "mice_neurons_oblique",
 46    "mouse_dodt": "mice_neurons_dodt",
 47    "rat_oblique": "rat_neurons_oblique",
 48    "rat_dodt": "rat_neurons_dodt",
 49}
 50
 51SPECIES = ["human", "mouse", "rat"]
 52MODALITIES = ["oblique", "dodt"]
 53
 54
 55def get_neurons_3d_multispecies_data(
 56    path: Union[os.PathLike, str],
 57    species: Optional[List[Literal["human", "mouse", "rat"]]] = None,
 58    modality: Optional[List[Literal["oblique", "dodt"]]] = None,
 59    download: bool = False,
 60) -> List[str]:
 61    """Download the multi-species 3D neuron segmentation dataset.
 62
 63    Args:
 64        path: Filepath to a folder where the downloaded data will be saved.
 65        species: The species subset(s) to download. Defaults to all species.
 66        modality: The imaging modality subset(s) to download. Defaults to all modalities.
 67        download: Whether to download the data if it is not present.
 68
 69    Returns:
 70        List of filepaths to the extracted data directories.
 71    """
 72    species = SPECIES if species is None else species
 73    modality = MODALITIES if modality is None else modality
 74
 75    os.makedirs(path, exist_ok=True)
 76
 77    data_dirs = []
 78    for sp in species:
 79        for mod in modality:
 80            key = f"{sp}_{mod}"
 81            data_dir = os.path.join(path, ARCHIVE_NAMES[key])
 82            if not os.path.exists(data_dir):
 83                zip_path = os.path.join(path, f"{ARCHIVE_NAMES[key]}.zip")
 84                util.download_source(zip_path, URLS[key], download, checksum=CHECKSUMS[key])
 85                util.unzip(zip_path, path)
 86            data_dirs.append(data_dir)
 87
 88    return data_dirs
 89
 90
 91def get_neurons_3d_multispecies_paths(
 92    path: Union[os.PathLike, str],
 93    species: Optional[List[Literal["human", "mouse", "rat"]]] = None,
 94    modality: Optional[List[Literal["oblique", "dodt"]]] = None,
 95    split: Optional[Literal["train", "val", "test"]] = None,
 96    download: bool = False,
 97) -> Tuple[List[str], List[str]]:
 98    """Get paths to the multi-species 3D neuron segmentation data.
 99
100    Args:
101        path: Filepath to a folder where the downloaded data will be saved.
102        species: The species subset(s) to use. Defaults to all species.
103        modality: The imaging modality subset(s) to use. Defaults to all modalities.
104        split: The data split to use. Either 'train', 'val' or 'test'. Defaults to using all splits.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        List of filepaths for the image data.
109        List of filepaths for the label data.
110    """
111    data_dirs = get_neurons_3d_multispecies_data(path, species, modality, download)
112
113    raw_paths, label_paths = [], []
114    for data_dir in data_dirs:
115        if split is None:
116            fnames = natsorted(os.path.basename(p) for p in glob(os.path.join(data_dir, "images", "*.tif")))
117        else:
118            split_file = os.path.join(data_dir, "split.json")
119            with open(split_file) as f:
120                fnames = natsorted(json.load(f)[split])
121
122        for fname in fnames:
123            raw_path = os.path.join(data_dir, "images", fname)
124            label_path = os.path.join(data_dir, "masks", fname)
125            if os.path.exists(raw_path) and os.path.exists(label_path):
126                raw_paths.append(raw_path)
127                label_paths.append(label_path)
128
129    if len(raw_paths) == 0:
130        raise RuntimeError(f"No image files found under {path}. Please check the dataset structure.")
131
132    return raw_paths, label_paths
133
134
135def get_neurons_3d_multispecies_dataset(
136    path: Union[os.PathLike, str],
137    patch_shape: Tuple[int, int, int],
138    species: Optional[List[Literal["human", "mouse", "rat"]]] = None,
139    modality: Optional[List[Literal["oblique", "dodt"]]] = None,
140    split: Optional[Literal["train", "val", "test"]] = None,
141    download: bool = False,
142    **kwargs,
143) -> Dataset:
144    """Get the multi-species 3D neuron segmentation dataset.
145
146    Args:
147        path: Filepath to a folder where the downloaded data will be saved.
148        patch_shape: The patch shape to use for training.
149        species: The species subset(s) to use. Defaults to all species.
150        modality: The imaging modality subset(s) to use. Defaults to all modalities.
151        split: The data split to use. Either 'train', 'val' or 'test'. Defaults to using all splits.
152        download: Whether to download the data if it is not present.
153        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
154
155    Returns:
156        The segmentation dataset.
157    """
158    raw_paths, label_paths = get_neurons_3d_multispecies_paths(path, species, modality, split, download)
159
160    return torch_em.default_segmentation_dataset(
161        raw_paths=raw_paths,
162        raw_key=None,
163        label_paths=label_paths,
164        label_key=None,
165        patch_shape=patch_shape,
166        **kwargs,
167    )
168
169
170def get_neurons_3d_multispecies_loader(
171    path: Union[os.PathLike, str],
172    batch_size: int,
173    patch_shape: Tuple[int, int, int],
174    species: Optional[List[Literal["human", "mouse", "rat"]]] = None,
175    modality: Optional[List[Literal["oblique", "dodt"]]] = None,
176    split: Optional[Literal["train", "val", "test"]] = None,
177    download: bool = False,
178    **kwargs,
179) -> DataLoader:
180    """Get the multi-species 3D neuron segmentation dataloader.
181
182    Args:
183        path: Filepath to a folder where the downloaded data will be saved.
184        batch_size: The batch size for training.
185        patch_shape: The patch shape to use for training.
186        species: The species subset(s) to use. Defaults to all species.
187        modality: The imaging modality subset(s) to use. Defaults to all modalities.
188        split: The data split to use. Either 'train', 'val' or 'test'. Defaults to using all splits.
189        download: Whether to download the data if it is not present.
190        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
191
192    Returns:
193        The DataLoader.
194    """
195    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
196    dataset = get_neurons_3d_multispecies_dataset(path, patch_shape, species, modality, split, download, **ds_kwargs)
197    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)