torch_em.data.datasets.light_microscopy.bbbc050

  1"""The BBBC050 dataset contains 3D time-lapse fluorescence microscopy images of early mouse embryos
  2(from fertilization to blastocyst formation) with nuclei labeled via histone H2B, and ground truth
  3annotations for nucleus segmentation.
  4
  5The dataset consists of 11 training embryos (121 volumes, imaged with an IX71 microscope,
  6voxel size 0.8 x 0.8 x 1.75 um) and 4 test embryos (44 volumes, imaged with a CV1000 microscope,
  7voxel size 0.8 x 0.8 x 2.0 um). Every embryo is sampled at 11 time points.
  8
  9Three types of ground truth are available for the training split (the test split only has 'QCANet'):
 10- 'QCANet': instance segmentation, every nucleus has its own id.
 11- 'NSN': semantic segmentation of the whole nuclear regions (foreground = 255).
 12- 'NDN': semantic segmentation of the nuclear center regions (foreground = 255).
 13
 14The dataset is located at https://bbbc.broadinstitute.org/BBBC050.
 15
 16This dataset is from the publication https://doi.org/10.1038/s41540-020-00152-8.
 17Please cite it if you use this dataset in your research.
 18"""
 19
 20import os
 21from glob import glob
 22from natsort import natsorted
 23from typing import Union, Tuple, Literal, List
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URLS = {
 33    "images": "https://data.broadinstitute.org/bbbc/BBBC050/Images.zip",
 34    "labels": "https://data.broadinstitute.org/bbbc/BBBC050/GroundTruth.zip",
 35}
 36
 37CHECKSUMS = {
 38    "images": "29f100abbfebfb1986b8e87eac091e86d8ec27cd8194f9a1c02c805e76b6dcd8",
 39    "labels": "1f19b308730dccf217c4d4dcf5745ad0fcde4eeb9f9c9306b2c8abd1fe73e5d1",
 40}
 41
 42
 43def get_bbbc050_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 44    """Download the BBBC050 dataset.
 45
 46    Args:
 47        path: Filepath to a folder where the data is downloaded for further processing.
 48        download: Whether to download the data if it is not present.
 49
 50    Returns:
 51        Filepath where the data is downloaded.
 52    """
 53    image_dir, label_dir = os.path.join(path, "Images"), os.path.join(path, "GroundTruth")
 54    if os.path.exists(image_dir) and os.path.exists(label_dir):
 55        return path
 56
 57    os.makedirs(path, exist_ok=True)
 58
 59    for name, url in URLS.items():
 60        zip_path = os.path.join(path, os.path.basename(url))
 61        util.download_source(path=zip_path, url=url, download=download, checksum=CHECKSUMS[name])
 62        util.unzip(zip_path=zip_path, dst=path)
 63
 64    return path
 65
 66
 67def get_bbbc050_paths(
 68    path: Union[os.PathLike, str],
 69    split: Literal["train", "test"],
 70    label_type: Literal["QCANet", "NSN", "NDN"] = "QCANet",
 71    download: bool = False,
 72) -> Tuple[List[str], List[str]]:
 73    """Get paths to the BBBC050 data.
 74
 75    Args:
 76        path: Filepath to a folder where the data is downloaded for further processing.
 77        split: The choice of data split. Either 'train' or 'test'.
 78        label_type: The choice of ground truth. Either 'QCANet' (instances), 'NSN' (nuclear regions)
 79            or 'NDN' (nuclear center regions). The 'test' split only provides 'QCANet' labels.
 80        download: Whether to download the data if it is not present.
 81
 82    Returns:
 83        List of filepaths for the image data.
 84        List of filepaths for the label data.
 85    """
 86    if split not in ("train", "test"):
 87        raise ValueError(f"'{split}' is not a valid split. Choose either 'train' or 'test'.")
 88    if label_type not in ("QCANet", "NSN", "NDN"):
 89        raise ValueError(f"'{label_type}' is not a valid label type. Choose one of 'QCANet', 'NSN' or 'NDN'.")
 90    if split == "test" and label_type != "QCANet":
 91        raise ValueError("The 'test' split only provides 'QCANet' labels.")
 92
 93    data_dir = get_bbbc050_data(path, download)
 94
 95    raw_paths = natsorted(glob(os.path.join(data_dir, "Images", split, "Images", "*.tif")))
 96    label_paths = [
 97        os.path.join(data_dir, "GroundTruth", split, f"GroundTruth_{label_type}", os.path.basename(p))
 98        for p in raw_paths
 99    ]
100
101    assert len(raw_paths) > 0, f"No volumes found for the '{split}' split at '{data_dir}'."
102    assert all(os.path.exists(p) for p in label_paths), "Some label volumes are missing."
103
104    return raw_paths, label_paths
105
106
107def get_bbbc050_dataset(
108    path: Union[os.PathLike, str],
109    patch_shape: Tuple[int, ...],
110    split: Literal["train", "test"],
111    label_type: Literal["QCANet", "NSN", "NDN"] = "QCANet",
112    resize_inputs: bool = False,
113    download: bool = False,
114    **kwargs
115) -> Dataset:
116    """Get the BBBC050 dataset for nucleus segmentation in 3D time-lapse images of mouse embryos.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        patch_shape: The patch shape to use for training.
121        split: The choice of data split. Either 'train' or 'test'.
122        label_type: The choice of ground truth. Either 'QCANet' (instances), 'NSN' (nuclear regions)
123            or 'NDN' (nuclear center regions). The 'test' split only provides 'QCANet' labels.
124        resize_inputs: Whether to resize inputs to the desired patch shape.
125        download: Whether to download the data if it is not present.
126        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
127
128    Returns:
129        The segmentation dataset.
130    """
131    raw_paths, label_paths = get_bbbc050_paths(path, split, label_type, download)
132
133    if resize_inputs:
134        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
135        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
136            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
137        )
138
139    return torch_em.default_segmentation_dataset(
140        raw_paths=raw_paths,
141        raw_key=None,
142        label_paths=label_paths,
143        label_key=None,
144        patch_shape=patch_shape,
145        is_seg_dataset=True,
146        **kwargs
147    )
148
149
150def get_bbbc050_loader(
151    path: Union[os.PathLike, str],
152    batch_size: int,
153    patch_shape: Tuple[int, ...],
154    split: Literal["train", "test"],
155    label_type: Literal["QCANet", "NSN", "NDN"] = "QCANet",
156    resize_inputs: bool = False,
157    download: bool = False,
158    **kwargs
159) -> DataLoader:
160    """Get the BBBC050 dataloader for nucleus segmentation in 3D time-lapse images of mouse embryos.
161
162    Args:
163        path: Filepath to a folder where the data is downloaded for further processing.
164        batch_size: The batch size for training.
165        patch_shape: The patch shape to use for training.
166        split: The choice of data split. Either 'train' or 'test'.
167        label_type: The choice of ground truth. Either 'QCANet' (instances), 'NSN' (nuclear regions)
168            or 'NDN' (nuclear center regions). The 'test' split only provides 'QCANet' labels.
169        resize_inputs: Whether to resize inputs to the desired patch shape.
170        download: Whether to download the data if it is not present.
171        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
172
173    Returns:
174        The DataLoader.
175    """
176    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
177    dataset = get_bbbc050_dataset(path, patch_shape, split, label_type, resize_inputs, download, **ds_kwargs)
178    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)