torch_em.data.datasets.electron_microscopy.parlakgul_fasting_liver

  1"""The Parlakgul fasting liver dataset contains FIB-SEM volumes of mouse liver in fasting and obesity
  2with dense organelle segmentations.
  3
  4Four FIB-SEM volumes (8 nm isotropic) are available:
  5- 9428 (lean fasted): 7984 x 6984 x 9501 voxels
  6- 8970 (obese fasted): 9402 x 6949 x 8450 voxels
  7- 2679 (obese, LacZ control): 6500 x 6500 x 7705 voxels
  8- 1858 (obese, RRBP1 overexpression): 9375 x 10625 x 5795 voxels (the first 65 slices contain resin and are removed)
  9
 10Segmentations are available via the `label_choice` parameter:
 11- "er": endoplasmic reticulum, binary (all samples)
 12- "mito": mitochondria, instance segmentation (all samples)
 13- "lipid_droplet": lipid droplets, binary (9428 and 8970)
 14- "nuclear_membrane": nuclear membrane, binary (9428 and 8970)
 15- "cells": cells, instance segmentation (9428 and 8970). Only a subset of the cells is annotated, the remaining
 16  cells are background, so these labels are not suitable for dense cell segmentation.
 17
 18The binary labels are converted to 0 (background) and 1 (foreground). The instance ids are made consecutive
 19within each requested subvolume. Note that the labels of a large bounding box are held in memory while streaming
 20(4 bytes per voxel for "mito").
 21
 22Data is streamed lazily from EMPIAR-12017 via HTTP range requests: only the rows of the raw TIFFs and the
 23segmentation slices (from the ZIP archives) of the requested bounding box are fetched and cached as zarr v3.
 24
 25Bounding boxes are specified as (x_min, x_max, y_min, y_max, z_min, z_max) in voxels.
 26
 27This dataset is from the publication https://doi.org/10.1038/s41467-024-48272-7.
 28Please cite it if you use this dataset in your research.
 29
 30The data is publicly available at https://www.ebi.ac.uk/empiar/EMPIAR-12017/.
 31"""
 32
 33import os
 34import shutil
 35import zipfile
 36from concurrent import futures
 37from typing import Dict, List, Literal, Tuple, Union
 38
 39import numpy as np
 40from torch.utils.data import DataLoader, Dataset
 41
 42import torch_em
 43from .. import util
 44from .parlakgul_liver import (
 45    _HttpFile, _bbox_to_str, _index_zip_slices, _read_raw_slice, _read_zip_slice, N_THREADS
 46)
 47
 48
 49EMPIAR_BASE = "https://ftp.ebi.ac.uk/empiar/world_availability/12017/data"
 50
 51# The raw file of slice z is 'raw_pattern' formatted with z + 'raw_z_offset'. The first 'z_offset' slices of
 52# a sample are removed (they contain resin), so the shape and all slice indices are relative to this offset.
 53# The segmentation slices are found by the slice number in their name, see '_index_zip_slices'.
 54PARLAKGUL_FASTING_SAMPLES: Dict[str, dict] = {
 55    "9428": {
 56        "condition": "lean_fasted",
 57        "dir": "9428%20Lean%20Fasted",
 58        "raw_dir": "Lean_Fasted_9428_raw_8-bit",
 59        "raw_pattern": "Lean_Fasted_9428_raw_T001_Z{z:03d}_C05.tif",
 60        "z_offset": 0,
 61        "shape": (9501, 6984, 7984),
 62        "seg_zips": {
 63            "er": "Lean_Fasted_9428_ER_8-bit.zip",
 64            "mito": "Lean_Fasted_9428_mitochondria_instance_32-bit.zip",
 65            "lipid_droplet": "Lean_Fasted_9428_LD_8-bit.zip",
 66            "nuclear_membrane": "Lean_Fasted_9428_nucleus_8-bit.zip",
 67            "cells": "Lean_Fasted_9428_plasma_membrane_8-bit.zip",
 68        },
 69    },
 70    "8970": {
 71        "condition": "obese_fasted",
 72        "dir": "8970%20Obese%20Fasted",
 73        "raw_dir": "Obese_Fasted_8970_raw_8-bit",
 74        "raw_pattern": "Obese_Fasted_8970_raw_T001_Z{z:03d}_C05.tif",
 75        "z_offset": 0,
 76        "shape": (8450, 6949, 9402),
 77        "seg_zips": {
 78            "er": "Obese_Fasted_8970_ER_8-bit.zip",
 79            "mito": "Obese_fasted_8970_mitochondria_instance_32-bit.zip",
 80            "lipid_droplet": "Obese_Fasted_8970_LD_8-bit.zip",
 81            "nuclear_membrane": "Obese_Fasted_8970_nucleus_8-bit.zip",
 82            "cells": "Obese_fasted_8970_plasma_membrane_8-bit.zip",
 83        },
 84    },
 85    "2679": {
 86        "condition": "obese_lacz",
 87        "dir": "Obese%20LacZ",
 88        "raw_dir": "ob_LacZ_2679_raw_8-bit",
 89        "raw_pattern": "2679_Ob_LacZ_T001_Z{z:03d}_C01.tif",
 90        "z_offset": 0,
 91        "shape": (7705, 6500, 6500),
 92        "seg_zips": {
 93            "er": "ob_LacZ_2679_ER_8-bit.zip",
 94            "mito": "ob_LacZ_2679_mitochondria-instance_32-bit.zip",
 95        },
 96    },
 97    "1858": {
 98        "condition": "obese_rrbp1",
 99        "dir": "Obese%20RRBP1",
100        "raw_dir": "1858_Ob_RRBP1-OE_Raw_8-bit",
101        "raw_pattern": "1858_Ob_RRBP1-OE_Raw_T001_Z{z:03d}_C01.tif",
102        "z_offset": 65,
103        "shape": (5795, 10625, 9375),
104        "seg_zips": {
105            "er": "ob_RRBP1_OE_1858_ER_8-bit.zip",
106            "mito": "ob_RRBP1_OE_1858_mito-instance_32-bit.zip",
107        },
108    },
109}
110
111INSTANCE_LABELS = {"mito": np.uint32, "cells": np.uint8}
112CHUNK_SHAPE = (64, 256, 256)
113LabelChoice = Literal["er", "mito", "lipid_droplet", "nuclear_membrane", "cells"]
114SampleChoice = Literal["9428", "8970", "2679", "1858"]
115
116
117def get_parlakgul_fasting_liver_data(
118    path: Union[os.PathLike, str],
119    bounding_box: Tuple[int, int, int, int, int, int],
120    sample: SampleChoice = "9428",
121    label_choice: LabelChoice = "mito",
122    download: bool = False,
123) -> str:
124    """Stream a subvolume from the Parlakgul fasting liver dataset and cache it as a zarr v3 store.
125
126    Args:
127        path: Filepath to a folder where the cached zarr store will be saved.
128        bounding_box: The region to fetch as (x_min, x_max, y_min, y_max, z_min, z_max)
129            in voxel coordinates at 8 nm isotropic resolution.
130        sample: Which liver sample to use. One of "9428" (lean fasted), "8970" (obese fasted),
131            "2679" (obese LacZ) or "1858" (obese RRBP1).
132        label_choice: Which segmentation to use as labels. One of "er", "mito", "lipid_droplet",
133            "nuclear_membrane" or "cells". The last three are only available for "9428" and "8970".
134        download: Whether to stream and cache the data if it is not present.
135
136    Returns:
137        The filepath to the cached zarr store.
138    """
139    import zarr
140    from skimage.segmentation import relabel_sequential
141    from zarr.codecs import BloscCodec
142
143    if sample not in PARLAKGUL_FASTING_SAMPLES:
144        raise ValueError(f"'{sample}' is not a valid sample. Choose from {list(PARLAKGUL_FASTING_SAMPLES)}.")
145    sample_info = PARLAKGUL_FASTING_SAMPLES[sample]
146    if label_choice not in sample_info["seg_zips"]:
147        raise ValueError(
148            f"label_choice='{label_choice}' is not available for sample='{sample}'. "
149            f"Choose from {list(sample_info['seg_zips'])}."
150        )
151
152    x_min, x_max, y_min, y_max, z_min, z_max = bounding_box
153    n_z, n_y, n_x = sample_info["shape"]
154    if not (0 <= x_min < x_max <= n_x and 0 <= y_min < y_max <= n_y and 0 <= z_min < z_max <= n_z):
155        raise ValueError(f"The bounding box {bounding_box} is not inside sample '{sample}' of shape (z, y, x) "
156                         f"{sample_info['shape']}.")
157
158    os.makedirs(str(path), exist_ok=True)
159    zarr_path = os.path.join(str(path), f"{sample}_{label_choice}_{_bbox_to_str(bounding_box)}.zarr")
160    if os.path.exists(os.path.join(zarr_path, "raw")) and os.path.exists(os.path.join(zarr_path, "labels")):
161        return zarr_path
162
163    if not download:
164        raise RuntimeError(
165            f"No cached data found at '{zarr_path}'. Set download=True to stream it from EMPIAR."
166        )
167
168    shape = (z_max - z_min, y_max - y_min, x_max - x_min)
169    label_dtype = INSTANCE_LABELS.get(label_choice, np.uint8)
170    raw_arr = np.zeros(shape, dtype=np.uint8)
171    lbl_arr = np.zeros(shape, dtype=label_dtype)
172
173    sample_base = f"{EMPIAR_BASE}/{sample_info['dir']}"
174    seg_zip_url = f"{sample_base}/{sample_info['seg_zips'][label_choice]}"
175    z_offset = sample_info["z_offset"]
176
177    print(f"Streaming Parlakgul fasting liver {sample} ({sample_info['condition']}) EM + {label_choice} ...")
178
179    # The ZIP directory is read once, the segmentation slices are then fetched one by one from it.
180    zf = zipfile.ZipFile(_HttpFile(seg_zip_url))
181    slices = _index_zip_slices(zf)
182    if len(slices) != n_z + z_offset:
183        raise RuntimeError(f"Expected {n_z + z_offset} segmentation slices in {seg_zip_url}, found {len(slices)}.")
184
185    def _fetch_raw(i):
186        # The raw files are numbered from 1.
187        name = sample_info["raw_pattern"].format(z=z_min + i + z_offset + 1)
188        raw_arr[i] = _read_raw_slice(f"{sample_base}/{sample_info['raw_dir']}/{name}", x_min, x_max, y_min, y_max)
189
190    with futures.ThreadPoolExecutor(N_THREADS) as tp:
191        raw_tasks = [tp.submit(_fetch_raw, i) for i in range(shape[0])]
192        for i, z in enumerate(range(z_min, z_max)):
193            seg = _read_zip_slice(zf, slices[z + z_offset], x_min, x_max, y_min, y_max)
194            lbl_arr[i] = seg if label_choice in INSTANCE_LABELS else seg > 0
195            if (i + 1) % 10 == 0:
196                print(f"  {i + 1}/{shape[0]} segmentation slices done")
197        for task in raw_tasks:
198            task.result()
199
200    # The instance ids are global to the whole volume, so they are made consecutive within the subvolume.
201    if label_choice in INSTANCE_LABELS:
202        lbl_arr = relabel_sequential(lbl_arr)[0].astype(label_dtype)
203
204    # The store is written under a temporary name and moved in place once it is complete,
205    # so that an interrupted download is not mistaken for a complete one.
206    tmp_path = f"{zarr_path}.incomplete"
207    if os.path.exists(tmp_path):
208        shutil.rmtree(tmp_path)
209    root = zarr.open_group(tmp_path, mode="w")
210    root.attrs.update({
211        "bounding_box": list(bounding_box), "sample": sample, "label_choice": label_choice, "resolution_nm": [8, 8, 8],
212    })
213    for name, data, shuffle in (("raw", raw_arr, "shuffle"), ("labels", lbl_arr, "bitshuffle")):
214        arr = root.create_array(
215            name, shape=data.shape, chunks=CHUNK_SHAPE, dtype=data.dtype,
216            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
217        )
218        arr[:] = data
219    if os.path.exists(zarr_path):
220        shutil.rmtree(zarr_path)
221    os.rename(tmp_path, zarr_path)
222
223    print(f"Cached to {zarr_path} (shape {shape})")
224    return zarr_path
225
226
227def get_parlakgul_fasting_liver_paths(
228    path: Union[os.PathLike, str],
229    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
230    sample: SampleChoice = "9428",
231    label_choice: LabelChoice = "mito",
232    download: bool = False,
233) -> List[str]:
234    """Get paths to Parlakgul fasting liver zarr stores.
235
236    Args:
237        path: Filepath to a folder where the cached zarr stores will be saved.
238        bounding_boxes: List of regions to fetch, each as
239            (x_min, x_max, y_min, y_max, z_min, z_max) in voxel coordinates.
240        sample: Which liver sample to use.
241        label_choice: Which segmentation to use as labels.
242        download: Whether to stream and cache the data if it is not present.
243
244    Returns:
245        List of filepaths to the cached zarr stores.
246    """
247    return [get_parlakgul_fasting_liver_data(path, bbox, sample, label_choice, download) for bbox in bounding_boxes]
248
249
250def get_parlakgul_fasting_liver_dataset(
251    path: Union[os.PathLike, str],
252    patch_shape: Tuple[int, int, int],
253    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
254    sample: SampleChoice = "9428",
255    label_choice: LabelChoice = "mito",
256    download: bool = False,
257    **kwargs,
258) -> Dataset:
259    """Get the Parlakgul fasting liver dataset for organelle segmentation.
260
261    Args:
262        path: Filepath to a folder where the cached zarr stores will be saved.
263        patch_shape: The patch shape (z, y, x) to use for training.
264        bounding_boxes: List of subvolumes to use, each as
265            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
266        sample: Which liver sample to use. One of "9428", "8970", "2679", "1858".
267        label_choice: Which segmentation to use. One of "er", "mito", "lipid_droplet",
268            "nuclear_membrane" or "cells".
269        download: Whether to stream and cache data if not already present.
270        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
271
272    Returns:
273        The segmentation dataset.
274    """
275    assert len(patch_shape) == 3
276
277    paths = get_parlakgul_fasting_liver_paths(path, bounding_boxes, sample, label_choice, download)
278
279    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
280
281    return torch_em.default_segmentation_dataset(
282        raw_paths=paths,
283        raw_key="raw",
284        label_paths=paths,
285        label_key="labels",
286        patch_shape=patch_shape,
287        **kwargs,
288    )
289
290
291def get_parlakgul_fasting_liver_loader(
292    path: Union[os.PathLike, str],
293    patch_shape: Tuple[int, int, int],
294    batch_size: int,
295    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
296    sample: SampleChoice = "9428",
297    label_choice: LabelChoice = "mito",
298    download: bool = False,
299    **kwargs,
300) -> DataLoader:
301    """Get the DataLoader for organelle segmentation in the Parlakgul fasting liver dataset.
302
303    Args:
304        path: Filepath to a folder where the cached zarr stores will be saved.
305        patch_shape: The patch shape (z, y, x) to use for training.
306        batch_size: The batch size for training.
307        bounding_boxes: List of subvolumes to use, each as
308            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
309        sample: Which liver sample to use. One of "9428", "8970", "2679", "1858".
310        label_choice: Which segmentation to use. One of "er", "mito", "lipid_droplet",
311            "nuclear_membrane" or "cells".
312        download: Whether to stream and cache data if not already present.
313        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
314            or for the PyTorch DataLoader.
315
316    Returns:
317        The DataLoader.
318    """
319    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
320    dataset = get_parlakgul_fasting_liver_dataset(
321        path, patch_shape, bounding_boxes, sample=sample, label_choice=label_choice,
322        download=download, **ds_kwargs
323    )
324    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)