torch_em.data.datasets.light_microscopy.bbbc027

  1"""The BBBC027 dataset contains synthetic 3D fluorescence microscopy images of colon tissue
  2with foreground segmentation ground truth for the cell nuclei.
  3
  4The images were generated with a virtual microscope (CytoPacq) imitating a Zeiss S100 confocal
  5microscope. The dataset contains 30 images of shape (129, 1030, 1300) voxels (ZYX), each provided
  6in a low SNR and a high SNR variant. The ground truth is a binary foreground / background mask
  7of all nuclei inside the tissue region (label 1 for nuclei, 0 for background); individual nuclei are
  8not separated into instances.
  9
 10The images and ground truth are distributed in the ICS format and are converted to HDF5 files
 11(with the keys 'raw' and 'labels') when the data is prepared.
 12
 13The dataset is located at https://bbbc.broadinstitute.org/BBBC027.
 14This dataset is from the publication https://doi.org/10.1007/978-3-642-21596-4_4.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19import shutil
 20import zlib
 21from glob import glob
 22from natsort import natsorted
 23from typing import List, Literal, Tuple, Union
 24
 25import numpy as np
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URL = "https://data.broadinstitute.org/bbbc/BBBC027/BBBC027_{snr}SNR_{kind}_part{part}.zip"
 35CHECKSUMS = {
 36    "low": {
 37        "images": [
 38            "bf537d5e63c63f3b86e31d0dfe1b888a5bac8df78d9562f7c4fb738d9ad4e5d1",
 39            "4e581c373a2656d50265345a0cdaa1ec32926028c56a46ac9a6cc33bd8fca295",
 40            "b0ca1814101ea7c8f71802c762195f3e69a639bf4ba24065d4efbb780ea94d63",
 41        ],
 42        "foreground": [
 43            "1530cd10f40479045e54c6a56e1ba0081f2b0383b7008376abcb9a774ebdf33e",
 44            "96a59ff604e515c084ecf8a8c45ac96685ff7259cb210759a06a8a4d4dff9cd3",
 45            "073c53019799ee40f9dfe2d5f93b841f8e0728f3c8419ae6f211b991ad2f3c84",
 46        ],
 47    },
 48    "high": {
 49        "images": [
 50            "f07d6aa56b990dfe564380e004e0b414336e610bfd8e93d3bfff1ea15843a1f0",
 51            "0ad770285041e53cc6a5950b22b82d23c50a2a4e053e31dcb9653b300fd03ab5",
 52            "1454e75882d7de39898e54cb02d6502d051cf2a75e7229031d30be6910090eb2",
 53        ],
 54        "foreground": [
 55            "0dd7c5e03fa216e71add8983e4c4a05c21b4208b800055d3c41596188308af5e",
 56            "d04115940d17d09918a0b2116a3389ced753c549ea792d0352e886fa7eeaab0c",
 57            "71dff9e5cc2fc202418db85084104ef729d36c61b5c79b8085c63cc186bacdd3",
 58        ],
 59    },
 60}
 61
 62
 63def _read_ics(ics_path):
 64    """Read an ICS 2.0 file (header and data in the same file, optionally gzip compressed).
 65    """
 66    with open(ics_path, "rb") as f:
 67        data = f.read()
 68
 69    # The first two bytes define the field and the line separator of the header.
 70    field_sep, line_sep = data[0:1], data[1:2]
 71    end_marker = b"end" + field_sep + line_sep
 72    header_end = data.index(end_marker)
 73
 74    header = {}
 75    for line in data[2:header_end].split(line_sep):
 76        fields = line.split(field_sep)
 77        if fields[0] in (b"layout", b"representation") and len(fields) > 2:
 78            header[(fields[0], fields[1])] = fields[2:]
 79
 80    order = [dim.decode() for dim in header[(b"layout", b"order")]]
 81    sizes = [int(size) for size in header[(b"layout", b"sizes")]]
 82    assert order[0] == "bits"
 83    bits, sizes, order = sizes[0], sizes[1:], order[1:]
 84
 85    fmt = header[(b"representation", b"format")][0]
 86    sign = header[(b"representation", b"sign")][0]
 87    if fmt == b"integer":
 88        dtype = np.dtype(f"{'u' if sign == b'unsigned' else 'i'}{bits // 8}")
 89    else:
 90        dtype = np.dtype(f"f{bits // 8}")
 91    byte_order = header[(b"representation", b"byte_order")]
 92    dtype = dtype.newbyteorder("<" if byte_order[0] == b"1" else ">")
 93
 94    raw = data[header_end + len(end_marker):]
 95    if header.get((b"representation", b"compression"), [b"uncompressed"])[0] == b"gzip":
 96        raw = zlib.decompress(raw, 16 + zlib.MAX_WBITS)
 97
 98    # The ICS layout lists the fastest varying axis first, so we reverse to get a C-ordered array.
 99    volume = np.frombuffer(raw, dtype=dtype).reshape(sizes[::-1])
100    # The volumes in BBBC027 are stored in the order x, y, z, i.e. we get a zyx array after reversing.
101    assert order[::-1] == ["z", "y", "x"], f"Unexpected axis order in {ics_path}: {order}"
102    return volume
103
104
105def _convert_to_h5(image_dir, label_dir, data_dir):
106    import h5py
107
108    image_paths = natsorted(glob(os.path.join(image_dir, "**", "image-final_*.ics"), recursive=True))
109    label_paths = natsorted(glob(os.path.join(label_dir, "**", "image-labels_*.ics"), recursive=True))
110    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
111
112    for image_path, label_path in zip(image_paths, label_paths):
113        image_id = os.path.basename(image_path).replace("image-final_", "").replace(".ics", "")
114        assert os.path.basename(label_path) == f"image-labels_{image_id}.ics"
115
116        raw = _read_ics(image_path)
117        labels = _read_ics(label_path)
118        assert raw.shape == labels.shape, f"{raw.shape}, {labels.shape}"
119
120        with h5py.File(os.path.join(data_dir, f"image_{image_id}.h5"), "w") as f:
121            f.create_dataset("raw", data=raw, compression="gzip")
122            f.create_dataset("labels", data=labels, compression="gzip")
123
124
125def get_bbbc027_data(
126    path: Union[os.PathLike, str], snr: Literal["low", "high"] = "high", download: bool = False
127) -> str:
128    """Download the BBBC027 dataset.
129
130    Args:
131        path: Filepath to a folder where the downloaded data will be saved.
132        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
133        download: Whether to download the data if it is not present.
134
135    Returns:
136        Filepath where the data for the chosen SNR level is stored.
137    """
138    if snr not in ("low", "high"):
139        raise ValueError(f"'{snr}' is not a valid SNR level. Choose from 'low' or 'high'.")
140
141    data_dir = os.path.join(path, "BBBC027", f"{snr}SNR")
142    if os.path.exists(data_dir):
143        return data_dir
144
145    os.makedirs(path, exist_ok=True)
146
147    tmp_dir = os.path.join(path, "BBBC027", f"{snr}SNR_ics")
148    for kind in ("images", "foreground"):
149        for part in (1, 2, 3):
150            zip_path = os.path.join(path, f"BBBC027_{snr}SNR_{kind}_part{part}.zip")
151            url = URL.format(snr=snr, kind=kind, part=part)
152            util.download_source(zip_path, url, download, CHECKSUMS[snr][kind][part - 1])
153            util.unzip(zip_path, os.path.join(tmp_dir, kind))
154
155    os.makedirs(data_dir, exist_ok=True)
156    _convert_to_h5(os.path.join(tmp_dir, "images"), os.path.join(tmp_dir, "foreground"), data_dir)
157    shutil.rmtree(tmp_dir)
158
159    return data_dir
160
161
162def get_bbbc027_paths(
163    path: Union[os.PathLike, str], snr: Literal["low", "high"] = "high", download: bool = False
164) -> List[str]:
165    """Get paths to the BBBC027 data.
166
167    Args:
168        path: Filepath to a folder where the downloaded data will be saved.
169        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
170        download: Whether to download the data if it is not present.
171
172    Returns:
173        List of filepaths for the HDF5 files, which contain the image data (key 'raw') and the labels (key 'labels').
174    """
175    data_dir = get_bbbc027_data(path, snr, download)
176    volume_paths = natsorted(glob(os.path.join(data_dir, "image_*.h5")))
177    assert len(volume_paths) > 0
178    return volume_paths
179
180
181def get_bbbc027_dataset(
182    path: Union[os.PathLike, str],
183    patch_shape: Tuple[int, ...],
184    snr: Literal["low", "high"] = "high",
185    resize_inputs: bool = False,
186    download: bool = False,
187    **kwargs
188) -> Dataset:
189    """Get the BBBC027 dataset for nucleus foreground segmentation.
190
191    Args:
192        path: Filepath to a folder where the downloaded data will be saved.
193        patch_shape: The patch shape to use for training.
194        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
195        resize_inputs: Whether to resize the inputs to the patch shape.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
198
199    Returns:
200        The segmentation dataset.
201    """
202    volume_paths = get_bbbc027_paths(path, snr, download)
203
204    if resize_inputs:
205        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
206        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
207            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
208        )
209
210    return torch_em.default_segmentation_dataset(
211        raw_paths=volume_paths,
212        raw_key="raw",
213        label_paths=volume_paths,
214        label_key="labels",
215        patch_shape=patch_shape,
216        is_seg_dataset=True,
217        **kwargs
218    )
219
220
221def get_bbbc027_loader(
222    path: Union[os.PathLike, str],
223    batch_size: int,
224    patch_shape: Tuple[int, ...],
225    snr: Literal["low", "high"] = "high",
226    resize_inputs: bool = False,
227    download: bool = False,
228    **kwargs
229) -> DataLoader:
230    """Get the BBBC027 dataloader for nucleus foreground segmentation.
231
232    Args:
233        path: Filepath to a folder where the downloaded data will be saved.
234        batch_size: The batch size for training.
235        patch_shape: The patch shape to use for training.
236        snr: The signal-to-noise ratio of the images. Either 'low' or 'high'.
237        resize_inputs: Whether to resize the inputs to the patch shape.
238        download: Whether to download the data if it is not present.
239        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
240
241    Returns:
242        The DataLoader.
243    """
244    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
245    dataset = get_bbbc027_dataset(path, patch_shape, snr, resize_inputs, download, **ds_kwargs)
246    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)