torch_em.data.datasets.light_microscopy.ovarian_reserve

  1"""The Ovarian Reserve dataset contains annotations for 3d oocyte segmentation in
  2light-sheet microscopy volumes of whole-mount mouse ovaries.
  3
  4The ovaries come from C57BL/6J mice between 5 and 60 weeks of age. They were cleared, stained for
  5DDX4 and imaged on a SPIM light-sheet microscope. Every volume holds 40 planes of 256 x 256 pixels
  6at a voxel size of 5.0 x 0.867 x 0.867 micrometer, together with hand curated instance labels of the
  7oocytes. The archive holds 66 volumes for training and 7 for validation.
  8
  9NOTE: This is the representative labeled subset of the study. The whole ovaries of the study carry no
 10labels. They are available at https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD3593 and at
 11https://doi.org/10.5281/zenodo.19085211 .
 12
 13NOTE: The masks come as uint16, int32, uint32 and float32, so this loader writes them all as uint32.
 14
 15The dataset is distributed with the segmentation tutorial at
 16https://biapy.readthedocs.io/en/latest/tutorials/instance_seg/ovarian-reserve.html , and the BioImage
 17Model Zoo lists it under the CC BY 4.0 license at https://bioimage.io/#/artifacts/splendid-falafel .
 18This dataset is from the publication https://doi.org/10.1038/s43587-026-01178-z .
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23from glob import glob
 24from natsort import natsorted
 25from typing import List, Literal, Optional, Tuple, Union
 26
 27from torch.utils.data import DataLoader, Dataset
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URL = (
 35    "https://upvehueus-my.sharepoint.com/:u:/g/personal/ignacio_arganda_ehu_eus/"
 36    "IQBlTg1-y8MlSqwgDpLZuPAgAU5oE0HOqc6vjDK7vVh_xBM?e=MMgzZf&download=1"
 37)
 38CHECKSUM = "d9774be882229451ddf415475a7a404caf190714a7091a679f8d5df97ae00a7c"
 39
 40SPLITS = ("train", "val")
 41
 42# The voxel size in micrometer, in the order of the axes of the arrays.
 43RESOLUTION = (5.0, 0.867, 0.867)
 44
 45
 46def _create_h5(data_dir: str, split: str) -> str:
 47    """Write one h5 file per volume, with the raw data, the labels and the voxel size."""
 48    import h5py
 49    import tifffile
 50    from tqdm import tqdm
 51
 52    output_dir = os.path.join(data_dir, "preprocessed", split)
 53    os.makedirs(output_dir, exist_ok=True)
 54
 55    raw_paths = natsorted(glob(os.path.join(data_dir, split, "raw", "*.tif")))
 56    if not raw_paths:
 57        raise RuntimeError(f"Could not find any raw data for the split '{split}' in {data_dir}.")
 58
 59    for raw_path in tqdm(raw_paths, desc=f"Preprocess '{split}'"):
 60        name = os.path.basename(raw_path)
 61        # The archive gives an image and its mask the same file name.
 62        label_path = os.path.join(data_dir, split, "label", name)
 63        if not os.path.exists(label_path):
 64            raise RuntimeError(f"Could not find the mask for the image '{name}' at {label_path}.")
 65
 66        output_path = os.path.join(output_dir, f"{os.path.splitext(name)[0]}.h5")
 67        if os.path.exists(output_path):
 68            continue
 69
 70        raw = tifffile.imread(raw_path)
 71        labels = tifffile.imread(label_path)
 72        if raw.shape != labels.shape:
 73            raise RuntimeError(
 74                f"The image '{name}' has the shape {raw.shape}, but its mask has the shape {labels.shape}."
 75            )
 76
 77        temporary_path = f"{output_path}.tmp"
 78        with h5py.File(temporary_path, "w") as f:
 79            f.attrs["modality"] = "selective plane illumination fluorescence microscopy"
 80            f.attrs["tissue"] = "whole-mount mouse ovary"
 81            f.attrs["stain"] = "DDX4"
 82            f.attrs["split"] = split
 83            f.attrs["resolution"] = RESOLUTION
 84            f.attrs["axes"] = "zyx"
 85            f.attrs["image_file"] = name
 86
 87            raw_dataset = f.create_dataset("raw", data=raw, compression="gzip")
 88            raw_dataset.attrs["resolution"] = RESOLUTION
 89            label_dataset = f.create_dataset("labels", data=labels.astype("uint32"), compression="gzip")
 90            label_dataset.attrs["resolution"] = RESOLUTION
 91        os.replace(temporary_path, output_path)
 92
 93    return output_dir
 94
 95
 96def get_ovarian_reserve_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 97    """Download the Ovarian Reserve dataset.
 98
 99    Args:
100        path: Filepath to a folder where the downloaded data will be saved.
101        download: Whether to download the data if it is not present.
102
103    Returns:
104        The filepath to the extracted data.
105    """
106    data_dir = os.path.join(path, "oocyte_training")
107    if os.path.exists(data_dir):
108        return data_dir
109
110    os.makedirs(path, exist_ok=True)
111    zip_path = os.path.join(path, "oocyte_training.zip")
112    util.download_source(zip_path, URL, download, CHECKSUM)
113    util.unzip(zip_path=zip_path, dst=path)
114
115    return data_dir
116
117
118def get_ovarian_reserve_paths(
119    path: Union[os.PathLike, str], split: Literal["train", "val"] = "train", download: bool = False,
120) -> List[str]:
121    """Get paths to the Ovarian Reserve data.
122
123    Args:
124        path: Filepath to a folder where the downloaded data will be saved.
125        split: The data split. Either 'train' or 'val'.
126        download: Whether to download the data if it is not present.
127
128    Returns:
129        List of filepaths for the h5 data.
130    """
131    if split not in SPLITS:
132        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
133
134    data_dir = get_ovarian_reserve_data(path, download)
135    output_dir = _create_h5(data_dir, split)
136    volume_paths = natsorted(glob(os.path.join(output_dir, "*.h5")))
137
138    if not volume_paths:
139        raise RuntimeError(f"Could not find any Ovarian Reserve data in {data_dir}.")
140
141    return volume_paths
142
143
144def get_ovarian_reserve_dataset(
145    path: Union[os.PathLike, str],
146    patch_shape: Tuple[int, int, int],
147    split: Literal["train", "val"] = "train",
148    offsets: Optional[List[List[int]]] = None,
149    boundaries: bool = False,
150    binary: bool = False,
151    download: bool = False,
152    **kwargs,
153) -> Dataset:
154    """Get the Ovarian Reserve dataset for 3d oocyte segmentation.
155
156    Args:
157        path: Filepath to a folder where the downloaded data will be saved.
158        patch_shape: The 3D patch shape to use for training.
159        split: The data split. Either 'train' or 'val'.
160        offsets: Offset values for affinity computation used as target.
161        boundaries: Whether to compute boundaries as the target.
162        binary: Whether to use a binary segmentation target.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
165
166    Returns:
167        The segmentation dataset.
168    """
169    if len(patch_shape) != 3:
170        raise ValueError(f"The Ovarian Reserve patch shape must be three-dimensional, got {patch_shape}.")
171
172    volume_paths = get_ovarian_reserve_paths(path, split, download)
173
174    kwargs, _ = util.add_instance_label_transform(
175        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
176    )
177    kwargs = util.ensure_transforms(ndim=3, **kwargs)
178
179    return torch_em.default_segmentation_dataset(
180        raw_paths=volume_paths,
181        raw_key="raw",
182        label_paths=volume_paths,
183        label_key="labels",
184        patch_shape=patch_shape,
185        ndim=3,
186        **kwargs,
187    )
188
189
190def get_ovarian_reserve_loader(
191    path: Union[os.PathLike, str],
192    batch_size: int,
193    patch_shape: Tuple[int, int, int],
194    split: Literal["train", "val"] = "train",
195    offsets: Optional[List[List[int]]] = None,
196    boundaries: bool = False,
197    binary: bool = False,
198    download: bool = False,
199    **kwargs,
200) -> DataLoader:
201    """Get the Ovarian Reserve dataloader for 3d oocyte segmentation.
202
203    Args:
204        path: Filepath to a folder where the downloaded data will be saved.
205        batch_size: The batch size for training.
206        patch_shape: The 3D patch shape to use for training.
207        split: The data split. Either 'train' or 'val'.
208        offsets: Offset values for affinity computation used as target.
209        boundaries: Whether to compute boundaries as the target.
210        binary: Whether to use a binary segmentation target.
211        download: Whether to download the data if it is not present.
212        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
213
214    Returns:
215        The DataLoader.
216    """
217    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
218    dataset = get_ovarian_reserve_dataset(
219        path=path,
220        patch_shape=patch_shape,
221        split=split,
222        offsets=offsets,
223        boundaries=boundaries,
224        binary=binary,
225        download=download,
226        **ds_kwargs,
227    )
228    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)