torch_em.data.datasets.electron_microscopy.cefa_hela

  1"""The Cefa-HeLa dataset contains SBF-SEM images of HeLa cells with semantic segmentation
  2annotations for nuclear envelope, nucleus, cell, and other cells.
  3
  4The dataset is from the publications:
  5- https://doi.org/10.3390/jimaging5090075
  6- https://doi.org/10.1371/journal.pone.0230605
  7
  8Please cite these if you use this dataset in your research.
  9
 10The ground truth labels are available at https://zenodo.org/records/4590903 (five classes, only the central
 11cell is annotated) and https://zenodo.org/records/6355622 (four classes, multiple nuclei are annotated).
 12The annotation can be selected via `label_version`. The multi-nuclei labels are semantic masks, not instances.
 13The raw electron microscopy images are available via EMPIAR-10094
 14(https://www.ebi.ac.uk/empiar/EMPIAR-10094/).
 15
 16Downloading from EMPIAR requires the aspera CLI. Install it via:
 17    conda install -c hcc aspera-cli
 18"""
 19
 20import os
 21from glob import glob
 22from typing import List, Literal, Optional, Sequence, Tuple, Union
 23
 24import numpy as np
 25from tqdm import tqdm
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34ZENODO_URLS = {
 35    "central": "https://zenodo.org/api/records/4590903/files-archive",
 36    "multinuclei": "https://zenodo.org/api/records/6355622/files-archive",
 37}
 38ZENODO_CHECKSUM = None
 39
 40LABEL_CLASSES_BY_VERSION = {
 41    "central": {
 42        "nuclear_envelope": 1,
 43        "nucleus": 2,
 44        "other_cells": 3,
 45        "background": 4,
 46        "cell": 5,
 47    },
 48    "multinuclei": {
 49        "nuclear_envelope": 1,
 50        "nucleus": 2,
 51        "cell": 3,
 52        "background": 4,
 53    },
 54}
 55# Kept for backwards compatibility.
 56LABEL_CLASSES = LABEL_CLASSES_BY_VERSION["central"]
 57
 58
 59def _check_label_version(label_version):
 60    if label_version not in LABEL_CLASSES_BY_VERSION:
 61        raise ValueError(f"Invalid label_version '{label_version}'. Choose from {list(LABEL_CLASSES_BY_VERSION)}.")
 62
 63
 64def _preprocess_data(empiar_path, gt_path, output_path):
 65    import h5py
 66
 67    try:
 68        import ncempy.io.dm as dm
 69    except ImportError:
 70        raise ImportError(
 71            "ncempy is required to read DM4 files from EMPIAR. "
 72            "Install it via 'pip install ncempy'."
 73        )
 74
 75    import scipy.io
 76
 77    gt_files = sorted(glob(os.path.join(gt_path, "GT_Slice_*.mat")))
 78    dm4_files = sorted(glob(os.path.join(empiar_path, "*.dm4")))
 79
 80    if not gt_files:
 81        raise RuntimeError(f"No GT_Slice_*.mat files found at {gt_path}.")
 82    if not dm4_files:
 83        raise RuntimeError(f"No DM4 files found at {empiar_path}.")
 84
 85    n_files = min(len(gt_files), len(dm4_files))
 86    os.makedirs(output_path, exist_ok=True)
 87
 88    for i in tqdm(range(n_files), desc="Preprocessing Cefa-HeLa"):
 89        h5_path = os.path.join(output_path, f"slice_{i + 1:03d}.h5")
 90        if os.path.exists(h5_path):
 91            continue
 92
 93        label = scipy.io.loadmat(gt_files[i])["groundTruth"]
 94
 95        with dm.fileDM(dm4_files[i]) as f:
 96            raw = np.array(f.getDataset(0)["data"], dtype=np.float32)
 97
 98        # Center-crop raw to match label size if the DM4 image is larger.
 99        if raw.shape[:2] != label.shape[:2]:
100            h, w = label.shape[:2]
101            cy, cx = raw.shape[0] // 2, raw.shape[1] // 2
102            raw = raw[cy - h // 2:cy - h // 2 + h, cx - w // 2:cx - w // 2 + w]
103
104        with h5py.File(h5_path, "w") as f:
105            f.create_dataset("raw", data=raw, compression="gzip")
106            f.create_dataset("labels", data=label, compression="gzip")
107
108
109def get_cefa_hela_data(
110    path: Union[os.PathLike, str],
111    download: bool = False,
112    label_version: Literal["central", "multinuclei"] = "central",
113) -> str:
114    """Download and prepare the Cefa-HeLa dataset.
115
116    Ground truth labels are downloaded from Zenodo and raw SBF-SEM images from
117    EMPIAR-10094. Both are paired by slice index and stored as H5 files.
118
119    Args:
120        path: Filepath to a folder where the downloaded data will be saved.
121        download: Whether to download the data if it is not present.
122            Downloading from EMPIAR requires the aspera CLI.
123        label_version: The annotation version. Either 'central' (five classes, only the central cell,
124            Zenodo 4590903) or 'multinuclei' (four classes, multiple nuclei, Zenodo 6355622).
125
126    Returns:
127        The filepath for the preprocessed H5 data.
128    """
129    _check_label_version(label_version)
130    suffix = "" if label_version == "central" else f"_{label_version}"
131    output_path = os.path.join(path, f"data{suffix}")
132    if os.path.exists(output_path) and len(glob(os.path.join(output_path, "*.h5"))) > 0:
133        return output_path
134
135    # Download GT labels from Zenodo.
136    gt_dir = os.path.join(path, f"gt_labels{suffix}")
137    if not os.path.exists(gt_dir) or len(glob(os.path.join(gt_dir, "GT_Slice_*.mat"))) == 0:
138        zip_path = os.path.join(path, f"cefa_hela_gt{suffix}.zip")
139        util.download_source(zip_path, ZENODO_URLS[label_version], download, ZENODO_CHECKSUM)
140        util.unzip(zip_path, gt_dir, remove=True)
141
142    # Download raw images from EMPIAR-10094.
143    empiar_root = util.download_source_empiar(path, "10094", download)
144    empiar_path = os.path.join(empiar_root, "data", "Micrographs_ROI_00")
145    if not os.path.exists(empiar_path):
146        raise RuntimeError(
147            f"Expected EMPIAR-10094 data at {empiar_path}. "
148            "Please ensure the download completed successfully."
149        )
150
151    _preprocess_data(empiar_path, gt_dir, output_path)
152    return output_path
153
154
155def get_cefa_hela_paths(
156    path: Union[os.PathLike, str],
157    download: bool = False,
158    label_version: Literal["central", "multinuclei"] = "central",
159) -> List[str]:
160    """Get paths to the Cefa-HeLa data.
161
162    Args:
163        path: Filepath to a folder where the downloaded data will be saved.
164        download: Whether to download the data if it is not present.
165        label_version: The annotation version, either 'central' or 'multinuclei'.
166
167    Returns:
168        The filepaths to the stored data.
169    """
170    data_root = get_cefa_hela_data(path, download, label_version)
171    paths = sorted(glob(os.path.join(data_root, "*.h5")))
172    return paths
173
174
175def get_cefa_hela_dataset(
176    path: Union[os.PathLike, str],
177    patch_shape: Tuple[int, int],
178    label_classes: Optional[Sequence[str]] = None,
179    download: bool = False,
180    label_version: Literal["central", "multinuclei"] = "central",
181    **kwargs,
182) -> Dataset:
183    """Get the Cefa-HeLa dataset for semantic segmentation of HeLa cell structures in SBF-SEM.
184
185    Args:
186        path: Filepath to a folder where the downloaded data will be saved.
187        patch_shape: The patch shape to use for training.
188        label_classes: The label classes to use for one-hot encoding.
189            For 'central' the classes are 'nuclear_envelope', 'nucleus', 'other_cells', 'background', 'cell'.
190            For 'multinuclei' they are 'nuclear_envelope', 'nucleus', 'cell', 'background'.
191            If None, returns the full label map
192            (central: 1=nuclear_envelope, 2=nucleus, 3=other_cells, 4=background, 5=cell;
193            multinuclei: 1=nuclear_envelope, 2=nucleus, 3=cell, 4=background).
194        download: Whether to download the data if it is not present.
195        label_version: The annotation version, either 'central' (only the central cell is annotated)
196            or 'multinuclei' (multiple nuclei are annotated, semantic labels only).
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
198
199    Returns:
200        The segmentation dataset.
201    """
202    _check_label_version(label_version)
203    label_classes_map = LABEL_CLASSES_BY_VERSION[label_version]
204    paths = get_cefa_hela_paths(path, download, label_version)
205
206    if label_classes is not None:
207        class_ids = []
208        for cls_name in label_classes:
209            if cls_name not in label_classes_map:
210                raise ValueError(
211                    f"Invalid class name: '{cls_name}'. Choose from {list(label_classes_map.keys())}."
212                )
213            class_ids.append(label_classes_map[cls_name])
214        label_transform = torch_em.transform.label.OneHotTransform(class_ids=class_ids)
215        msg = "'label_classes' is set, but 'label_transform' is in kwargs. It will be overridden."
216        kwargs = util.update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
217
218    return torch_em.default_segmentation_dataset(
219        raw_paths=paths,
220        raw_key="raw",
221        label_paths=paths,
222        label_key="labels",
223        patch_shape=patch_shape,
224        **kwargs,
225    )
226
227
228def get_cefa_hela_loader(
229    path: Union[os.PathLike, str],
230    patch_shape: Tuple[int, int],
231    batch_size: int,
232    label_classes: Optional[Sequence[str]] = None,
233    download: bool = False,
234    label_version: Literal["central", "multinuclei"] = "central",
235    **kwargs,
236) -> DataLoader:
237    """Get the DataLoader for the Cefa-HeLa dataset.
238
239    Args:
240        path: Filepath to a folder where the downloaded data will be saved.
241        patch_shape: The patch shape to use for training.
242        batch_size: The batch size for training.
243        label_classes: The label classes to use for one-hot encoding.
244            For 'central' the classes are 'nuclear_envelope', 'nucleus', 'other_cells', 'background', 'cell'.
245            For 'multinuclei' they are 'nuclear_envelope', 'nucleus', 'cell', 'background'.
246            If None, returns the full label map.
247        download: Whether to download the data if it is not present.
248        label_version: The annotation version, either 'central' or 'multinuclei'.
249        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
250            or for the PyTorch DataLoader.
251
252    Returns:
253        The DataLoader.
254    """
255    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
256    dataset = get_cefa_hela_dataset(path, patch_shape, label_classes, download, label_version, **ds_kwargs)
257    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)