torch_em.data.datasets.electron_microscopy.fafb_nuclei

  1"""This dataset contains crops of EM data with annotated nuclei from the adult fruit fly brain.
  2
  3The data is extracted from https://doi.org/10.1101/2021.11.04.46719, which contains a segmentation
  4of all nuclei in the fruit fly brain. Please cite it if you use this dataset for a publication.
  5"""
  6
  7import os
  8from glob import glob
  9from typing import Tuple, Union, Literal, List
 10
 11import torch_em
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15from .. import util
 16
 17
 18URL = "https://owncloud.gwdg.de/index.php/s/PDrkO02w7FWLrsh/download"
 19CHECKSUM = "7f245866fa8ffb473f7bda1fbbea23d265183e4ada5e34a36ecd60475809dca3"
 20
 21
 22def get_fafb_nuclei_data(path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool) -> str:
 23    """Download the FAFB Nucleus data.
 24
 25    Args:
 26        path: Filepath to a folder where the downloaded data will be saved.
 27        split: The split to use. One of 'train', 'val', 'test'.
 28        download: Whether to download the data if it is not present.
 29
 30    Returns:
 31        The filepath to the downloaded data.
 32    """
 33    assert split in ("train", "val", "test")
 34    split_folder = os.path.join(path, split)
 35    if not os.path.exists(split_folder):
 36        os.makedirs(path, exist_ok=True)
 37        zip_path = os.path.join(path, "fafb_nucleus_data.zip")
 38        util.download_source(zip_path, URL, download, CHECKSUM)
 39        util.unzip(zip_path, path, remove=True)
 40    return split_folder
 41
 42
 43def get_fafb_nuclei_paths(
 44    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool
 45) -> List[str]:
 46    """Get paths to the FAFB Nucleus data.
 47
 48    Args:
 49        path: Filepath to a folder where the downloaded data will be saved.
 50        split: The split to use. One of 'train', 'val', 'test'.
 51        download: Whether to download the data if it is not present.
 52
 53    Returns:
 54        The filepaths to the stored data.
 55    """
 56    get_fafb_nuclei_data(path, split, download)
 57    split_folder = os.path.join(path, split)
 58    paths = sorted(glob(os.path.join(split_folder, "*.h5")))
 59    return paths
 60
 61
 62def get_fafb_nuclei_dataset(
 63    path: Union[os.PathLike, str],
 64    split: Literal["train", "val", "test"],
 65    patch_shape: Tuple[int, int, int],
 66    download: bool = False,
 67    **kwargs
 68) -> Dataset:
 69    """Get the FAFB nucleus dataset for the segmentation of nuclei in EM.
 70
 71    Args:
 72        path: Filepath to a folder where the downloaded data will be saved.
 73        split: The split for the dataset, either 'train, 'val', or 'test'.
 74        patch_shape: The patch shape to use for training.
 75        download: Whether to download the data if it is not present.
 76        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 77
 78    Returns:
 79       The segmentation dataset.
 80    """
 81    paths = get_fafb_nuclei_paths(path, split, download)
 82    return torch_em.default_segmentation_dataset(
 83        raw_paths=paths,
 84        raw_key="raw",
 85        label_paths=paths,
 86        label_key="labels/nuclei",
 87        patch_shape=patch_shape,
 88        is_seg_dataset=True,
 89        **kwargs
 90    )
 91
 92
 93def get_fafb_nuclei_loader(
 94    path: Union[os.PathLike, str],
 95    split: Literal["train", "val", "test"],
 96    patch_shape: Tuple[int, int, int],
 97    batch_size: int,
 98    download: bool = False,
 99    **kwargs
100) -> DataLoader:
101    """Get the FAFB nucleus dataloader for the segmentation of nuclei in EM.
102
103    Args:
104        path: Filepath to a folder where the downloaded data will be saved.
105        split: The split for the dataset, either 'train', 'val', or 'test'.
106        patch_shape: The patch shape to use for training.
107        batch_size: The batch size for training.
108        download: Whether to download the data if it is not present.
109        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
110
111    Returns:
112       The segmentation dataset.
113    """
114    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
115    ds = get_fafb_nuclei_dataset(path, split, patch_shape, download, **ds_kwargs)
116    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)