torch_em.data.datasets.light_microscopy.cochleanet

  1"""The CochleaNet dataset contains annotations for 3D instance segmentation of spiral ganglion neurons (SGNs)
  2and inner hair cells (IHCs) in light-sheet microscopy volumes of cleared cochleae.
  3
  4The data are 3D crops with isotropic voxel size of 0.38 um. The SGN data has 48 crops of mouse cochleae with
  5parvalbumin stain and the IHC data has 89 crops of mouse and gerbil cochleae with Vglut3 stain. Some crops contain
  6no annotated cell. The official train and val splits are the ones used for training the published CochleaNet models.
  7
  8The data is located at https://doi.org/10.5281/zenodo.23478266, released under the MIT license.
  9The dataset is from the publication https://doi.org/10.1101/2025.11.16.688700.
 10Please cite it if you use this dataset in your research.
 11"""
 12
 13import os
 14from glob import glob
 15from natsort import natsorted
 16from typing import Union, Tuple, List, Literal, Optional
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URLS = {
 26    "sgn": "https://zenodo.org/records/23478267/files/SGN_v1.zip",
 27    "ihc": "https://zenodo.org/records/23478267/files/IHC_v1.zip",
 28}
 29
 30CHECKSUMS = {
 31    "sgn": "3d80bce3b7bd635609f15fb64703a1ab4480bef396bfe3d0f11c850091064f10",
 32    "ihc": "26e0576ede1f1730d77636da0ac75066f5c5dbd7cc3d6ff3be3be9ad973d13ae",
 33}
 34
 35FOLDERS = {"sgn": "SGN_v1", "ihc": "IHC_v1"}
 36
 37
 38def get_cochleanet_data(
 39    path: Union[os.PathLike, str], target: Literal["sgn", "ihc"] = "sgn", download: bool = False
 40) -> str:
 41    """Download the CochleaNet data.
 42
 43    Args:
 44        path: Filepath to a folder where the downloaded data will be saved.
 45        target: The annotated cell type. Either 'sgn' for spiral ganglion neurons or 'ihc' for inner hair cells.
 46        download: Whether to download the data if it is not present.
 47
 48    Returns:
 49        The filepath to the folder with the data of the chosen target.
 50    """
 51    if target not in URLS:
 52        raise ValueError(f"'{target}' is not a valid target. Choose from {list(URLS)}.")
 53
 54    data_dir = os.path.join(path, FOLDERS[target])
 55    if os.path.exists(data_dir):
 56        return data_dir
 57
 58    os.makedirs(path, exist_ok=True)
 59    zip_path = os.path.join(path, f"{FOLDERS[target]}.zip")
 60    util.download_source(path=zip_path, url=URLS[target], download=download, checksum=CHECKSUMS[target])
 61    util.unzip(zip_path=zip_path, dst=path)
 62
 63    return data_dir
 64
 65
 66def get_cochleanet_paths(
 67    path: Union[os.PathLike, str],
 68    split: Literal["train", "val"],
 69    target: Literal["sgn", "ihc"] = "sgn",
 70    download: bool = False,
 71) -> Tuple[List[str], List[str]]:
 72    """Get paths to the CochleaNet data.
 73
 74    Args:
 75        path: Filepath to a folder where the downloaded data will be saved.
 76        split: The data split to use. Either 'train' or 'val'.
 77        target: The annotated cell type. Either 'sgn' for spiral ganglion neurons or 'ihc' for inner hair cells.
 78        download: Whether to download the data if it is not present.
 79
 80    Returns:
 81        List of filepaths for the image data.
 82        List of filepaths for the label data.
 83    """
 84    if split not in ("train", "val"):
 85        raise ValueError(f"'{split}' is not a valid split. Choose from 'train' or 'val'.")
 86
 87    data_dir = get_cochleanet_data(path, target, download)
 88
 89    label_paths = natsorted(glob(os.path.join(data_dir, split, "*_annotations.tif")))
 90    raw_paths = [p.replace("_annotations.tif", ".tif") for p in label_paths]
 91    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths)
 92
 93    return raw_paths, label_paths
 94
 95
 96def get_cochleanet_dataset(
 97    path: Union[os.PathLike, str],
 98    patch_shape: Tuple[int, int, int],
 99    split: Literal["train", "val"],
100    target: Literal["sgn", "ihc"] = "sgn",
101    offsets: Optional[List[List[int]]] = None,
102    boundaries: bool = False,
103    binary: bool = False,
104    download: bool = False,
105    **kwargs
106) -> Dataset:
107    """Get the CochleaNet dataset for SGN or IHC segmentation in light-sheet microscopy of cochleae.
108
109    Args:
110        path: Filepath to a folder where the downloaded data will be saved.
111        patch_shape: The patch shape to use for training.
112        split: The data split to use. Either 'train' or 'val'.
113        target: The annotated cell type. Either 'sgn' for spiral ganglion neurons or 'ihc' for inner hair cells.
114        offsets: Offset values for affinity computation used as target.
115        boundaries: Whether to compute boundaries as the target.
116        binary: Whether to use a binary segmentation target.
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_cochleanet_paths(path, split, target, download)
124
125    kwargs = util.ensure_transforms(ndim=3, **kwargs)
126    kwargs, _ = util.add_instance_label_transform(
127        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
128    )
129
130    return torch_em.default_segmentation_dataset(
131        raw_paths=raw_paths,
132        raw_key=None,
133        label_paths=label_paths,
134        label_key=None,
135        patch_shape=patch_shape,
136        ndim=3,
137        **kwargs
138    )
139
140
141def get_cochleanet_loader(
142    path: Union[os.PathLike, str],
143    batch_size: int,
144    patch_shape: Tuple[int, int, int],
145    split: Literal["train", "val"],
146    target: Literal["sgn", "ihc"] = "sgn",
147    offsets: Optional[List[List[int]]] = None,
148    boundaries: bool = False,
149    binary: bool = False,
150    download: bool = False,
151    **kwargs
152) -> DataLoader:
153    """Get the CochleaNet dataloader for SGN or IHC segmentation in light-sheet microscopy of cochleae.
154
155    Args:
156        path: Filepath to a folder where the downloaded data will be saved.
157        batch_size: The batch size for training.
158        patch_shape: The patch shape to use for training.
159        split: The data split to use. Either 'train' or 'val'.
160        target: The annotated cell type. Either 'sgn' for spiral ganglion neurons or 'ihc' for inner hair cells.
161        offsets: Offset values for affinity computation used as target.
162        boundaries: Whether to compute boundaries as the target.
163        binary: Whether to use a binary segmentation target.
164        download: Whether to download the data if it is not present.
165        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
166
167    Returns:
168        The DataLoader.
169    """
170    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
171    dataset = get_cochleanet_dataset(
172        path, patch_shape, split, target, offsets=offsets, boundaries=boundaries, binary=binary,
173        download=download, **ds_kwargs
174    )
175    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)