torch_em.data.datasets.light_microscopy.flywing

  1"""The FlyWing dataset contains annotated fluorescence microscopy images of Drosophila wing epithelia.
  2
  3This loader uses the zero-noise segmentation release published with DenoiSeg. It contains 1428 training and
  4252 validation patches of shape 128 x 128, as well as 42 test images of shape 512 x 512. The data is derived from
  5the epithelial cell tracking benchmark introduced in https://doi.org/10.1007/978-3-030-11024-6_33.
  6
  7The dataset is located at https://doi.org/10.5281/zenodo.5156991 under the CC BY 4.0 license. This release does
  8not contain the eight complete time-lapse movies used by the original tracking benchmark or its movie-level splits.
  9It is from the publication https://doi.org/10.1007/978-3-030-66415-2_21. Please cite the dataset and publications
 10if you use this dataset in your research.
 11"""
 12
 13import os
 14from glob import glob
 15from typing import List, Literal, Optional, Tuple, Union
 16
 17import imageio.v3 as imageio
 18import numpy as np
 19from natsort import natsorted
 20from torch.utils.data import DataLoader, Dataset
 21from tqdm import tqdm
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://zenodo.org/api/records/5156991/files/Flywing_n0.zip/content"
 29CHECKSUM = "3fb49ba44e7e3e20b4fc3c77754f1bbff7184af7f343f23653f258d50e5d5aca"
 30
 31SPLIT_INFO = {
 32    "train": ("train/train_data.npz", "X_train", "Y_train", 1428),
 33    "val": ("train/train_data.npz", "X_val", "Y_val", 252),
 34    "test": ("test/test_data.npz", "X_test", "Y_test", 42),
 35}
 36
 37
 38def _get_split_paths(data_dir: str, split: str) -> Tuple[List[str], List[str]]:
 39    raw_paths = natsorted(glob(os.path.join(data_dir, split, "images", "*.tif")))
 40    label_paths = natsorted(glob(os.path.join(data_dir, split, "labels", "*.tif")))
 41    return raw_paths, label_paths
 42
 43
 44def _is_complete(data_dir: str, split: str) -> bool:
 45    raw_paths, label_paths = _get_split_paths(data_dir, split)
 46    expected_images = SPLIT_INFO[split][3]
 47    return (
 48        len(raw_paths) == expected_images
 49        and len(label_paths) == expected_images
 50        and [os.path.basename(path) for path in raw_paths] == [os.path.basename(path) for path in label_paths]
 51    )
 52
 53
 54def _preprocess_split(data_dir: str, split: str) -> None:
 55    npz_name, raw_key, label_key, expected_images = SPLIT_INFO[split]
 56    npz_path = os.path.join(data_dir, npz_name)
 57    if not os.path.exists(npz_path):
 58        raise RuntimeError(f"Could not find the FlyWing source data at '{npz_path}'.")
 59
 60    raw_dir = os.path.join(data_dir, split, "images")
 61    label_dir = os.path.join(data_dir, split, "labels")
 62    os.makedirs(raw_dir, exist_ok=True)
 63    os.makedirs(label_dir, exist_ok=True)
 64
 65    with np.load(npz_path) as data:
 66        raw = data[raw_key]
 67        labels = data[label_key]
 68        if raw.shape != labels.shape or len(raw) != expected_images:
 69            raise RuntimeError(
 70                f"Unexpected FlyWing data for split '{split}': raw={raw.shape}, labels={labels.shape}."
 71            )
 72
 73        for index, (image, instances) in tqdm(
 74            enumerate(zip(raw, labels)), total=expected_images, desc=f"Preprocessing FlyWing '{split}' split"
 75        ):
 76            filename = f"image_{index:04d}.tif"
 77            raw_path = os.path.join(raw_dir, filename)
 78            label_path = os.path.join(label_dir, filename)
 79            if not os.path.exists(raw_path):
 80                imageio.imwrite(raw_path, image, compression="zlib")
 81            if not os.path.exists(label_path):
 82                imageio.imwrite(label_path, instances, compression="zlib")
 83
 84
 85def get_flywing_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 86    """Download and preprocess the FlyWing segmentation dataset.
 87
 88    Args:
 89        path: Filepath to a folder where the downloaded data will be saved.
 90        download: Whether to download the data if it is not present.
 91
 92    Returns:
 93        The filepath to the preprocessed data.
 94    """
 95    data_dir = os.path.join(path, "Flywing_n0")
 96    if all(_is_complete(data_dir, split) for split in SPLIT_INFO):
 97        return data_dir
 98
 99    os.makedirs(path, exist_ok=True)
100    if not all(os.path.exists(os.path.join(data_dir, info[0])) for info in SPLIT_INFO.values()):
101        zip_path = os.path.join(path, "Flywing_n0.zip")
102        util.download_source(zip_path, URL, download, CHECKSUM)
103        util.unzip(zip_path, path)
104
105    for split in SPLIT_INFO:
106        if not _is_complete(data_dir, split):
107            _preprocess_split(data_dir, split)
108
109    incomplete_splits = [split for split in SPLIT_INFO if not _is_complete(data_dir, split)]
110    if incomplete_splits:
111        raise RuntimeError(f"FlyWing preprocessing failed for splits: {incomplete_splits}.")
112    return data_dir
113
114
115def get_flywing_paths(
116    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False,
117) -> Tuple[List[str], List[str]]:
118    """Get paths to the FlyWing fluorescence images and cell instance labels.
119
120    Args:
121        path: Filepath to a folder where the downloaded data will be saved.
122        split: The data split. Either 'train', 'val' or 'test'.
123        download: Whether to download the data if it is not present.
124
125    Returns:
126        The image paths and corresponding label paths.
127    """
128    if split not in SPLIT_INFO:
129        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLIT_INFO)}.")
130
131    data_dir = get_flywing_data(path, download)
132    raw_paths, label_paths = _get_split_paths(data_dir, split)
133    return raw_paths, label_paths
134
135
136def get_flywing_dataset(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int],
139    split: Literal["train", "val", "test"],
140    offsets: Optional[List[List[int]]] = None,
141    boundaries: bool = False,
142    binary: bool = False,
143    download: bool = False,
144    **kwargs,
145) -> Dataset:
146    """Get the FlyWing dataset for epithelial cell instance segmentation.
147
148    Args:
149        path: Filepath to a folder where the downloaded data will be saved.
150        patch_shape: The 2D patch shape to use for training.
151        split: The data split. Either 'train', 'val' or 'test'.
152        offsets: Offset values for affinity computation used as target.
153        boundaries: Whether to compute boundaries as the target.
154        binary: Whether to use a binary segmentation target.
155        download: Whether to download the data if it is not present.
156        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
157
158    Returns:
159        The segmentation dataset.
160    """
161    if len(patch_shape) != 2:
162        raise ValueError(f"The FlyWing patch shape must be two-dimensional, got {patch_shape}.")
163
164    raw_paths, label_paths = get_flywing_paths(path, split, download)
165    kwargs, _ = util.add_instance_label_transform(
166        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
167    )
168    kwargs = util.ensure_transforms(ndim=2, **kwargs)
169
170    return torch_em.default_segmentation_dataset(
171        raw_paths=raw_paths,
172        raw_key=None,
173        label_paths=label_paths,
174        label_key=None,
175        patch_shape=patch_shape,
176        is_seg_dataset=False,
177        **kwargs,
178    )
179
180
181def get_flywing_loader(
182    path: Union[os.PathLike, str],
183    batch_size: int,
184    patch_shape: Tuple[int, int],
185    split: Literal["train", "val", "test"],
186    offsets: Optional[List[List[int]]] = None,
187    boundaries: bool = False,
188    binary: bool = False,
189    download: bool = False,
190    **kwargs,
191) -> DataLoader:
192    """Get the FlyWing dataloader for epithelial cell instance segmentation.
193
194    Args:
195        path: Filepath to a folder where the downloaded data will be saved.
196        batch_size: The batch size for training.
197        patch_shape: The 2D patch shape to use for training.
198        split: The data split. Either 'train', 'val' or 'test'.
199        offsets: Offset values for affinity computation used as target.
200        boundaries: Whether to compute boundaries as the target.
201        binary: Whether to use a binary segmentation target.
202        download: Whether to download the data if it is not present.
203        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
204
205    Returns:
206        The DataLoader.
207    """
208    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
209    dataset = get_flywing_dataset(
210        path=path,
211        patch_shape=patch_shape,
212        split=split,
213        offsets=offsets,
214        boundaries=boundaries,
215        binary=binary,
216        download=download,
217        **ds_kwargs,
218    )
219    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)