torch_em.data.datasets.light_microscopy.oocount

  1"""This dataset contains annotations for 3d instance segmentation of oocytes in whole-mount
  2immunofluorescence imaging of the mouse ovary, at adult and perinatal developmental timepoints.
  3
  4The dataset is located at https://doi.org/10.5061/dryad.nk98sf81r.
  5This dataset is from the publication https://doi.org/10.1093/biolre/ioaf023.
  6Please cite it if you use this dataset in your research.
  7"""
  8
  9import os
 10from glob import glob
 11from natsort import natsorted
 12from typing import List, Literal, Tuple, Union
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21def get_oocount_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 22    """Instructions to obtain the OoCount training data.
 23
 24    NOTE: Dryad blocks automated downloads for this dataset. Please manually download
 25    'TrainingDatasets.zip' from https://datadryad.org/dataset/doi:10.5061/dryad.nk98sf81r
 26    and place it at the given 'path'.
 27
 28    Args:
 29        path: Filepath to a folder where the manually downloaded zip file is placed.
 30        download: Whether to download the data if it is not present.
 31
 32    Returns:
 33        The filepath to the extracted data directory.
 34    """
 35    data_dir = os.path.join(path, "TrainingDatasets")
 36    if os.path.exists(data_dir):
 37        return data_dir
 38
 39    if download:
 40        raise NotImplementedError(
 41            "The OoCount dataset cannot be downloaded automatically because Dryad blocks automated requests. "
 42            "Please manually download 'TrainingDatasets.zip' from "
 43            "https://datadryad.org/dataset/doi:10.5061/dryad.nk98sf81r and place it at the given 'path'."
 44        )
 45
 46    zip_path = os.path.join(path, "TrainingDatasets.zip")
 47    if not os.path.exists(zip_path):
 48        raise RuntimeError(
 49            f"The manually downloaded zip file should be placed at '{zip_path}'. Please download "
 50            "'TrainingDatasets.zip' from https://datadryad.org/dataset/doi:10.5061/dryad.nk98sf81r."
 51        )
 52
 53    util.unzip(zip_path=zip_path, dst=path, remove=False)
 54    assert os.path.exists(data_dir), data_dir
 55
 56    return data_dir
 57
 58
 59def get_oocount_paths(
 60    path: Union[os.PathLike, str],
 61    timepoint: Literal["adult", "perinatal"] = "adult",
 62    split: Literal["train", "val"] = "train",
 63    download: bool = False,
 64) -> Tuple[List[str], List[str]]:
 65    """Get paths to the OoCount data.
 66
 67    Args:
 68        path: Filepath to a folder where the data is stored.
 69        timepoint: The ovary developmental timepoint. Either 'adult' or 'perinatal'.
 70        split: The data split to use. 'train' uses the training images, 'val' uses the held-out QC images.
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        List of filepaths for the image data.
 75        List of filepaths for the label data.
 76    """
 77    if timepoint not in ("adult", "perinatal"):
 78        raise ValueError(f"'{timepoint}' is not a valid timepoint. Choose 'adult' or 'perinatal'.")
 79    if split not in ("train", "val"):
 80        raise ValueError(f"'{split}' is not a valid split. Choose 'train' or 'val'.")
 81
 82    data_dir = get_oocount_data(path, download)
 83
 84    sample_dir = os.path.join(data_dir, f"Vasa-{timepoint.capitalize()}")
 85    image_dir = os.path.join(sample_dir, "Images" if split == "train" else "QC Images")
 86    label_dir = os.path.join(sample_dir, "Masks" if split == "train" else "QC Masks")
 87
 88    raw_paths = natsorted(glob(os.path.join(image_dir, "*.tif")))
 89    label_paths = natsorted(glob(os.path.join(label_dir, "*.tif")))
 90
 91    if len(raw_paths) == 0:
 92        raise RuntimeError(f"No image files found in {image_dir}.")
 93    if len(raw_paths) != len(label_paths):
 94        raise RuntimeError(
 95            f"Mismatch between images ({len(raw_paths)}) and masks ({len(label_paths)}) in {sample_dir}."
 96        )
 97
 98    return raw_paths, label_paths
 99
100
101def get_oocount_dataset(
102    path: Union[os.PathLike, str],
103    patch_shape: Tuple[int, int],
104    timepoint: Literal["adult", "perinatal"] = "adult",
105    split: Literal["train", "val"] = "train",
106    download: bool = False,
107    **kwargs,
108) -> Dataset:
109    """Get the OoCount dataset for oocyte instance segmentation in 3d fluorescence microscopy.
110
111    Args:
112        path: Filepath to a folder where the data is stored.
113        patch_shape: The patch shape to use for training.
114        timepoint: The ovary developmental timepoint. Either 'adult' or 'perinatal'.
115        split: The data split to use. 'train' uses the training images, 'val' uses the held-out QC images.
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    raw_paths, label_paths = get_oocount_paths(path, timepoint, split, download)
123
124    return torch_em.default_segmentation_dataset(
125        raw_paths=raw_paths,
126        raw_key=None,
127        label_paths=label_paths,
128        label_key=None,
129        patch_shape=patch_shape,
130        **kwargs,
131    )
132
133
134def get_oocount_loader(
135    path: Union[os.PathLike, str],
136    batch_size: int,
137    patch_shape: Tuple[int, int],
138    timepoint: Literal["adult", "perinatal"] = "adult",
139    split: Literal["train", "val"] = "train",
140    download: bool = False,
141    **kwargs,
142) -> DataLoader:
143    """Get the OoCount dataloader for oocyte instance segmentation in 3d fluorescence microscopy.
144
145    Args:
146        path: Filepath to a folder where the data is stored.
147        batch_size: The batch size for training.
148        patch_shape: The patch shape to use for training.
149        timepoint: The ovary developmental timepoint. Either 'adult' or 'perinatal'.
150        split: The data split to use. 'train' uses the training images, 'val' uses the held-out QC images.
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
153
154    Returns:
155        The DataLoader.
156    """
157    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
158    dataset = get_oocount_dataset(path, patch_shape, timepoint, split, download, **ds_kwargs)
159    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)