torch_em.data.datasets.electron_microscopy.u2os_fibsem

  1"""The U2OS FIB-SEM dataset contains training data for segmenting organelles in a FIB-SEM volume
  2of a human U2-OS (bone osteosarcoma) cell.
  3
  4It contains four sub-datasets with binary semantic annotations of 2D slices:
  5- ne: Nuclear envelope (44 slices of 1235 x 530 pixels).
  6- mitochondria: Mitochondria (45 slices of 692 x 1697 pixels).
  7- er: Endoplasmic reticulum (45 slices of 1385 x 3394 pixels).
  8- golgi: Golgi apparatus (45 slices of 1385 x 3394 pixels).
  9
 10For mitochondria, ER and Golgi the data is distributed as small stacks (9 or 5 slices), where only the central slice
 11is annotated. Only the annotated slices are used here.
 12
 13The dataset is located at https://doi.org/10.5281/zenodo.10043461 and is distributed under the CC BY 4.0 license.
 14The full FIB-SEM volume is available at EMPIAR-11746. The annotations were created for training
 15the networks in DeepMIB (Belevich & Jokitalo, https://doi.org/10.1371/journal.pcbi.1008374).
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20import re
 21from glob import glob
 22from natsort import natsorted
 23from typing import Union, Tuple, List, Literal
 24
 25import numpy as np
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URLS = {
 35    "ne": "https://zenodo.org/records/10043461/files/1_NE.zip?download=1",
 36    "mitochondria": "https://zenodo.org/records/10043461/files/1_Mitochondria.zip?download=1",
 37    "er": "https://zenodo.org/records/10043461/files/1_ER.zip?download=1",
 38    "golgi": "https://zenodo.org/records/10043461/files/1_Golgi.zip?download=1",
 39}
 40CHECKSUMS = {
 41    "ne": "4ae915ce64632d1f27328cf3a79e85806d6b899cc51e0e85be619f3c161233bc",
 42    "mitochondria": "6725fb41ed7c14061ee96ab6110209aa0a6f9bd9fdb3a9aca2adb82a408f5680",
 43    "er": "1c5df4101b5f3f376124fc0b87da03198158d039ce146536542bc90c943feffc",
 44    "golgi": "594bf6a51dc9c85b2af504063759e8e46b39e604e3eb506728f8ba9254e31102",
 45}
 46# The folders with the training images and labels inside of the extracted archives.
 47FOLDERS = {
 48    "ne": ("1_NE", "1_Train_NE", "Images", "Labels"),
 49    "mitochondria": ("1_Mitochondria", "1_Train_Mito", "Images", "Labels"),
 50    "er": ("1_ER", "1_Train_ER", "Images", "Labels"),
 51    "golgi": ("1_Golgi", "1_Train_Golgi", "Images_for_patches", "Labels_for_patches"),
 52}
 53
 54
 55def _read_amira(am_path):
 56    with open(am_path, "rb") as f:
 57        content = f.read()
 58
 59    # The header is followed by the raw (uncompressed) voxel data, which starts after the '@1' marker.
 60    marker = content.find(b"\n@1\n")
 61    header = content[:marker].decode("latin1")
 62    assert "Lattice { byte Data } @1" in header, f"Unsupported Amira data type in {am_path}"
 63    x, y, z = map(int, re.search(r"define Lattice (\d+) (\d+) (\d+)", header).groups())
 64
 65    start = marker + len(b"\n@1\n")
 66    return np.frombuffer(content[start:start + x * y * z], dtype="uint8").reshape(z, y, x)
 67
 68
 69def _convert_to_annotated_slice(am_path, label_path, image_tif_path, label_tif_path):
 70    import tifffile
 71
 72    image = _read_amira(am_path)
 73    label = tifffile.imread(label_path)
 74    label = label[None] if label.ndim == 2 else label
 75    assert image.shape == label.shape, f"Image and label shapes differ for {am_path}: {image.shape}, {label.shape}"
 76
 77    # Only a single slice per stack is annotated.
 78    z = int(np.argmax((label > 0).reshape(len(label), -1).sum(axis=1)))
 79    tifffile.imwrite(image_tif_path, image[z], compression="zlib")
 80    tifffile.imwrite(label_tif_path, label[z], compression="zlib")
 81
 82
 83def get_u2os_fibsem_data(
 84    path: Union[os.PathLike, str], name: Literal["ne", "mitochondria", "er", "golgi"], download: bool = False
 85) -> str:
 86    """Download the U2OS FIB-SEM dataset.
 87
 88    The images are distributed in the Amira format and are converted to tif files after downloading.
 89    Only the annotated slice of each image / label stack is kept for training.
 90
 91    Args:
 92        path: Filepath to a folder where the downloaded data will be saved.
 93        name: The name of the sub-dataset. Either 'ne', 'mitochondria', 'er' or 'golgi'.
 94        download: Whether to download the data if it is not present.
 95
 96    Returns:
 97        The filepath to the folder with the data.
 98    """
 99    assert name in URLS, f"'{name}' is not a valid sub-dataset name. Choose from {list(URLS)}."
100    root, train_dir, _, _ = FOLDERS[name]
101
102    data_dir = os.path.join(path, root)
103    if os.path.exists(os.path.join(data_dir, train_dir)):
104        return data_dir
105
106    os.makedirs(path, exist_ok=True)
107    zip_path = os.path.join(path, f"{root}.zip")
108    util.download_source(path=zip_path, url=URLS[name], download=download, checksum=CHECKSUMS[name])
109    util.unzip(zip_path=zip_path, dst=path)
110
111    _, _, image_dir, label_dir = FOLDERS[name]
112    image_tif_dir = os.path.join(data_dir, train_dir, "images_tif")
113    label_tif_dir = os.path.join(data_dir, train_dir, "labels_tif")
114    os.makedirs(image_tif_dir, exist_ok=True)
115    os.makedirs(label_tif_dir, exist_ok=True)
116
117    am_paths = natsorted(glob(os.path.join(data_dir, train_dir, image_dir, "*.am")))
118    label_paths = natsorted(glob(os.path.join(data_dir, train_dir, label_dir, "*.tif")))
119    assert len(am_paths) > 0 and len(am_paths) == len(label_paths)
120    for am_path, label_path in zip(am_paths, label_paths):
121        fname = os.path.splitext(os.path.basename(am_path))[0]
122        _convert_to_annotated_slice(
123            am_path,
124            label_path,
125            os.path.join(image_tif_dir, f"{fname}.tif"),
126            os.path.join(label_tif_dir, f"{fname}.tif"),
127        )
128
129    return data_dir
130
131
132def get_u2os_fibsem_paths(
133    path: Union[os.PathLike, str],
134    name: Literal["ne", "mitochondria", "er", "golgi"],
135    download: bool = False,
136) -> Tuple[List[str], List[str]]:
137    """Get paths to the U2OS FIB-SEM data.
138
139    Args:
140        path: Filepath to a folder where the downloaded data will be saved.
141        name: The name of the sub-dataset. Either 'ne', 'mitochondria', 'er' or 'golgi'.
142        download: Whether to download the data if it is not present.
143
144    Returns:
145        List of filepaths for the image data.
146        List of filepaths for the label data.
147    """
148    data_dir = get_u2os_fibsem_data(path, name, download)
149    train_dir = FOLDERS[name][1]
150
151    image_paths = natsorted(glob(os.path.join(data_dir, train_dir, "images_tif", "*.tif")))
152    label_paths = natsorted(glob(os.path.join(data_dir, train_dir, "labels_tif", "*.tif")))
153    assert len(image_paths) > 0 and len(image_paths) == len(label_paths)
154    return image_paths, label_paths
155
156
157def get_u2os_fibsem_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, int],
160    name: Literal["ne", "mitochondria", "er", "golgi"],
161    download: bool = False,
162    **kwargs
163) -> Dataset:
164    """Get the U2OS FIB-SEM dataset for organelle segmentation.
165
166    The labels are binary (1 is the organelle).
167
168    Args:
169        path: Filepath to a folder where the downloaded data will be saved.
170        patch_shape: The patch shape to use for training.
171        name: The name of the sub-dataset. Either 'ne', 'mitochondria', 'er' or 'golgi'.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
174
175    Returns:
176        The segmentation dataset.
177    """
178    image_paths, label_paths = get_u2os_fibsem_paths(path, name, download)
179
180    kwargs = util.update_kwargs(kwargs, "label_transform", torch_em.transform.label.labels_to_binary)
181    kwargs = util.update_kwargs(kwargs, "ndim", 2)
182
183    return torch_em.default_segmentation_dataset(
184        raw_paths=image_paths,
185        raw_key=None,
186        label_paths=label_paths,
187        label_key=None,
188        patch_shape=patch_shape,
189        **kwargs
190    )
191
192
193def get_u2os_fibsem_loader(
194    path: Union[os.PathLike, str],
195    patch_shape: Tuple[int, int],
196    batch_size: int,
197    name: Literal["ne", "mitochondria", "er", "golgi"],
198    download: bool = False,
199    **kwargs
200) -> DataLoader:
201    """Get the DataLoader for organelle segmentation in the U2OS FIB-SEM dataset.
202
203    Args:
204        path: Filepath to a folder where the downloaded data will be saved.
205        patch_shape: The patch shape to use for training.
206        batch_size: The batch size for training.
207        name: The name of the sub-dataset. Either 'ne', 'mitochondria', 'er' or 'golgi'.
208        download: Whether to download the data if it is not present.
209        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
210
211    Returns:
212        The DataLoader.
213    """
214    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
215    dataset = get_u2os_fibsem_dataset(path, patch_shape, name, download, **ds_kwargs)
216    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)