torch_em.data.datasets.electron_microscopy.human_liver_em

  1"""The human liver EM dataset contains multiscale SBF-SEM images of human liver tissue
  2with semantic segmentations of cellular and organelle structures.
  3
  4The dataset covers a reconstructed human periportal liver volume (597 z-slices,
  520000 x 20000 pixels per slice) with 9 binary semantic segmentation classes
  6(0=background, 255=foreground - not instance segmentation).
  7The segmentations were produced automatically by the authors' multiscale deep learning approach,
  8they are not manual annotations. The mitochondria labels only cover part of each slice
  9(mainly the upper part of the tissue), so unlabeled tissue is not necessarily background.
 10
 11Currently recommended label choices for training: "er", "mito", "nucleus".
 12These are the classes that are well-represented and have clear biological meaning
 13at the imaging resolution of this dataset.
 14
 15All 9 available classes:
 16- "bile_duct": bile duct (sparse, often absent in a given crop)
 17- "cell_boundary": cell boundary region (coarse, marks interior near cell edge)
 18- "cholangiocyte": cholangiocyte cells (sparse)
 19- "endothelial": endothelial cells (sparse)
 20- "er": endoplasmic reticulum (recommended)
 21- "hepatocyte": hepatocyte cell interior - tissue-level mask (large filled regions)
 22- "mito": mitochondria (recommended)
 23- "nucleus": nucleus (recommended)
 24- "sinusoid": sinusoidal capillary (sparse)
 25
 26NOTE (on other organelles): the 9 classes above are tissue/cell-level annotations.
 27Organelle-level segmentations for additional structures (lipid droplets, Golgi, etc.)
 28are not available in EMPIAR-13356. The Parlakgul liver dataset (EMPIAR-10791) provides
 29richer organelle annotations (ER sheets/tubules, lipid droplets, nuclear membrane,
 30plasma membrane) at higher FIB-SEM resolution for mouse liver.
 31
 32NOTE (on resolution): the pixel size is not documented in EMPIAR-13356. The bioRxiv abstract (linked below)
 33reports a volume of 152 x 140 x 33 um at 8 nm pixel size, i.e. 8 nm in xy and about
 3455 nm section thickness in z (33 um over 597 slices).
 35
 36Data is streamed lazily from EMPIAR-13356 via HTTP: raw 16-bit TIFFs and binary PNG
 37masks are fetched per z-slice and cached in a single zarr v3 store per bounding box.
 38All 9 label classes are stored together (raw, bile_duct, cell_boundary, ..., sinusoid).
 39The `label_choice` parameter in the loader selects which array to use as labels.
 40
 41Bounding boxes are specified as (x_min, x_max, y_min, y_max, z_min, z_max) in voxels.
 42The full volume is (597, 20000, 20000) voxels (z, y, x). Tissue spans roughly
 43x=[1195, 18890], y=[469, 19570] - the volume edges are empty.
 44
 45This dataset is from the publication https://www.biorxiv.org/content/10.64898/2026.04.22.719970v1.
 46Please cite it if you use this dataset in your research.
 47
 48The data is publicly available at https://www.ebi.ac.uk/empiar/EMPIAR-13356/.
 49"""
 50
 51import hashlib
 52import io
 53import os
 54from typing import List, Literal, Optional, Tuple, Union
 55
 56import numpy as np
 57from torch.utils.data import DataLoader, Dataset
 58
 59import torch_em
 60from .. import util
 61
 62
 63EMPIAR_BASE = "https://ftp.ebi.ac.uk/empiar/world_availability/13356/data"
 64
 65HUMAN_LIVER_EM_LABEL_DIRS = {
 66    "er": "humanliver_er_mask",
 67    "mito": "humanliver_mito_mask",
 68    "nucleus": "humanliver_nucleus_mask",
 69}
 70
 71HUMAN_LIVER_EM_SHAPE = (597, 20000, 20000)
 72# Tissue spans x=[1195,18890], y=[469,19570] - edges are empty background.
 73HUMAN_LIVER_EM_TISSUE_BBOX = (1195, 18890, 469, 19570, 0, 597)
 74
 75# Zarr layout for bbox crops.
 76HUMAN_LIVER_EM_CHUNK_SHAPE = (64, 256, 256)
 77# Zarr layout for full-volume sharded store.
 78# Shards: (64, 4096, 4096) outer; chunks: (8, 256, 256) inner.
 79HUMAN_LIVER_EM_SHARD_SHAPE = (64, 4096, 4096)
 80HUMAN_LIVER_EM_INNER_CHUNK = (8, 256, 256)
 81
 82LabelChoice = Literal["er", "mito", "nucleus"]
 83
 84
 85def _bbox_to_str(bbox):
 86    return hashlib.md5("_".join(str(v) for v in bbox).encode()).hexdigest()[:12]
 87
 88
 89class _HttpFile:
 90    """Seekable file-like object backed by HTTP range requests for efficient partial TIFF reading."""
 91
 92    def __init__(self, url):
 93        import requests
 94        self.url = url
 95        self._pos = 0
 96        r = requests.head(url, timeout=30)
 97        r.raise_for_status()
 98        self._size = int(r.headers["Content-Length"])
 99
100    def read(self, n=-1):
101        import requests
102        end = (self._size - 1) if n == -1 else min(self._pos + n - 1, self._size - 1)
103        if self._pos > end:
104            return b""
105        r = requests.get(self.url, headers={"Range": f"bytes={self._pos}-{end}"}, timeout=120)
106        data = r.content
107        self._pos += len(data)
108        return data
109
110    def seek(self, pos, whence=0):
111        if whence == 0:
112            self._pos = pos
113        elif whence == 1:
114            self._pos += pos
115        elif whence == 2:
116            self._pos = self._size + pos
117        self._pos = max(0, min(self._pos, self._size))
118        return self._pos
119
120    def tell(self):
121        return self._pos
122
123    def seekable(self):
124        return True
125
126    def readable(self):
127        return True
128
129    def __enter__(self):
130        return self
131
132    def __exit__(self, *args):
133        pass
134
135
136def _read_raw_slice(z, x_min, x_max, y_min, y_max):
137    """Read a cropped region from a remote TIFF using a single HTTP range request covering
138    only the strips needed for the y range, avoiding downloading the full ~880 MB file."""
139    import requests
140    import tifffile
141
142    url = f"{EMPIAR_BASE}/humanliver_raw_images/humanliver_raw_{z:03d}.tif"
143
144    # Read TIFF metadata to get strip offsets and byte counts.
145    with tifffile.TiffFile(_HttpFile(url)) as tif:
146        page = tif.pages[0]
147        offsets = page.dataoffsets
148        bytecounts = page.databytecounts
149        width = page.imagewidth
150        dtype = np.dtype(page.dtype)
151
152    # One strip per row - request only strips y_min..y_max in one range request.
153    from imagecodecs import lzw_decode
154    start_byte = offsets[y_min]
155    end_byte = offsets[y_max - 1] + bytecounts[y_max - 1] - 1
156    r = requests.get(url, headers={"Range": f"bytes={start_byte}-{end_byte}"}, timeout=120)
157    r.raise_for_status()
158    raw_bytes = r.content
159
160    # Predictor=2 (horizontal differencing) is used - need cumsum after LZW decode.
161    from imagecodecs import delta_decode
162    rows = []
163    for i in range(y_min, y_max):
164        strip_start = offsets[i] - start_byte
165        strip_data = raw_bytes[strip_start:strip_start + bytecounts[i]]
166        decoded = np.frombuffer(lzw_decode(strip_data), dtype=dtype)
167        row = delta_decode(decoded, axis=-1, dist=1, out=decoded).reshape(width)
168        rows.append(row[x_min:x_max])
169
170    return np.stack(rows, axis=0)
171
172
173def _read_mask_slice(label_key, z, x_min, x_max, y_min, y_max):
174    import PIL.Image
175    import requests
176    PIL.Image.MAX_IMAGE_PIXELS = None
177    label_dir = HUMAN_LIVER_EM_LABEL_DIRS[label_key]
178    url = f"{EMPIAR_BASE}/{label_dir}/{label_dir}_{z:03d}.png"
179    r = requests.get(url, timeout=120)
180    r.raise_for_status()
181    img = np.array(PIL.Image.open(io.BytesIO(r.content)))
182    if img.ndim == 3:
183        img = img[..., 0]
184    return img[y_min:y_max, x_min:x_max]
185
186
187def get_human_liver_em_data(
188    path: Union[os.PathLike, str],
189    bounding_box: Tuple[int, int, int, int, int, int],
190    download: bool = False,
191) -> str:
192    """Stream a subvolume from the human liver EM dataset and cache it as a zarr v3 store.
193
194    All 9 label classes are stored in a single zarr alongside the raw EM, so each
195    bounding box is only downloaded once regardless of which label_choice is used.
196
197    Args:
198        path: Filepath to a folder where the cached zarr store will be saved.
199        bounding_box: The region to fetch as (x_min, x_max, y_min, y_max, z_min, z_max)
200            in voxel coordinates. Tissue spans roughly x=[1195,18890], y=[469,19570].
201        download: Whether to stream and cache the data if it is not present.
202
203    Returns:
204        The filepath to the cached zarr store.
205    """
206    import zarr
207    from zarr.codecs import BloscCodec
208
209    os.makedirs(str(path), exist_ok=True)
210    zarr_path = os.path.join(str(path), f"{_bbox_to_str(bounding_box)}.zarr")
211
212    root = zarr.open_group(zarr_path, mode="a")
213    all_keys = ["raw"] + list(HUMAN_LIVER_EM_LABEL_DIRS.keys())
214    if all(k in root for k in all_keys):
215        return zarr_path
216
217    if not download:
218        raise RuntimeError(
219            f"No cached data found at '{zarr_path}'. Set download=True to stream it from EMPIAR."
220        )
221
222    x_min, x_max, y_min, y_max, z_min, z_max = bounding_box
223    shape = (z_max - z_min, y_max - y_min, x_max - x_min)
224
225    print(f"Streaming Human Liver EM + all labels for bbox {bounding_box} ...")
226    raw_vol = np.zeros(shape, dtype=np.uint16)
227    label_vols = {k: np.zeros(shape, dtype=np.uint8) for k in HUMAN_LIVER_EM_LABEL_DIRS}
228
229    for i, z in enumerate(range(z_min, z_max)):
230        raw_vol[i] = _read_raw_slice(z, x_min, x_max, y_min, y_max)
231        for label_key in HUMAN_LIVER_EM_LABEL_DIRS:
232            label_vols[label_key][i] = _read_mask_slice(label_key, z, x_min, x_max, y_min, y_max)
233        if (i + 1) % 5 == 0:
234            print(f"  {i + 1}/{z_max - z_min} slices done")
235
236    def _make_array(name, data, is_label):
237        shuffle = "bitshuffle" if is_label else "shuffle"
238        arr = root.create_array(
239            name, shape=data.shape, chunks=HUMAN_LIVER_EM_CHUNK_SHAPE, dtype=data.dtype,
240            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
241        )
242        arr[:] = data
243
244    root.attrs["bounding_box"] = list(bounding_box)
245
246    if "raw" not in root:
247        _make_array("raw", raw_vol, is_label=False)
248    for label_key, vol in label_vols.items():
249        if label_key not in root:
250            _make_array(label_key, vol, is_label=True)
251
252    print(f"Cached to {zarr_path} (shape {shape})")
253    return zarr_path
254
255
256def get_human_liver_em_full_volume(
257    path: Union[os.PathLike, str],
258    download: bool = False,
259) -> str:
260    """Download the full human liver EM tissue volume into a sharded zarr v3 store.
261
262    Downloads all 597 z-slices for the tissue region x=[1195,18890], y=[469,19570]
263    with raw EM + er/mito/nucleus labels. Data is written slice by slice to avoid
264    memory issues. Estimated storage: ~100-150 GB compressed. Estimated download
265    time: ~12 hours (one-time cost).
266
267    The sharded zarr uses shard shape (64, 4096, 4096) with inner chunks (8, 256, 256),
268    enabling efficient random crop access during training without loading the full volume.
269
270    Args:
271        path: Filepath to a folder where the zarr store will be saved.
272        download: Whether to stream and cache the data if it is not present.
273
274    Returns:
275        The filepath to the full-volume zarr store.
276    """
277    import zarr
278    from zarr.codecs import BloscCodec, ShardingCodec
279
280    os.makedirs(str(path), exist_ok=True)
281    zarr_path = os.path.join(str(path), "full_volume.zarr")
282
283    x_min, x_max, y_min, y_max, z_min, z_max = HUMAN_LIVER_EM_TISSUE_BBOX
284    shape = (z_max - z_min, y_max - y_min, x_max - x_min)
285
286    root = zarr.open_group(zarr_path, mode="a")
287    all_keys = ["raw"] + list(HUMAN_LIVER_EM_LABEL_DIRS.keys())
288    if all(k in root for k in all_keys):
289        return zarr_path
290
291    if not download:
292        raise RuntimeError(
293            f"Full-volume zarr not found at '{zarr_path}'. Set download=True to stream from EMPIAR."
294            " Note: download takes ~12 hours and requires ~100-150 GB disk space."
295        )
296
297    def _make_sharded(name, dtype, is_label):
298        shuffle = "bitshuffle" if is_label else "shuffle"
299        return root.create_array(
300            name, shape=shape, chunks=HUMAN_LIVER_EM_SHARD_SHAPE, dtype=dtype,
301            compressors=ShardingCodec(
302                chunk_shape=HUMAN_LIVER_EM_INNER_CHUNK,
303                codecs=[BloscCodec(cname="zstd", clevel=6, shuffle=shuffle)],
304            ),
305        )
306
307    if "raw" not in root:
308        _make_sharded("raw", np.dtype("uint16"), is_label=False)
309    for label_key in HUMAN_LIVER_EM_LABEL_DIRS:
310        if label_key not in root:
311            _make_sharded(label_key, np.dtype("uint8"), is_label=True)
312
313    root.attrs["tissue_bbox"] = list(HUMAN_LIVER_EM_TISSUE_BBOX)
314    n_slices = z_max - z_min
315    print(f"Streaming full Human Liver EM volume ({shape}) - this will take several hours ...")
316
317    for i, z in enumerate(range(z_min, z_max)):
318        raw_slice = _read_raw_slice(z, x_min, x_max, y_min, y_max)
319        root["raw"][i] = raw_slice
320        for label_key in HUMAN_LIVER_EM_LABEL_DIRS:
321            mask_slice = _read_mask_slice(label_key, z, x_min, x_max, y_min, y_max)
322            root[label_key][i] = mask_slice
323        if (i + 1) % 10 == 0:
324            print(f"  {i + 1}/{n_slices} slices done")
325
326    print(f"Full volume cached to {zarr_path} (shape {shape})")
327    return zarr_path
328
329
330def get_human_liver_em_paths(
331    path: Union[os.PathLike, str],
332    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
333    download: bool = False,
334    full_volume: bool = False,
335) -> List[str]:
336    """Get paths to human liver EM zarr stores.
337
338    Args:
339        path: Filepath to a folder where the cached zarr stores will be saved.
340        bounding_boxes: List of regions to fetch, each as
341            (x_min, x_max, y_min, y_max, z_min, z_max) in voxel coordinates.
342            Ignored when full_volume=True.
343        download: Whether to stream and cache the data if it is not present.
344        full_volume: If True, download/use the full tissue volume as a single sharded
345            zarr v3 store (~12h download, ~100-150 GB). Supersedes bounding_boxes.
346
347    Returns:
348        List of filepaths to the cached zarr stores.
349    """
350    if full_volume:
351        return [get_human_liver_em_full_volume(path, download)]
352    if bounding_boxes is None:
353        raise ValueError("Provide bounding_boxes or set full_volume=True.")
354    return [get_human_liver_em_data(path, bbox, download) for bbox in bounding_boxes]
355
356
357def get_human_liver_em_dataset(
358    path: Union[os.PathLike, str],
359    patch_shape: Tuple[int, int, int],
360    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
361    label_choice: LabelChoice = "mito",
362    download: bool = False,
363    full_volume: bool = False,
364    **kwargs,
365) -> Dataset:
366    """Get the human liver EM dataset for semantic segmentation.
367
368    Args:
369        path: Filepath to a folder where the cached zarr stores will be saved.
370        patch_shape: The patch shape (z, y, x) to use for training. The pixel
371            resolution is unconfirmed (estimated ~50-100 nm/px xy). Check the
372            paper for the exact values when selecting isotropic patch shapes.
373        bounding_boxes: List of subvolumes to use, each as
374            (x_min, x_max, y_min, y_max, z_min, z_max) in voxel coordinates.
375            Ignored when full_volume=True.
376        label_choice: Which structure to segment.
377        download: Whether to stream and cache data if not already present.
378        full_volume: If True, use the full sharded tissue volume (~12h one-time download).
379        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
380
381    Returns:
382        The segmentation dataset.
383    """
384    assert len(patch_shape) == 3
385
386    paths = get_human_liver_em_paths(path, bounding_boxes, download, full_volume)
387
388    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
389
390    return torch_em.default_segmentation_dataset(
391        raw_paths=paths,
392        raw_key="raw",
393        label_paths=paths,
394        label_key=label_choice,
395        patch_shape=patch_shape,
396        **kwargs,
397    )
398
399
400def get_human_liver_em_loader(
401    path: Union[os.PathLike, str],
402    patch_shape: Tuple[int, int, int],
403    batch_size: int,
404    bounding_boxes: Optional[List[Tuple[int, int, int, int, int, int]]] = None,
405    label_choice: LabelChoice = "mito",
406    download: bool = False,
407    full_volume: bool = False,
408    **kwargs,
409) -> DataLoader:
410    """Get the DataLoader for semantic segmentation in the human liver EM dataset.
411
412    Args:
413        path: Filepath to a folder where the cached zarr stores will be saved.
414        patch_shape: The patch shape (z, y, x) to use for training.
415        batch_size: The batch size for training.
416        bounding_boxes: List of subvolumes to use, each as
417            (x_min, x_max, y_min, y_max, z_min, z_max) in voxel coordinates.
418            Ignored when full_volume=True.
419        label_choice: Which structure to segment.
420        download: Whether to stream and cache data if not already present.
421        full_volume: If True, use the full sharded tissue volume (~12h one-time download).
422        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
423            or for the PyTorch DataLoader.
424
425    Returns:
426        The DataLoader.
427    """
428    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
429    dataset = get_human_liver_em_dataset(
430        path, patch_shape, bounding_boxes, label_choice=label_choice,
431        download=download, full_volume=full_volume, **ds_kwargs
432    )
433    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)