torch_em.data.datasets.electron_microscopy.mitoemv2

  1"""MitoEM v2 is a benchmark collection for 3D mitochondria instance segmentation in electron microscopy.
  2
  3It contains eight challenging datasets with expert-verified labels, covering biologically complex
  4scenarios such as dense mitochondrial packing, hyperfused networks, thin-necked morphologies,
  5and ultrastructurally ambiguous boundaries.
  6
  7The data is located at https://doi.org/10.5281/zenodo.17635006.
  8The dataset is from the publication https://doi.org/10.1101/2025.11.12.687478.
  9Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13import json
 14from glob import glob
 15from typing import Union, Literal, Optional, Tuple, List
 16
 17import numpy as np
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26BASE_URL = "https://zenodo.org/records/17635006/files"
 27
 28DATASETS = {
 29    "beta": "Dataset001_ME2-Beta",
 30    "jurkat": "Dataset002_ME2-Jurkat",
 31    "macro": "Dataset003_ME2-Macro",
 32    "mossy": "Dataset004_ME2-Mossy",
 33    "podo": "Dataset005_ME2-Podo",
 34    "pyra": "Dataset006_ME2-Pyra",
 35    "sperm": "Dataset007_ME2-Sperm",
 36    "stem": "Dataset008_ME2-Stem",
 37}
 38
 39DATASET_NAMES = list(DATASETS.keys())
 40
 41
 42CHUNKS = (32, 256, 256)
 43
 44
 45def _read_split_info(dataset_dir):
 46    # The split file is called 'split.json' in most datasets, but 'splits.json' in the pyra dataset.
 47    for name in ("split.json", "splits.json"):
 48        split_path = os.path.join(dataset_dir, name)
 49        if os.path.exists(split_path):
 50            with open(split_path) as f:
 51                return json.load(f)[0]
 52    raise FileNotFoundError(f"Cannot find the split file in {dataset_dir}.")
 53
 54
 55def _write_nifti(f, key, nifti_path, transpose):
 56    """Write a NIfTI volume into the n5 file as (z, y, x).
 57
 58    NIfTI stores (x, y, z) in Fortran order, so a slab along the last axis is contiguous on disk and the volume can
 59    be converted slab by slab, without loading it into memory as a whole.
 60    """
 61    import nibabel as nib
 62
 63    nii = nib.load(nifti_path)
 64    dtype = nii.get_data_dtype()
 65    if not transpose:
 66        f.create_dataset(key, data=np.asarray(nii.dataobj).astype(dtype), chunks=CHUNKS, compression="gzip")
 67        return
 68
 69    shape = nii.shape[::-1]
 70    ds = f.create_dataset(key, shape=shape, dtype=dtype, chunks=CHUNKS, compression="gzip")
 71    for z0 in range(0, shape[0], CHUNKS[0]):
 72        z1 = min(z0 + CHUNKS[0], shape[0])
 73        ds[z0:z1] = np.asarray(nii.dataobj[:, :, z0:z1]).astype(dtype).transpose(2, 1, 0)
 74
 75
 76def _preprocess_dataset(path, dataset_name, dataset_dir):
 77    """Preprocess a single dataset: convert NIfTI to n5."""
 78    import shutil
 79    import z5py
 80
 81    n5_dir = os.path.join(path, "n5_data", dataset_name)
 82    os.makedirs(n5_dir, exist_ok=True)
 83    split_info = _read_split_info(dataset_dir)
 84
 85    processed = {}
 86    for split_name, split_key in [("train", "train"), ("val", "val"), ("test", "test")]:
 87        samples = split_info.get(split_key, [])
 88        if not samples:
 89            continue
 90
 91        for sample in samples:
 92            # Determine source directories based on split
 93            if split_name == "test":
 94                img_dir = "imagesTs"
 95                lbl_dir = "labelsTs"
 96            else:
 97                img_dir = "imagesTr"
 98                lbl_dir = "labelsTr"
 99
100            img_nifti = os.path.join(dataset_dir, img_dir, f"{sample}_0000.nii.gz")
101            lbl_nifti = os.path.join(dataset_dir, lbl_dir, f"{sample}.nii.gz")
102
103            if not os.path.exists(img_nifti) or not os.path.exists(lbl_nifti):
104                continue
105
106            n5_path = os.path.join(n5_dir, f"{sample}.n5")
107
108            if not os.path.exists(n5_path):
109                print(f"Converting {sample} to n5...")
110                # The volume is written under a temporary name and moved in place once it is complete,
111                # so that an interrupted conversion is not mistaken for a complete one.
112                tmp_path = os.path.join(n5_dir, f"{sample}_incomplete.n5")
113                if os.path.exists(tmp_path):
114                    shutil.rmtree(tmp_path)
115
116                # The labels of this sample are stored in (z, y, x) order in the NIfTI file already.
117                # NOTE: Even then they do not match the raw data (the mitochondria are drawn in empty resin).
118                transpose_labels = sample != "me2-jurkat_train02"
119                with z5py.File(tmp_path, "a") as f:
120                    _write_nifti(f, "raw", img_nifti, transpose=True)
121                    _write_nifti(f, "labels", lbl_nifti, transpose=transpose_labels)
122                    if f["raw"].shape != f["labels"].shape:
123                        raise RuntimeError(f"There is a shape mismatch between raw and labels for {sample}.")
124                os.rename(tmp_path, n5_path)
125
126            if split_name not in processed:
127                processed[split_name] = []
128            processed[split_name].append(n5_path)
129
130    return processed
131
132
133def get_mitoemv2_data(
134    path: Union[os.PathLike, str],
135    dataset: str,
136    download: bool = False,
137) -> str:
138    """Download and preprocess a MitoEM v2 dataset.
139
140    Args:
141        path: Filepath to a folder where the downloaded data will be saved.
142        dataset: The dataset to download. One of 'beta', 'jurkat', 'macro', 'mossy',
143            'podo', 'pyra', 'sperm', or 'stem'.
144        download: Whether to download the data if it is not present.
145
146    Returns:
147        The filepath to the preprocessed n5 data directory.
148    """
149    assert dataset in DATASETS, f"'{dataset}' is not valid. Choose from {DATASET_NAMES}."
150
151    dataset_folder = DATASETS[dataset]
152    n5_dir = os.path.join(path, "n5_data", dataset)
153    zip_path = os.path.join(path, f"{dataset_folder}.zip")
154    dataset_dir = os.path.join(path, dataset_folder)
155
156    # Without the NIfTI data there is nothing left to convert, so the converted data is used as it is.
157    if not os.path.exists(dataset_dir) and len(glob(os.path.join(n5_dir, "*.n5"))) > 0:
158        return n5_dir
159
160    # Download if needed
161    if not os.path.exists(dataset_dir):
162        os.makedirs(path, exist_ok=True)
163        url = f"{BASE_URL}/{dataset_folder}.zip"
164        util.download_source(path=zip_path, url=url, download=download, checksum=None)
165        util.unzip(zip_path=zip_path, dst=path)
166
167    # Preprocess
168    _preprocess_dataset(path, dataset, dataset_dir)
169
170    return n5_dir
171
172
173def get_mitoemv2_paths(
174    path: Union[os.PathLike, str],
175    dataset: Optional[Union[str, List[str]]] = None,
176    split: Literal["train", "val", "test"] = "train",
177    download: bool = False,
178) -> List[str]:
179    """Get paths to the MitoEM v2 data.
180
181    Args:
182        path: Filepath to a folder where the downloaded data will be saved.
183        dataset: The dataset(s) to use. One of 'beta', 'jurkat', 'macro', 'mossy',
184            'podo', 'pyra', 'sperm', or 'stem'. Can also be a list of dataset names.
185            If None, all datasets will be used.
186        split: The data split to use. One of 'train', 'val', or 'test'.
187        download: Whether to download the data if it is not present.
188
189    Returns:
190        List of filepaths for the n5 data.
191    """
192    from natsort import natsorted
193
194    assert split in ("train", "val", "test"), f"'{split}' is not a valid split."
195
196    if dataset is None:
197        dataset = DATASET_NAMES
198    elif isinstance(dataset, str):
199        dataset = [dataset]
200
201    all_n5_paths = []
202    for ds in dataset:
203        n5_dir = get_mitoemv2_data(path, ds, download)
204
205        # Read split info to get correct samples
206        dataset_folder = DATASETS[ds]
207        dataset_dir = os.path.join(path, dataset_folder)
208        split_info = _read_split_info(dataset_dir)
209
210        samples = split_info.get(split, [])
211        n5_paths = [os.path.join(n5_dir, f"{sample}.n5") for sample in samples]
212        n5_paths = [p for p in n5_paths if os.path.exists(p)]
213        all_n5_paths.extend(n5_paths)
214
215    assert len(all_n5_paths) > 0, f"No data found for {dataset}/{split}"
216
217    return natsorted(all_n5_paths)
218
219
220def get_mitoemv2_dataset(
221    path: Union[os.PathLike, str],
222    patch_shape: Tuple[int, int, int],
223    dataset: Optional[Union[str, List[str]]] = None,
224    split: Literal["train", "val", "test"] = "train",
225    download: bool = False,
226    offsets: Optional[List[List[int]]] = None,
227    boundaries: bool = False,
228    binary: bool = False,
229    **kwargs
230) -> Dataset:
231    """Get the MitoEM v2 dataset for mitochondria segmentation in EM.
232
233    Args:
234        path: Filepath to a folder where the downloaded data will be saved.
235        patch_shape: The patch shape to use for training.
236        dataset: The dataset(s) to use. One of 'beta', 'jurkat', 'macro', 'mossy',
237            'podo', 'pyra', 'sperm', or 'stem'. Can also be a list of dataset names.
238            If None, all datasets will be used.
239        split: The data split to use. One of 'train', 'val', or 'test'.
240        download: Whether to download the data if it is not present.
241        offsets: Offset values for affinity computation used as target.
242        boundaries: Whether to compute boundaries as the target.
243        binary: Whether to return a binary segmentation target.
244        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
245
246    Returns:
247        The segmentation dataset.
248    """
249    assert len(patch_shape) == 3
250
251    n5_paths = get_mitoemv2_paths(path, dataset, split, download)
252
253    kwargs, _ = util.add_instance_label_transform(
254        kwargs, add_binary_target=True, binary=binary, boundaries=boundaries, offsets=offsets
255    )
256
257    return torch_em.default_segmentation_dataset(
258        raw_paths=n5_paths,
259        raw_key="raw",
260        label_paths=n5_paths,
261        label_key="labels",
262        patch_shape=patch_shape,
263        **kwargs
264    )
265
266
267def get_mitoemv2_loader(
268    path: Union[os.PathLike, str],
269    batch_size: int,
270    patch_shape: Tuple[int, int, int],
271    dataset: Optional[Union[str, List[str]]] = None,
272    split: Literal["train", "val", "test"] = "train",
273    download: bool = False,
274    offsets: Optional[List[List[int]]] = None,
275    boundaries: bool = False,
276    binary: bool = False,
277    **kwargs
278) -> DataLoader:
279    """Get the MitoEM v2 dataloader for mitochondria segmentation in EM.
280
281    Args:
282        path: Filepath to a folder where the downloaded data will be saved.
283        batch_size: The batch size for training.
284        patch_shape: The patch shape to use for training.
285        dataset: The dataset(s) to use. One of 'beta', 'jurkat', 'macro', 'mossy',
286            'podo', 'pyra', 'sperm', or 'stem'. Can also be a list of dataset names.
287            If None, all datasets will be used.
288        split: The data split to use. One of 'train', 'val', or 'test'.
289        download: Whether to download the data if it is not present.
290        offsets: Offset values for affinity computation used as target.
291        boundaries: Whether to compute boundaries as the target.
292        binary: Whether to return a binary segmentation target.
293        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
294
295    Returns:
296        The DataLoader.
297    """
298    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
299    dataset_obj = get_mitoemv2_dataset(
300        path=path,
301        patch_shape=patch_shape,
302        dataset=dataset,
303        split=split,
304        download=download,
305        offsets=offsets,
306        boundaries=boundaries,
307        binary=binary,
308        **ds_kwargs,
309    )
310    return torch_em.get_data_loader(dataset=dataset_obj, batch_size=batch_size, **loader_kwargs)