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 warnings
 12import subprocess
 13from glob import glob
 14from natsort import natsorted
 15from typing import Union, Tuple, List, Optional
 16
 17import torch_em
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21from .. import util
 22from .neurips_cell_seg import to_rgb
 23
 24
 25DOWNLOAD_SCRIPT = 'wget -c -nH -np -r -R "index.html*" --cut-dirs 4 ftp://ftp.cngb.org/pub/CNSA/data5/CNP0006370/Other/'
 26
 27# Files that are corrupted in the source archive (5 of 1044 pairs): one truncated ssDNA image and four
 28# instance masks that are not TIFF files. Their image-label pairs are skipped.
 29CORRUPTED_FILES = {
 30    "HH799999864GO_W8_36_62-img.tif",
 31    "X98668W8-x10151_y9176_w256_h256-instancemask.tif",
 32    "Z98801V5-x15153_y13890_w256_h256-instancemask.tif",
 33    "HH799999356_M1-x1536_y9216_w512_h512-instancemask.tif",
 34    "X97754Z3-x17024_y19696_w512_h512-instancemask.tif",
 35}
 36
 37CHOICES = ["10×Genomics_DAPI", "10×Genomics_HE", "DAPI", "HE", "mIF", "ssDNA"]
 38
 39
 40def get_cellbindb_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 41    """Download the CellBinDB dataset.
 42
 43    Args:
 44        path: Filepath to a folder where the data is downloaded.
 45        download: Whether to download the data if it is not present.
 46
 47    Returns:
 48        The filepath to the data.
 49    """
 50    data_dir = os.path.join(path, "Other")
 51    if os.path.exists(data_dir):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55
 56    if not download:
 57        raise AssertionError("The dataset is not found and download is set to 'False'.")
 58
 59    print(
 60        "Downloading the dataset takes several hours and is extremely (like very very) slow. "
 61        "Make sure you have consistent internet connection or run it in background over a cluster."
 62    )
 63    splits = DOWNLOAD_SCRIPT.split(" ")
 64    subprocess.run([*splits[:-1], "-P", os.path.abspath(path), splits[-1]])
 65    return data_dir
 66
 67
 68def get_cellbindb_paths(
 69    path: Union[os.PathLike, str], data_choice: Optional[Union[str, List[str]]] = None, download: bool = False
 70) -> Tuple[List[str], List[str]]:
 71    """Get paths to the CellBinDB data.
 72
 73    Args:
 74        path: Filepath to a folder where the data is downloaded.
 75        data_choice: The choice of datasets.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image data.
 80        List of filepaths for the label data.
 81    """
 82    data_dir = get_cellbindb_data(path, download)
 83
 84    if data_choice is None:
 85        data_choice = CHOICES
 86    else:
 87        if isinstance(data_choice, str):
 88            data_choice = [data_choice]
 89
 90    raw_paths, label_paths = [], []
 91    for dchoice in data_choice:
 92        assert dchoice in CHOICES, f"'{dchoice}' is not a valid data choice."
 93        raw_paths.extend(natsorted(glob(os.path.join(data_dir, dchoice, "*", "*-img.tif"))))
 94        label_paths.extend(natsorted(glob(os.path.join(data_dir, dchoice, "*", "*-instancemask.tif"))))
 95
 96    # NOTE: Some files are corrupted from source. Since it's just a few of them, let's bump them out.
 97    valid_paired_images = [
 98        (rp, lp) for rp, lp in zip(raw_paths, label_paths)
 99        if not {os.path.basename(rp), os.path.basename(lp)} & CORRUPTED_FILES
100        and _is_valid_image(rp) and _is_valid_image(lp)
101    ]
102    raw_paths, label_paths = zip(*valid_paired_images)
103    raw_paths, label_paths = list(raw_paths), list(label_paths)
104
105    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
106
107    return raw_paths, label_paths
108
109
110def _is_valid_image(im_path):
111    import tifffile
112
113    try:
114        _ = tifffile.imread(im_path)
115        return True
116    except Exception as e:
117        warnings.warn(f"Skipping the corrupted CellBinDB file '{im_path}': {type(e).__name__}: {e}")
118        return False
119
120
121def get_cellbindb_dataset(
122    path: Union[os.PathLike, str],
123    patch_shape: Tuple[int, int],
124    data_choice: Optional[Union[str, List[str]]] = None,
125    download: bool = False,
126    **kwargs
127) -> Dataset:
128    """Get the CellBinDB dataset for cell segmentation.
129
130    Args:
131        path: Filepath to a folder where the data is downloaded.
132        patch_shape: The patch shape to use for training.
133        data_choice: The choice of datasets.
134        download: Whether to download the data if it is not present.
135        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
136
137    Returns:
138        The segmentation dataset.
139    """
140    raw_paths, label_paths = get_cellbindb_paths(path, data_choice, download)
141
142    if "raw_transform" not in kwargs:
143        kwargs["raw_transform"] = torch_em.transform.get_raw_transform(augmentation2=to_rgb)
144
145    return torch_em.default_segmentation_dataset(
146        raw_paths=raw_paths,
147        raw_key=None,
148        label_paths=label_paths,
149        label_key=None,
150        is_seg_dataset=False,
151        ndim=2,
152        patch_shape=patch_shape,
153        **kwargs
154    )
155
156
157def get_cellbindb_loader(
158    path: Union[os.PathLike, str],
159    batch_size: int,
160    patch_shape: Tuple[int, int],
161    data_choice: Optional[Union[str, List[str]]] = None,
162    download: bool = False,
163    **kwargs
164) -> DataLoader:
165    """Get the CellBinDB dataloader for cell segmentation.
166
167    Args:
168        path: Filepath to a folder where the data is downloaded.
169        patch_shape: The patch shape to use for training.
170        data_choice: The choice of datasets.
171        download: Whether to download the data if it is not present.
172        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
173
174    Returns:
175        The DataLoader.
176    """
177    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
178    dataset = get_cellbindb_dataset(path, patch_shape, data_choice, download, **ds_kwargs)
179    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)