torch_em.data.datasets.light_microscopy.cell_acdc

  1"""This dataset contains phase-contrast time-lapse images of budding yeast (S. cerevisiae) with
  2per-frame single-cell instance segmentation, lineage tracking (mother-bud relationships) and
  3cell-cycle stage annotations, from the Cell-ACDC software test data.
  4
  5The dataset is hosted on Zenodo at https://doi.org/10.5281/zenodo.6795124.
  6The dataset is from the publication https://doi.org/10.1186/s12915-022-01372-6.
  7
  8Please cite it if you use this dataset for your research.
  9"""
 10
 11import os
 12from glob import glob
 13from typing import List, 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
 25URL = "https://zenodo.org/records/6795124/files/test_data_budding_yeast.zip"
 26CHECKSUM = "391b774888946ccd201be2ad0719a9ddd966b2d2215712ea3f4df9e52fbf9cc6"
 27
 28
 29def get_cell_acdc_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the Cell-ACDC budding yeast dataset.
 31
 32    Args:
 33        path: Filepath to a folder where the downloaded data will be saved.
 34        download: Whether to download the data if it is not present.
 35
 36    Returns:
 37        Filepath where the dataset is stored.
 38    """
 39    data_dir = os.path.join(path, "test_data_budding_yeast")
 40    if os.path.exists(data_dir):
 41        return data_dir
 42
 43    os.makedirs(path, exist_ok=True)
 44
 45    zip_path = os.path.join(path, "test_data_budding_yeast.zip")
 46    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 47    util.unzip(zip_path=zip_path, dst=path)
 48
 49    return data_dir
 50
 51
 52def _prepare_position(segm_path):
 53    # 'last_tracked_i.txt' records the last annotated frame; raw stacks always have more frames.
 54    last_tracked_path = segm_path.replace("_segm.npz", "_last_tracked_i.txt")
 55    with open(last_tracked_path) as f:
 56        n_frames = int(f.read().strip()) + 1
 57
 58    label_path = segm_path.replace(".npz", ".tif")
 59    if not os.path.exists(label_path):
 60        labels = np.load(segm_path)["arr_0"][:n_frames]
 61        tifffile.imwrite(label_path, labels)
 62
 63    raw_path = segm_path.replace("_segm.npz", "_phase_contr.tif")
 64    if tifffile.TiffFile(raw_path).series[0].shape[0] != n_frames:
 65        matched_raw_path = segm_path.replace("_segm.npz", "_phase_contr_matched.tif")
 66        if not os.path.exists(matched_raw_path):
 67            raw = tifffile.imread(raw_path)[:n_frames]
 68            tifffile.imwrite(matched_raw_path, raw)
 69        raw_path = matched_raw_path
 70
 71    return raw_path, label_path
 72
 73
 74def get_cell_acdc_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 75    """Get paths for the Cell-ACDC budding yeast dataset.
 76
 77    Args:
 78        path: Filepath to a folder where the downloaded data will be saved.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        List of filepaths for the raw phase-contrast images.
 83        List of filepaths for the instance segmentation and tracking labels.
 84    """
 85    data_dir = get_cell_acdc_data(path, download)
 86
 87    segm_paths = sorted(glob(os.path.join(data_dir, "TimeLapse_2D", "*_labeled", "Position_*", "Images", "*_segm.npz")))
 88    assert segm_paths, f"No labeled positions found at {data_dir}."
 89
 90    raw_paths, label_paths = [], []
 91    for segm_path in segm_paths:
 92        raw_path, label_path = _prepare_position(segm_path)
 93        raw_paths.append(raw_path)
 94        label_paths.append(label_path)
 95
 96    return raw_paths, label_paths
 97
 98
 99def get_cell_acdc_dataset(
100    path: Union[os.PathLike, str], patch_shape: Tuple[int, int, int], download: bool = False, **kwargs
101) -> Dataset:
102    """Get the Cell-ACDC dataset for budding yeast segmentation and tracking.
103
104    Args:
105        path: Filepath to a folder where the downloaded data will be saved.
106        patch_shape: The patch shape to use for training.
107        download: Whether to download the data if it is not present.
108        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
109
110    Returns:
111        The segmentation dataset.
112    """
113    raw_paths, label_paths = get_cell_acdc_paths(path, download)
114
115    kwargs = util.update_kwargs(kwargs, "ndim", 2)
116
117    return torch_em.default_segmentation_dataset(
118        raw_paths=raw_paths,
119        raw_key=None,
120        label_paths=label_paths,
121        label_key=None,
122        patch_shape=patch_shape,
123        is_seg_dataset=True,
124        **kwargs
125    )
126
127
128def get_cell_acdc_loader(
129    path: Union[os.PathLike, str],
130    batch_size: int,
131    patch_shape: Tuple[int, int, int],
132    download: bool = False,
133    **kwargs
134) -> DataLoader:
135    """Get the Cell-ACDC dataloader for budding yeast segmentation and tracking.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        batch_size: The batch size for training.
140        patch_shape: The patch shape to use for training.
141        download: Whether to download the data if it is not present.
142        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
143
144    Returns:
145        The DataLoader.
146    """
147    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
148    dataset = get_cell_acdc_dataset(path, patch_shape, download, **ds_kwargs)
149    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)