torch_em.data.datasets.light_microscopy.bbbc032

  1"""The BBBC032 dataset contains a 3D fluorescence microscopy image of a mouse embryo blastocyst
  2with instance segmentation ground truth for the nuclei.
  3
  4The volume was acquired with a spinning disk confocal microscope and has a shape of (172, 1344, 1024) voxels (ZYX)
  5with a voxel size of 0.5 x 0.101 x 0.101 micrometer. It contains four channels, which are stored as separate volumes:
  6- channel 0: BMP4 transcripts (647 nm)
  7- channel 1: GAPDH transcripts (568 nm)
  8- channel 2: WGA, wheat germ agglutinin membrane stain (488 nm)
  9- channel 3: Hoechst nuclear stain (405 nm)
 10The ground truth contains 56 manually annotated nuclei as a labeled 16-bit volume (one id per nucleus, 0 background).
 11NOTE: The annotations are sparse, only a subset of the nuclei visible in the volume is annotated.
 12
 13The dataset is located at https://bbbc.broadinstitute.org/BBBC032.
 14This dataset is from the publication https://doi.org/10.1038/s41586-018-0051-0.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19import shutil
 20from typing import List, Tuple, Union
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URL = "https://data.broadinstitute.org/bbbc/BBBC032/BBBC032_v1_dataset.zip"
 30CHECKSUM = "02df5ca7cdc9afb751c63161cabc0e7967310ed1911edaa8572764fb44455321"
 31
 32GT_URL = "https://data.broadinstitute.org/bbbc/BBBC032/BBBC032_v1_DatasetGroundTruth.tif"
 33GT_CHECKSUM = "7ef577da64e1f95038d7eb40c03b3bce56ca2d5cd358cc9ea0aeba2be4526f3b"
 34
 35
 36def get_bbbc032_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 37    """Download the BBBC032 dataset.
 38
 39    Args:
 40        path: Filepath to a folder where the downloaded data will be saved.
 41        download: Whether to download the data if it is not present.
 42
 43    Returns:
 44        Filepath where the data is stored.
 45    """
 46    data_dir = os.path.join(path, "BBBC032")
 47    if os.path.exists(data_dir):
 48        return data_dir
 49
 50    os.makedirs(path, exist_ok=True)
 51
 52    zip_path = os.path.join(path, "BBBC032_v1_dataset.zip")
 53    util.download_source(zip_path, URL, download, CHECKSUM)
 54    util.unzip(zip_path, data_dir)
 55    # The zip file contains macOS metadata files, which we remove.
 56    shutil.rmtree(os.path.join(data_dir, "__MACOSX"), ignore_errors=True)
 57
 58    gt_path = os.path.join(data_dir, "BBBC032_v1_DatasetGroundTruth.tif")
 59    util.download_source(gt_path, GT_URL, download, GT_CHECKSUM)
 60
 61    return data_dir
 62
 63
 64def get_bbbc032_paths(
 65    path: Union[os.PathLike, str], channel: int = 3, download: bool = False
 66) -> Tuple[List[str], List[str]]:
 67    """Get paths to the BBBC032 data.
 68
 69    Args:
 70        path: Filepath to a folder where the downloaded data will be saved.
 71        channel: The channel to use as raw input. 0: BMP4, 1: GAPDH, 2: WGA (membranes), 3: Hoechst (nuclei).
 72        download: Whether to download the data if it is not present.
 73
 74    Returns:
 75        List of filepaths for the image data.
 76        List of filepaths for the label data.
 77    """
 78    if channel not in (0, 1, 2, 3):
 79        raise ValueError(f"'{channel}' is not a valid channel. Choose from 0, 1, 2 or 3.")
 80
 81    data_dir = get_bbbc032_data(path, download)
 82    raw_path = os.path.join(data_dir, f"BMP4blastocystC{channel}.tif")
 83    label_path = os.path.join(data_dir, "BBBC032_v1_DatasetGroundTruth.tif")
 84    assert os.path.exists(raw_path) and os.path.exists(label_path)
 85
 86    return [raw_path], [label_path]
 87
 88
 89def get_bbbc032_dataset(
 90    path: Union[os.PathLike, str],
 91    patch_shape: Tuple[int, ...],
 92    channel: int = 3,
 93    resize_inputs: bool = False,
 94    download: bool = False,
 95    **kwargs
 96) -> Dataset:
 97    """Get the BBBC032 dataset for nucleus segmentation.
 98
 99    Args:
100        path: Filepath to a folder where the downloaded data will be saved.
101        patch_shape: The patch shape to use for training.
102        channel: The channel to use as raw input. 0: BMP4, 1: GAPDH, 2: WGA (membranes), 3: Hoechst (nuclei).
103        resize_inputs: Whether to resize the inputs to the patch shape.
104        download: Whether to download the data if it is not present.
105        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
106
107    Returns:
108        The segmentation dataset.
109    """
110    raw_paths, label_paths = get_bbbc032_paths(path, channel, download)
111
112    if resize_inputs:
113        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
114        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
115            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
116        )
117
118    return torch_em.default_segmentation_dataset(
119        raw_paths=raw_paths,
120        raw_key=None,
121        label_paths=label_paths,
122        label_key=None,
123        patch_shape=patch_shape,
124        is_seg_dataset=True,
125        **kwargs
126    )
127
128
129def get_bbbc032_loader(
130    path: Union[os.PathLike, str],
131    batch_size: int,
132    patch_shape: Tuple[int, ...],
133    channel: int = 3,
134    resize_inputs: bool = False,
135    download: bool = False,
136    **kwargs
137) -> DataLoader:
138    """Get the BBBC032 dataloader for nucleus segmentation.
139
140    Args:
141        path: Filepath to a folder where the downloaded data will be saved.
142        batch_size: The batch size for training.
143        patch_shape: The patch shape to use for training.
144        channel: The channel to use as raw input. 0: BMP4, 1: GAPDH, 2: WGA (membranes), 3: Hoechst (nuclei).
145        resize_inputs: Whether to resize the inputs to the patch shape.
146        download: Whether to download the data if it is not present.
147        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
148
149    Returns:
150        The DataLoader.
151    """
152    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
153    dataset = get_bbbc032_dataset(path, patch_shape, channel, resize_inputs, download, **ds_kwargs)
154    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)