torch_em.data.datasets.electron_microscopy.wafer4

  1"""The Wafer4 dataset contains annotations for neuron segmentation
  2in serial section electron microscopy of the mouse medial entorhinal cortex.
  3
  4The dataset is one volume of 125 x 1250 x 1250 voxels with a voxel level instance annotation of the
  5neurons. It was imaged with multi beam scanning electron microscopy at a resolution of 8 x 8 x 35
  6nanometer. The authors split it along the z axis, into 100 sections for training and 25 sections for
  7testing, and this loader follows that split.
  8
  9NOTE: The volume covers layer six of the allocortex, which the common electron microscopy neuron
 10datasets do not cover.
 11
 12The dataset is located at https://github.com/liuxy1103/CAD under the CC BY-NC 4.0 license.
 13This dataset is from the publication https://doi.org/10.1109/CVPR52733.2024.01056.
 14Please cite it if you use this dataset in your research.
 15"""
 16
 17import os
 18from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 19
 20import numpy as np
 21
 22from torch.utils.data import DataLoader, Dataset
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URLS = {
 30    "raw": "https://drive.usercontent.google.com/download?id=1l8Lhk-icIWyDb3fDt_2dDIvvx7S_9vLS&confirm=xxx",
 31    "labels": "https://drive.usercontent.google.com/download?id=1yyr3eo3-IQsEVvdZIgLsf_QrzaqXfUUX&confirm=xxx",
 32}
 33
 34CHECKSUMS = {
 35    "raw": "5bb64ae54d5d89a501b6942a999a4210eb178671293b5053ad63e12211df603a",
 36    "labels": "7d3eab21447a0efe327b5192303cc8df8f74ca280eb04e020cbe7dc27beb59c4",
 37}
 38
 39FILE_NAMES = {"raw": "wafer4_inputs.h5", "labels": "wafer4_labels.h5"}
 40
 41N_SECTIONS = 125
 42N_TRAIN_SECTIONS = 100
 43
 44
 45def get_wafer4_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 46    """Download the Wafer4 dataset.
 47
 48    Args:
 49        path: Filepath to a folder where the downloaded data will be saved.
 50        download: Whether to download the data if it is not present.
 51
 52    Returns:
 53        The filepath to the folder that holds the data.
 54    """
 55    os.makedirs(path, exist_ok=True)
 56
 57    for key, file_name in FILE_NAMES.items():
 58        file_path = os.path.join(path, file_name)
 59        if os.path.exists(file_path):
 60            continue
 61        util.download_source(file_path, URLS[key], download, CHECKSUMS[key])
 62
 63    return path
 64
 65
 66def get_wafer4_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[str, str]:
 67    """Get paths to the Wafer4 data.
 68
 69    Args:
 70        path: Filepath to a folder where the downloaded data will be saved.
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        The filepath for the image data.
 75        The filepath for the label data.
 76    """
 77    data_dir = get_wafer4_data(path, download)
 78    raw_path = os.path.join(data_dir, FILE_NAMES["raw"])
 79    label_path = os.path.join(data_dir, FILE_NAMES["labels"])
 80    return raw_path, label_path
 81
 82
 83def _get_split_roi(split: Optional[str]) -> Any:
 84    """Get the region of interest of a split. The authors split the volume along the z axis."""
 85    if split is None:
 86        return np.s_[:, :, :]
 87    if split == "train":
 88        return np.s_[:N_TRAIN_SECTIONS, :, :]
 89    if split == "test":
 90        return np.s_[N_TRAIN_SECTIONS:, :, :]
 91    raise ValueError(f"'{split}' is not a valid split. Choose 'train' or 'test', or None for the full volume.")
 92
 93
 94def get_wafer4_dataset(
 95    path: Union[os.PathLike, str],
 96    patch_shape: Tuple[int, int, int],
 97    split: Optional[Literal["train", "test"]] = "train",
 98    offsets: Optional[List[List[int]]] = None,
 99    boundaries: bool = False,
100    binary: bool = False,
101    rois: Optional[Dict[str, Any]] = None,
102    download: bool = False,
103    **kwargs,
104) -> Dataset:
105    """Get the Wafer4 dataset for the segmentation of neurons in EM.
106
107    Args:
108        path: Filepath to a folder where the downloaded data will be saved.
109        patch_shape: The 3D patch shape to use for training.
110        split: The data split. Either 'train' for the first 100 sections, 'test' for the last 25
111            sections, or None for the full volume.
112        offsets: Offset values for affinity computation used as target.
113        boundaries: Whether to compute boundaries as the target.
114        binary: Whether to use a binary segmentation target.
115        rois: The region of interest to use. Overrides the region of interest of the split.
116        download: Whether to download the data if it is not present.
117        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
118
119    Returns:
120        The segmentation dataset.
121    """
122    if len(patch_shape) != 3:
123        raise ValueError(f"The Wafer4 patch shape must be three-dimensional, got {patch_shape}.")
124
125    raw_path, label_path = get_wafer4_paths(path, download)
126    roi = _get_split_roi(split) if rois is None else rois
127
128    kwargs, _ = util.add_instance_label_transform(
129        kwargs, add_binary_target=False, offsets=offsets, boundaries=boundaries, binary=binary,
130    )
131
132    return torch_em.default_segmentation_dataset(
133        raw_paths=raw_path,
134        raw_key="main",
135        label_paths=label_path,
136        label_key="main",
137        patch_shape=patch_shape,
138        rois=roi,
139        ndim=3,
140        **kwargs,
141    )
142
143
144def get_wafer4_loader(
145    path: Union[os.PathLike, str],
146    batch_size: int,
147    patch_shape: Tuple[int, int, int],
148    split: Optional[Literal["train", "test"]] = "train",
149    offsets: Optional[List[List[int]]] = None,
150    boundaries: bool = False,
151    binary: bool = False,
152    rois: Optional[Dict[str, Any]] = None,
153    download: bool = False,
154    **kwargs,
155) -> DataLoader:
156    """Get the Wafer4 dataloader for the segmentation of neurons in EM.
157
158    Args:
159        path: Filepath to a folder where the downloaded data will be saved.
160        batch_size: The batch size for training.
161        patch_shape: The 3D patch shape to use for training.
162        split: The data split. Either 'train' for the first 100 sections, 'test' for the last 25
163            sections, or None for the full volume.
164        offsets: Offset values for affinity computation used as target.
165        boundaries: Whether to compute boundaries as the target.
166        binary: Whether to use a binary segmentation target.
167        rois: The region of interest to use. Overrides the region of interest of the split.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
170
171    Returns:
172        The DataLoader.
173    """
174    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
175    dataset = get_wafer4_dataset(
176        path=path,
177        patch_shape=patch_shape,
178        split=split,
179        offsets=offsets,
180        boundaries=boundaries,
181        binary=binary,
182        rois=rois,
183        download=download,
184        **ds_kwargs,
185    )
186    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)