torch_em.data.datasets.light_microscopy.bbbc024

  1"""The BBBC024 dataset contains synthetic 3D fluorescence microscopy images of HL60 cell nuclei
  2with instance segmentation ground truth.
  3
  4The images were generated with a virtual microscope (CytoPacq) imitating a Zeiss S100 confocal
  5microscope. Each image contains 20 nuclei and has a shape of (129, 565, 807) voxels (ZYX).
  6The dataset is organized in 8 subsets of 30 images each: the nuclei cluster with a probability
  7of 0%, 25%, 50% or 75%, and every clustering level is provided in a low SNR and a high SNR variant.
  8The ground truth is a 16-bit labeled volume with the ids 1 to 20 for the individual nuclei and 0 for background.
  9
 10The dataset is located at https://bbbc.broadinstitute.org/BBBC024.
 11This dataset is from the publication https://doi.org/10.1002/cyto.a.20714.
 12Please cite it if you use this dataset in your research.
 13"""
 14
 15import os
 16from glob import glob
 17from natsort import natsorted
 18from typing import List, Literal, Tuple, Union
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://data.broadinstitute.org/bbbc/BBBC024/BBBC024_v1_{subset}_images_TIFF.zip"
 28CHECKSUMS = {
 29    "c00_lowSNR": "c725828c267c347245e94f2dd4bfd65222fbbb526eb020fd5143cf7c7400e495",
 30    "c00_highSNR": "2a18179503dbd00172c6ea62a7e55752a5e86816b44faf8d1abc8f4cc9346c7b",
 31    "c25_lowSNR": "7a6a302ab8335657dfb98d8d129691204f333b87f00abdedc5737cc27c80101a",
 32    "c25_highSNR": "83662cad96f01d6956eba40a3b4b1d1e05f90400ff4fb63406d6bcebbac63a6b",
 33    "c50_lowSNR": "80f1b11f985fae6052619083c2a8bff0dc6a4a75c1beb2d7b2d90257f59c69f7",
 34    "c50_highSNR": "ba8fe5d26adec8d6e7ab224b06c1261218825ab1b3bf260a30515d6ebf77fa21",
 35    "c75_lowSNR": "b0e2f19d70012e2be4ccf7a7f1d0502858a45d60d1af01d2e32fad7c86263af3",
 36    "c75_highSNR": "2d56055a09d7dd22911f593dcfdef93be074ae27d58d7f7d32d089a25217fc09",
 37}
 38
 39
 40def _get_subset_name(clustering, snr):
 41    if clustering not in (0, 25, 50, 75):
 42        raise ValueError(f"'{clustering}' is not a valid clustering probability. Choose from 0, 25, 50 or 75.")
 43    if snr not in ("low", "high"):
 44        raise ValueError(f"'{snr}' is not a valid SNR level. Choose from 'low' or 'high'.")
 45    return f"c{clustering:02d}_{snr}SNR"
 46
 47
 48def get_bbbc024_data(
 49    path: Union[os.PathLike, str],
 50    clustering: Literal[0, 25, 50, 75] = 0,
 51    snr: Literal["low", "high"] = "high",
 52    download: bool = False,
 53) -> str:
 54    """Download the BBBC024 dataset.
 55
 56    Args:
 57        path: Filepath to a folder where the downloaded data will be saved.
 58        clustering: The clustering probability of the nuclei in percent. One of 0, 25, 50 or 75.
 59        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
 60        download: Whether to download the data if it is not present.
 61
 62    Returns:
 63        Filepath where the data for the chosen subset is stored.
 64    """
 65    subset = _get_subset_name(clustering, snr)
 66    data_dir = os.path.join(path, "BBBC024", subset)
 67    if os.path.exists(data_dir):
 68        return data_dir
 69
 70    os.makedirs(path, exist_ok=True)
 71    zip_path = os.path.join(path, f"BBBC024_v1_{subset}_images_TIFF.zip")
 72    util.download_source(zip_path, URL.format(subset=subset), download, CHECKSUMS[subset])
 73    util.unzip(zip_path, data_dir)
 74
 75    return data_dir
 76
 77
 78def get_bbbc024_paths(
 79    path: Union[os.PathLike, str],
 80    clustering: Literal[0, 25, 50, 75] = 0,
 81    snr: Literal["low", "high"] = "high",
 82    download: bool = False,
 83) -> Tuple[List[str], List[str]]:
 84    """Get paths to the BBBC024 data.
 85
 86    Args:
 87        path: Filepath to a folder where the downloaded data will be saved.
 88        clustering: The clustering probability of the nuclei in percent. One of 0, 25, 50 or 75.
 89        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
 90        download: Whether to download the data if it is not present.
 91
 92    Returns:
 93        List of filepaths for the image data.
 94        List of filepaths for the label data.
 95    """
 96    data_dir = get_bbbc024_data(path, clustering, snr, download)
 97
 98    raw_paths = natsorted(glob(os.path.join(data_dir, "image-final_*.tif")))
 99    label_paths = natsorted(glob(os.path.join(data_dir, "image-labels_*.tif")))
100    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
101
102    return raw_paths, label_paths
103
104
105def get_bbbc024_dataset(
106    path: Union[os.PathLike, str],
107    patch_shape: Tuple[int, ...],
108    clustering: Literal[0, 25, 50, 75] = 0,
109    snr: Literal["low", "high"] = "high",
110    resize_inputs: bool = False,
111    download: bool = False,
112    **kwargs
113) -> Dataset:
114    """Get the BBBC024 dataset for nucleus segmentation.
115
116    Args:
117        path: Filepath to a folder where the downloaded data will be saved.
118        patch_shape: The patch shape to use for training.
119        clustering: The clustering probability of the nuclei in percent. One of 0, 25, 50 or 75.
120        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
121        resize_inputs: Whether to resize the inputs to the patch shape.
122        download: Whether to download the data if it is not present.
123        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
124
125    Returns:
126        The segmentation dataset.
127    """
128    raw_paths, label_paths = get_bbbc024_paths(path, clustering, snr, download)
129
130    if resize_inputs:
131        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
132        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
133            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
134        )
135
136    return torch_em.default_segmentation_dataset(
137        raw_paths=raw_paths,
138        raw_key=None,
139        label_paths=label_paths,
140        label_key=None,
141        patch_shape=patch_shape,
142        is_seg_dataset=True,
143        **kwargs
144    )
145
146
147def get_bbbc024_loader(
148    path: Union[os.PathLike, str],
149    batch_size: int,
150    patch_shape: Tuple[int, ...],
151    clustering: Literal[0, 25, 50, 75] = 0,
152    snr: Literal["low", "high"] = "high",
153    resize_inputs: bool = False,
154    download: bool = False,
155    **kwargs
156) -> DataLoader:
157    """Get the BBBC024 dataloader for nucleus segmentation.
158
159    Args:
160        path: Filepath to a folder where the downloaded data will be saved.
161        batch_size: The batch size for training.
162        patch_shape: The patch shape to use for training.
163        clustering: The clustering probability of the nuclei in percent. One of 0, 25, 50 or 75.
164        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
165        resize_inputs: Whether to resize the inputs to the patch shape.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
168
169    Returns:
170        The DataLoader.
171    """
172    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
173    dataset = get_bbbc024_dataset(path, patch_shape, clustering, snr, resize_inputs, download, **ds_kwargs)
174    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)