torch_em.data.datasets.light_microscopy.brain_organoids

  1"""The Brain Organoids dataset contains annotations for organoid segmentation in
  22d brightfield images of brain organoids.
  3
  4This dataset is from the publication https://doi.org/10.1038/s41597-024-03330-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 natsort import natsorted
 12from typing import Union, Tuple, List
 13
 14import numpy as np
 15import imageio.v3 as imageio
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22from .neurips_cell_seg import to_rgb
 23
 24
 25URL = "https://zenodo.org/records/10301912/files/data.zip"
 26CHECKSUM = "bc2ed56717a65ccd49e27bac92c3b714ca4bb245299698b68baa599c9d510a26"
 27
 28
 29def get_brain_organoids_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the Brain Organoids 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 data is downloaded.
 38    """
 39    data_dir = os.path.join(path, "data")
 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, "data.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 get_brain_organoids_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 53    """Get the Brain Organoids data.
 54
 55    Args:
 56        path: Filepath to a folder where the downloaded data will be saved.
 57        download: Whether to download the data if it is not present.
 58
 59    Returns:
 60        List of filepaths for the image data.
 61        List of filepaths for the label data.
 62    """
 63    data_dir = get_brain_organoids_data(path, download)
 64
 65    raw_paths = natsorted(glob(os.path.join(data_dir, "imgs", "*")))
 66    label_paths = natsorted(glob(os.path.join(data_dir, "labels", "*.npy")))
 67
 68    preprocessed_dir = os.path.join(data_dir, "labels_preprocessed")
 69    os.makedirs(preprocessed_dir, exist_ok=True)
 70
 71    neu_label_paths = []
 72    for lpath in tqdm(label_paths, desc="Preprocessing labels"):
 73        neu_lpath = lpath.replace("labels", "labels_preprocessed").replace(".npy", ".tif")
 74        neu_label_paths.append(neu_lpath)
 75        if os.path.exists(neu_lpath):
 76            continue
 77
 78        label = np.load(lpath)
 79        imageio.imwrite(neu_lpath, label)
 80
 81    return raw_paths, neu_label_paths
 82
 83
 84def get_brain_organoids_dataset(
 85    path: Union[os.PathLike, str], patch_shape: Tuple[int, int], download: bool = False, **kwargs
 86) -> Dataset:
 87    """Get the Brain Organoids dataset for organoid segmentation.
 88
 89    Args:
 90        path: Filepath to a folder where the downloaded data will be saved.
 91        patch_shape: The patch shape to use for training.
 92        download: Whether to download the data if it is not present.
 93        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 94
 95    Returns:
 96        The segmentation dataset.
 97    """
 98    raw_paths, label_paths = get_brain_organoids_paths(path, download)
 99
100    if "raw_transform" not in kwargs:
101        kwargs["raw_transform"] = torch_em.transform.get_raw_transform(augmentation2=to_rgb)
102
103    return torch_em.default_segmentation_dataset(
104        raw_paths=raw_paths,
105        raw_key=None,
106        label_paths=label_paths,
107        label_key=None,
108        is_seg_dataset=False,
109        with_channels=True,
110        patch_shape=patch_shape,
111        **kwargs
112    )
113
114
115def get_brain_organoids_loader(
116    path: Union[os.PathLike, str], batch_size: int, patch_shape: Tuple[int, int], download: bool = False, **kwargs
117) -> DataLoader:
118    """Get the Brain Organoids dataloader for organoid segmentation.
119
120    Args:
121        path: Filepath to a folder where the downloaded data will be saved.
122        patch_shape: The patch shape to use for training.
123        download: Whether to download the data if it is not present.
124        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
125
126    Returns:
127        The DataLoader.
128    """
129    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
130    dataset = get_brain_organoids_dataset(path, patch_shape, download, **ds_kwargs)
131    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)