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
 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], 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        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        List of filepaths for the image data.
 69        List of filepaths for the label data.
 70    """
 71    data_dir = get_phmamm_data(path, download)
 72
 73    raw_paths = natsorted(glob(os.path.join(data_dir, "Inputs", "*.tiff")))
 74    label_paths = natsorted(glob(os.path.join(data_dir, "ASTEC_Ground_truth", "*.tiff")))
 75    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 76
 77    return raw_paths, label_paths
 78
 79
 80def get_phmamm_dataset(
 81    path: Union[os.PathLike, str],
 82    patch_shape: Tuple[int, int, int],
 83    offsets: Optional[List[List[int]]] = None,
 84    boundaries: bool = False,
 85    binary: bool = False,
 86    download: bool = False,
 87    **kwargs
 88) -> Dataset:
 89    """Get the PhMamm dataset for cell segmentation in light-sheet microscopy.
 90
 91    Args:
 92        path: Filepath to a folder where the downloaded data will be saved.
 93        patch_shape: The patch shape to use for training.
 94        offsets: Offset values for affinity computation used as target.
 95        boundaries: Whether to compute boundaries as the target.
 96        binary: Whether to use a binary segmentation target.
 97        download: Whether to download the data if it is not present.
 98        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 99
100    Returns:
101        The segmentation dataset.
102    """
103    raw_paths, label_paths = get_phmamm_paths(path, download)
104
105    kwargs, _ = util.add_instance_label_transform(
106        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
107    )
108
109    return torch_em.default_segmentation_dataset(
110        raw_paths=raw_paths,
111        raw_key=None,
112        label_paths=label_paths,
113        label_key=None,
114        patch_shape=patch_shape,
115        **kwargs
116    )
117
118
119def get_phmamm_loader(
120    path: Union[os.PathLike, str],
121    batch_size: int,
122    patch_shape: Tuple[int, int, int],
123    offsets: Optional[List[List[int]]] = None,
124    boundaries: bool = False,
125    binary: bool = False,
126    download: bool = False,
127    **kwargs
128) -> DataLoader:
129    """Get the PhMamm dataloader for cell segmentation in light-sheet microscopy.
130
131    Args:
132        path: Filepath to a folder where the downloaded data will be saved.
133        batch_size: The batch size for training.
134        patch_shape: The patch shape to use for training.
135        offsets: Offset values for affinity computation used as target.
136        boundaries: Whether to compute boundaries as the target.
137        binary: Whether to use a binary segmentation target.
138        download: Whether to download the data if it is not present.
139        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
140
141    Returns:
142        The DataLoader.
143    """
144    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
145    dataset = get_phmamm_dataset(
146        path=path,
147        patch_shape=patch_shape,
148        offsets=offsets,
149        boundaries=boundaries,
150        binary=binary,
151        download=download,
152        **ds_kwargs,
153    )
154    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)