torch_em.data.datasets.electron_microscopy.hela_metaphase

  1"""The HeLa Metaphase dataset contains 3D SBF-SEM volumes of two metaphase HeLa cells with
  2semantic segmentation annotations of organelles.
  3
  4It contains two sub-datasets, each showing one quadrant of a metaphase cell (30 nm z-resolution):
  5- control: A control cell (578 slices of 1477 x 1453 pixels) with the labels
  6  0 (exterior), 1 (mitochondria), 2 (ER), 3 (ER in close proximity to chromosomes, within 320 nm),
  7  4 (chromosomes), 5 (centrioles) and 6 (peripheral ER). It also has a binary cell outline annotation.
  8- reep3_4_kd: A REEP3/4 knock-down cell (471 slices of 1586 x 1600 pixels) with the labels
  9  0 (exterior), 1 (mitochondria), 2 (lipid droplets), 3 (ER) and 4 (chromosomes).
 10
 11Note that the label indices differ between the two sub-datasets. The annotations were created
 12with Microscopy Image Browser and cover one quadrant of the cell and the insides of the spindle area.
 13
 14The data is available at EMPIAR-10994 (https://doi.org/10.6019/EMPIAR-10994) and is distributed under
 15the CC0 license. It was published in https://doi.org/10.1091/mbc.e18-11-0698.
 16Please cite this publication if you use the dataset in your research.
 17"""
 18
 19import os
 20import hashlib
 21import time
 22from warnings import warn
 23from concurrent.futures import ThreadPoolExecutor
 24from typing import List, Literal, Tuple, Union
 25
 26import numpy as np
 27import requests
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35BASE_URL = "https://ftp.ebi.ac.uk/empiar/world_availability/10994/data"
 36# The URL patterns of the (processed) raw slices and the label slices for each sub-dataset.
 37# NOTE: Only the processed tif slices are downloaded (not the raw dm4 files, nor the full 8.7 GB entry).
 38URLS = {
 39    "control": {
 40        "raw": BASE_URL + "/160621_HeLa_Control/02_processed_dataset/160621_HeLa_R01_{i:03d}.tif",
 41        "organelles": BASE_URL + "/160621_HeLa_Control/03_models/OrganelleModel_TIF/Labels_160621_HeLa_R01_combined3_{i:03d}.tif",  # noqa
 42        "cell_outline": BASE_URL + "/160621_HeLa_Control/03_models/CellOutlineModel_TIF/Labels_cell_outlines_{i:03d}.tif",  # noqa
 43    },
 44    "reep3_4_kd": {
 45        "raw": BASE_URL + "/161107_HeLa_REEP3-4_KD_LBR/02_processed_dataset/161107_HeLa_REEP3-4_KD_LBR_{i:03d}.tif",
 46        "organelles": BASE_URL + "/161107_HeLa_REEP3-4_KD_LBR/03_models/Model_TIF/Labels_161107_HeLa_REEP3-4_KD_LBR_{i:03d}.tif",  # noqa
 47    },
 48}
 49NUM_SLICES = {"control": 578, "reep3_4_kd": 471}
 50# The sha256 checksums of the stacked raw volumes (the bytes of the uint8 array in zyx order).
 51CHECKSUMS = {
 52    "control": "882bc5f1794a35b96617964edb294aee243b30ed57758b964738f244bd2b481f",
 53    "reep3_4_kd": "9bc5f9689f0e96d3aafd0bd97cfb5f50dff2e82fb0d7dbbe2fc8f09832c0597c",
 54}
 55
 56# The organelle label of slice 455 of the control cell is an empty folder in the EMPIAR entry.
 57# It is filled with the label of the previous slice (the z resolution is 30 nm).
 58MISSING_LABELS = {"control": {"organelles": [455]}}
 59
 60NAMES = list(URLS.keys())
 61LABEL_TYPES = {"control": ["organelles", "cell_outline"], "reep3_4_kd": ["organelles"]}
 62
 63
 64def _download_slice(url, tmp_dir, key, i):
 65    import tifffile
 66
 67    out_path = os.path.join(tmp_dir, f"{key}_{i:03d}.tif")
 68    if not os.path.exists(out_path):
 69        for attempt in range(5):
 70            try:
 71                with requests.get(url, timeout=120) as r:
 72                    r.raise_for_status()
 73                    content = r.content
 74                # The server may answer with an html error page instead of the tif under load.
 75                if content[:2] not in (b"II", b"MM"):
 76                    raise requests.ConnectionError(f"Did not receive a tif file from {url}.")
 77                break
 78            except (requests.ConnectionError, requests.Timeout):
 79                if attempt == 4:
 80                    raise
 81                time.sleep(2 ** attempt)
 82        with open(out_path + ".incomplete", "wb") as f:
 83            f.write(content)
 84        os.replace(out_path + ".incomplete", out_path)
 85    return tifffile.imread(out_path)
 86
 87
 88def _create_volume(name, h5_path, tmp_dir):
 89    import h5py
 90
 91    n_slices = NUM_SLICES[name]
 92    sha = hashlib.sha256()
 93    tmp_h5 = h5_path + ".incomplete"
 94
 95    with h5py.File(tmp_h5, "w") as f, ThreadPoolExecutor(max_workers=4) as pool:
 96        datasets = {}
 97        for key, url in URLS[name].items():
 98            # Download in blocks of slices, so that tifs are read in order but fetched in parallel.
 99            ds_name = "raw" if key == "raw" else f"labels/{key}"
100            for start in range(1, n_slices + 1, 32):
101                ids = range(start, min(start + 32, n_slices + 1))
102                missing = MISSING_LABELS.get(name, {}).get(key, [])
103                slices = list(pool.map(
104                    lambda i: None if i in missing else _download_slice(url.format(i=i), tmp_dir, key, i), ids
105                ))
106                if ds_name not in datasets:
107                    shape = (n_slices,) + slices[0].shape
108                    datasets[ds_name] = f.create_dataset(
109                        ds_name, shape=shape, dtype="uint8", chunks=(1,) + slices[0].shape, compression="gzip"
110                    )
111                for i, slice_ in zip(ids, slices):
112                    if slice_ is None:
113                        warn(f"The {key} label of slice {i} is missing in EMPIAR-10994, using the previous slice.")
114                        slice_ = datasets[ds_name][i - 2]
115                    assert slice_.shape == datasets[ds_name].shape[1:], f"Unexpected shape of slice {i} for {key}."
116                    datasets[ds_name][i - 1] = slice_.astype("uint8")
117                    if key == "raw":
118                        sha.update(np.ascontiguousarray(slice_, dtype="uint8").tobytes())
119
120    checksum = sha.hexdigest()
121    if CHECKSUMS[name] is not None and checksum != CHECKSUMS[name]:
122        os.remove(tmp_h5)
123        raise RuntimeError(
124            f"The checksum of the downloaded data does not match the expected checksum. "
125            f"Expected: {CHECKSUMS[name]}, got: {checksum}"
126        )
127    os.replace(tmp_h5, h5_path)
128
129
130def get_hela_metaphase_data(
131    path: Union[os.PathLike, str], name: Literal["control", "reep3_4_kd"], download: bool = False
132) -> str:
133    """Download the HeLa Metaphase dataset.
134
135    The slices are downloaded from EMPIAR-10994 and stacked into a single hdf5 file per sub-dataset.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        name: The name of the sub-dataset. Either 'control' or 'reep3_4_kd'.
140        download: Whether to download the data if it is not present.
141
142    Returns:
143        The filepath to the hdf5 file with the data.
144    """
145    if name not in URLS:
146        raise ValueError(f"'{name}' is not a valid sub-dataset name. Please choose from {NAMES}.")
147
148    h5_path = os.path.join(path, f"hela_metaphase_{name}.h5")
149    if os.path.exists(h5_path):
150        return h5_path
151    if not download:
152        raise RuntimeError(f"Cannot find the data at {h5_path}, but download was set to False")
153
154    tmp_dir = os.path.join(path, f"slices_{name}")
155    os.makedirs(tmp_dir, exist_ok=True)
156    _create_volume(name, h5_path, tmp_dir)
157
158    # The individual slices are not needed anymore.
159    for fname in os.listdir(tmp_dir):
160        os.remove(os.path.join(tmp_dir, fname))
161    os.rmdir(tmp_dir)
162
163    return h5_path
164
165
166def get_hela_metaphase_paths(
167    path: Union[os.PathLike, str],
168    names: Union[Literal["control", "reep3_4_kd"], List[str], None] = None,
169    download: bool = False,
170) -> List[str]:
171    """Get paths to the HeLa Metaphase data.
172
173    Args:
174        path: Filepath to a folder where the downloaded data will be saved.
175        names: The name(s) of the sub-dataset(s). Either 'control' or 'reep3_4_kd', or a list of them.
176            By default, both are used.
177        download: Whether to download the data if it is not present.
178
179    Returns:
180        List of filepaths for the data.
181    """
182    if names is None:
183        names = NAMES
184    elif isinstance(names, str):
185        names = [names]
186
187    return [get_hela_metaphase_data(path, name, download) for name in names]
188
189
190def get_hela_metaphase_dataset(
191    path: Union[os.PathLike, str],
192    patch_shape: Tuple[int, int, int],
193    names: Union[Literal["control", "reep3_4_kd"], List[str], None] = None,
194    label_type: Literal["organelles", "cell_outline"] = "organelles",
195    download: bool = False,
196    **kwargs
197) -> Dataset:
198    """Get the HeLa Metaphase dataset for organelle segmentation in SBF-SEM.
199
200    The 'organelles' labels are semantic (multi-class) labels, see the module description for the label indices.
201    The 'cell_outline' labels are binary and only available for the 'control' sub-dataset.
202
203    Args:
204        path: Filepath to a folder where the downloaded data will be saved.
205        patch_shape: The patch shape to use for training.
206        names: The name(s) of the sub-dataset(s). Either 'control' or 'reep3_4_kd', or a list of them.
207            By default, both are used (only possible for 'organelles' labels).
208        label_type: The type of labels. Either 'organelles' or 'cell_outline'.
209        download: Whether to download the data if it is not present.
210        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
211
212    Returns:
213        The segmentation dataset.
214    """
215    if label_type not in ("organelles", "cell_outline"):
216        raise ValueError(f"'{label_type}' is not a valid label type. Please choose 'organelles' or 'cell_outline'.")
217    if label_type == "cell_outline":
218        if names is None:
219            names = "control"
220        if names not in ("control", ["control"]):
221            raise ValueError("The 'cell_outline' labels are only available for the 'control' sub-dataset.")
222
223    paths = get_hela_metaphase_paths(path, names, download)
224
225    return torch_em.default_segmentation_dataset(
226        raw_paths=paths,
227        raw_key="raw",
228        label_paths=paths,
229        label_key=f"labels/{label_type}",
230        patch_shape=patch_shape,
231        is_seg_dataset=True,
232        **kwargs
233    )
234
235
236def get_hela_metaphase_loader(
237    path: Union[os.PathLike, str],
238    patch_shape: Tuple[int, int, int],
239    batch_size: int,
240    names: Union[Literal["control", "reep3_4_kd"], List[str], None] = None,
241    label_type: Literal["organelles", "cell_outline"] = "organelles",
242    download: bool = False,
243    **kwargs
244) -> DataLoader:
245    """Get the DataLoader for organelle segmentation in the HeLa Metaphase dataset.
246
247    Args:
248        path: Filepath to a folder where the downloaded data will be saved.
249        patch_shape: The patch shape to use for training.
250        batch_size: The batch size for training.
251        names: The name(s) of the sub-dataset(s). Either 'control' or 'reep3_4_kd', or a list of them.
252            By default, both are used (only possible for 'organelles' labels).
253        label_type: The type of labels. Either 'organelles' or 'cell_outline'.
254        download: Whether to download the data if it is not present.
255        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
256
257    Returns:
258        The DataLoader.
259    """
260    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
261    dataset = get_hela_metaphase_dataset(path, patch_shape, names, label_type, download, **ds_kwargs)
262    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)