torch_em.data.datasets.electron_microscopy.mitoem

  1"""MitoEM is a dataset for segmenting mitochondria in electron microscopy.
  2It contains two large annotated volumes, one from rat cortex, the other from human cortex.
  3This dataset was used for a segmentation challenge at ISBI 2022.
  4
  5If you use it in your research then please cite https://doi.org/10.1007/978-3-030-59722-1_7.
  6"""
  7
  8import os
  9from tqdm import tqdm
 10import multiprocessing
 11from shutil import rmtree
 12from concurrent import futures
 13from typing import List, Optional, Sequence, Tuple, Union
 14
 15import imageio
 16import numpy as np
 17
 18import torch_em
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22from .. import util
 23
 24
 25URLS = {
 26    "raw": {
 27        "human": "https://www.dropbox.com/s/z41qtu4y735j95e/EM30-H-im.zip?dl=1",
 28        "rat": "https://huggingface.co/datasets/pytc/EM30/resolve/main/EM30-R-im.zip"
 29    },
 30    "labels": {
 31        "human": "https://www.dropbox.com/s/dhf89bc14kemw4e/EM30-H-mito-train-val-v2.zip?dl=1",
 32        "rat": "https://huggingface.co/datasets/pytc/MitoEM/resolve/main/EM30-R-mito-train-val-v2.zip"
 33    }
 34}
 35CHECKSUMS = {
 36    "raw": {
 37        "human": "98fe259f36a7d8d43f99981b7a0ef8cdeba2ce2615ff91595f428ae57207a041",
 38        "rat": "6a2cac68adde5d01984542d3ee1d7753d1fa3e6eb2a042ce15ce297c95885bbe"
 39    },
 40    "labels": {
 41        "human": "0e8ed292cfcd0c58701d9f4299244a1b66d6aeb506c85754c34f98a4eda0ef1b",
 42        "rat": "c56380ac575428a818bd293ca3509d1249999846c3702ccbf11d308acdd2ae86"
 43    }
 44}
 45
 46
 47def _check_data(path, sample):
 48    splits = ["train", "val", "test"]
 49    expected_paths = [os.path.join(path, f"{sample}_{split}.n5") for split in splits]
 50    return all(os.path.exists(pp) for pp in expected_paths)
 51
 52
 53def get_slices(folder):
 54    files = os.listdir(folder)
 55    files.sort()
 56    files = [os.path.splitext(ff)[0] for ff in files]
 57    slice_ids = [int(ff[2:]) if ff.startswith('im') else int(ff[3:]) for ff in files]
 58    return slice_ids
 59
 60
 61def _load_vol(pattern, slice_ids, desc, n_threads, dtype=None):
 62    im0 = pattern % slice_ids[0]
 63    im0 = imageio.imread(im0)
 64
 65    shape = (len(slice_ids),) + im0.shape
 66
 67    dtype = im0.dtype if dtype is None else dtype
 68    out = np.zeros(shape, dtype=dtype)
 69    out[0] = im0
 70
 71    def load_slice(z, slice_id):
 72        out[z] = imageio.imread(pattern % slice_id)
 73
 74    zs = list(range(1, len(slice_ids)))
 75    assert len(zs) == len(slice_ids) - 1
 76    with futures.ThreadPoolExecutor(n_threads) as tp:
 77        list(tqdm(tp.map(load_slice, zs, slice_ids[1:]), total=len(slice_ids) - 1, desc=desc))
 78
 79    return out
 80
 81
 82def _create_volume(out_path, im_folder, label_folder=None, z_start=None):
 83    import z5py
 84
 85    if label_folder is None:
 86        assert z_start is not None
 87        n_slices = len(get_slices(im_folder))
 88        slices = list(range(z_start, n_slices))
 89    else:
 90        assert z_start is None
 91        slices = get_slices(label_folder)
 92
 93    n_threads = min(16, multiprocessing.cpu_count())
 94    raw = _load_vol(os.path.join(im_folder, "im%04i.png"), slices, "load raw", n_threads)
 95    if label_folder is not None:
 96        labels = _load_vol(os.path.join(label_folder, "seg%04i.tif"), slices, "load labels", n_threads, dtype="uint64")
 97
 98    print("Write volume to", out_path)
 99    chunks = (32, 256, 256)
100    with z5py.File(out_path, "a") as f:
101        f.create_dataset("raw", data=raw, chunks=chunks, compression="gzip", n_threads=n_threads)
102        if label_folder is not None:
103            ds = f.create_dataset("labels", data=labels, chunks=chunks, compression="gzip", n_threads=n_threads)
104            ds.attrs["maxId"] = int(labels.max()) + 1
105
106    return slices[-1]
107
108
109def _require_mitoem_sample(path, sample, download):
110    os.makedirs(path, exist_ok=True)
111
112    for name in ("raw", "labels"):
113        url = URLS[name][sample]
114        checksum = CHECKSUMS[name][sample]
115        zip_path = os.path.join(path, f"{sample}.zip")
116        util.download_source(zip_path, url, download, checksum)
117        util.unzip(zip_path, path, remove=True)
118
119    im_folder = os.path.join(path, "im")
120    train_folder = os.path.join(path, "mito-train-v2")
121    val_folder = os.path.join(path, "mito-val-v2")
122
123    print("Create train volume")
124    train_path = os.path.join(path, f"{sample}_train.n5")
125    _create_volume(train_path, im_folder, train_folder)
126
127    print("Create validation volume")
128    val_path = os.path.join(path, f"{sample}_val.n5")
129    z = _create_volume(val_path, im_folder, val_folder)
130
131    print("Create test volume")
132    test_path = os.path.join(path, f"{sample}_test.n5")
133    _create_volume(test_path, im_folder, z_start=z)
134
135    rmtree(im_folder)
136    rmtree(train_folder)
137    rmtree(val_folder)
138
139
140def get_mitoem_data(path: Union[os.PathLike, str], samples: Sequence[str], splits: Sequence[str], download: bool):
141    """Download the MitoEM training data.
142
143    Args:
144        path: Filepath to a folder where the downloaded data will be saved.
145        samples: The samples to download. The available samples are 'human' and 'rat'.
146        splits: The data splits to download. The available splits are 'train', 'val' and 'test'.
147        download: Whether to download the data if it is not present.
148    """
149    assert len(set(splits) - {"train", "val"}) == 0, f"{splits}"
150    assert len(set(samples) - {"human", "rat"}) == 0, f"{samples}"
151    os.makedirs(path, exist_ok=True)
152
153    for sample in samples:
154        if not _check_data(path, sample):
155            print("The MitoEM data for sample", sample, "is not available yet and will be downloaded and created.")
156            print("Note that this dataset is large, so this step can take several hours (depending on your internet).")
157            _require_mitoem_sample(path, sample, download)
158            print("The MitoEM data for sample", sample, "has been created.")
159
160        for split in splits:
161            split_path = os.path.join(path, f"{sample}_{split}.n5")
162            assert os.path.exists(split_path), split_path
163
164
165def get_mitoem_paths(
166    path: Union[os.PathLike, str],
167    splits: Sequence[str],
168    samples: Sequence[str] = ("human", "rat"),
169    download: bool = False,
170) -> List[str]:
171    """Get paths for MitoEM data.
172
173    Args:
174        path: Filepath to a folder where the downloaded data will be saved.
175        samples: The samples to download. The available samples are 'human' and 'rat'.
176        splits: The data splits to download. The available splits are 'train', 'val' and 'test'.
177        download: Whether to download the data if it is not present.
178
179    Returns:
180        The filepaths for the stored data.
181    """
182    if isinstance(splits, str):
183        splits = [splits]
184
185    if isinstance(samples, str):
186        samples = [samples]
187
188    get_mitoem_data(path, samples, splits, download)
189    data_paths = [os.path.join(path, f"{sample}_{split}.n5") for split in splits for sample in samples]
190
191    return data_paths
192
193
194def get_mitoem_dataset(
195    path: Union[os.PathLike, str],
196    splits: Sequence[str],
197    patch_shape: Tuple[int, int, int],
198    samples: Sequence[str] = ("human", "rat"),
199    download: bool = False,
200    offsets: Optional[List[List[int]]] = None,
201    boundaries: bool = False,
202    binary: bool = False,
203    **kwargs,
204) -> Dataset:
205    """Get the MitoEM dataset for the segmentation of mitochondria in EM.
206
207    Args:
208        path: Filepath to a folder where the downloaded data will be saved.
209        splits: The splits to use for the dataset. Available values are 'train', 'val' and 'test'.
210        patch_shape: The patch shape to use for training.
211        samples: The samples to use for the dataset. The available samples are 'human' and 'rat'.
212        download: Whether to download the data if it is not present.
213        offsets: Offset values for affinity computation used as target.
214        boundaries: Whether to compute boundaries as the target.
215        binary: Whether to return a binary segmentation target.
216        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
217
218    Returns:
219       The segmentation dataset.
220    """
221    assert len(patch_shape) == 3
222
223    data_paths = get_mitoem_paths(path, splits, samples, download)
224
225    kwargs, _ = util.add_instance_label_transform(
226        kwargs, add_binary_target=True, binary=binary, boundaries=boundaries, offsets=offsets
227    )
228
229    return torch_em.default_segmentation_dataset(
230        raw_paths=data_paths,
231        raw_key="raw",
232        label_paths=data_paths,
233        label_key="labels",
234        patch_shape=patch_shape,
235        **kwargs
236    )
237
238
239def get_mitoem_loader(
240    path: Union[os.PathLike, str],
241    splits: Sequence[str],
242    patch_shape: Tuple[int, int, int],
243    batch_size: int,
244    samples: Sequence[str] = ("human", "rat"),
245    download: bool = False,
246    offsets: Optional[List[List[int]]] = None,
247    boundaries: bool = False,
248    binary: bool = False,
249    **kwargs,
250) -> DataLoader:
251    """Get the MitoEM dataloader for the segmentation of mitochondria in EM.
252
253    Args:
254        path: Filepath to a folder where the downloaded data will be saved.
255        splits: The splits to use for the dataset. Available values are 'train', 'val' and 'test'.
256        patch_shape: The patch shape to use for training.
257        batch_size: The batch size for training.
258        samples: The samples to use for the dataset. The available samples are 'human' and 'rat'.
259        download: Whether to download the data if it is not present.
260        offsets: Offset values for affinity computation used as target.
261        boundaries: Whether to compute boundaries as the target.
262        binary: Whether to return a binary segmentation target.
263        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
264
265    Returns:
266       The DataLoader.
267    """
268    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
269    dataset = get_mitoem_dataset(path, splits, patch_shape, samples, download, offsets, boundaries, binary, **ds_kwargs)
270    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)