torch_em.data.datasets.light_microscopy.cellbindb

  1"""CellBinDB contains annotations for cell segmentation in multi-modal images.
  2- Consists of DAPI, ssDNA, H&E, and mIF staining.
  3- Covers more than 30 normal and diseased tissue types from human and mouse samples.
  4
  5The dataset is located at https://db.cngb.org/search/project/CNP0006370/.
  6This dataset is from the publication https://doi.org/10.1101/2024.11.20.619750.
  7Please cite it if you use this dataset for your research.
  8"""
  9
 10import os
 11import subprocess
 12from glob import glob
 13from natsort import natsorted
 14from typing import Union, Tuple, List, Optional
 15
 16import torch_em
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20from .. import util
 21from .neurips_cell_seg import to_rgb
 22
 23
 24DOWNLOAD_SCRIPT = 'wget -c -nH -np -r -R "index.html*" --cut-dirs 4 ftp://ftp.cngb.org/pub/CNSA/data5/CNP0006370/Other/'
 25
 26CHOICES = ["10×Genomics_DAPI", "10×Genomics_HE", "DAPI", "HE", "mIF", "ssDNA"]
 27
 28
 29def get_cellbindb_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the CellBinDB dataset.
 31
 32    Args:
 33        path: Filepath to a folder where the data is downloaded.
 34        download: Whether to download the data if it is not present.
 35
 36    Returns:
 37        The filepath to the data.
 38    """
 39    data_dir = os.path.join(path, "Other")
 40    if os.path.exists(data_dir):
 41        return data_dir
 42
 43    os.makedirs(path, exist_ok=True)
 44
 45    if not download:
 46        raise AssertionError("The dataset is not found and download is set to 'False'.")
 47
 48    print(
 49        "Downloading the dataset takes several hours and is extremely (like very very) slow. "
 50        "Make sure you have consistent internet connection or run it in background over a cluster."
 51    )
 52    splits = DOWNLOAD_SCRIPT.split(" ")
 53    subprocess.run([*splits[:-1], "-P", os.path.abspath(path), splits[-1]])
 54    return data_dir
 55
 56
 57def get_cellbindb_paths(
 58    path: Union[os.PathLike, str], data_choice: Optional[Union[str, List[str]]] = None, download: bool = False
 59) -> Tuple[List[str], List[str]]:
 60    """Get paths to the CellBinDB data.
 61
 62    Args:
 63        path: Filepath to a folder where the data is downloaded.
 64        data_choice: The choice of datasets.
 65        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        List of filepaths for the image data.
 69        List of filepaths for the label data.
 70    """
 71    data_dir = get_cellbindb_data(path, download)
 72
 73    if data_choice is None:
 74        data_choice = CHOICES
 75    else:
 76        if isinstance(data_choice, str):
 77            data_choice = [data_choice]
 78
 79    raw_paths, label_paths = [], []
 80    for dchoice in data_choice:
 81        assert dchoice in CHOICES, f"'{dchoice}' is not a valid data choice."
 82        raw_paths.extend(natsorted(glob(os.path.join(data_dir, dchoice, "*", "*-img.tif"))))
 83        label_paths.extend(natsorted(glob(os.path.join(data_dir, dchoice, "*", "*-instancemask.tif"))))
 84
 85    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 86
 87    return raw_paths, label_paths
 88
 89
 90def get_cellbindb_dataset(
 91    path: Union[os.PathLike, str],
 92    patch_shape: Tuple[int, int],
 93    data_choice: Optional[Union[str, List[str]]] = None,
 94    download: bool = False,
 95    **kwargs
 96) -> Dataset:
 97    """Get the CellBinDB dataset for cell segmentation.
 98
 99    Args:
100        path: Filepath to a folder where the data is downloaded.
101        patch_shape: The patch shape to use for training.
102        data_choice: The choice of datasets.
103        download: Whether to download the data if it is not present.
104        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
105
106    Returns:
107        The segmentation dataset.
108    """
109    raw_paths, label_paths = get_cellbindb_paths(path, data_choice, download)
110
111    if "raw_transform" not in kwargs:
112        kwargs["raw_transform"] = torch_em.transform.get_raw_transform(augmentation2=to_rgb)
113
114    return torch_em.default_segmentation_dataset(
115        raw_paths=raw_paths,
116        raw_key=None,
117        label_paths=label_paths,
118        label_key=None,
119        is_seg_dataset=False,
120        ndim=2,
121        patch_shape=patch_shape,
122        **kwargs
123    )
124
125
126def get_cellbindb_loader(
127    path: Union[os.PathLike, str],
128    batch_size: int,
129    patch_shape: Tuple[int, int],
130    data_choice: Optional[Union[str, List[str]]] = None,
131    download: bool = False,
132    **kwargs
133) -> DataLoader:
134    """Get the CellBinDB dataloader for cell segmentation.
135
136    Args:
137        path: Filepath to a folder where the data is downloaded.
138        patch_shape: The patch shape to use for training.
139        data_choice: The choice of datasets.
140        download: Whether to download the data if it is not present.
141        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
142
143    Returns:
144        The DataLoader.
145    """
146    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
147    dataset = get_cellbindb_dataset(path, patch_shape, data_choice, download, **ds_kwargs)
148    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)