torch_em.data.datasets.light_microscopy.mouse_embryo

  1"""This dataset contains confocal microscopy stacks of a mouse embryo
  2with annotations for cell and nucleus segmentation.
  3
  4This dataset is part of the publication https://doi.org/10.15252/embj.2022113280.
  5Please cite it if you use this data in your research.
  6"""
  7
  8import os
  9from glob import glob
 10from typing import List, Optional, Tuple, Union
 11
 12from torch.utils.data import Dataset, DataLoader
 13
 14import torch_em
 15
 16from .. import util
 17
 18
 19URL = "https://zenodo.org/record/6546550/files/MouseEmbryos.zip?download=1"
 20CHECKSUM = "bf24df25e5f919489ce9e674876ff27e06af84445c48cf2900f1ab590a042622"
 21
 22
 23def get_mouse_embryo_data(path: Union[os.PathLike, str], download: bool) -> str:
 24    """Download the mouse embryo dataset.
 25
 26    Args:
 27        path: Filepath to a folder where the downloaded data will be saved.
 28        download: Whether to download the data if it is not present.
 29
 30    Returns:
 31        The filepath for the downloaded data.
 32    """
 33    if os.path.exists(path):
 34        return path
 35    os.makedirs(path, exist_ok=True)
 36    tmp_path = os.path.join(path, "mouse_embryo.zip")
 37    util.download_source(tmp_path, URL, download, CHECKSUM)
 38    util.unzip(tmp_path, path, remove=True)
 39    # Remove empty volume.
 40    os.remove(os.path.join(path, "Membrane", "train", "fused_paral_stack0_chan2_tp00073_raw_crop_bg_noise.h5"))
 41    return path
 42
 43
 44def get_mouse_embryo_paths(path: Union[os.PathLike, str], name: str, split: str, download: bool = False) -> List[str]:
 45    """Get paths to the Mouse Embryo data.
 46
 47    Args:
 48        path: Filepath to a folder where the downloaded data will be saved.
 49        name: The name of the segmentation task. Either 'membrane' or 'nuclei'.
 50        split: The split to use for the dataset. Either 'train' or 'val'.
 51        download: Whether to download the data if it is not present.
 52
 53    Returns:
 54        List of filepaths for the stored data.
 55    """
 56    get_mouse_embryo_data(path, download)
 57
 58    # the naming of the data is inconsistent: membrane has val, nuclei has test;
 59    # we treat nuclei:test as val
 60    split_ = "test" if name == "nuclei" and split == "val" else split
 61    file_paths = glob(os.path.join(path, name.capitalize(), split_, "*.h5"))
 62    file_paths.sort()
 63
 64    return file_paths
 65
 66
 67def get_mouse_embryo_dataset(
 68    path: Union[os.PathLike, str],
 69    name: str,
 70    split: str,
 71    patch_shape: Tuple[int, int],
 72    download: bool = False,
 73    offsets: Optional[List[List[int]]] = None,
 74    boundaries: bool = False,
 75    binary: bool = False,
 76    **kwargs,
 77) -> Dataset:
 78    """Get the mouse embryo dataset for cell or nucleus segmentation.
 79
 80    Args:
 81        path: Filepath to a folder where the downloaded data will be saved.
 82        name: The name of the segmentation task. Either 'membrane' or 'nuclei'.
 83        split: The split to use for the dataset. Either 'train' or 'val'.
 84        patch_shape: The patch shape to use for training.
 85        download: Whether to download the data if it is not present.
 86        offsets: Offset values for affinity computation used as target.
 87        boundaries: Whether to compute boundaries as the target.
 88        binary: Whether to use a binary segmentation target.
 89        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 90
 91    Returns:
 92       The segmentation dataset.
 93    """
 94    assert name in ("membrane", "nuclei")
 95    assert split in ("train", "val")
 96    assert len(patch_shape) == 3
 97
 98    file_paths = get_mouse_embryo_paths(path, name, split, download)
 99
100    kwargs, _ = util.add_instance_label_transform(
101        kwargs,
102        add_binary_target=binary,
103        binary=binary,
104        boundaries=boundaries,
105        offsets=offsets,
106        binary_is_exclusive=False
107    )
108
109    return torch_em.default_segmentation_dataset(
110        raw_paths=file_paths,
111        raw_key="raw",
112        label_paths=file_paths,
113        label_key="label",
114        patch_shape=patch_shape,
115        **kwargs
116    )
117
118
119def get_mouse_embryo_loader(
120    path: Union[os.PathLike, str],
121    name: str,
122    split: str,
123    patch_shape: Tuple[int, int, int],
124    batch_size: int,
125    download: bool = False,
126    offsets: Optional[List[List[int]]] = None,
127    boundaries: bool = False,
128    binary: bool = False,
129    **kwargs,
130) -> DataLoader:
131    """Get the mouse embryo dataset for cell or nucleus segmentation.
132
133    Args:
134        path: Filepath to a folder where the downloaded data will be saved.
135        name: The name of the segmentation task. Either 'membrane' or 'nuclei'.
136        split: The split to use for the dataset. Either 'train' or 'val'.
137        patch_shape: The patch shape to use for training.
138        batch_size: The batch size for training.
139        download: Whether to download the data if it is not present.
140        offsets: Offset values for affinity computation used as target.
141        boundaries: Whether to compute boundaries as the target.
142        binary: Whether to use a binary segmentation target.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
144
145    Returns:
146        The DataLoader.
147    """
148    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
149    dataset = get_mouse_embryo_dataset(
150        path, name, split, patch_shape, download=download, offsets=offsets,
151        boundaries=boundaries, binary=binary, **ds_kwargs
152    )
153    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)