torch_em.data.datasets.electron_microscopy.skeletal_muscle_fibsem

  1"""The Skeletal Muscle FIB-SEM dataset contains FIB-SEM volumes of human skeletal muscle biopsies
  2with semi-automated segmentations of mitochondria.
  3
  4The data was acquired from biopsies of 10 individuals (aged 22 to 81) and the standardized volumes were
  5segmented in a semi-automated manner. In total, 32 standardized volumes of 9 individuals have mitochondria
  6segmentations.
  7The dataset is located at https://zenodo.org/records/5873883 (part 1 of 6; the other parts contain the raw data)
  8and is distributed under the CC BY 4.0 license. It is part of the publication
  9https://doi.org/10.1016/j.jsb.2019.03.008. Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13import re
 14import warnings
 15from glob import glob
 16from shutil import rmtree
 17from typing import Union, Tuple, List, Literal
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26URL = "https://zenodo.org/api/records/5873883/files/Human%20Skeletal%20Muscle%201%20of%206.zip/content"
 27# The md5 checksum of the zip file is 9c43bfbe8971c279b080d0ab59e6afe1, but only sha256 checksums are supported.
 28CHECKSUM = None
 29
 30# The volumes of these individuals are used for validation, all other volumes for training.
 31VAL_DONORS = ("045", "110")
 32
 33
 34def _volume_key(fname):
 35    # The naming of the files is not consistent, so we normalize it to match the raw volumes and segmentations.
 36    stem = os.path.splitext(os.path.basename(fname))[0]
 37    stem = re.sub(r"[\s_]*\(?mito\)?$", "", stem, flags=re.IGNORECASE)
 38    return re.sub(r"[^a-z0-9]", "", stem.lower())
 39
 40
 41def _find_volume_pairs(names):
 42    raws, segs = {}, {}
 43    for name in names:
 44        parts = name.split("/")
 45        if parts[0] == "__MACOSX" or not name.lower().endswith(".tif"):
 46            continue
 47        donor = re.search(r"GT_?(\d+)", parts[1])
 48        if donor is None:
 49            continue
 50        key = (donor.group(1), _volume_key(name))
 51        if "Standardised Files" in parts:
 52            raws[key] = name
 53        elif "Segmentations" in parts and "z-disk" not in parts[-1].lower():
 54            segs[key] = name
 55    return {key: (raws[key], segs[key]) for key in sorted(raws) if key in segs}
 56
 57
 58def _convert_to_h5(zf, pairs, path):
 59    import h5py
 60    import tifffile
 61
 62    tmp_dir = os.path.join(path, "tmp")
 63    for (donor, key), (raw_name, seg_name) in pairs.items():
 64        h5_path = os.path.join(path, f"{key}.h5")
 65        if os.path.exists(h5_path):
 66            continue
 67
 68        raw = tifffile.imread(zf.extract(raw_name, tmp_dir))
 69        labels = tifffile.imread(zf.extract(seg_name, tmp_dir))
 70        if raw.shape != labels.shape:
 71            warnings.warn(f"Skipping {key}: shape of the raw data {raw.shape} and labels {labels.shape} do not match.")
 72            continue
 73
 74        chunks = tuple(min(c, s) for c, s in zip((32, 256, 256), raw.shape))
 75        # NOTE: The segmentations are binary masks, with foreground values 244 and 255.
 76        with h5py.File(h5_path, "w") as f:
 77            f.create_dataset("raw", data=raw, compression="gzip", chunks=chunks)
 78            f.create_dataset("labels/mitochondria", data=(labels > 0).astype("uint8"), compression="gzip",
 79                             chunks=chunks)
 80
 81        rmtree(tmp_dir)
 82
 83
 84def get_skeletal_muscle_fibsem_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 85    """Download the Skeletal Muscle FIB-SEM dataset.
 86
 87    Note: The data is only available as a 35.9 GB zip archive, which also contains binned volumes, volumes without
 88    segmentations and Z-disk segmentations that are not used. Only the standardized volumes with mitochondria
 89    segmentations are extracted and converted to hdf5 files (one per volume), the zip archive is removed afterwards.
 90
 91    Args:
 92        path: Filepath to a folder where the downloaded data will be saved.
 93        download: Whether to download the data if it is not present.
 94
 95    Returns:
 96        The filepath to the folder with the hdf5 files.
 97    """
 98    import zipfile
 99
100    if len(glob(os.path.join(path, "*.h5"))) > 0:
101        return str(path)
102
103    os.makedirs(path, exist_ok=True)
104    zip_path = os.path.join(path, "skeletal_muscle_fibsem.zip")
105    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
106
107    with zipfile.ZipFile(zip_path) as zf:
108        pairs = _find_volume_pairs(zf.namelist())
109        _convert_to_h5(zf, pairs, path)
110
111    os.remove(zip_path)
112    return str(path)
113
114
115def get_skeletal_muscle_fibsem_paths(
116    path: Union[os.PathLike, str], split: Literal["train", "val"] = "train", download: bool = False
117) -> List[str]:
118    """Get the paths to the Skeletal Muscle FIB-SEM data.
119
120    Args:
121        path: Filepath to a folder where the downloaded data will be saved.
122        split: The data split to use. Either 'train' or 'val'.
123        download: Whether to download the data if it is not present.
124
125    Returns:
126        List of the filepaths to the hdf5 files.
127    """
128    assert split in ("train", "val"), f"'{split}' is not a valid data split."
129    get_skeletal_muscle_fibsem_data(path, download)
130
131    h5_paths = sorted(glob(os.path.join(path, "*.h5")))
132    is_val = [re.match(r"gt(\d+)", os.path.basename(p)).group(1) in VAL_DONORS for p in h5_paths]
133    return [p for p, val in zip(h5_paths, is_val) if val == (split == "val")]
134
135
136def get_skeletal_muscle_fibsem_dataset(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int, int],
139    split: Literal["train", "val"] = "train",
140    download: bool = False,
141    **kwargs
142) -> Dataset:
143    """Get the Skeletal Muscle FIB-SEM dataset for mitochondria segmentation.
144
145    Args:
146        path: Filepath to a folder where the downloaded data will be saved.
147        patch_shape: The patch shape to use for training.
148        split: The data split to use. Either 'train' or 'val'. The split is done by individual.
149        download: Whether to download the data if it is not present.
150        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
151
152    Returns:
153        The segmentation dataset.
154    """
155    paths = get_skeletal_muscle_fibsem_paths(path, split, download)
156
157    kwargs = util.update_kwargs(kwargs, "ndim", 3)
158
159    return torch_em.default_segmentation_dataset(
160        raw_paths=paths,
161        raw_key="raw",
162        label_paths=paths,
163        label_key="labels/mitochondria",
164        patch_shape=patch_shape,
165        **kwargs
166    )
167
168
169def get_skeletal_muscle_fibsem_loader(
170    path: Union[os.PathLike, str],
171    patch_shape: Tuple[int, int, int],
172    batch_size: int,
173    split: Literal["train", "val"] = "train",
174    download: bool = False,
175    **kwargs
176) -> DataLoader:
177    """Get the DataLoader for mitochondria segmentation in the Skeletal Muscle FIB-SEM dataset.
178
179    Args:
180        path: Filepath to a folder where the downloaded data will be saved.
181        patch_shape: The patch shape to use for training.
182        batch_size: The batch size for training.
183        split: The data split to use. Either 'train' or 'val'.
184        download: Whether to download the data if it is not present.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    dataset = get_skeletal_muscle_fibsem_dataset(path, patch_shape, split, download, **ds_kwargs)
192    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)