torch_em.data.datasets.light_microscopy.bbbc033

  1"""The BBBC033 dataset contains a 3D fluorescence microscopy image of a clustered monolayer of
  2mouse trophoblast stem cells with instance segmentation ground truth for the nuclei.
  3
  4The volume was acquired with a spinning disk confocal microscope and has a shape of (32, 1344, 1024) voxels (ZYX)
  5with a z-spacing of 0.5 micrometer. It contains two channels, which are distributed as separate 8-bit RGB volumes
  6(named 'C0' and 'C2'). The BBBC page does not document the stains. From the image content, 'C0' is a membrane /
  7cytoplasm stain (stored as grayscale RGB) and 'C2' is the nuclear stain (stored as a blue-tinted RGB rendering).
  8Both are converted to single-channel 8-bit volumes and stored, together with the labels, in a HDF5 file
  9(keys 'raw/c0', 'raw/c2' and 'labels') when the data is prepared.
 10The ground truth contains 15 manually annotated nuclei as a labeled 16-bit volume (one id per nucleus, 0 background).
 11
 12The dataset is located at https://bbbc.broadinstitute.org/BBBC033.
 13This dataset is from the publication https://doi.org/10.1038/s41586-018-0051-0.
 14Please cite it if you use this dataset in your research.
 15"""
 16
 17import os
 18import shutil
 19from typing import Literal, Tuple, Union
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://data.broadinstitute.org/bbbc/BBBC033/BBBC033_v1_dataset.zip"
 29CHECKSUM = "232e73c781f8658cef2abe2ea41ac1e352f8032e9292ec9618e2eb2aa5dc760b"
 30
 31GT_URL = "https://data.broadinstitute.org/bbbc/BBBC033/BBBC033_v1_DatasetGroundTruth.tif"
 32GT_CHECKSUM = "c9bef6906ee450fac0d8fbcf588dad13d941ed71620010cf92a137984368c21f"
 33
 34
 35def _convert_to_h5(tmp_dir, gt_path, volume_path):
 36    import h5py
 37    import tifffile
 38
 39    with h5py.File(volume_path, "w") as f:
 40        for channel in ("C0", "C2"):
 41            # The channels are stored as RGB volumes, we reduce them to a single intensity channel.
 42            raw = tifffile.imread(os.path.join(tmp_dir, f"{channel}.tif"))
 43            assert raw.ndim == 4 and raw.shape[-1] == 3, f"Unexpected shape for {channel}: {raw.shape}"
 44            f.create_dataset(f"raw/{channel.lower()}", data=raw.max(axis=-1), compression="gzip")
 45
 46        labels = tifffile.imread(gt_path)
 47        f.create_dataset("labels", data=labels, compression="gzip")
 48
 49
 50def get_bbbc033_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 51    """Download the BBBC033 dataset.
 52
 53    Args:
 54        path: Filepath to a folder where the downloaded data will be saved.
 55        download: Whether to download the data if it is not present.
 56
 57    Returns:
 58        Filepath to the HDF5 file with the data.
 59    """
 60    data_dir = os.path.join(path, "BBBC033")
 61    volume_path = os.path.join(data_dir, "BBBC033.h5")
 62    if os.path.exists(volume_path):
 63        return volume_path
 64
 65    os.makedirs(data_dir, exist_ok=True)
 66
 67    zip_path = os.path.join(path, "BBBC033_v1_dataset.zip")
 68    util.download_source(zip_path, URL, download, CHECKSUM)
 69    tmp_dir = os.path.join(data_dir, "tif")
 70    util.unzip(zip_path, tmp_dir)
 71
 72    gt_path = os.path.join(path, "BBBC033_v1_DatasetGroundTruth.tif")
 73    util.download_source(gt_path, GT_URL, download, GT_CHECKSUM)
 74
 75    _convert_to_h5(tmp_dir, gt_path, volume_path)
 76    shutil.rmtree(tmp_dir)
 77    os.remove(gt_path)
 78
 79    return volume_path
 80
 81
 82def get_bbbc033_paths(
 83    path: Union[os.PathLike, str], channel: Literal[0, 2] = 2, download: bool = False
 84) -> Tuple[str, str]:
 85    """Get paths to the BBBC033 data.
 86
 87    Args:
 88        path: Filepath to a folder where the downloaded data will be saved.
 89        channel: The channel to use as raw input. 0: membrane / cytoplasm stain, 2: nuclear stain.
 90        download: Whether to download the data if it is not present.
 91
 92    Returns:
 93        Filepath to the HDF5 file with the data.
 94        The key of the raw data for the chosen channel within the HDF5 file.
 95    """
 96    if channel not in (0, 2):
 97        raise ValueError(f"'{channel}' is not a valid channel. Choose from 0 or 2.")
 98    volume_path = get_bbbc033_data(path, download)
 99    return volume_path, f"raw/c{channel}"
100
101
102def get_bbbc033_dataset(
103    path: Union[os.PathLike, str],
104    patch_shape: Tuple[int, ...],
105    channel: Literal[0, 2] = 2,
106    resize_inputs: bool = False,
107    download: bool = False,
108    **kwargs
109) -> Dataset:
110    """Get the BBBC033 dataset for nucleus segmentation.
111
112    Args:
113        path: Filepath to a folder where the downloaded data will be saved.
114        patch_shape: The patch shape to use for training.
115        channel: The channel to use as raw input. 0: membrane / cytoplasm stain, 2: nuclear stain.
116        resize_inputs: Whether to resize the inputs to the patch shape.
117        download: Whether to download the data if it is not present.
118        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
119
120    Returns:
121        The segmentation dataset.
122    """
123    volume_path, raw_key = get_bbbc033_paths(path, channel, download)
124
125    if resize_inputs:
126        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
127        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
128            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
129        )
130
131    return torch_em.default_segmentation_dataset(
132        raw_paths=volume_path,
133        raw_key=raw_key,
134        label_paths=volume_path,
135        label_key="labels",
136        patch_shape=patch_shape,
137        is_seg_dataset=True,
138        **kwargs
139    )
140
141
142def get_bbbc033_loader(
143    path: Union[os.PathLike, str],
144    batch_size: int,
145    patch_shape: Tuple[int, ...],
146    channel: Literal[0, 2] = 2,
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> DataLoader:
151    """Get the BBBC033 dataloader for nucleus segmentation.
152
153    Args:
154        path: Filepath to a folder where the downloaded data will be saved.
155        batch_size: The batch size for training.
156        patch_shape: The patch shape to use for training.
157        channel: The channel to use as raw input. 0: membrane / cytoplasm stain, 2: nuclear stain.
158        resize_inputs: Whether to resize the inputs to the patch shape.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_bbbc033_dataset(path, patch_shape, channel, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)