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