torch_em.data.datasets.electron_microscopy.aive

  1"""The AIVE dataset contains FIB-SEM volumes of cells and tissue with dense organelle annotations.
  2
  3It was released with the AI-directed voxel extraction (AIVE) approach for volume EM. The overview volumes have a
  4voxel size of 6.5 x 6.5 x 10 nm and contain semantic annotations of several organelles:
  5- Dataset 1: 724 x 720 x 2960 voxels (z, y, x), with the classes of `CLASS_IDS` except 'chromatin'.
  6- Dataset 2: 858 x 806 x 2930 voxels, with the classes of `CLASS_IDS` except 'golgi'.
  7- Dataset 3: 618 x 830 x 2948 voxels, with all classes of `CLASS_IDS`.
  8- Dataset 4 (muscle tissue): 311 x 500 x 500 voxels, with mitochondria annotations only.
  9
 10The per-dataset class ids of the original annotations differ. They are harmonized to the ids in `CLASS_IDS`
 11(0 is background) when the data is converted to hdf5 files after downloading.
 12Note that class 11 of dataset 3 is not described in the EMPIAR entry, it is assumed to be 'chromatin' (as in
 13dataset 2). Also, 'plasma_membrane' covers a large area, i.e. it is not a thin membrane annotation only.
 14
 15The data is available at https://www.ebi.ac.uk/empiar/EMPIAR-12885/ (doi: 10.6019/EMPIAR-12885).
 16The licence of the EMPIAR entry could not be confirmed (EMPIAR entries are CC0 by default).
 17The corresponding publication is https://doi.org/10.1083/jcb.202411138. Please cite it if you use this dataset
 18in your research.
 19"""
 20
 21import os
 22from typing import List, Literal, Optional, Tuple, Union
 23
 24import numpy as np
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32BASE_URL = "https://ftp.ebi.ac.uk/empiar/world_availability/12885/data"
 33URLS = {
 34    1: {
 35        "raw": f"{BASE_URL}/Set1_-_AIVE_SOURCE_DATA/Dataset1_overview.tif",
 36        "labels": f"{BASE_URL}/Set1_-_AIVE_SOURCE_DATA/Dataset1_overview%20-%20ORGANELLE%20CLASS%20LABELS.tif",
 37    },
 38    2: {
 39        "raw": f"{BASE_URL}/Set2_-_AIVE_SOURCE_DATA/Dataset2_overview.tif",
 40        "labels": f"{BASE_URL}/Set2_-_AIVE_SOURCE_DATA/Dataset2_overview%20-%20ORGANELLE%20CLASS%20LABELS.tif",
 41    },
 42    3: {
 43        "raw": f"{BASE_URL}/Set3_-_AIVE_SOURCE_DATA/Dataset3_overview.tif",
 44        "labels": f"{BASE_URL}/Set3_-_AIVE_SOURCE_DATA/Dataset3_overview%20-%20ORGANELLE%20CLASS%20LABELS.tif",
 45    },
 46    4: {
 47        "raw": f"{BASE_URL}/Set4_-_AIVE_SOURCE_DATA/Dataset4_overview.tif",
 48        "labels": f"{BASE_URL}/Set4_-_AIVE_SOURCE_DATA/Dataset4_MitoClasss.tif",
 49    },
 50}
 51CHECKSUMS = {
 52    1: {
 53        "raw": "925d6e36fdc268408b41f85563f78627c26a6c1dffabb8eec1d2c05acffe8b41",
 54        "labels": "5d343e525757143ac92db47f1be2c70b5c51e5bba3120ca1b83c732dba909eba",
 55    },
 56    2: {
 57        "raw": "71d39cc101f5d1754ab6e8950561e1cb322750c7af89e0c540be7b74bd2ac3e5",
 58        "labels": "b805978d995b7fad3a172a225f4553ceb5957fbd1493390d9162d9d5a32caf62",
 59    },
 60    3: {
 61        "raw": "682c54ef1f945f565a50577d94c334cde36ff71e520ebb22d5360a2d2b76e58a",
 62        "labels": "1546d16b4057ee6c409e52bf1a49beface2a2af8780f7a2dda8997c68b0cfa4a",
 63    },
 64    4: {
 65        "raw": "565a74783f8f9769c0054359728b37eea333fc6c05fe148ebeed4818c2454355",
 66        "labels": "be7dc8749c33a32f0d9f5a64853268e07d37418f458edc0b021aedeff32fb962",
 67    },
 68}
 69
 70# The harmonized class ids, which are used for the labels after the conversion.
 71CLASS_IDS = {
 72    "mitochondria": 1,
 73    "lipid_droplet": 2,
 74    "early_endosome": 3,
 75    "golgi": 4,
 76    "late_endosome": 5,
 77    "nucleus": 6,
 78    "cytoskeleton": 7,
 79    "er": 8,
 80    "plasma_membrane": 9,
 81    "vesicle": 10,
 82    "chromatin": 11,
 83}
 84# The class ids in the original annotations of the datasets.
 85ORIGINAL_CLASS_IDS = {
 86    1: {
 87        "mitochondria": 1, "lipid_droplet": 2, "early_endosome": 3, "golgi": 4, "late_endosome": 5, "nucleus": 6,
 88        "cytoskeleton": 7, "er": 8, "plasma_membrane": 9, "vesicle": 10,
 89    },
 90    2: {
 91        "nucleus": 1, "mitochondria": 2, "lipid_droplet": 3, "plasma_membrane": 4, "chromatin": 5, "er": 6,
 92        "cytoskeleton": 7, "late_endosome": 8, "early_endosome": 9, "vesicle": 10,
 93    },
 94    3: {
 95        "mitochondria": 1, "lipid_droplet": 2, "late_endosome": 3, "early_endosome": 4, "nucleus": 5,
 96        "cytoskeleton": 6, "golgi": 7, "er": 8, "plasma_membrane": 9, "vesicle": 10, "chromatin": 11,
 97    },
 98    4: {"mitochondria": 255},
 99}
100DATASET_IDS = [1, 2, 3, 4]
101Organelle = Literal[
102    "mitochondria", "lipid_droplet", "early_endosome", "golgi", "late_endosome", "nucleus", "cytoskeleton", "er",
103    "plasma_membrane", "vesicle", "chromatin"
104]
105
106
107class _OrganelleTransform:
108    """Convert the harmonized labels to a binary mask for one organelle (picklable for multiprocessing)."""
109    def __init__(self, class_id: int):
110        self.class_id = class_id
111
112    def __call__(self, labels):
113        return (labels == self.class_id).astype("uint8")
114
115
116def _convert_to_hdf5(raw_path, label_path, out_path, dataset_id):
117    import h5py
118    import tifffile
119
120    # Lookup table from the original to the harmonized class ids.
121    lut = np.zeros(256, dtype="uint8")
122    for name, original_id in ORIGINAL_CLASS_IDS[dataset_id].items():
123        lut[original_id] = CLASS_IDS[name]
124
125    chunks = (32, 256, 256)
126    with tifffile.TiffFile(raw_path) as f_raw, tifffile.TiffFile(label_path) as f_label:
127        shape = f_raw.series[0].shape
128        assert shape == f_label.series[0].shape, f"Raw and label shapes differ: {shape}, {f_label.series[0].shape}"
129        n_slices = shape[0]
130        tmp_path = out_path + ".tmp"
131        with h5py.File(tmp_path, "w") as f:
132            kw = dict(shape=shape, dtype="uint8", chunks=chunks, compression="gzip", compression_opts=1)
133            ds_raw = f.create_dataset("raw", **kw)
134            ds_label = f.create_dataset("labels", **kw)
135            for start in range(0, n_slices, chunks[0]):
136                stop = min(start + chunks[0], n_slices)
137                ds_raw[start:stop] = np.stack([f_raw.pages[i].asarray() for i in range(start, stop)])
138                ds_label[start:stop] = lut[np.stack([f_label.pages[i].asarray() for i in range(start, stop)])]
139    os.rename(tmp_path, out_path)
140
141
142def get_aive_data(
143    path: Union[os.PathLike, str], dataset_ids: Optional[List[int]] = None, download: bool = False
144) -> List[str]:
145    """Download the AIVE dataset.
146
147    The volumes are converted to hdf5 files with the keys 'raw' and 'labels' after downloading.
148
149    Args:
150        path: Filepath to a folder where the downloaded data will be saved.
151        dataset_ids: The ids of the datasets to download. Choose from 1, 2, 3 and 4. By default all are used.
152        download: Whether to download the data if it is not present.
153
154    Returns:
155        The filepaths to the hdf5 files with the data.
156    """
157    dataset_ids = DATASET_IDS if dataset_ids is None else dataset_ids
158    for dataset_id in dataset_ids:
159        if dataset_id not in DATASET_IDS:
160            raise ValueError(f"'{dataset_id}' is not a valid dataset id. Choose from {DATASET_IDS}.")
161
162    os.makedirs(path, exist_ok=True)
163    h5_paths = []
164    for dataset_id in dataset_ids:
165        h5_path = os.path.join(path, f"aive_dataset{dataset_id}.h5")
166        h5_paths.append(h5_path)
167        if os.path.exists(h5_path):
168            continue
169
170        tif_paths = {}
171        for key in ("raw", "labels"):
172            tif_paths[key] = os.path.join(path, f"dataset{dataset_id}_{key}.tif")
173            util.download_source(
174                tif_paths[key], URLS[dataset_id][key], download, checksum=CHECKSUMS[dataset_id][key]
175            )
176        _convert_to_hdf5(tif_paths["raw"], tif_paths["labels"], h5_path, dataset_id)
177        for tif_path in tif_paths.values():
178            os.remove(tif_path)
179
180    return h5_paths
181
182
183def get_aive_paths(
184    path: Union[os.PathLike, str],
185    dataset_ids: Optional[List[int]] = None,
186    organelle: Optional[Organelle] = None,
187    download: bool = False,
188) -> List[str]:
189    """Get paths to the AIVE data.
190
191    Args:
192        path: Filepath to a folder where the downloaded data will be saved.
193        dataset_ids: The ids of the datasets to use. Choose from 1, 2, 3 and 4. By default all datasets
194            that contain annotations for the organelle (or all datasets if no organelle is given) are used.
195        organelle: The organelle that needs to be annotated in the datasets. Used to select the default datasets.
196        download: Whether to download the data if it is not present.
197
198    Returns:
199        The filepaths to the hdf5 files with the data.
200    """
201    if organelle is not None and organelle not in CLASS_IDS:
202        raise ValueError(f"'{organelle}' is not a valid organelle. Choose from {list(CLASS_IDS)}.")
203
204    if dataset_ids is None:
205        dataset_ids = DATASET_IDS if organelle is None else [
206            i for i in DATASET_IDS if organelle in ORIGINAL_CLASS_IDS[i]
207        ]
208    elif organelle is not None:
209        missing = [i for i in dataset_ids if organelle not in ORIGINAL_CLASS_IDS[i]]
210        if missing:
211            raise ValueError(f"The datasets {missing} do not contain annotations for '{organelle}'.")
212
213    return get_aive_data(path, dataset_ids, download)
214
215
216def get_aive_dataset(
217    path: Union[os.PathLike, str],
218    patch_shape: Tuple[int, int, int],
219    dataset_ids: Optional[List[int]] = None,
220    organelle: Optional[Organelle] = "mitochondria",
221    download: bool = False,
222    **kwargs
223) -> Dataset:
224    """Get the AIVE dataset for organelle segmentation in FIB-SEM.
225
226    Args:
227        path: Filepath to a folder where the downloaded data will be saved.
228        patch_shape: The patch shape to use for training.
229        dataset_ids: The ids of the datasets to use. Choose from 1, 2, 3 and 4. By default all datasets
230            that contain annotations for the organelle are used.
231        organelle: The organelle to segment, which gives a binary segmentation task. One of the names in `CLASS_IDS`.
232            If None, the multi-class labels with the harmonized class ids of `CLASS_IDS` are returned.
233        download: Whether to download the data if it is not present.
234        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
235
236    Returns:
237        The segmentation dataset.
238    """
239    paths = get_aive_paths(path, dataset_ids, organelle, download)
240
241    if organelle is not None:
242        kwargs = util.update_kwargs(kwargs, "label_transform", _OrganelleTransform(CLASS_IDS[organelle]))
243    kwargs = util.update_kwargs(kwargs, "ndim", 3)
244
245    return torch_em.default_segmentation_dataset(
246        raw_paths=paths,
247        raw_key="raw",
248        label_paths=paths,
249        label_key="labels",
250        patch_shape=patch_shape,
251        **kwargs
252    )
253
254
255def get_aive_loader(
256    path: Union[os.PathLike, str],
257    patch_shape: Tuple[int, int, int],
258    batch_size: int,
259    dataset_ids: Optional[List[int]] = None,
260    organelle: Optional[Organelle] = "mitochondria",
261    download: bool = False,
262    **kwargs
263) -> DataLoader:
264    """Get the DataLoader for organelle segmentation in the AIVE dataset.
265
266    Args:
267        path: Filepath to a folder where the downloaded data will be saved.
268        patch_shape: The patch shape to use for training.
269        batch_size: The batch size for training.
270        dataset_ids: The ids of the datasets to use. Choose from 1, 2, 3 and 4. By default all datasets
271            that contain annotations for the organelle are used.
272        organelle: The organelle to segment, which gives a binary segmentation task. One of the names in `CLASS_IDS`.
273            If None, the multi-class labels with the harmonized class ids of `CLASS_IDS` are returned.
274        download: Whether to download the data if it is not present.
275        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
276
277    Returns:
278        The DataLoader.
279    """
280    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
281    dataset = get_aive_dataset(path, patch_shape, dataset_ids, organelle, download, **ds_kwargs)
282    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)