torch_em.data.datasets.electron_microscopy.snemi_mito

  1"""The SNEMI Mito dataset contains annotations for mitochondria segmentation in electron microscopy (EM)
  2images of mouse neocortex.
  3
  4It provides mitochondria labels for the sections of the SNEMI3D challenge volume, which comes from the publication
  5https://doi.org/10.1016/j.cell.2015.06.054. The data consists of 100 images (1024 x 1024 pixels, 8-bit, CLAHE contrast
  6enhanced) that are split into 50 train, 20 validation and 30 test images.
  7The images are consecutive sections of the volume (6 x 6 x 30 nm), so they can also be used as one 3D volume.
  8The volume is stored as `snemi_mito.h5` with 'raw' and binary 'labels', and the splits are its ROIs in `VOLUME_ROIS`.
  9
 10The data is available at https://github.com/sumiya-kuroda/EMsegmentation-mito/tree/master/dataset/SNEMI3D_mito.
 11The repository does not state a license. The mitochondria labels were prepared by members of the Kubota lab
 12(see the repository for details). Please cite the SNEMI3D / Kasthuri et al. publication and the repository
 13if you use this dataset in your research.
 14"""
 15
 16import os
 17import shutil
 18from glob import glob
 19from typing import List, Tuple, Union
 20
 21import h5py
 22import numpy as np
 23import imageio.v3 as imageio
 24from natsort import natsorted
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33# The archive is pinned to a commit, so that the checksum does not change with new commits to the repository.
 34URL = "https://github.com/sumiya-kuroda/EMsegmentation-mito/archive/44134f6dc27b1ad7f04293398e482fec854bb658.zip"
 35CHECKSUM = "1d560ad05d0faab8e6fedd9347131c2a0f395dcdba079f8699a66c3e39d80a25"
 36SPLITS = {"train": "train", "val": "valid", "test": "test"}
 37# The splits are contiguous ranges of sections in the volume.
 38VOLUME_ROIS = {"train": np.s_[0:50, :, :], "val": np.s_[50:70, :, :], "test": np.s_[70:100, :, :]}
 39
 40
 41def _convert_rgba_labels(data_dir):
 42    """Convert the label images that are stored as RGBA, with the gray values in every channel, to one channel."""
 43    from PIL import Image
 44
 45    for label_path in glob(os.path.join(data_dir, "*", "labels", "*.png")):
 46        with Image.open(label_path) as image:
 47            if image.mode != "RGBA":
 48                continue
 49            labels = image.getchannel(0)
 50        labels.save(label_path)
 51
 52
 53def _create_volume(data_dir, volume_path):
 54    raw_paths, label_paths = [], []
 55    for split in SPLITS.values():
 56        raw_paths.extend(natsorted(glob(os.path.join(data_dir, split, "images", "*.png"))))
 57        label_paths.extend(natsorted(glob(os.path.join(data_dir, split, "labels", "*.png"))))
 58    names = [os.path.basename(raw_path) for raw_path in raw_paths]
 59    assert names == [os.path.basename(label_path) for label_path in label_paths]
 60    assert names == [f"{i:04}.png" for i in range(len(names))], "The sections are not consecutive."
 61
 62    raw = np.stack([imageio.imread(raw_path) for raw_path in raw_paths])
 63    labels = np.stack([imageio.imread(label_path) > 0 for label_path in label_paths]).astype("uint8")
 64
 65    tmp_path = f"{volume_path}.incomplete"
 66    with h5py.File(tmp_path, "w") as f:
 67        f.create_dataset("raw", data=raw, compression="gzip", chunks=(8, 256, 256))
 68        f.create_dataset("labels", data=labels, compression="gzip", chunks=(8, 256, 256))
 69        f.attrs["voxel_size_nm"] = [30, 6, 6]
 70    os.replace(tmp_path, volume_path)
 71
 72
 73def _download_data(path, data_dir, download):
 74    os.makedirs(path, exist_ok=True)
 75    zip_path = os.path.join(path, "snemi_mito.zip")
 76    util.download_source(zip_path, URL, download, checksum=CHECKSUM)
 77
 78    # The archive contains the whole repository, we only keep the SNEMI3D_mito folder.
 79    tmp_dir = os.path.join(path, "tmp")
 80    util.unzip(zip_path, tmp_dir, remove=True)
 81    extracted = glob(os.path.join(tmp_dir, "*", "dataset", "SNEMI3D_mito"))
 82    assert len(extracted) == 1, "Unexpected structure of the downloaded archive."
 83    shutil.move(extracted[0], data_dir)
 84    shutil.rmtree(tmp_dir)
 85
 86
 87def get_snemi_mito_data(path: Union[os.PathLike, str], download: bool = False, as_volume: bool = False) -> str:
 88    """Download the SNEMI Mito dataset.
 89
 90    Args:
 91        path: Filepath to a folder where the downloaded data will be saved.
 92        download: Whether to download the data if it is not present.
 93        as_volume: Whether to stack the sections into one 3D volume.
 94
 95    Returns:
 96        The filepath to the folder with the data, or the filepath to the volume if `as_volume` is set.
 97    """
 98    data_dir = os.path.join(path, "SNEMI3D_mito")
 99    volume_path = os.path.join(path, "snemi_mito.h5")
100    if as_volume and os.path.exists(volume_path):
101        return volume_path
102
103    if not os.path.exists(data_dir):
104        _download_data(path, data_dir, download)
105    _convert_rgba_labels(data_dir)
106
107    if as_volume:
108        _create_volume(data_dir, volume_path)
109        return volume_path
110    return data_dir
111
112
113def get_snemi_mito_paths(
114    path: Union[os.PathLike, str], split: str, download: bool = False, as_volume: bool = False
115) -> Union[Tuple[List[str], List[str]], Tuple[str, Tuple[slice, ...]]]:
116    """Get paths to the SNEMI Mito data.
117
118    Args:
119        path: Filepath to a folder where the downloaded data will be saved.
120        split: The data split to use. Either 'train', 'val' or 'test'.
121        download: Whether to download the data if it is not present.
122        as_volume: Whether to return the 3D volume and the ROI of the split instead of the 2D images.
123
124    Returns:
125        List of filepaths for the image data, or the filepath to the volume if `as_volume` is set.
126        List of filepaths for the label data, or the ROI of the split in the volume if `as_volume` is set.
127    """
128    if split not in SPLITS:
129        raise ValueError(f"'{split}' is not a valid split. Please choose from {list(SPLITS.keys())}.")
130
131    if as_volume:
132        return get_snemi_mito_data(path, download, as_volume=True), VOLUME_ROIS[split]
133
134    data_dir = get_snemi_mito_data(path, download)
135    raw_paths = natsorted(glob(os.path.join(data_dir, SPLITS[split], "images", "*.png")))
136    label_paths = natsorted(glob(os.path.join(data_dir, SPLITS[split], "labels", "*.png")))
137    assert len(raw_paths) > 0 and len(raw_paths) == len(label_paths)
138
139    return raw_paths, label_paths
140
141
142def _binary_labels(labels):
143    return (labels > 0).astype("uint8")
144
145
146def get_snemi_mito_dataset(
147    path: Union[os.PathLike, str],
148    patch_shape: Tuple[int, int],
149    split: str,
150    binary: bool = True,
151    download: bool = False,
152    as_volume: bool = False,
153    **kwargs
154) -> Dataset:
155    """Get the SNEMI Mito dataset for mitochondria segmentation in EM.
156
157    Args:
158        path: Filepath to a folder where the downloaded data will be saved.
159        patch_shape: The patch shape to use for training.
160        split: The data split to use. Either 'train', 'val' or 'test'.
161        binary: Whether to return a binary foreground mask. Otherwise the mask values are returned as they are.
162            The labels of the volume are always binary.
163        download: Whether to download the data if it is not present.
164        as_volume: Whether to use the split of the 3D volume instead of the 2D images.
165        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
166
167    Returns:
168        The segmentation dataset.
169    """
170    if as_volume:
171        volume_path, roi = get_snemi_mito_paths(path, split, download, as_volume=True)
172        return torch_em.default_segmentation_dataset(
173            raw_paths=volume_path,
174            raw_key="raw",
175            label_paths=volume_path,
176            label_key="labels",
177            patch_shape=patch_shape,
178            rois=roi,
179            is_seg_dataset=True,
180            **kwargs
181        )
182
183    raw_paths, label_paths = get_snemi_mito_paths(path, split, download)
184
185    if binary:
186        kwargs = util.update_kwargs(kwargs, "label_transform", _binary_labels)
187
188    return torch_em.default_segmentation_dataset(
189        raw_paths=raw_paths,
190        raw_key=None,
191        label_paths=label_paths,
192        label_key=None,
193        patch_shape=patch_shape,
194        is_seg_dataset=False,
195        **kwargs
196    )
197
198
199def get_snemi_mito_loader(
200    path: Union[os.PathLike, str],
201    batch_size: int,
202    patch_shape: Tuple[int, int],
203    split: str,
204    binary: bool = True,
205    download: bool = False,
206    as_volume: bool = False,
207    **kwargs
208) -> DataLoader:
209    """Get the SNEMI Mito dataloader for mitochondria segmentation in EM.
210
211    Args:
212        path: Filepath to a folder where the downloaded data will be saved.
213        batch_size: The batch size for training.
214        patch_shape: The patch shape to use for training.
215        split: The data split to use. Either 'train', 'val' or 'test'.
216        binary: Whether to return a binary foreground mask. Otherwise the mask values are returned as they are.
217            The labels of the volume are always binary.
218        download: Whether to download the data if it is not present.
219        as_volume: Whether to use the split of the 3D volume instead of the 2D images.
220        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
221
222    Returns:
223        The DataLoader.
224    """
225    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
226    dataset = get_snemi_mito_dataset(path, patch_shape, split, binary, download, as_volume, **ds_kwargs)
227    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)