torch_em.data.datasets.electron_microscopy.parlakgul_liver

  1"""The Parlakgul liver dataset contains FIB-SEM volumes of mouse liver with dense
  2semantic segmentation of 7 organelle classes. All labels are binary semantic masks
  3(0=background, 255=foreground) - not instance segmentation.
  4
  5Four FIB-SEM volumes are available across lean and obese conditions:
  6- 6461 (lean): 12000 x 8000 x 5638 voxels, 8 nm isotropic
  7- 6464 (obese 1): 9112 x 10200 x 7896 voxels, 8 nm isotropic
  8- 9430 (obese 2): 8000 x 8050 x 8501 voxels, 8 nm isotropic
  9- 1857 (obese Climp63): 9700 x 9650 x 3629 voxels, 8 nm isotropic
 10
 11Seven semantic segmentation classes are available via the `label_choice` parameter:
 12- "er": endoplasmic reticulum
 13- "er_sheets": ER sheets
 14- "er_tubules": ER tubules
 15- "mito": mitochondria
 16- "lipid_droplet": lipid droplets
 17- "nuclear_membrane": nuclear membrane (not available for 1857)
 18- "plasma_membrane": plasma membrane
 19
 20Data is streamed lazily from EMPIAR-10791 via HTTP range requests: only the rows of the raw TIFFs and the
 21segmentation slices (from the ZIP archives) of the requested bounding box are fetched and cached as zarr v3.
 22
 23Bounding boxes are specified as (x_min, x_max, y_min, y_max, z_min, z_max) in voxels.
 24
 25This dataset is from the publication https://doi.org/10.1038/s41586-022-04518-2.
 26Please cite it if you use this dataset in your research.
 27
 28The data is publicly available at https://www.ebi.ac.uk/empiar/EMPIAR-10791/.
 29"""
 30
 31import hashlib
 32import io
 33import os
 34import re
 35import shutil
 36import zipfile
 37from concurrent import futures
 38from typing import Dict, List, Literal, Tuple, Union
 39
 40import numpy as np
 41from torch.utils.data import DataLoader, Dataset
 42
 43import torch_em
 44from .. import util
 45
 46
 47EMPIAR_BASE = "https://ftp.ebi.ac.uk/empiar/world_availability/10791/data"
 48PARLAKGUL_PAPER_BASE = (
 49    "Parlakgul%20et%20al%20-%20Regulation%20of%20liver%20subcellular%20architecture%20"
 50    "controls%20metabolic%20homeostasis/FIB-SEM%20Raw%20and%20Segmentation%20Data"
 51)
 52
 53# The raw file of slice z is 'raw_pattern' formatted with z + 'raw_z_offset' (the raw files of 6461 are numbered
 54# from 0, the others from 1). The segmentation slices are found by the slice number in their name, see
 55# '_index_zip_slices'.
 56PARLAKGUL_SAMPLES: Dict[str, dict] = {
 57    "6461": {
 58        "condition": "lean",
 59        "raw_dir": "6461%20-%20Lean%20Liver/6461%20Lean%20Liver%20-%20Raw",
 60        "seg_dir": "6461%20-%20Lean%20Liver/6461%20Lean%20Liver%20-%20Segmentation",
 61        "raw_pattern": "Gunes_WT1_8x8x8nm_3MHz.{z:04d}.tif",
 62        "raw_z_offset": 0,
 63        "shape": (5638, 8000, 12000),
 64        "seg_zips": {
 65            "er": "6461%20Lean%20ER.zip",
 66            "er_sheets": "6461%20Lean%20ER%20Sheets.zip",
 67            "er_tubules": "6461%20Lean%20ER%20Tubules.zip",
 68            "mito": "6461%20Lean%20Mitochondria.zip",
 69            "lipid_droplet": "6461%20Lean%20Lipid%20Droplet.zip",
 70            "nuclear_membrane": "6461%20Lean%20Nuclear%20membrane.zip",
 71            "plasma_membrane": "6461%20Lean%20Plasma%20Membrane.zip",
 72        },
 73    },
 74    "6464": {
 75        "condition": "obese1",
 76        "raw_dir": "6464%20-%20Obese1%20Liver/6464%20Obese1%20Liver%20-%20Raw",
 77        "seg_dir": "6464%20-%20Obese1%20Liver/6464%20Obese1%20Liver%20-%20Segmentation",
 78        "raw_pattern": "6464%20Ob%20Fed%20Raw%20with%20All_T001_Z{z:03d}_C09.tif",
 79        "raw_z_offset": 1,
 80        "shape": (7896, 10200, 9112),
 81        "seg_zips": {
 82            "er": "6464%20Obese1%20ER.zip",
 83            "er_sheets": "6464%20Obese1%20ER%20Sheets.zip",
 84            "er_tubules": "6464%20Obese1%20ER%20Tubules.zip",
 85            "mito": "6464%20Obese1%20Mitochondria.zip",
 86            "lipid_droplet": "6464%20Obese1%20Lipid%20Droplet.zip",
 87            "nuclear_membrane": "6464%20Obese1%20Nuclear%20membrane.zip",
 88            "plasma_membrane": "6464%20Obese1%20Plasma%20membrane.zip",
 89        },
 90    },
 91    "9430": {
 92        "condition": "obese2",
 93        "raw_dir": "9430%20-%20Obese2%20Liver/9430%20Obese2%20Liver%20-%20Raw/9430%20raw",
 94        "seg_dir": "9430%20-%20Obese2%20Liver/9430%20Obese2%20Liver%20-%20Segmentation",
 95        "raw_pattern": "9430%20Ob%20Fed%20Raw_T001_Z{z:03d}_C07.tif",
 96        "raw_z_offset": 1,
 97        "shape": (8501, 8050, 8000),
 98        "seg_zips": {
 99            "er": "9430%20Obese2%20ER.zip",
100            "er_sheets": "9430%20Obese2%20ER%20Sheets.zip",
101            "er_tubules": "9430%20Obese2%20ER%20Tubules.zip",
102            "mito": "9430%20Obese2%20Mitochondria.zip",
103            "lipid_droplet": "9430%20Obese2%20Lipid%20Droplet.zip",
104            "nuclear_membrane": "9430%20Obese2%20nuclear%20membrane.zip",
105            "plasma_membrane": "9430%20Obese2%20Plasma%20membrane.zip",
106        },
107    },
108    "1857": {
109        "condition": "obese_climp63",
110        "raw_dir": "1857%20-%20Obese%20Climp-63%20Liver/1857%20Obese%20Climp-63%20Raw",
111        "seg_dir": "1857%20-%20Obese%20Climp-63%20Liver/1857%20Obese%20Climp-63%20Segmentation",
112        "raw_pattern": "1857_Obese_Climp63_raw_T001_Z{z:03d}_C07.tif",
113        "raw_z_offset": 1,
114        "shape": (3629, 9650, 9700),
115        "seg_zips": {
116            "er": "1857_Obese_Climp63_ER.zip",
117            "er_sheets": "1857_Obese_Climp63_ER-sheets.zip",
118            "er_tubules": "1857_Obese_Climp63_ER-tubules.zip",
119            "mito": "1857_Obese_Climp63_Mitochondria.zip",
120            "lipid_droplet": "1857_Obese_Climp63_Lipid-droplet.zip",
121            "plasma_membrane": "1857_Obese_Climp63_Plasma-membrane.zip",
122        },
123    },
124}
125
126PARLAKGUL_CHUNK_SHAPE = (64, 256, 256)
127N_THREADS = 4
128LabelChoice = Literal[
129    "er", "er_sheets", "er_tubules", "mito", "lipid_droplet", "nuclear_membrane", "plasma_membrane"
130]
131
132
133def _bbox_to_str(bbox):
134    return hashlib.md5("_".join(str(v) for v in bbox).encode()).hexdigest()[:12]
135
136
137class _HttpFile:
138    """Seekable file-like object backed by HTTP range requests.
139
140    Reading a TIFF or ZIP header issues many small reads, so small reads are served from cached 1 MB blocks.
141    Large reads are fetched with a single range request.
142    """
143    block_size = 1 << 20
144    n_cached_blocks = 8
145
146    def __init__(self, url):
147        import requests
148        self.url = url
149        self._pos = 0
150        self._blocks = {}
151        r = requests.head(url, timeout=30)
152        r.raise_for_status()
153        self._size = int(r.headers["Content-Length"])
154
155    def _fetch(self, start, stop):
156        import time
157        import requests
158        for attempt in range(5):
159            try:
160                r = requests.get(self.url, headers={"Range": f"bytes={start}-{stop - 1}"}, timeout=300)
161                r.raise_for_status()
162                if len(r.content) != stop - start:
163                    raise IOError(f"Expected {stop - start} bytes from {self.url}, got {len(r.content)}.")
164                return r.content
165            except Exception:
166                if attempt == 4:
167                    raise
168                time.sleep(2 ** attempt)
169
170    def _block(self, index):
171        if index not in self._blocks:
172            if len(self._blocks) >= self.n_cached_blocks:
173                self._blocks.pop(next(iter(self._blocks)))
174            start = index * self.block_size
175            self._blocks[index] = self._fetch(start, min(start + self.block_size, self._size))
176        return self._blocks[index]
177
178    def read(self, n=-1):
179        stop = self._size if n is None or n < 0 else min(self._pos + n, self._size)
180        if self._pos >= stop:
181            return b""
182        if stop - self._pos > self.block_size:
183            data = self._fetch(self._pos, stop)
184        else:
185            first, last = self._pos // self.block_size, (stop - 1) // self.block_size
186            data = b"".join(self._block(index) for index in range(first, last + 1))
187            data = data[self._pos - first * self.block_size:stop - first * self.block_size]
188        self._pos = stop
189        return data
190
191    def seek(self, pos, whence=0):
192        if whence == 0:
193            self._pos = pos
194        elif whence == 1:
195            self._pos += pos
196        elif whence == 2:
197            self._pos = self._size + pos
198        self._pos = max(0, min(self._pos, self._size))
199        return self._pos
200
201    def tell(self):
202        return self._pos
203
204    def seekable(self):
205        return True
206
207    def readable(self):
208        return True
209
210    def __enter__(self):
211        return self
212
213    def __exit__(self, *args):
214        pass
215
216
217def _index_zip_slices(zf):
218    """Map the z index of each segmentation slice to its member in the ZIP archive.
219
220    The archives number their slices differently ('0000.tiff', '000' to '7895' without consistent zero padding,
221    or '..._Z001_C01.tif' starting at 1), and do not always store them in z order. So the z index is parsed from
222    the slice number in the member name, relative to the smallest slice number.
223    """
224    numbers = {}
225    for info in zf.infolist():
226        if not info.filename.lower().endswith((".tif", ".tiff")):
227            continue
228        stem = os.path.splitext(os.path.basename(info.filename))[0]
229        match = re.search(r"Z(\d+)_C\d+$", stem) or re.search(r"(\d+)$", stem)
230        numbers[int(match.group(1))] = info
231    first = min(numbers)
232    return {number - first: info for number, info in numbers.items()}
233
234
235def _read_zip_slice(zf, info, x_min, x_max, y_min, y_max):
236    """Extract one segmentation slice from the remote ZIP archive and crop it to the requested region."""
237    import tifffile
238    return tifffile.imread(io.BytesIO(zf.read(info)))[y_min:y_max, x_min:x_max]
239
240
241def _read_raw_slice(raw_url, x_min, x_max, y_min, y_max):
242    """Read the rows y_min:y_max of one raw TIFF slice and crop them to the requested region.
243
244    The raw TIFFs are uncompressed with contiguous strips, so the rows are fetched with a single range request.
245    Other layouts fall back to reading the whole slice.
246    """
247    import tifffile
248
249    fh = _HttpFile(raw_url)
250    with tifffile.TiffFile(fh) as tif:
251        page = tif.pages[0]
252        width, rows_per_strip = page.shape[1], page.rowsperstrip
253        offsets = np.asarray(page.dataoffsets)
254        contiguous = page.compression == 1 and page.dtype == np.uint8 and page.samplesperpixel == 1 and bool(
255            np.all(np.diff(offsets) == rows_per_strip * width)
256        )
257    if not contiguous:
258        fh.seek(0)
259        return tifffile.imread(io.BytesIO(fh.read()))[y_min:y_max, x_min:x_max]
260
261    fh.seek(int(offsets[y_min // rows_per_strip]) + (y_min % rows_per_strip) * width)
262    rows = np.frombuffer(fh.read((y_max - y_min) * width), dtype="uint8").reshape(y_max - y_min, width)
263    return rows[:, x_min:x_max]
264
265
266def get_parlakgul_liver_data(
267    path: Union[os.PathLike, str],
268    bounding_box: Tuple[int, int, int, int, int, int],
269    sample: Literal["6461", "6464", "9430", "1857"] = "6461",
270    label_choice: LabelChoice = "mito",
271    download: bool = False,
272) -> str:
273    """Stream a subvolume from the Parlakgul liver dataset and cache it as a zarr v3 store.
274
275    Args:
276        path: Filepath to a folder where the cached zarr store will be saved.
277        bounding_box: The region to fetch as (x_min, x_max, y_min, y_max, z_min, z_max)
278            in voxel coordinates at 8 nm isotropic resolution.
279        sample: Which liver sample to use. One of "6461" (lean), "6464" (obese 1),
280            "9430" (obese 2), "1857" (obese Climp63).
281        label_choice: Which organelle segmentation to use as labels.
282        download: Whether to stream and cache the data if it is not present.
283
284    Returns:
285        The filepath to the cached zarr store.
286    """
287    import zarr
288    from zarr.codecs import BloscCodec
289
290    if sample not in PARLAKGUL_SAMPLES:
291        raise ValueError(f"'{sample}' is not a valid sample. Choose from {list(PARLAKGUL_SAMPLES)}.")
292    sample_info = PARLAKGUL_SAMPLES[sample]
293    if label_choice not in sample_info["seg_zips"]:
294        raise ValueError(f"label_choice='{label_choice}' not available for sample='{sample}'")
295
296    x_min, x_max, y_min, y_max, z_min, z_max = bounding_box
297    n_z, n_y, n_x = sample_info["shape"]
298    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):
299        raise ValueError(f"The bounding box {bounding_box} is not inside sample '{sample}' of shape (z, y, x) "
300                         f"{sample_info['shape']}.")
301
302    os.makedirs(str(path), exist_ok=True)
303    zarr_path = os.path.join(str(path), f"{sample}_{label_choice}_{_bbox_to_str(bounding_box)}.zarr")
304    if os.path.exists(os.path.join(zarr_path, "raw")) and os.path.exists(os.path.join(zarr_path, "labels")):
305        return zarr_path
306
307    if not download:
308        raise RuntimeError(
309            f"No cached data found at '{zarr_path}'. Set download=True to stream it from EMPIAR."
310        )
311
312    shape = (z_max - z_min, y_max - y_min, x_max - x_min)
313    raw_arr = np.zeros(shape, dtype=np.uint8)
314    lbl_arr = np.zeros(shape, dtype=np.uint8)
315
316    raw_base = f"{EMPIAR_BASE}/{PARLAKGUL_PAPER_BASE}/{sample_info['raw_dir']}"
317    zip_name = sample_info["seg_zips"][label_choice]
318    seg_zip_url = f"{EMPIAR_BASE}/{PARLAKGUL_PAPER_BASE}/{sample_info['seg_dir']}/{zip_name}"
319
320    print(f"Streaming Parlakgul {sample} ({sample_info['condition']}) EM + {label_choice} ...")
321
322    # The ZIP directory is read once, the segmentation slices are then fetched one by one from it.
323    zf = zipfile.ZipFile(_HttpFile(seg_zip_url))
324    slices = _index_zip_slices(zf)
325    if len(slices) != n_z:
326        raise RuntimeError(f"Expected {n_z} segmentation slices in {zip_name}, found {len(slices)}.")
327
328    def _fetch_raw(i):
329        url = f"{raw_base}/{sample_info['raw_pattern'].format(z=z_min + i + sample_info['raw_z_offset'])}"
330        raw_arr[i] = _read_raw_slice(url, x_min, x_max, y_min, y_max)
331
332    with futures.ThreadPoolExecutor(N_THREADS) as tp:
333        raw_tasks = [tp.submit(_fetch_raw, i) for i in range(shape[0])]
334        for i, z in enumerate(range(z_min, z_max)):
335            lbl_arr[i] = _read_zip_slice(zf, slices[z], x_min, x_max, y_min, y_max)
336            if (i + 1) % 10 == 0:
337                print(f"  {i + 1}/{shape[0]} segmentation slices done")
338        for task in raw_tasks:
339            task.result()
340
341    # The store is written under a temporary name and moved in place once it is complete,
342    # so that an interrupted download is not mistaken for a complete one.
343    tmp_path = f"{zarr_path}.incomplete"
344    if os.path.exists(tmp_path):
345        shutil.rmtree(tmp_path)
346    root = zarr.open_group(tmp_path, mode="w")
347    root.attrs.update({
348        "bounding_box": list(bounding_box), "sample": sample, "label_choice": label_choice, "resolution_nm": [8, 8, 8],
349    })
350    for name, data, shuffle in (("raw", raw_arr, "shuffle"), ("labels", lbl_arr, "bitshuffle")):
351        arr = root.create_array(
352            name, shape=data.shape, chunks=PARLAKGUL_CHUNK_SHAPE, dtype=data.dtype,
353            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
354        )
355        arr[:] = data
356    if os.path.exists(zarr_path):
357        shutil.rmtree(zarr_path)
358    os.rename(tmp_path, zarr_path)
359
360    print(f"Cached to {zarr_path} (shape {shape})")
361    return zarr_path
362
363
364def get_parlakgul_liver_paths(
365    path: Union[os.PathLike, str],
366    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
367    sample: Literal["6461", "6464", "9430", "1857"] = "6461",
368    label_choice: LabelChoice = "mito",
369    download: bool = False,
370) -> List[str]:
371    """Get paths to Parlakgul liver zarr stores.
372
373    Args:
374        path: Filepath to a folder where the cached zarr stores will be saved.
375        bounding_boxes: List of regions to fetch, each as
376            (x_min, x_max, y_min, y_max, z_min, z_max) in voxel coordinates.
377        sample: Which liver sample to use.
378        label_choice: Which organelle segmentation to use as labels.
379        download: Whether to stream and cache the data if it is not present.
380
381    Returns:
382        List of filepaths to the cached zarr stores.
383    """
384    return [get_parlakgul_liver_data(path, bbox, sample, label_choice, download) for bbox in bounding_boxes]
385
386
387def get_parlakgul_liver_dataset(
388    path: Union[os.PathLike, str],
389    patch_shape: Tuple[int, int, int],
390    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
391    sample: Literal["6461", "6464", "9430", "1857"] = "6461",
392    label_choice: LabelChoice = "mito",
393    download: bool = False,
394    **kwargs,
395) -> Dataset:
396    """Get the Parlakgul liver dataset for organelle segmentation.
397
398    Args:
399        path: Filepath to a folder where the cached zarr stores will be saved.
400        patch_shape: The patch shape (z, y, x) to use for training.
401        bounding_boxes: List of subvolumes to use, each as
402            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
403        sample: Which liver sample to use. One of "6461", "6464", "9430", "1857".
404        label_choice: Which organelle to segment.
405        download: Whether to stream and cache data if not already present.
406        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
407
408    Returns:
409        The segmentation dataset.
410    """
411    assert len(patch_shape) == 3
412
413    paths = get_parlakgul_liver_paths(path, bounding_boxes, sample, label_choice, download)
414
415    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
416
417    return torch_em.default_segmentation_dataset(
418        raw_paths=paths,
419        raw_key="raw",
420        label_paths=paths,
421        label_key="labels",
422        patch_shape=patch_shape,
423        **kwargs,
424    )
425
426
427def get_parlakgul_liver_loader(
428    path: Union[os.PathLike, str],
429    patch_shape: Tuple[int, int, int],
430    batch_size: int,
431    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
432    sample: Literal["6461", "6464", "9430", "1857"] = "6461",
433    label_choice: LabelChoice = "mito",
434    download: bool = False,
435    **kwargs,
436) -> DataLoader:
437    """Get the DataLoader for organelle segmentation in the Parlakgul liver dataset.
438
439    Args:
440        path: Filepath to a folder where the cached zarr stores will be saved.
441        patch_shape: The patch shape (z, y, x) to use for training.
442        batch_size: The batch size for training.
443        bounding_boxes: List of subvolumes to use, each as
444            (x_min, x_max, y_min, y_max, z_min, z_max) in 8 nm voxel coordinates.
445        sample: Which liver sample to use. One of "6461", "6464", "9430", "1857".
446        label_choice: Which organelle to segment.
447        download: Whether to stream and cache data if not already present.
448        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
449            or for the PyTorch DataLoader.
450
451    Returns:
452        The DataLoader.
453    """
454    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
455    dataset = get_parlakgul_liver_dataset(
456        path, patch_shape, bounding_boxes, sample=sample, label_choice=label_choice,
457        download=download, **ds_kwargs
458    )
459    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)