torch_em.data.datasets.light_microscopy.phmamm

  1"""The PhMamm dataset contains 3D light-sheet microscopy volumes of Phallusia mammillata
  2embryos with cell membrane segmentations.
  3
  4The dataset is located at https://figshare.com/articles/dataset/3D_Mask_R-CNN_data/26973085.
  5The original data is from the publication https://doi.org/10.1126/science.aar5663.
  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 Union, Tuple, Optional, List, Sequence
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URLS = {
 22    "inputs": "https://ndownloader.figshare.com/files/51130115",
 23    "ground_truth": "https://ndownloader.figshare.com/files/51130100",
 24}
 25CHECKSUMS = {
 26    "inputs": None,
 27    "ground_truth": None,
 28}
 29
 30
 31def get_phmamm_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the PhMamm dataset.
 33
 34    Args:
 35        path: Filepath to a folder where the downloaded data will be saved.
 36        download: Whether to download the data if it is not present.
 37
 38    Returns:
 39        Filepath where the dataset is stored.
 40    """
 41    data_dir = os.path.join(path, "data")
 42    if os.path.exists(data_dir):
 43        return data_dir
 44
 45    os.makedirs(path, exist_ok=True)
 46
 47    inputs_zip = os.path.join(path, "Inputs.zip")
 48    util.download_source(inputs_zip, URLS["inputs"], download, checksum=CHECKSUMS["inputs"])
 49    util.unzip(inputs_zip, data_dir, remove=True)
 50
 51    gt_zip = os.path.join(path, "ASTEC_Ground_truth.zip")
 52    util.download_source(gt_zip, URLS["ground_truth"], download, checksum=CHECKSUMS["ground_truth"])
 53    util.unzip(gt_zip, data_dir, remove=True)
 54
 55    return data_dir
 56
 57
 58def get_phmamm_paths(
 59    path: Union[os.PathLike, str], timepoints: Optional[Sequence[int]] = None, download: bool = False,
 60) -> Tuple[List[str], List[str]]:
 61    """Get paths to the PhMamm data.
 62
 63    Args:
 64        path: Filepath to a folder where the downloaded data will be saved.
 65        timepoints: The timepoints of the time-lapse to restrict to, e.g. range(1, 81). By default all 100 are used.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        List of filepaths for the image data.
 70        List of filepaths for the label data.
 71    """
 72    data_dir = get_phmamm_data(path, download)
 73
 74    raw_paths = natsorted(glob(os.path.join(data_dir, "Inputs", "*.tiff")))
 75    label_paths = natsorted(glob(os.path.join(data_dir, "ASTEC_Ground_truth", "*.tiff")))
 76    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 77
 78    if timepoints is not None:
 79        timepoints = set(timepoints)
 80        keep = [int(os.path.basename(p).split("_t")[1][:3]) in timepoints for p in raw_paths]
 81        raw_paths = [p for p, k in zip(raw_paths, keep) if k]
 82        label_paths = [p for p, k in zip(label_paths, keep) if k]
 83
 84    return raw_paths, label_paths
 85
 86
 87def get_phmamm_dataset(
 88    path: Union[os.PathLike, str],
 89    patch_shape: Tuple[int, int, int],
 90    offsets: Optional[List[List[int]]] = None,
 91    boundaries: bool = False,
 92    binary: bool = False,
 93    timepoints: Optional[Sequence[int]] = None,
 94    download: bool = False,
 95    **kwargs
 96) -> Dataset:
 97    """Get the PhMamm dataset for cell segmentation in light-sheet microscopy.
 98
 99    Args:
100        path: Filepath to a folder where the downloaded data will be saved.
101        patch_shape: The patch shape to use for training.
102        offsets: Offset values for affinity computation used as target.
103        boundaries: Whether to compute boundaries as the target.
104        binary: Whether to use a binary segmentation target.
105        timepoints: The timepoints of the time-lapse to restrict to, e.g. range(1, 81). By default all 100 are used.
106        download: Whether to download the data if it is not present.
107        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
108
109    Returns:
110        The segmentation dataset.
111    """
112    raw_paths, label_paths = get_phmamm_paths(path, timepoints, download)
113
114    kwargs, _ = util.add_instance_label_transform(
115        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
116    )
117
118    return torch_em.default_segmentation_dataset(
119        raw_paths=raw_paths,
120        raw_key=None,
121        label_paths=label_paths,
122        label_key=None,
123        patch_shape=patch_shape,
124        **kwargs
125    )
126
127
128def get_phmamm_loader(
129    path: Union[os.PathLike, str],
130    batch_size: int,
131    patch_shape: Tuple[int, int, int],
132    offsets: Optional[List[List[int]]] = None,
133    boundaries: bool = False,
134    binary: bool = False,
135    timepoints: Optional[Sequence[int]] = None,
136    download: bool = False,
137    **kwargs
138) -> DataLoader:
139    """Get the PhMamm dataloader for cell segmentation in light-sheet microscopy.
140
141    Args:
142        path: Filepath to a folder where the downloaded data will be saved.
143        batch_size: The batch size for training.
144        patch_shape: The patch shape to use for training.
145        offsets: Offset values for affinity computation used as target.
146        boundaries: Whether to compute boundaries as the target.
147        binary: Whether to use a binary segmentation target.
148        timepoints: The timepoints of the time-lapse to restrict to, e.g. range(1, 81). By default all 100 are used.
149        download: Whether to download the data if it is not present.
150        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
151
152    Returns:
153        The DataLoader.
154    """
155    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
156    dataset = get_phmamm_dataset(
157        path=path,
158        patch_shape=patch_shape,
159        offsets=offsets,
160        boundaries=boundaries,
161        binary=binary,
162        timepoints=timepoints,
163        download=download,
164        **ds_kwargs,
165    )
166    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)