torch_em.data.datasets.electron_microscopy.perez

  1"""The Perez dataset contains annotations for semantic segmentation of mitochondria, lysosomes, nuclei and nucleoli
  2in serial block-face scanning electron microscopy (SBF-SEM) images of the mouse suprachiasmatic nucleus.
  3
  4Each organelle is a separate binary segmentation task, with its own set of 2D tiles.
  5
  6The original data is from the publication https://doi.org/10.3389/fnana.2014.00126.
  7The preprocessed version used here was released with the RETINA benchmark: https://github.com/BaderLab/RETINA
  8(publication: https://doi.org/10.1371/journal.pcbi.1013115).
  9Please cite both publications if you use this dataset in your research.
 10"""
 11
 12import os
 13from glob import glob
 14from zipfile import ZipFile
 15from typing import List, Literal, Tuple, Union
 16
 17from natsort import natsorted
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25# NOTE: This is the full RETINA benchmark archive (~1.2 GB), which contains four other datasets besides Perez.
 26# Only the Perez files are extracted.
 27URL = "https://drive.google.com/uc?id=1T0fShblGYgKxTRwPlbMSvMbHQqmEk8YJ"
 28CHECKSUM = "e65555a039af5cacd64d72209f34a31c2af7afcdb887b5420e3ef38401ba3f45"
 29
 30ORGANELLES = ["mito", "lyso", "nuclei", "nucleoli"]
 31
 32
 33def get_perez_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the Perez dataset.
 35
 36    Args:
 37        path: Filepath to a folder where the downloaded data will be saved.
 38        download: Whether to download the data if it is not present.
 39
 40    Returns:
 41        The filepath to the extracted data.
 42    """
 43    data_dir = os.path.join(path, "perez")
 44    if os.path.exists(data_dir):
 45        return data_dir
 46
 47    os.makedirs(path, exist_ok=True)
 48
 49    zip_path = os.path.join(path, "benchmark_dataset.zip")
 50    util.download_source_gdrive(zip_path, URL, download, checksum=CHECKSUM)
 51
 52    prefix = "benchdata/perez/"
 53    with ZipFile(zip_path) as f:
 54        members = [m for m in f.namelist() if m.startswith(prefix) and not m.endswith(".DS_Store")]
 55        f.extractall(path, members=members)
 56
 57    os.rename(os.path.join(path, "benchdata", "perez"), data_dir)
 58    os.rmdir(os.path.join(path, "benchdata"))
 59    os.remove(zip_path)
 60
 61    return data_dir
 62
 63
 64def get_perez_paths(
 65    path: Union[os.PathLike, str],
 66    split: Literal["train", "test"],
 67    organelle: Literal["mito", "lyso", "nuclei", "nucleoli"] = "mito",
 68    download: bool = False,
 69) -> Tuple[List[str], List[str]]:
 70    """Get paths to the Perez data.
 71
 72    Args:
 73        path: Filepath to a folder where the downloaded data will be saved.
 74        split: The data split. Either 'train' or 'test'.
 75        organelle: The organelle to segment. Either 'mito', 'lyso', 'nuclei' or 'nucleoli'.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image data.
 80        List of filepaths for the label data.
 81    """
 82    if split not in ["train", "test"]:
 83        raise ValueError(f"'{split}' is not a valid split. Please choose from 'train' / 'test'.")
 84
 85    if organelle not in ORGANELLES:
 86        raise ValueError(f"'{organelle}' is not a valid organelle. Please choose from {ORGANELLES}.")
 87
 88    data_dir = get_perez_data(path, download)
 89
 90    raw_paths = natsorted(glob(os.path.join(data_dir, organelle, split, "images", "*.png")))
 91    label_paths = natsorted(glob(os.path.join(data_dir, organelle, split, "masks", "*.png")))
 92    assert raw_paths and len(raw_paths) == len(label_paths)
 93
 94    return raw_paths, label_paths
 95
 96
 97def get_perez_dataset(
 98    path: Union[os.PathLike, str],
 99    patch_shape: Tuple[int, int],
100    split: Literal["train", "test"],
101    organelle: Literal["mito", "lyso", "nuclei", "nucleoli"] = "mito",
102    download: bool = False,
103    **kwargs
104) -> Dataset:
105    """Get the Perez dataset for organelle segmentation in SBF-SEM.
106
107    Args:
108        path: Filepath to a folder where the downloaded data will be saved.
109        patch_shape: The patch shape to use for training.
110        split: The data split. Either 'train' or 'test'.
111        organelle: The organelle to segment. Either 'mito', 'lyso', 'nuclei' or 'nucleoli'.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
114
115    Returns:
116        The segmentation dataset.
117    """
118    raw_paths, label_paths = get_perez_paths(path, split, organelle, download)
119
120    return torch_em.default_segmentation_dataset(
121        raw_paths=raw_paths,
122        raw_key=None,
123        label_paths=label_paths,
124        label_key=None,
125        patch_shape=patch_shape,
126        is_seg_dataset=False,
127        **kwargs
128    )
129
130
131def get_perez_loader(
132    path: Union[os.PathLike, str],
133    batch_size: int,
134    patch_shape: Tuple[int, int],
135    split: Literal["train", "test"],
136    organelle: Literal["mito", "lyso", "nuclei", "nucleoli"] = "mito",
137    download: bool = False,
138    **kwargs
139) -> DataLoader:
140    """Get the Perez dataloader for organelle segmentation in SBF-SEM.
141
142    Args:
143        path: Filepath to a folder where the downloaded data will be saved.
144        batch_size: The batch size for training.
145        patch_shape: The patch shape to use for training.
146        split: The data split. Either 'train' or 'test'.
147        organelle: The organelle to segment. Either 'mito', 'lyso', 'nuclei' or 'nucleoli'.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
150
151    Returns:
152        The DataLoader.
153    """
154    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
155    dataset = get_perez_dataset(path, patch_shape, split, organelle, download, **ds_kwargs)
156    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)