torch_em.data.datasets.light_microscopy.cvz_fluo

  1"""The CVZ-Fluo dataset contains annotations for cell and nuclei segmentation in
  2fluorescence microscopy images.
  3
  4The dataset is from the publication https://doi.org/10.1038/s41597-023-02108-z.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9from glob import glob
 10from tqdm import tqdm
 11from pathlib import Path
 12from natsort import natsorted
 13from typing import Union, Literal, Tuple, Optional, List
 14
 15import numpy as np
 16import imageio.v3 as imageio
 17from bioimage_cpp.segmentation import label as connected_components
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24from .neurips_cell_seg import to_rgb
 25
 26
 27URL = "https://www.synapse.org/Synapse:syn27624812/"
 28
 29
 30def get_cvz_fluo_data(path: Union[os.PathLike, str], download: bool = False):
 31    """Download the CVZ-Fluo dataset.
 32
 33    Args:
 34        path: Filepath to a folder where the downloaded data is saved.
 35        download: Whether to download the data if it is not present.
 36    """
 37    data_dir = os.path.join(path, r"Annotation Panel Table.xlsx")
 38    if not os.path.exists(data_dir):
 39        os.makedirs(path, exist_ok=True)
 40        # Download the dataset from 'synapse'.
 41        util.download_source_synapse(path=path, entity="syn27624812", download=download)
 42
 43    return
 44
 45
 46def _preprocess_labels(label_paths):
 47    neu_label_paths, to_process = [], []
 48
 49    # First, make simple checks to avoid redundant progress bar runs.
 50    for lpath in label_paths:
 51        neu_lpath = lpath.replace(".png", ".tif")
 52        neu_label_paths.append(neu_lpath)
 53
 54        if not os.path.exists(neu_lpath):
 55            to_process.append((lpath, neu_lpath))
 56
 57    if to_process:  # Next, process valid inputs.
 58        for lpath, neu_lpath in tqdm(to_process, desc="Preprocessing labels"):
 59            if not os.path.exists(lpath):  # HACK: Some paths have weird spacing nomenclature.
 60                lpath = Path(lpath).parent / rf" {os.path.basename(lpath)}"
 61
 62            label = imageio.imread(lpath)
 63            # The source masks are binary uint8. The connected components must not be cast back to that dtype:
 64            # crops have several hundred cells, so the ids would wrap at 255.
 65            instances = connected_components(label)
 66            dtype = "uint16" if instances.max() < np.iinfo("uint16").max else "uint32"
 67            imageio.imwrite(neu_lpath, instances.astype(dtype), compression="zlib")
 68
 69    return neu_label_paths
 70
 71
 72def get_cvz_fluo_paths(
 73    path: Union[os.PathLike, str],
 74    stain_choice: Literal["cell", "dapi"],
 75    data_choice: Optional[Literal["CODEX", "Vectra", "Zeiss"]] = None,
 76    download: bool = False,
 77) -> Tuple[List[str], List[str]]:
 78    """Get paths to the CVZ-Fluo data.
 79
 80    Args:
 81        path: Filepath to a folder where the downloaded data will be saved.
 82        download: Whether to download the data if it is not present.
 83
 84    Returns:
 85        List of filepaths for the image data.
 86        List of filepaths for the label data.
 87    """
 88    get_cvz_fluo_data(path, download)
 89
 90    if data_choice is None:
 91        data_choice = "**"
 92    else:
 93        if data_choice == "Zeiss" and stain_choice == "dapi":
 94            raise ValueError("'Zeiss' data does not have DAPI stained images.")
 95
 96        data_choice = f"{data_choice}/**"
 97
 98    if stain_choice not in ["cell", "dapi"]:
 99        raise ValueError(f"'{stain_choice}' is not a valid stain choice.")
100
101    raw_paths = natsorted(
102        glob(os.path.join(path, data_choice, f"*-Crop_{stain_choice.title()}_Png.png"), recursive=True)
103    )
104    label_paths = [p.replace("_Png.png", "_Mask_Png.png") for p in raw_paths]
105    label_paths = _preprocess_labels(label_paths)
106
107    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
108
109    return raw_paths, label_paths
110
111
112def get_cvz_fluo_dataset(
113    path: Union[os.PathLike, str],
114    patch_shape: Tuple[int, int],
115    stain_choice: Literal["cell", "dapi"],
116    data_choice: Optional[Literal["CODEX", "Vectra", "Zeiss"]] = None,
117    download: bool = False,
118    **kwargs
119) -> Dataset:
120    """Get the CVZ-Fluo dataset for cell and nucleus segmentation.
121
122    Args:
123        path: Filepath to a folder where the downloaded data will be saved.
124        patch_shape: The patch shape to use for training.
125        stain_choice: Decides for annotations based on staining. Either "cell" (for cells) or "dapi" (for nuclei).
126        data_choice: The choice of dataset.
127        download: Whether to download the data if it is not present.
128        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
129
130    Returns:
131        The segmentation dataset.
132    """
133    raw_paths, label_paths = get_cvz_fluo_paths(path, stain_choice, data_choice, download)
134
135    if "raw_transform" not in kwargs:
136        kwargs["raw_transform"] = torch_em.transform.get_raw_transform(augmentation2=to_rgb)
137
138    if "transform" not in kwargs:
139        kwargs["transform"] = torch_em.transform.get_augmentations(ndim=2)
140
141    return torch_em.default_segmentation_dataset(
142        raw_paths=raw_paths,
143        raw_key=None,
144        label_paths=label_paths,
145        label_key=None,
146        is_seg_dataset=False,
147        patch_shape=patch_shape,
148        **kwargs
149    )
150
151
152def get_cvz_fluo_loader(
153    path: Union[os.PathLike, str],
154    batch_size: int,
155    patch_shape: Tuple[int, int],
156    stain_choice: Literal["cell", "dapi"],
157    data_choice: Optional[Literal["CODEX", "Vectra", "Zeiss"]] = None,
158    download: bool = False,
159    **kwargs
160) -> DataLoader:
161    """Get the CVZ-Fluo dataloader for cell and nucleus segmentation.
162
163    Args:
164        path: Filepath to a folder where the downloaded data will be saved.
165        batch_size: The batch size for training
166        patch_shape: The patch shape to use for training.
167        stain_choice: Decides for annotations based on staining. Either "cell" (for cells) or "dapi" (for nuclei).
168        data_choice: The choice of dataset.
169        download: Whether to download the data if it is not present.
170        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
171
172    Returns:
173        The DataLoader.
174    """
175    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
176    dataset = get_cvz_fluo_dataset(path, patch_shape, stain_choice, data_choice, download, **ds_kwargs)
177    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)