torch_em.data.datasets.electron_microscopy.axondeepseg

  1"""AxonDeepSeg is a dataset for the segmentation of myelinated axons in EM.
  2It contains two different data types: TEM and SEM.
  3The dataset was published in https://doi.org/10.1038/s41598-018-22181-4.
  4Please cite this publication if you use the dataset in your research.
  5"""
  6
  7import os
  8from glob import glob
  9from shutil import rmtree
 10from typing import Optional, Tuple, Union
 11
 12import imageio
 13import h5py
 14import numpy as np
 15import torch_em
 16
 17from torch.utils.data import Dataset, DataLoader
 18from .. import util
 19
 20URLS = {
 21    "sem": "https://github.com/axondeepseg/data_axondeepseg_sem/archive/refs/heads/master.zip",
 22    "tem": "https://osf.io/download/uewd9"
 23}
 24CHECKSUMS = {
 25    "sem": "d334cbacf548f78ce8dd4a597bf86b884bd15a47a230a0ccc46e1ffa94d58426",
 26    "tem": "e4657280808f3b80d3bf1fba87d1cbbf2455f519baf1a7b16d2ddf2e54739a95"
 27}
 28
 29
 30def _preprocess_sem_data(out_path):
 31    # preprocess the data to get it to a better data format
 32    data_root = os.path.join(out_path, "data_axondeepseg_sem-master")
 33    assert os.path.exists(data_root)
 34
 35    # get the raw data paths
 36    raw_folders = glob(os.path.join(data_root, "sub-rat*"))
 37    raw_folders.sort()
 38    raw_paths = []
 39    for folder in raw_folders:
 40        paths = glob(os.path.join(folder, "micr", "*.png"))
 41        paths.sort()
 42        raw_paths.extend(paths)
 43
 44    # get the label paths
 45    label_folders = glob(os.path.join(
 46        data_root, "derivatives", "labels", "sub-rat*"
 47    ))
 48    label_folders.sort()
 49    label_paths = []
 50    for folder in label_folders:
 51        paths = glob(os.path.join(folder, "micr", "*axonmyelin-manual.png"))
 52        paths.sort()
 53        label_paths.extend(paths)
 54    assert len(raw_paths) == len(label_paths), f"{len(raw_paths)}, {len(label_paths)}"
 55
 56    # process raw data and labels
 57    for i, (rp, lp) in enumerate(zip(raw_paths, label_paths)):
 58        outp = os.path.join(out_path, f"sem_data_{i}.h5")
 59        with h5py.File(outp, "w") as f:
 60
 61            # raw data: invert to match tem em intensities
 62            raw = imageio.imread(rp)
 63            assert np.dtype(raw.dtype) == np.dtype("uint8")
 64            if raw.ndim == 3:  # (one of the images is RGBA)
 65                raw = np.mean(raw[..., :-3], axis=-1)
 66            raw = 255 - raw
 67            f.create_dataset("raw", data=raw, compression="gzip")
 68
 69            # labels: map from
 70            # 0 -> 0
 71            # 127, 128 -> 1
 72            # 255 -> 2
 73            labels = imageio.imread(lp)
 74            assert labels.shape == raw.shape, f"{labels.shape}, {raw.shape}"
 75            label_vals = np.unique(labels)
 76            # 127, 128: both myelin labels, 130, 233: noise
 77            assert len(np.setdiff1d(label_vals, [0, 127, 128, 130, 233, 255])) == 0, f"{label_vals}"
 78            new_labels = np.zeros_like(labels)
 79            new_labels[labels == 127] = 1
 80            new_labels[labels == 128] = 1
 81            new_labels[labels == 255] = 2
 82            f.create_dataset("labels", data=new_labels, compression="gzip")
 83
 84    # clean up
 85    rmtree(data_root)
 86
 87
 88def _preprocess_tem_data(out_path):
 89    data_root = os.path.join(out_path, "TEM_dataset")
 90    folder_names = os.listdir(data_root)
 91    folders = [os.path.join(data_root, fname) for fname in folder_names
 92               if os.path.isdir(os.path.join(data_root, fname))]
 93    for i, folder in enumerate(folders):
 94        data_out = os.path.join(out_path, f"tem_{i}.h5")
 95        with h5py.File(data_out, "w") as f:
 96            im = imageio.imread(os.path.join(folder, "image.png"))
 97            f.create_dataset("raw", data=im, compression="gzip")
 98
 99            # labels: map from
100            # 0 -> 0
101            # 128 -> 1
102            # 255 -> 2
103            # the rest are noise
104            labels = imageio.imread(os.path.join(folder, "mask.png"))
105            new_labels = np.zeros_like(labels)
106            new_labels[labels == 128] = 1
107            new_labels[labels == 255] = 2
108            f.create_dataset("labels", data=new_labels, compression="gzip")
109
110    # clean up
111    rmtree(data_root)
112
113
114def get_axondeepseg_data(path: Union[str, os.PathLike], name: str, download: bool) -> str:
115    """Download the axondeepseg data.
116
117    Args:
118        path: Filepath to a folder where the downloaded data will be saved.
119        name: The name of the dataset to download. Can be either 'sem' or 'tem'.
120        download: Whether to download the data if it is not present.
121
122    Returns:
123        The filepath for the downloaded data.
124    """
125
126    # download and unzip the data
127    url, checksum = URLS[name], CHECKSUMS[name]
128    os.makedirs(path, exist_ok=True)
129    out_path = os.path.join(path, name)
130    if os.path.exists(out_path):
131        return out_path
132
133    tmp_path = os.path.join(path, f"{name}.zip")
134    util.download_source(tmp_path, url, download, checksum=checksum)
135    util.unzip(tmp_path, out_path, remove=True)
136
137    if name == "sem":
138        _preprocess_sem_data(out_path)
139    elif name == "tem":
140        _preprocess_tem_data(out_path)
141    else:
142        raise ValueError(f"Invalid dataset name for axondeepseg, expected 'sem' or 'tem', got {name}.")
143
144    return out_path
145
146
147def get_axondeepseg_dataset(
148    path: Union[str, os.PathLike],
149    name: str,
150    patch_shape: Tuple[int, int],
151    download: bool = False,
152    one_hot_encoding: bool = False,
153    val_fraction: Optional[float] = None,
154    split: Optional[str] = None,
155    **kwargs,
156) -> Dataset:
157    """Get dataset for segmnetation of myelinated axons.
158
159    Args:
160        path: Filepath to a folder where the downloaded data will be saved.
161        name: The name of the dataset to download. Can be either 'sem' or 'tem'.
162        patch_shape: The patch shape to use for training.
163        download: Whether to download the data if it is not present.
164        one_hot_encoding: Whether to return the labels as one hot encoding.
165        val_fraction: The fraction of the data to use for validation.
166        split: The data split. Either 'train' or 'val'.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
168
169    Returns:
170        The segmentation dataset.
171    """
172    if isinstance(name, str):
173        name = [name]
174    assert isinstance(name, (tuple, list))
175
176    all_paths = []
177    for nn in name:
178        data_root = get_axondeepseg_data(path, nn, download)
179        paths = glob(os.path.join(data_root, "*.h5"))
180        paths.sort()
181        if val_fraction is not None:
182            assert split is not None
183            n_samples = int(len(paths) * (1 - val_fraction))
184            paths = paths[:n_samples] if split == "train" else paths[n_samples:]
185        all_paths.extend(paths)
186
187    if one_hot_encoding:
188        if isinstance(one_hot_encoding, bool):
189            # add transformation to go from [0, 1, 2] to one hot encoding
190            class_ids = [0, 1, 2]
191        elif isinstance(one_hot_encoding, int):
192            class_ids = list(range(one_hot_encoding))
193        elif isinstance(one_hot_encoding, (list, tuple)):
194            class_ids = list(one_hot_encoding)
195        else:
196            raise ValueError(
197                f"Invalid value {one_hot_encoding} passed for 'one_hot_encoding', expect bool, int or list."
198            )
199        label_transform = torch_em.transform.label.OneHotTransform(class_ids=class_ids)
200        msg = "'one_hot' is set to True, but 'label_transform' is in the kwargs. It will be over-ridden."
201        kwargs = util.update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
202
203    raw_key, label_key = "raw", "labels"
204    return torch_em.default_segmentation_dataset(all_paths, raw_key, all_paths, label_key, patch_shape, **kwargs)
205
206
207def get_axondeepseg_loader(
208    path: Union[str, os.PathLike],
209    name: str,
210    patch_shape: Tuple[int, int],
211    batch_size: int,
212    download: bool = False,
213    one_hot_encoding: bool = False,
214    val_fraction: Optional[float] = None,
215    split: Optional[str] = None,
216    **kwargs
217) -> DataLoader:
218    """Get dataloader for the segmentation of myelinated axons.
219
220    Args:
221        path: Filepath to a folder where the downloaded data will be saved.
222        name: The name of the dataset to download. Can be either 'sem' or 'tem'.
223        patch_shape: The patch shape to use for training.
224        batch_size: The batch size for training.
225        download: Whether to download the data if it is not present.
226        one_hot_encoding: Whether to return the labels as one hot encoding.
227        val_fraction: The fraction of the data to use for validation.
228        split: The data split. Either 'train' or 'val'.
229        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
230
231    Returns:
232        The PyTorch DataLoader.
233    """
234    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
235    dataset = get_axondeepseg_dataset(
236        path, name, patch_shape, download=download, one_hot_encoding=one_hot_encoding,
237        val_fraction=val_fraction, split=split, **ds_kwargs
238    )
239    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)