torch_em.data.datasets.electron_microscopy.etchacell

  1"""The Etch a Cell dataset contains annotations for nuclear envelope segmentation in serial blockface
  2scanning electron microscopy (SBF-SEM) volumes of HeLa cells.
  3
  4The labels are consensus segmentations, aggregated from the annotations of volunteers of the citizen science project
  5'Etch a Cell' (https://www.zooniverse.org/projects/dwright04/etch-a-cell-nuclear-envelope).
  6The data consists of 18 cropped regions of interest (ROIs) with a voxel size of 10 x 10 x 50 nm.
  7Each ROI has up to 300 slices of 2000 x 2000 pixels, some of the ROIs are smaller.
  8
  9The labels are available at https://www.ebi.ac.uk/biostudies/studies/S-BSST448.
 10The raw data is available at https://www.ebi.ac.uk/empiar/EMPIAR-10478/ (~15.6 GB for all ROIs).
 11Use `rois` to download only a subset of the ROIs.
 12The licence of the BioStudies entry is not stated. Please check the licence of EMPIAR-10478 before using the data.
 13
 14The dataset is from the publication 'Citizen science, cells and CNNs: deep learning for automatic segmentation of the
 15nuclear envelope in electron microscopy data, trained with volunteer segmentations' (Spiers et al.).
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20from typing import List, Optional, Tuple, Union
 21
 22import numpy as np
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30LABEL_URL = "https://www.ebi.ac.uk/biostudies/files/S-BSST448/Aggregations.zip"
 31RAW_URL = "https://ftp.ebi.ac.uk/empiar/world_availability/10478/data/ROIs"
 32LABEL_CHECKSUM = None
 33# NOTE: The raw data has no checksums available, so the downloaded files are not verified.
 34
 35ROIS = [
 36    "ROI_1416-1932-171", "ROI_1536-3456-213", "ROI_1584-6996-1", "ROI_1608-912-1", "ROI_1656-6756-329",
 37    "ROI_1716-7800-517", "ROI_2052-5784-112", "ROI_2448-4704-271", "ROI_2820-6780-468", "ROI_2832-1692-1",
 38    "ROI_3000-3264-393", "ROI_3516-5712-314", "ROI_3576-5232-35", "ROI_3588-3972-1", "ROI_3624-2712-201",
 39    "ROI_3768-7248-143", "ROI_3972-1956-438", "ROI_4320-1260-95",
 40]
 41
 42
 43def _check_rois(rois):
 44    rois = ROIS if rois is None else rois
 45    for roi in rois:
 46        if roi not in ROIS:
 47            raise ValueError(f"'{roi}' is not a valid ROI. Please choose from {ROIS}.")
 48    return list(rois)
 49
 50
 51def get_etchacell_data(
 52    path: Union[os.PathLike, str], rois: Optional[List[str]] = None, download: bool = False
 53) -> str:
 54    """Download the Etch a Cell dataset.
 55
 56    Args:
 57        path: Filepath to a folder where the downloaded data will be saved.
 58        rois: The ROIs to download. By default all 18 ROIs are used, e.g. 'ROI_1416-1932-171'.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        The filepath to the folder with the data.
 63    """
 64    rois = _check_rois(rois)
 65
 66    label_dir = os.path.join(path, "labels")
 67    raw_dir = os.path.join(path, "raw")
 68
 69    if not all(os.path.exists(os.path.join(label_dir, f"{roi}.tiff")) for roi in rois):
 70        os.makedirs(path, exist_ok=True)
 71        zip_path = os.path.join(path, "Aggregations.zip")
 72        util.download_source(zip_path, LABEL_URL, download, checksum=LABEL_CHECKSUM)
 73        util.unzip(zip_path, label_dir, remove=True)
 74
 75    os.makedirs(raw_dir, exist_ok=True)
 76    for roi in rois:
 77        raw_path = os.path.join(raw_dir, f"{roi}.tiff")
 78        if not os.path.exists(raw_path):
 79            util.download_source(raw_path, f"{RAW_URL}/{roi}.tiff", download, checksum=None)
 80
 81    return path
 82
 83
 84def get_etchacell_paths(
 85    path: Union[os.PathLike, str], rois: Optional[List[str]] = None, download: bool = False
 86) -> Tuple[List[str], List[str]]:
 87    """Get paths to the Etch a Cell data.
 88
 89    Args:
 90        path: Filepath to a folder where the downloaded data will be saved.
 91        rois: The ROIs to use. By default all 18 ROIs are used.
 92        download: Whether to download the data if it is not present.
 93
 94    Returns:
 95        List of filepaths for the image data.
 96        List of filepaths for the label data.
 97    """
 98    rois = _check_rois(rois)
 99    get_etchacell_data(path, rois, download)
100
101    raw_paths = [os.path.join(path, "raw", f"{roi}.tiff") for roi in rois]
102    label_paths = [os.path.join(path, "labels", f"{roi}.tiff") for roi in rois]
103    return raw_paths, label_paths
104
105
106def _binary_labels(labels):
107    return (labels > 0).astype(np.uint8)
108
109
110def get_etchacell_dataset(
111    path: Union[os.PathLike, str],
112    patch_shape: Tuple[int, int, int],
113    rois: Optional[List[str]] = None,
114    download: bool = False,
115    **kwargs
116) -> Dataset:
117    """Get the Etch a Cell dataset for nuclear envelope segmentation in SBF-SEM.
118
119    Args:
120        path: Filepath to a folder where the downloaded data will be saved.
121        patch_shape: The patch shape to use for training.
122        rois: The ROIs to use. By default all 18 ROIs are used.
123        download: Whether to download the data if it is not present.
124        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
125
126    Returns:
127        The segmentation dataset.
128    """
129    raw_paths, label_paths = get_etchacell_paths(path, rois, download)
130
131    kwargs = util.update_kwargs(kwargs, "label_transform", _binary_labels)
132
133    return torch_em.default_segmentation_dataset(
134        raw_paths=raw_paths,
135        raw_key=None,
136        label_paths=label_paths,
137        label_key=None,
138        patch_shape=patch_shape,
139        is_seg_dataset=True,
140        **kwargs
141    )
142
143
144def get_etchacell_loader(
145    path: Union[os.PathLike, str],
146    batch_size: int,
147    patch_shape: Tuple[int, int, int],
148    rois: Optional[List[str]] = None,
149    download: bool = False,
150    **kwargs
151) -> DataLoader:
152    """Get the Etch a Cell dataloader for nuclear envelope segmentation in SBF-SEM.
153
154    Args:
155        path: Filepath to a folder where the downloaded data will be saved.
156        batch_size: The batch size for training.
157        patch_shape: The patch shape to use for training.
158        rois: The ROIs to use. By default all 18 ROIs are used.
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_etchacell_dataset(path, patch_shape, rois, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)