torch_em.data.datasets.electron_microscopy.parlakgul_liver_eval

  1"""The Parlakgul liver evaluation dataset contains manually annotated FIB-SEM crops of mouse liver
  2for mitochondria and lipid droplets. They are the independent evaluation volumes of the model-generated
  3segmentations of the Parlakgul liver dataset (see `parlakgul_liver.py`).
  4
  5For each of the four liver samples (6461, 6464, 9430, 1857) there are three crops of 25 x 150 x 150 voxels
  6(z, y, x) per organelle, at 8 nm isotropic resolution. Every crop comes with:
  7- the raw image,
  8- the manual annotation (instance ids, 0 is background; for a few crops the annotation is semantic),
  9- the instance segmentation predicted by the authors' model for the same crop (ids are arbitrary).
 10
 11This dataset is from the publication https://doi.org/10.1038/s41586-022-04518-2.
 12Please cite it if you use this dataset in your research.
 13
 14The data is publicly available at https://www.ebi.ac.uk/empiar/EMPIAR-10791/.
 15"""
 16
 17import io
 18import os
 19import re
 20import shutil
 21import zipfile
 22from typing import List, Literal, Tuple, Union
 23
 24import numpy as np
 25from torch.utils.data import DataLoader, Dataset
 26
 27import torch_em
 28from .. import util
 29
 30
 31URL = (
 32    "https://ftp.ebi.ac.uk/empiar/world_availability/10791/data/"
 33    "Parlakgul%20et%20al%20-%20Regulation%20of%20liver%20subcellular%20architecture%20controls%20metabolic%20"
 34    "homeostasis/FIB-SEM%20Raw%20and%20Segmentation%20Data/Segmentation%20Evaluation/"
 35    "Manual%20annotation%20-%20Prediction%20-%20Raw%20Slices.zip"
 36)
 37CHECKSUM = "bfb94ec5818a5b25e7ae380e624a29df366e32a66a674ad71030312951c9775f"
 38ZIP_NAME = "Manual annotation - Prediction - Raw Slices.zip"
 39
 40SAMPLES = ("6461", "6464", "9430", "1857")
 41LABEL_CHOICES = {"mito": "mitochondria", "lipid_droplet": "lipid droplets"}
 42N_EVALUATION_VOLUMES = 3
 43
 44
 45def _read_tiff_stack(archive, name):
 46    """Read all slices of a nested zip archive as a (z, y, x) array, ordered by the slice number."""
 47    from PIL import Image
 48
 49    with zipfile.ZipFile(io.BytesIO(archive.read(name))) as zf:
 50        members = [
 51            info.filename for info in zf.infolist()
 52            if info.filename.lower().endswith((".tif", ".tiff")) and "__MACOSX" not in info.filename
 53        ]
 54        members.sort(key=lambda f: int(re.search(r"(\d+)\.tiff?$", f.lower()).group(1)))
 55        return np.stack([np.array(Image.open(io.BytesIO(zf.read(f)))) for f in members])
 56
 57
 58def _find_member(names, sample, organelle, index, kind):
 59    # The archives are not named consistently: the manual annotation of 6464 is called 'manual evaluation'.
 60    pattern = {
 61        "raw": r"raw", "labels": r"manual (annotation|evaluation)", "prediction": r"prediction",
 62    }[kind]
 63    regex = re.compile(rf"^{sample}[ _].*{organelle} evaluation 00{index} {pattern}\.zip$")
 64    matches = [n for n in names if regex.match(os.path.basename(n))]
 65    if len(matches) != 1:
 66        raise RuntimeError(f"Expected one archive for {sample}, {organelle}, {index}, {kind}, found {matches}.")
 67    return matches[0]
 68
 69
 70def get_parlakgul_liver_eval_data(
 71    path: Union[os.PathLike, str],
 72    sample: Literal["6461", "6464", "9430", "1857"] = "6461",
 73    label_choice: Literal["mito", "lipid_droplet"] = "mito",
 74    download: bool = False,
 75) -> List[str]:
 76    """Download the Parlakgul liver evaluation crops and convert them to zarr stores.
 77
 78    Args:
 79        path: Filepath to a folder where the downloaded data will be saved.
 80        sample: Which liver sample to use. One of "6461" (lean), "6464" (obese 1),
 81            "9430" (obese 2), "1857" (obese Climp63).
 82        label_choice: The organelle. One of "mito" or "lipid_droplet".
 83        download: Whether to download the data if it is not present.
 84
 85    Returns:
 86        The filepaths to the zarr stores of the three evaluation crops.
 87    """
 88    import zarr
 89    from zarr.codecs import BloscCodec
 90
 91    if sample not in SAMPLES:
 92        raise ValueError(f"'{sample}' is not a valid sample. Choose from {SAMPLES}.")
 93    if label_choice not in LABEL_CHOICES:
 94        raise ValueError(f"'{label_choice}' is not a valid label_choice. Choose from {list(LABEL_CHOICES)}.")
 95    organelle = LABEL_CHOICES[label_choice]
 96
 97    os.makedirs(path, exist_ok=True)
 98    zarr_paths = [
 99        os.path.join(path, f"{sample}_{label_choice}_{i:03d}.zarr") for i in range(1, N_EVALUATION_VOLUMES + 1)
100    ]
101    if all(os.path.exists(p) for p in zarr_paths):
102        return zarr_paths
103
104    zip_path = os.path.join(path, ZIP_NAME)
105    util.download_source(zip_path, URL, download, checksum=CHECKSUM)
106
107    with zipfile.ZipFile(zip_path) as archive:
108        names = [n for n in archive.namelist() if n.endswith(".zip") and "__MACOSX" not in n]
109        for index, zarr_path in enumerate(zarr_paths, start=1):
110            if os.path.exists(zarr_path):
111                continue
112
113            stacks = {
114                kind: _read_tiff_stack(archive, _find_member(names, sample, organelle, index, kind))
115                for kind in ("raw", "labels", "prediction")
116            }
117            # One crop (6461 mito 003) is annotated with 255 instead of 1.
118            if stacks["labels"].max() == 255:
119                stacks["labels"] = (stacks["labels"] > 0).astype("uint8")
120
121            tmp_path = f"{zarr_path}.incomplete"
122            if os.path.exists(tmp_path):
123                shutil.rmtree(tmp_path)
124            root = zarr.open_group(tmp_path, mode="w")
125            root.attrs.update({"sample": sample, "label_choice": label_choice, "resolution_nm": [8, 8, 8]})
126            for name, data in stacks.items():
127                arr = root.create_array(
128                    name, shape=data.shape, chunks=data.shape, dtype=data.dtype,
129                    compressors=BloscCodec(cname="zstd", clevel=6),
130                )
131                arr[:] = data
132            os.rename(tmp_path, zarr_path)
133
134    return zarr_paths
135
136
137def get_parlakgul_liver_eval_paths(
138    path: Union[os.PathLike, str],
139    samples: Union[str, List[str]] = "6461",
140    label_choice: Literal["mito", "lipid_droplet"] = "mito",
141    download: bool = False,
142) -> List[str]:
143    """Get paths to the Parlakgul liver evaluation zarr stores.
144
145    Args:
146        path: Filepath to a folder where the downloaded data will be saved.
147        samples: The liver sample(s) to use. One or several of "6461", "6464", "9430", "1857".
148        label_choice: The organelle. One of "mito" or "lipid_droplet".
149        download: Whether to download the data if it is not present.
150
151    Returns:
152        The filepaths to the zarr stores, three per sample.
153    """
154    samples = [samples] if isinstance(samples, str) else samples
155    return [p for sample in samples for p in get_parlakgul_liver_eval_data(path, sample, label_choice, download)]
156
157
158def get_parlakgul_liver_eval_dataset(
159    path: Union[os.PathLike, str],
160    patch_shape: Tuple[int, int, int],
161    samples: Union[str, List[str]] = "6461",
162    label_choice: Literal["mito", "lipid_droplet"] = "mito",
163    semantic: bool = False,
164    return_prediction: bool = False,
165    download: bool = False,
166    **kwargs,
167) -> Dataset:
168    """Get the Parlakgul liver evaluation dataset for mitochondria or lipid droplet segmentation.
169
170    Args:
171        path: Filepath to a folder where the downloaded data will be saved.
172        patch_shape: The patch shape (z, y, x). The crops have the shape (25, 150, 150).
173        samples: The liver sample(s) to use. One or several of "6461", "6464", "9430", "1857".
174        label_choice: The organelle. One of "mito" or "lipid_droplet".
175        semantic: Whether to return the binary foreground mask instead of the instance ids.
176        return_prediction: Whether to return the model prediction instead of the manual annotation.
177            This is meant to compute the evaluation scores of the original segmentations.
178        download: Whether to download the data if it is not present.
179        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
180
181    Returns:
182        The segmentation dataset.
183    """
184    assert len(patch_shape) == 3
185
186    paths = get_parlakgul_liver_eval_paths(path, samples, label_choice, download)
187
188    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
189    if semantic:
190        kwargs = util.update_kwargs(kwargs, "label_transform", _binarize)
191
192    return torch_em.default_segmentation_dataset(
193        raw_paths=paths,
194        raw_key="raw",
195        label_paths=paths,
196        label_key="prediction" if return_prediction else "labels",
197        patch_shape=patch_shape,
198        **kwargs,
199    )
200
201
202def _binarize(labels):
203    return (labels > 0).astype("float32")
204
205
206def get_parlakgul_liver_eval_loader(
207    path: Union[os.PathLike, str],
208    patch_shape: Tuple[int, int, int],
209    batch_size: int,
210    samples: Union[str, List[str]] = "6461",
211    label_choice: Literal["mito", "lipid_droplet"] = "mito",
212    semantic: bool = False,
213    return_prediction: bool = False,
214    download: bool = False,
215    **kwargs,
216) -> DataLoader:
217    """Get the DataLoader for the Parlakgul liver evaluation dataset.
218
219    Args:
220        path: Filepath to a folder where the downloaded data will be saved.
221        patch_shape: The patch shape (z, y, x). The crops have the shape (25, 150, 150).
222        batch_size: The batch size.
223        samples: The liver sample(s) to use. One or several of "6461", "6464", "9430", "1857".
224        label_choice: The organelle. One of "mito" or "lipid_droplet".
225        semantic: Whether to return the binary foreground mask instead of the instance ids.
226        return_prediction: Whether to return the model prediction instead of the manual annotation.
227        download: Whether to download the data if it is not present.
228        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
229            or for the PyTorch DataLoader.
230
231    Returns:
232        The DataLoader.
233    """
234    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
235    dataset = get_parlakgul_liver_eval_dataset(
236        path, patch_shape, samples=samples, label_choice=label_choice, semantic=semantic,
237        return_prediction=return_prediction, download=download, **ds_kwargs,
238    )
239    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)