torch_em.data.datasets.electron_microscopy.hela_stard3

  1"""The HeLa STARD3 dataset contains FIB-SEM volumes of HeLa cells with manual annotations of
  2mitochondria, endoplasmic reticulum (ER) and endosomes / lysosomes.
  3
  4It contains three volumes: control HeLa cells ('ctrl'), cells expressing wild-type STARD3 ('stard3')
  5and cells expressing the STARD3 FA/YA mutant ('stard3_fa_ya').
  6Each volume has a region of interest (ROI) with dense annotations. Only this ROI is used here:
  7the raw data is cropped to the ROI and stored together with the binary labels of the three organelles.
  8The annotations outside of the ROI are not used.
  9
 10NOTE: Four slices of the 'ctrl' volume (slices 149, 151 and 152 of the original stack are truncated and
 11slice 154 is missing in the EMPIAR archive). They are replaced by the closest valid slice (and a warning is raised).
 12
 13The dataset is located at https://www.ebi.ac.uk/empiar/EMPIAR-13156 (https://doi.org/10.6019/EMPIAR-13156).
 14The license is not stated in the EMPIAR metadata (EMPIAR entries are usually released under CC0),
 15please check the entry before using the data.
 16The data is part of the publication https://doi.org/10.1038/s44318-026-00705-3.
 17Please cite it if you use this dataset in your research.
 18"""
 19
 20import os
 21import time
 22import warnings
 23from concurrent.futures import ThreadPoolExecutor
 24from typing import Union, Tuple, List, Literal, Optional
 25
 26import numpy as np
 27import requests
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36BASE_URL = "https://ftp.ebi.ac.uk/empiar/world_availability/13156/data"
 37ORGANELLES = ["mitochondria", "er", "endosomes"]
 38FILE_NAMES = {"mitochondria": "Mitochondria", "er": "ER", "endosomes": "Endosomes"}
 39
 40# The EMPIAR folder, the number of slices, the first slices of the sub-folders, the pattern of the slice filenames
 41# and the region of interest (ROI) of the volumes. The ROI is given as start and shape in (z, y, x).
 42# The start is zero-based. NOTE: The label volumes of 'stard3' have fewer rows (2112) than the raw slices (2748),
 43# they are aligned to the top of the slices.
 44VOLUMES = {
 45    "ctrl": {
 46        "folder": "HeLaWT_Raw",
 47        "n_slices": 994,
 48        "chunk_starts": [1, 201, 401, 601, 801],
 49        "pattern": "ROI_CICGFA_{:03d}.tif",
 50        "roi_start": (5, 64, 1536),
 51        "roi_shape": (444, 900, 1200),
 52    },
 53    "stard3": {
 54        "folder": "HeLaSTARD3_WT_Raw",
 55        "n_slices": 810,
 56        "chunk_starts": [1, 201, 401, 601],
 57        "pattern": "image_{:03d}.tif",
 58        "roi_start": (428, 151, 16),
 59        "roi_shape": (382, 1100, 1400),
 60    },
 61    "stard3_fa_ya": {
 62        "folder": "HeLaSTARD3_FA_YA_Raw",
 63        "n_slices": 600,
 64        "chunk_starts": [1, 101, 201, 301, 401, 501],
 65        "pattern": "image{:03d}.tif",
 66        "roi_start": (120, 50, 386),
 67        "roi_shape": (449, 700, 700),
 68    },
 69}
 70# The checksums of the label files with the annotations inside of the ROI.
 71CHECKSUMS = {
 72    "ctrl": {
 73        "mitochondria": "1c05a133f4dbb763db5a9affb1afe3a54622beae55dab1d32a318ed3cd36aec7",
 74        "er": "3c3ea1b4c9a91aba8b982b712ed48e5f323427fc7d7f2951e05b542577a1dfa6",
 75        "endosomes": "984201582bd00f307ed2af627a33665168d4f5c35435569e850b979266fff049",
 76    },
 77    "stard3": {
 78        "mitochondria": "32117b402c2d92bafb9a1c7d999af5f05d5e359674dc0f308da788fe448c5d71",
 79        "er": "f54d5c7ea87d710e9e53ea8d01d062a92c975684237d14a9d336fb61423c8069",
 80        "endosomes": "b33b88a09c46abdb7786ffd474e1c6a52578efcdc1ed20bac166637ef36a6b2f",
 81    },
 82    "stard3_fa_ya": {
 83        "mitochondria": "b6278389178f881e2ecb0f65713746096b0ef207bc6233a4e63ffbd973695c21",
 84        "er": "51df8cc2d5d1f45571828d16dd61b53623edd9ba92e32481b4f8bb1188e4bafb",
 85        "endosomes": "2ce30154604aa751860ae9711c1c49626b6bf3122c97e30dfa8991fe122ea5c4",
 86    },
 87}
 88
 89
 90def _label_url(volume, organelle):
 91    return f"{BASE_URL}/{VOLUMES[volume]['folder']}/{FILE_NAMES[organelle]}_ROI.tif"
 92
 93
 94def _slice_url(volume, idx):
 95    # The slices are stored in sub-folders with a fixed number of slices. The index is one-based.
 96    info = VOLUMES[volume]
 97    starts = info["chunk_starts"]
 98    i = max(j for j, s in enumerate(starts) if s <= idx)
 99    start = starts[i]
100    stop = starts[i + 1] - 1 if i + 1 < len(starts) else info["n_slices"]
101    return f"{BASE_URL}/{info['folder']}/{start}-{stop}/{info['pattern'].format(idx)}"
102
103
104def _download_slices(volume, tmp_dir, download):
105    """Download the slices of the ROI. The slice with (zero-based) index z is stored as the file z + 1."""
106    info = VOLUMES[volume]
107    z0 = info["roi_start"][0]
108    indices = range(z0 + 1, z0 + info["roi_shape"][0] + 1)
109
110    def _download(idx):
111        path = os.path.join(tmp_dir, f"{idx:04d}.tif")
112        with warnings.catch_warnings():
113            # There are no checksums for the individual slices.
114            warnings.simplefilter("ignore")
115            for attempt in range(3):
116                try:
117                    util.download_source(path, _slice_url(volume, idx), download)
118                    break
119                except requests.HTTPError as e:
120                    # Some slices are missing in the archive.
121                    if e.response is not None and e.response.status_code == 404:
122                        return None
123                    if attempt == 2:
124                        raise
125                except requests.ConnectionError:  # Retry on transient connection problems.
126                    if attempt == 2:
127                        raise
128                time.sleep(5 * (attempt + 1))
129        return path
130
131    with ThreadPoolExecutor(4) as pool:
132        return list(pool.map(_download, indices))
133
134
135def _create_volume(volume, h5_path, download):
136    import h5py
137    import tifffile
138
139    info = VOLUMES[volume]
140    (z0, y0, x0), (sz, sy, sx) = info["roi_start"], info["roi_shape"]
141    roi = (slice(z0, z0 + sz), slice(y0, y0 + sy), slice(x0, x0 + sx))
142
143    tmp_dir = os.path.join(os.path.dirname(h5_path), f"{volume}_tmp")
144    os.makedirs(tmp_dir, exist_ok=True)
145
146    # Download the labels and check that they agree with the ROI.
147    labels = {}
148    for organelle in ORGANELLES:
149        label_path = os.path.join(tmp_dir, f"{FILE_NAMES[organelle]}_ROI.tif")
150        util.download_source(label_path, _label_url(volume, organelle), download, CHECKSUMS[volume][organelle])
151        label = tifffile.imread(label_path)
152        # The label volume has the shape of the full stack, but it is only annotated inside of the ROI.
153        assert label.shape[0] == info["n_slices"], f"Unexpected label shape {label.shape} for {label_path}"
154        assert int(label.sum()) == int(label[roi].sum()), f"Found annotations outside of the ROI in {label_path}"
155        labels[organelle] = label[roi].astype("uint8")
156        del label
157
158    slice_paths = _download_slices(volume, tmp_dir, download)
159
160    raw = np.zeros((sz, sy, sx), dtype="uint8")
161    valid = np.zeros(sz, dtype=bool)
162    for i, path in enumerate(slice_paths):
163        if path is None:
164            continue
165        try:
166            image = tifffile.imread(path)
167        except (ValueError, tifffile.TiffFileError):  # Truncated files can't be read.
168            continue
169        raw[i] = image[y0:y0 + sy, x0:x0 + sx]
170        valid[i] = True
171
172    if not valid.any():
173        raise RuntimeError(f"Could not read any of the slices for {volume}.")
174    if not valid.all():
175        bad = (np.where(~valid)[0] + z0 + 1).tolist()
176        warnings.warn(
177            f"The slices {bad} of '{volume}' are corrupted or missing and are replaced by the closest valid slice."
178        )
179        valid_ids = np.where(valid)[0]
180        for i in np.where(~valid)[0]:
181            raw[i] = raw[valid_ids[np.argmin(np.abs(valid_ids - i))]]
182
183    tmp_h5 = f"{h5_path}.incomplete"
184    with h5py.File(tmp_h5, "w") as f:
185        f.create_dataset("raw", data=raw, compression="gzip", chunks=(1, min(sy, 512), min(sx, 512)))
186        for organelle, label in labels.items():
187            f.create_dataset(
188                f"labels/{organelle}", data=label, compression="gzip", chunks=(1, min(sy, 512), min(sx, 512))
189            )
190    os.replace(tmp_h5, h5_path)
191
192    # Remove the downloaded files.
193    for fname in os.listdir(tmp_dir):
194        os.remove(os.path.join(tmp_dir, fname))
195    os.rmdir(tmp_dir)
196
197
198def get_hela_stard3_data(
199    path: Union[os.PathLike, str],
200    volume: Literal["ctrl", "stard3", "stard3_fa_ya"],
201    download: bool = False,
202) -> str:
203    """Download the HeLa STARD3 dataset.
204
205    Only the slices of the annotated region of interest (ROI) and the corresponding labels are downloaded.
206    They are stored in a single hdf5 file per volume.
207
208    Args:
209        path: Filepath to a folder where the downloaded data will be saved.
210        volume: The volume. Either 'ctrl', 'stard3' or 'stard3_fa_ya'.
211        download: Whether to download the data if it is not present.
212
213    Returns:
214        The filepath to the hdf5 file with the data.
215    """
216    if volume not in VOLUMES:
217        raise ValueError(f"'{volume}' is not a valid volume. Please choose from {list(VOLUMES)}.")
218
219    h5_path = os.path.join(path, f"{volume}.h5")
220    if os.path.exists(h5_path):
221        return h5_path
222
223    if not download:
224        raise RuntimeError(f"Cannot find the data at {h5_path}, but download was set to False")
225
226    os.makedirs(path, exist_ok=True)
227    _create_volume(volume, h5_path, download)
228    return h5_path
229
230
231def get_hela_stard3_paths(
232    path: Union[os.PathLike, str],
233    volumes: Optional[Union[str, List[str]]] = None,
234    download: bool = False,
235) -> List[str]:
236    """Get paths to the HeLa STARD3 data.
237
238    Args:
239        path: Filepath to a folder where the downloaded data will be saved.
240        volumes: The volumes to use. Either 'ctrl', 'stard3', 'stard3_fa_ya' or a list of them.
241            By default all volumes are used.
242        download: Whether to download the data if it is not present.
243
244    Returns:
245        List of filepaths for the data.
246    """
247    if volumes is None:
248        volumes = list(VOLUMES)
249    elif isinstance(volumes, str):
250        volumes = [volumes]
251
252    return [get_hela_stard3_data(path, volume, download) for volume in volumes]
253
254
255def get_hela_stard3_dataset(
256    path: Union[os.PathLike, str],
257    patch_shape: Tuple[int, int, int],
258    organelle: Literal["mitochondria", "er", "endosomes"] = "mitochondria",
259    volumes: Optional[Union[str, List[str]]] = None,
260    download: bool = False,
261    **kwargs
262) -> Dataset:
263    """Get the HeLa STARD3 dataset for organelle segmentation in FIB-SEM.
264
265    The labels are binary (1 is the organelle).
266
267    Args:
268        path: Filepath to a folder where the downloaded data will be saved.
269        patch_shape: The patch shape to use for training.
270        organelle: The organelle to segment. Either 'mitochondria', 'er' or 'endosomes'.
271        volumes: The volumes to use. Either 'ctrl', 'stard3', 'stard3_fa_ya' or a list of them.
272            By default all volumes are used.
273        download: Whether to download the data if it is not present.
274        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
275
276    Returns:
277        The segmentation dataset.
278    """
279    if organelle not in ORGANELLES:
280        raise ValueError(f"'{organelle}' is not a valid organelle. Please choose from {ORGANELLES}.")
281
282    paths = get_hela_stard3_paths(path, volumes, download)
283
284    return torch_em.default_segmentation_dataset(
285        raw_paths=paths,
286        raw_key="raw",
287        label_paths=paths,
288        label_key=f"labels/{organelle}",
289        patch_shape=patch_shape,
290        **kwargs
291    )
292
293
294def get_hela_stard3_loader(
295    path: Union[os.PathLike, str],
296    batch_size: int,
297    patch_shape: Tuple[int, int, int],
298    organelle: Literal["mitochondria", "er", "endosomes"] = "mitochondria",
299    volumes: Optional[Union[str, List[str]]] = None,
300    download: bool = False,
301    **kwargs
302) -> DataLoader:
303    """Get the HeLa STARD3 dataloader for organelle segmentation in FIB-SEM.
304
305    Args:
306        path: Filepath to a folder where the downloaded data will be saved.
307        batch_size: The batch size for training.
308        patch_shape: The patch shape to use for training.
309        organelle: The organelle to segment. Either 'mitochondria', 'er' or 'endosomes'.
310        volumes: The volumes to use. Either 'ctrl', 'stard3', 'stard3_fa_ya' or a list of them.
311            By default all volumes are used.
312        download: Whether to download the data if it is not present.
313        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
314
315    Returns:
316        The DataLoader.
317    """
318    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
319    dataset = get_hela_stard3_dataset(path, patch_shape, organelle, volumes, download, **ds_kwargs)
320    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)