torch_em.data.datasets.light_microscopy.epicure

  1"""This dataset contains fluorescence microscopy movies of developing epithelial tissue with per-frame
  2instance segmentation and full-movie lineage tracking, curated with the EpiCure napari plugin. It covers
  3four model systems: the Drosophila notum, Drosophila abdomen histoblasts, the zebrafish telencephalon,
  4and the gastrulating quail.
  5
  6The dataset is hosted on Zenodo at https://doi.org/10.5281/zenodo.20607705 under the CC BY 4.0 license.
  7It is from the publication https://doi.org/10.1242/dev.205701.
  8
  9Please cite it if you use this dataset for your research.
 10"""
 11
 12import os
 13from typing import List, Literal, Optional, Tuple, Union
 14
 15import numpy as np
 16import tifffile
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URLS = {
 26    "notum": "https://zenodo.org/records/20607705/files/notumMovie.zip",
 27    "generalization": "https://zenodo.org/records/20607705/files/MovieEpitheliumFigure5.zip",
 28}
 29
 30CHECKSUMS = {
 31    "notum": "9cbd2399469ea6b267a706490acc3224f4a8f9e23f3bac4020ac39943a4e6dee",
 32    "generalization": "677c2d39931933a05d5b4aed33891c2edc6f893e0b9444c15559ed6a1e082fe9",
 33}
 34
 35SOURCES = ["notum", "histoblast", "telencephalon", "quail_gastrula"]
 36
 37# Per source: which archive holds it, its raw / label file relative to the extracted archive root,
 38# and the channel to keep for movies where the raw file has more than one channel.
 39SOURCE_INFO = {
 40    "notum": {
 41        "archive": "notum",
 42        "raw": "notumMovie/Ecad.tif",
 43        "label": "notumMovie/epics/Ecad_labels.tif",
 44        "channel_axis": None,
 45    },
 46    "histoblast": {
 47        "archive": "generalization",
 48        "raw": "data_generalisations/movie2/abdomen_maxz_z15-24_t1-60_crop.tif",
 49        "label": "data_generalisations/movie2/epics_corrected/abdomen_maxz_z15-24_t1-60_crop_labels.tif",
 50        "channel_axis": 1,
 51        "main_channel": 1,
 52    },
 53    "telencephalon": {
 54        "archive": "generalization",
 55        "raw": "data_generalisations/movie3/moji_merged_3to13_crop.tif",
 56        "label": "data_generalisations/movie3/epics_corrected/moji_merged_3to13_crop_labels.tif",
 57        "channel_axis": 1,
 58        "main_channel": 0,
 59    },
 60    "quail_gastrula": {
 61        "archive": "generalization",
 62        "raw": "data_generalisations/movie4/Composite_cropped.tif",
 63        "label": "data_generalisations/movie4/epics_correctedWithTA/Composite_cropped_labels.tif",
 64        "channel_axis": 0,
 65        "main_channel": 1,  # Channel 0 of the composite holds only noise, the membranes are in channel 1.
 66        "single_frame": True,
 67    },
 68}
 69
 70
 71def get_epicure_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 72    """Download the EpiCure dataset.
 73
 74    Args:
 75        path: Filepath to a folder where the downloaded data will be saved.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        Filepath where the dataset is stored.
 80    """
 81    os.makedirs(path, exist_ok=True)
 82
 83    for archive, marker in [("notum", "notumMovie"), ("generalization", "data_generalisations")]:
 84        if os.path.exists(os.path.join(path, marker)):
 85            continue
 86        zip_path = os.path.join(path, f"{archive}.zip")
 87        util.download_source(path=zip_path, url=URLS[archive], download=download, checksum=CHECKSUMS[archive])
 88        util.unzip(zip_path=zip_path, dst=path)
 89
 90    return path
 91
 92
 93def _prepare_raw(data_dir, source, info):
 94    raw_path = os.path.join(data_dir, info["raw"])
 95    if info["channel_axis"] is None:
 96        return raw_path
 97
 98    prepared_path = os.path.join(data_dir, "prepared", f"{source}_raw.tif")
 99    if os.path.exists(prepared_path):
100        return prepared_path
101
102    os.makedirs(os.path.dirname(prepared_path), exist_ok=True)
103    raw = tifffile.imread(raw_path)
104    raw = np.take(raw, info["main_channel"], axis=info["channel_axis"])
105    if info.get("single_frame"):
106        raw = raw[None]  # add a singleton frame axis so the array matches the movie sources.
107    tifffile.imwrite(prepared_path, raw)
108    return prepared_path
109
110
111def _prepare_label(data_dir, source, info):
112    label_path = os.path.join(data_dir, info["label"])
113    if not info.get("single_frame"):
114        return label_path
115
116    prepared_path = os.path.join(data_dir, "prepared", f"{source}_labels.tif")
117    if os.path.exists(prepared_path):
118        return prepared_path
119
120    os.makedirs(os.path.dirname(prepared_path), exist_ok=True)
121    label = tifffile.imread(label_path)[None]  # add a singleton frame axis, see '_prepare_raw'.
122    tifffile.imwrite(prepared_path, label)
123    return prepared_path
124
125
126def get_epicure_paths(
127    path: Union[os.PathLike, str], sources: Optional[List[str]] = None, download: bool = False,
128) -> Tuple[List[str], List[str]]:
129    f"""Get paths for the EpiCure dataset.
130
131    Args:
132        path: Filepath to a folder where the downloaded data will be saved.
133        sources: The model systems to use. By default uses all of them.
134            The available sources are: {', '.join(SOURCES)}.
135        download: Whether to download the data if it is not present.
136
137    Returns:
138        List of filepaths for the raw movies.
139        List of filepaths for the instance segmentation and tracking labels.
140    """
141    sources = SOURCES if sources is None else sources
142    for source in sources:
143        if source not in SOURCES:
144            raise ValueError(f"'{source}' is not a valid source, choose one of {SOURCES}.")
145
146    data_dir = get_epicure_data(path, download)
147
148    raw_paths, label_paths = [], []
149    for source in sources:
150        info = SOURCE_INFO[source]
151        raw_paths.append(_prepare_raw(data_dir, source, info))
152        label_paths.append(_prepare_label(data_dir, source, info))
153
154    return raw_paths, label_paths
155
156
157def get_epicure_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, int, int],
160    sources: Optional[List[Literal["notum", "histoblast", "telencephalon", "quail_gastrula"]]] = None,
161    download: bool = False,
162    **kwargs
163) -> Dataset:
164    """Get the EpiCure dataset for epithelial cell segmentation and tracking.
165
166    Args:
167        path: Filepath to a folder where the downloaded data will be saved.
168        patch_shape: The patch shape to use for training.
169        sources: The model systems to use. By default uses all of them.
170        download: Whether to download the data if it is not present.
171        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
172
173    Returns:
174        The segmentation dataset.
175    """
176    raw_paths, label_paths = get_epicure_paths(path, sources, download)
177
178    kwargs = util.update_kwargs(kwargs, "ndim", 2)
179
180    return torch_em.default_segmentation_dataset(
181        raw_paths=raw_paths,
182        raw_key=None,
183        label_paths=label_paths,
184        label_key=None,
185        patch_shape=patch_shape,
186        is_seg_dataset=True,
187        **kwargs
188    )
189
190
191def get_epicure_loader(
192    path: Union[os.PathLike, str],
193    batch_size: int,
194    patch_shape: Tuple[int, int, int],
195    sources: Optional[List[Literal["notum", "histoblast", "telencephalon", "quail_gastrula"]]] = None,
196    download: bool = False,
197    **kwargs
198) -> DataLoader:
199    """Get the EpiCure dataloader for epithelial cell segmentation and tracking.
200
201    Args:
202        path: Filepath to a folder where the downloaded data will be saved.
203        batch_size: The batch size for training.
204        patch_shape: The patch shape to use for training.
205        sources: The model systems to use. By default uses all of them.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
208
209    Returns:
210        The DataLoader.
211    """
212    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
213    dataset = get_epicure_dataset(path, patch_shape, sources, download, **ds_kwargs)
214    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)