torch_em.data.datasets.light_microscopy.vibrio_cholerae

  1"""The Vibrio Cholerae dataset contains 3D confocal fluorescence microscopy images
  2of Vibrio cholerae biofilms with instance segmentation annotations for single-cell
  3segmentation.
  4
  5The dataset provides two annotation types for 5 biofilm volumes:
  6- semi-manual-annotation: all 5 volumes labeled via automated segmentation + manual correction.
  7- fully-manual-annotation: 1 cropped volume (biofilm_1) with fully manual annotations —
  8  intended as a held-out evaluation set.
  9
 10NOTE: The semi-manual labels are used by default for training. Whether all cells in each
 11volume are annotated should be verified against the paper before assuming dense coverage.
 12
 13The dataset is located at https://zenodo.org/records/7704410.
 14This dataset is from the publication https://doi.org/10.1111/mmi.15064.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import List, Tuple, Union
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://zenodo.org/records/7704410/files/ZENODO.zip"
 31CHECKSUM = "31edb3edbbd308261ead96fa6ec201aff4daf6a0fa8624462c0384e61d67d4c8"
 32
 33
 34def get_vibrio_cholerae_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the Vibrio Cholerae dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the downloaded data will be saved.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        The filepath to the training data directory.
 43    """
 44    data_dir = os.path.join(path, "training-data-from-experimentally-acquired-images")
 45    if os.path.exists(data_dir):
 46        return data_dir
 47
 48    os.makedirs(path, exist_ok=True)
 49    zip_path = os.path.join(path, "ZENODO.zip")
 50    util.download_source(zip_path, URL, download, checksum=CHECKSUM)
 51    util.unzip(zip_path, path)
 52
 53    return data_dir
 54
 55
 56def get_vibrio_cholerae_paths(
 57    path: Union[os.PathLike, str], download: bool = False,
 58) -> Tuple[List[str], List[str]]:
 59    """Get paths to the Vibrio Cholerae data.
 60
 61    Args:
 62        path: Filepath to a folder where the downloaded data will be saved.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    data_dir = get_vibrio_cholerae_data(path, download)
 70
 71    raw_paths = natsorted(glob(os.path.join(data_dir, "raw-data", "*_raw.tif")))
 72    label_paths = natsorted(glob(os.path.join(data_dir, "semi-manual-annotation", "*_labels.tif")))
 73
 74    if len(raw_paths) == 0:
 75        raise RuntimeError(
 76            f"No image files found in {os.path.join(data_dir, 'raw-data')}. "
 77            "Please check the dataset structure."
 78        )
 79    if len(raw_paths) != len(label_paths):
 80        raise RuntimeError(
 81            f"Number of images ({len(raw_paths)}) and labels ({len(label_paths)}) do not match."
 82        )
 83
 84    return raw_paths, label_paths
 85
 86
 87def get_vibrio_cholerae_dataset(
 88    path: Union[os.PathLike, str],
 89    patch_shape: Tuple[int, ...],
 90    download: bool = False,
 91    **kwargs,
 92) -> Dataset:
 93    """Get the Vibrio Cholerae dataset for 3D cell instance segmentation.
 94
 95    Args:
 96        path: Filepath to a folder where the downloaded data will be saved.
 97        patch_shape: The patch shape to use for training.
 98        download: Whether to download the data if it is not present.
 99        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
100
101    Returns:
102        The segmentation dataset.
103    """
104    raw_paths, label_paths = get_vibrio_cholerae_paths(path, download)
105
106    return torch_em.default_segmentation_dataset(
107        raw_paths=raw_paths,
108        raw_key=None,
109        label_paths=label_paths,
110        label_key=None,
111        patch_shape=patch_shape,
112        **kwargs,
113    )
114
115
116def get_vibrio_cholerae_loader(
117    path: Union[os.PathLike, str],
118    batch_size: int,
119    patch_shape: Tuple[int, ...],
120    download: bool = False,
121    **kwargs,
122) -> DataLoader:
123    """Get the Vibrio Cholerae dataloader for 3D cell instance segmentation.
124
125    Args:
126        path: Filepath to a folder where the downloaded data will be saved.
127        batch_size: The batch size for training.
128        patch_shape: The patch shape to use for training.
129        download: Whether to download the data if it is not present.
130        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
131
132    Returns:
133        The DataLoader.
134    """
135    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
136    dataset = get_vibrio_cholerae_dataset(path, patch_shape, download, **ds_kwargs)
137    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)