torch_em.data.datasets.electron_microscopy.macrophage_a431

  1"""The Macrophage / A431 dataset contains FIB-SEM and SBF-SEM volumes of human primary macrophages and A431 cells
  2with semantic annotations of organelles (nucleus, lipid droplets and, for one volume, also mitochondria,
  3Golgi apparatus and endoplasmic reticulum).
  4
  5It contains seven volumes:
  6- fib_tnf: FIB-SEM of a macrophage treated with TNFa for 1 h (2719 slices; nucleus, lipid droplets).
  7- fib_control: FIB-SEM of a macrophage in control medium (2888 slices; mitochondria, golgi, er, nucleus,
  8  lipid droplets).
  9- sbf_a431: SBF-SEM of an A431 cell (162 slices; nucleus, lipid droplets).
 10- sbf_oa_chol_7t, sbf_oa_chol_av1, sbf_oa_chol_av3: SBF-SEM of macrophages loaded with oleic acid / cholesterol
 11  (143 / 154 / 207 slices; nucleus, lipid droplets touching the nucleus).
 12- sbf_no_load: SBF-SEM of a macrophage in control medium (257 slices; nucleus, lipid droplets touching the nucleus).
 13
 14The annotations are quick semantic segmentations (created with Microscopy Image Browser) and are not necessarily
 15dense: the lipid droplet labels of the SBF-SEM volumes only cover the droplets close to the nucleus.
 16
 17The data is located at https://www.ebi.ac.uk/empiar/EMPIAR-13420/ (https://doi.org/10.6019/EMPIAR-13420).
 18Please cite the EMPIAR entry if you use this dataset in your research.
 19"""
 20
 21import os
 22from concurrent.futures import ThreadPoolExecutor
 23from typing import List, Literal, Tuple, Union
 24
 25import numpy as np
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33BASE_URL = "https://ftp.ebi.ac.uk/empiar/world_availability/13420/data"
 34
 35# For each volume: the folder, the folders and file patterns of the image / label slices, the number of slices
 36# and the label ids of the organelles. NOTE: The label ids differ between the volumes. If an organelle has
 37# several ids they are merged into a single binary mask.
 38VOLUMES = {
 39    "fib_tnf": {
 40        "folder": "FIB/240729_HumanMP_TNFa_1U",
 41        "raw_dir": "TIFs", "raw": "240729_HumanMP_TNFa_1U_{:04d}.tif",
 42        "label_dir": "Labels_NucleusLDs_TIFs", "label": "Labels_240729_HumanMP_TNFa_1U_nucleus_LD_{:04d}.tif",
 43        "n_slices": 2719,
 44        "organelles": {"nucleus": [1], "lipid_droplets": [2]},
 45    },
 46    "fib_control": {
 47        "folder": "FIB/240801_HumanMP_Control",
 48        "raw_dir": "TIFs", "raw": "240801_HumanMP_Control_{:04d}.tif",
 49        "label_dir": "Labels_Combined_model_TIFs", "label": "Labels_240801_HumanMP_Control_combined_{:04d}.tif",
 50        "n_slices": 2888,
 51        "organelles": {"mitochondria": [1], "golgi": [2], "er": [3], "nucleus": [4], "lipid_droplets": [5]},
 52    },
 53    "sbf_a431": {
 54        "folder": "SBFSEM/231106_A431_CLEM_CM1_3J",
 55        "raw_dir": "TIFs", "raw": "231106_A431_CLEM_CM1_3J_R01_{:03d}.tif",
 56        "label_dir": "Labels_TIF", "label": "Labels_231106_A431_CLEM_CM1_3J_R01_LDs_localization_{:03d}.tif",
 57        "n_slices": 162,
 58        "organelles": {"nucleus": [3], "lipid_droplets": [1, 2]},
 59    },
 60    "sbf_oa_chol_7t": {
 61        "folder": "SBFSEM/240306_Microphages_OA-Chol_7T",
 62        "raw_dir": "TIFs", "raw": "240306_OA-Chol_7T_R02_{:03d}.tif",
 63        "label_dir": "Labels_TIF", "label": "Labels_240306_OA-Chol_7T_R02_{:03d}.tif",
 64        "n_slices": 143,
 65        "organelles": {"nucleus": [1], "lipid_droplets": [2]},
 66    },
 67    "sbf_oa_chol_av1": {
 68        "folder": "SBFSEM/240307_Microphages_OA-Chol_AV-1",
 69        "raw_dir": "TIF", "raw": "240307_OA-Chol_AV_R02_AV-1_{:03d}.tif",
 70        "label_dir": "Labels_TIF", "label": "Labels_240307_OA-Chol_AV_R02_AV-1_{:03d}.tif",
 71        "n_slices": 154,
 72        "organelles": {"nucleus": [1], "lipid_droplets": [2]},
 73    },
 74    "sbf_oa_chol_av3": {
 75        "folder": "SBFSEM/240307_Microphages_OA-Chol_AV-3",
 76        "raw_dir": "TIFs", "raw": "240307_OA-Chol_AV_R01_AV-3_{:03d}.tif",
 77        "label_dir": "Labels_TIF", "label": "Labels_240307_OA-Chol_AV_R01_AV-3_{:03d}.tif",
 78        "n_slices": 207,
 79        "organelles": {"nucleus": [1], "lipid_droplets": [2]},
 80    },
 81    "sbf_no_load": {
 82        "folder": "SBFSEM/240424_Microphages_no_load",
 83        "raw_dir": "TIFs", "raw": "240424_Microphages_no_load_R02_{:03d}.tif",
 84        "label_dir": "Labels_TIF", "label": "Labels_240424_Microphages_no_load_R02_Nuclei_{:03d}.tif",
 85        "n_slices": 257,
 86        "organelles": {"nucleus": [1], "lipid_droplets": [2, 3]},
 87    },
 88}
 89# The urls of the first image slice of each volume. The urls of all other slices follow the file patterns above.
 90URLS = {
 91    name: f"{BASE_URL}/{v['folder']}/dataset/{v['raw_dir']}/{v['raw'].format(1)}" for name, v in VOLUMES.items()
 92}
 93# NOTE: There are no checksums, since the volumes are distributed as thousands of individual tif files.
 94CHECKSUMS = None
 95
 96
 97def _download_slice(args):
 98    import tifffile
 99
100    url, tmp_path, download = args
101    util.download_source(tmp_path, url, download, checksum=CHECKSUMS)
102    data = tifffile.imread(tmp_path)
103    os.remove(tmp_path)
104    return data
105
106
107def _slice_args(info, kind, tmp_dir, download):
108    if kind == "raw":
109        sub, folder, pattern = "dataset", info["raw_dir"], info["raw"]
110    else:
111        sub, folder, pattern = "labels", info["label_dir"], info["label"]
112    return [
113        (f"{BASE_URL}/{info['folder']}/{sub}/{folder}/{pattern.format(i + 1)}",
114         os.path.join(tmp_dir, f"{kind}_{i}.tif"), download)
115        for i in range(info["n_slices"])
116    ]
117
118
119def _download_volume(volume_path, name, download):
120    import h5py
121
122    info = VOLUMES[name]
123    n_slices = info["n_slices"]
124    organelles = info["organelles"]
125
126    tmp_dir = f"{volume_path}.tmp"
127    os.makedirs(tmp_dir, exist_ok=True)
128    raw_args = _slice_args(info, "raw", tmp_dir, download)
129    label_args = _slice_args(info, "label", tmp_dir, download)
130
131    tmp_h5 = f"{volume_path}.incomplete"
132    ds_raw, ds_labels = None, {}
133    # Download and convert the data in blocks of slices to keep the memory footprint low.
134    step = 16
135    with h5py.File(tmp_h5, "w") as f, ThreadPoolExecutor(4) as pool:
136        for start in range(0, n_slices, step):
137            raw = np.stack(list(pool.map(_download_slice, raw_args[start:start + step])))
138            labels = np.stack(list(pool.map(_download_slice, label_args[start:start + step])))
139            assert raw.shape == labels.shape, f"Shape mismatch of raw and labels: {raw.shape}, {labels.shape}"
140
141            if ds_raw is None:
142                shape = (n_slices,) + raw.shape[1:]
143                chunks = (1, min(512, shape[1]), min(512, shape[2]))
144                ds_raw = f.create_dataset("raw", shape=shape, dtype=raw.dtype, chunks=chunks, compression="gzip")
145                for organelle in organelles:
146                    ds_labels[organelle] = f.create_dataset(
147                        f"labels/{organelle}", shape=shape, dtype="uint8", chunks=chunks, compression="gzip"
148                    )
149
150            stop = start + len(raw)
151            ds_raw[start:stop] = raw
152            for organelle, ids in organelles.items():
153                ds_labels[organelle][start:stop] = np.isin(labels, ids).astype("uint8")
154
155    os.replace(tmp_h5, volume_path)
156    os.rmdir(tmp_dir)
157
158
159def get_macrophage_a431_data(
160    path: Union[os.PathLike, str],
161    name: Literal[
162        "fib_tnf", "fib_control", "sbf_a431", "sbf_oa_chol_7t", "sbf_oa_chol_av1", "sbf_oa_chol_av3", "sbf_no_load"
163    ],
164    download: bool = False,
165) -> str:
166    """Download a volume of the Macrophage / A431 dataset.
167
168    The volumes are large (the FIB-SEM volumes have ~3 GB, the SBF-SEM volumes ~4-12 GB). They are downloaded
169    slice by slice and converted to an hdf5 file with the raw data and one binary label volume per organelle.
170
171    Args:
172        path: Filepath to a folder where the downloaded data will be saved.
173        name: The name of the volume.
174        download: Whether to download the data if it is not present.
175
176    Returns:
177        The filepath to the hdf5 file with the volume.
178    """
179    if name not in VOLUMES:
180        raise ValueError(f"'{name}' is not a valid volume name. Please choose from {list(VOLUMES)}.")
181
182    volume_path = os.path.join(path, f"{name}.h5")
183    if os.path.exists(volume_path):
184        return volume_path
185
186    if not download:
187        raise RuntimeError(f"Cannot find the data at {volume_path}, but download was set to False")
188
189    os.makedirs(path, exist_ok=True)
190    _download_volume(volume_path, name, download)
191    return volume_path
192
193
194def get_macrophage_a431_paths(
195    path: Union[os.PathLike, str],
196    names: Union[str, List[str], None] = None,
197    download: bool = False,
198) -> List[str]:
199    """Get paths to the Macrophage / A431 data.
200
201    Args:
202        path: Filepath to a folder where the downloaded data will be saved.
203        names: The names of the volumes. By default the 'fib_control' volume is used.
204            Choose from 'fib_tnf', 'fib_control', 'sbf_a431', 'sbf_oa_chol_7t', 'sbf_oa_chol_av1',
205            'sbf_oa_chol_av3' and 'sbf_no_load'.
206        download: Whether to download the data if it is not present.
207
208    Returns:
209        List of filepaths for the hdf5 files.
210    """
211    if names is None:
212        names = ["fib_control"]
213    elif isinstance(names, str):
214        names = [names]
215
216    return [get_macrophage_a431_data(path, name, download) for name in names]
217
218
219def get_macrophage_a431_dataset(
220    path: Union[os.PathLike, str],
221    patch_shape: Tuple[int, int, int],
222    names: Union[str, List[str], None] = None,
223    organelle: Literal["nucleus", "lipid_droplets", "mitochondria", "golgi", "er"] = "nucleus",
224    download: bool = False,
225    **kwargs
226) -> Dataset:
227    """Get the Macrophage / A431 dataset for organelle segmentation in FIB-SEM and SBF-SEM.
228
229    Args:
230        path: Filepath to a folder where the downloaded data will be saved.
231        patch_shape: The patch shape to use for training.
232        names: The names of the volumes. By default the 'fib_control' volume is used.
233            Choose from 'fib_tnf', 'fib_control', 'sbf_a431', 'sbf_oa_chol_7t', 'sbf_oa_chol_av1',
234            'sbf_oa_chol_av3' and 'sbf_no_load'.
235        organelle: The organelle to segment. Either 'nucleus', 'lipid_droplets', 'mitochondria', 'golgi' or 'er'.
236            Mitochondria, golgi and er are only available for the 'fib_control' volume.
237        download: Whether to download the data if it is not present.
238        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
239
240    Returns:
241        The segmentation dataset.
242    """
243    all_organelles = {o for v in VOLUMES.values() for o in v["organelles"]}
244    if organelle not in all_organelles:
245        raise ValueError(f"'{organelle}' is not a valid organelle. Please choose from {sorted(all_organelles)}.")
246
247    if names is None:
248        names = ["fib_control"]
249    elif isinstance(names, str):
250        names = [names]
251    for name in names:
252        if name in VOLUMES and organelle not in VOLUMES[name]["organelles"]:
253            raise ValueError(f"The organelle '{organelle}' is not annotated in '{name}'.")
254
255    volume_paths = get_macrophage_a431_paths(path, names, download)
256
257    kwargs = util.update_kwargs(kwargs, "ndim", 3)
258
259    return torch_em.default_segmentation_dataset(
260        raw_paths=volume_paths,
261        raw_key="raw",
262        label_paths=volume_paths,
263        label_key=f"labels/{organelle}",
264        patch_shape=patch_shape,
265        **kwargs
266    )
267
268
269def get_macrophage_a431_loader(
270    path: Union[os.PathLike, str],
271    batch_size: int,
272    patch_shape: Tuple[int, int, int],
273    names: Union[str, List[str], None] = None,
274    organelle: Literal["nucleus", "lipid_droplets", "mitochondria", "golgi", "er"] = "nucleus",
275    download: bool = False,
276    **kwargs
277) -> DataLoader:
278    """Get the Macrophage / A431 dataloader for organelle segmentation in FIB-SEM and SBF-SEM.
279
280    Args:
281        path: Filepath to a folder where the downloaded data will be saved.
282        batch_size: The batch size for training.
283        patch_shape: The patch shape to use for training.
284        names: The names of the volumes. By default the 'fib_control' volume is used.
285            Choose from 'fib_tnf', 'fib_control', 'sbf_a431', 'sbf_oa_chol_7t', 'sbf_oa_chol_av1',
286            'sbf_oa_chol_av3' and 'sbf_no_load'.
287        organelle: The organelle to segment. Either 'nucleus', 'lipid_droplets', 'mitochondria', 'golgi' or 'er'.
288            Mitochondria, golgi and er are only available for the 'fib_control' volume.
289        download: Whether to download the data if it is not present.
290        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
291
292    Returns:
293        The DataLoader.
294    """
295    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
296    dataset = get_macrophage_a431_dataset(path, patch_shape, names, organelle, download, **ds_kwargs)
297    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)