torch_em.data.datasets.electron_microscopy.mitoverse

  1"""MitoVerse is a collection of 3D electron microscopy volumes with mitochondria instance labels.
  2
  3It brings 212 volumes from 13 source datasets (FIB-SEM, ssSEM, ssTEM and SBF-SEM, covering many species, tissues
  4and resolutions) into one format: each volume is a zarr store `data/<dataset>/<volume>.zarr` with the arrays
  5`img` (uint8, ZYX) and `mito` (instance labels, 0 is background), and the per-volume metadata (voxel size in nm
  6as XYZ, modality, species, tissue, provenance) is listed in `catalog.json`. The volumes of the 'han24' dataset
  7also have a `mask` array, which marks the annotated mossy fiber boutons: mitochondria outside of it are not labeled.
  8
  9Benchmarks are defined as split files that assign whole volumes to 'train', 'val' and 'test',
 10see `BENCHMARKS`. The 'mitoem2.0' benchmark is the MitoEM 2.0 benchmark (see also `mitoemv2.py`) as an overlay
 11over the MitoVerse volumes.
 12
 13The data is hosted at https://huggingface.co/datasets/pytc/MitoVerse.
 14The code to build it is at https://github.com/PytorchConnectomics/mitoverse.
 15The license differs per source dataset (see the 'source' of each volume in `catalog.json`), and the MitoVerse
 16authors mark several of them as to be verified, so please check them before reusing or redistributing the data.
 17Please cite the publications of the source datasets that you use.
 18"""
 19
 20import os
 21import json
 22import time
 23from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32REPO_ID = "pytc/MitoVerse"
 33# The data repository is updated in place, so the downloads are pinned to one revision.
 34REVISION = "d0af9a6fd11faaf16ce0822233b7bbf72553ce7d"
 35BENCHMARKS = ["mitoem2.0", "cellmap", "guay21"]
 36COMPLETE_MARKER = ".download_complete"
 37# The volumes consist of many small chunk files, so downloading many of them hits the Hugging Face rate limit.
 38# The download then retries after an increasing wait time. Setting the HF_TOKEN environment variable raises the limit.
 39N_RETRIES = 5
 40RETRY_WAIT = 60
 41
 42
 43def _is_transient(error):
 44    from huggingface_hub.errors import HfHubHTTPError
 45
 46    if isinstance(error, (ConnectionError, TimeoutError)) or "429" in str(error):
 47        return True
 48    status = getattr(getattr(error, "response", None), "status_code", None)
 49    return isinstance(error, HfHubHTTPError) and status is not None and status >= 500
 50
 51
 52def _download(path, patterns):
 53    try:
 54        from huggingface_hub import snapshot_download
 55    except ImportError:
 56        raise ImportError("Downloading MitoVerse requires the 'huggingface_hub' package: pip install huggingface_hub")
 57
 58    for attempt in range(N_RETRIES):
 59        try:
 60            snapshot_download(
 61                repo_id=REPO_ID, repo_type="dataset", revision=REVISION, local_dir=path, allow_patterns=patterns,
 62                max_workers=4,
 63            )
 64            return
 65        except Exception as error:
 66            if attempt == N_RETRIES - 1 or not _is_transient(error):
 67                raise
 68            wait = RETRY_WAIT * 2 ** attempt
 69            print(f"Downloading from {REPO_ID} failed ({type(error).__name__}), retrying in {wait} s.")
 70            time.sleep(wait)
 71
 72
 73def _download_file(path, file_name, download):
 74    file_path = os.path.join(path, file_name)
 75    if not os.path.exists(file_path):
 76        if not download:
 77            raise RuntimeError(f"Cannot find '{file_path}', but 'download' is set to False.")
 78        _download(path, [file_name])
 79    return file_path
 80
 81
 82def get_mitoverse_catalog(path: Union[os.PathLike, str], download: bool = False) -> Dict[str, dict]:
 83    """Get the MitoVerse catalog with the metadata of all volumes.
 84
 85    Args:
 86        path: Filepath to a folder where the downloaded data will be saved.
 87        download: Whether to download the catalog if it is not present.
 88
 89    Returns:
 90        The catalog, mapping each volume id to its metadata (dataset_id, zarr, shape_zyx, voxel_nm, ...).
 91    """
 92    with open(_download_file(path, "catalog.json", download)) as f:
 93        return json.load(f)
 94
 95
 96def get_mitoverse_benchmark(
 97    path: Union[os.PathLike, str], benchmark: str, download: bool = False
 98) -> Dict[str, Union[str, list]]:
 99    """Get a MitoVerse benchmark split file.
100
101    Args:
102        path: Filepath to a folder where the downloaded data will be saved.
103        benchmark: The benchmark name, one of `BENCHMARKS`.
104        download: Whether to download the split file if it is not present.
105
106    Returns:
107        The benchmark, with the list of its 'volumes' (id, zarr and split of each).
108    """
109    if benchmark not in BENCHMARKS:
110        raise ValueError(f"'{benchmark}' is not a valid benchmark. Choose from {BENCHMARKS}.")
111    with open(_download_file(path, f"splits/{benchmark}.json", download)) as f:
112        return json.load(f)
113
114
115def _select_volume_ids(path, catalog, datasets, volume_ids, benchmark, split, download):
116    if volume_ids is not None:
117        volume_ids = [volume_ids] if isinstance(volume_ids, str) else list(volume_ids)
118        invalid = [volume_id for volume_id in volume_ids if volume_id not in catalog]
119        if invalid:
120            raise ValueError(f"The volume ids {invalid} are not in the MitoVerse catalog.")
121    else:
122        volume_ids = list(catalog)
123
124    if datasets is not None:
125        datasets = [datasets] if isinstance(datasets, str) else list(datasets)
126        invalid = sorted(set(datasets) - {meta["dataset_id"] for meta in catalog.values()})
127        if invalid:
128            raise ValueError(f"The datasets {invalid} are not in the MitoVerse catalog.")
129        volume_ids = [volume_id for volume_id in volume_ids if catalog[volume_id]["dataset_id"] in datasets]
130
131    if benchmark is not None:
132        volumes = get_mitoverse_benchmark(path, benchmark, download)["volumes"]
133        in_benchmark = {volume["id"] for volume in volumes if split is None or volume["split"] == split}
134        volume_ids = [volume_id for volume_id in volume_ids if volume_id in in_benchmark]
135    elif split is not None:
136        raise ValueError("A split can only be chosen together with a benchmark.")
137
138    if not volume_ids:
139        raise ValueError("No MitoVerse volume matches the selection.")
140    return sorted(volume_ids)
141
142
143def get_mitoverse_data(
144    path: Union[os.PathLike, str], volume_ids: Sequence[str], download: bool = False
145) -> List[str]:
146    """Download MitoVerse volumes.
147
148    Only the requested volumes are downloaded. A volume counts as present once its download has completed,
149    so an interrupted download is resumed instead of silently leaving missing chunks.
150
151    Args:
152        path: Filepath to a folder where the downloaded data will be saved.
153        volume_ids: The ids of the volumes to download, see `get_mitoverse_catalog`.
154        download: Whether to download the data if it is not present.
155
156    Returns:
157        The filepaths to the zarr stores of the volumes.
158    """
159    catalog = get_mitoverse_catalog(path, download)
160    zarr_paths = [os.path.join(path, catalog[volume_id]["zarr"]) for volume_id in volume_ids]
161
162    missing = [p for p in zarr_paths if not os.path.exists(os.path.join(p, COMPLETE_MARKER))]
163    if missing and not download:
164        raise RuntimeError(f"Cannot find {len(missing)} of the MitoVerse volumes, but 'download' is set to False.")
165
166    # The volumes are downloaded per source dataset, so that the progress is kept if a later download fails.
167    by_dataset = {}
168    for zarr_path in missing:
169        by_dataset.setdefault(os.path.dirname(zarr_path), []).append(zarr_path)
170    for dataset_dir, dataset_paths in by_dataset.items():
171        print(f"Downloading {len(dataset_paths)} MitoVerse volumes of '{os.path.basename(dataset_dir)}'.")
172        _download(path, [f"{os.path.relpath(p, path)}/*" for p in dataset_paths])
173        for zarr_path in dataset_paths:
174            open(os.path.join(zarr_path, COMPLETE_MARKER), "w").close()
175
176    return zarr_paths
177
178
179def get_mitoverse_paths(
180    path: Union[os.PathLike, str],
181    datasets: Optional[Union[str, Sequence[str]]] = None,
182    volume_ids: Optional[Union[str, Sequence[str]]] = None,
183    benchmark: Optional[str] = None,
184    split: Optional[Literal["train", "val", "test"]] = None,
185    download: bool = False,
186) -> List[str]:
187    """Get paths to the MitoVerse data.
188
189    The volumes can be selected by source dataset, by volume id and by benchmark split; the selections combine.
190    Without any selection all volumes are used, which downloads the whole collection (about 28 GB).
191
192    Args:
193        path: Filepath to a folder where the downloaded data will be saved.
194        datasets: The source datasets to use, e.g. 'casser20' or 'han24'. By default all datasets are used.
195        volume_ids: The volumes to use, see `get_mitoverse_catalog`. By default all volumes are used.
196        benchmark: A benchmark that restricts the volumes to its own, one of `BENCHMARKS`.
197        split: The split of the benchmark to use, one of 'train', 'val' or 'test'. Requires a benchmark.
198        download: Whether to download the data if it is not present.
199
200    Returns:
201        The filepaths to the zarr stores of the selected volumes.
202    """
203    catalog = get_mitoverse_catalog(path, download)
204    volume_ids = _select_volume_ids(path, catalog, datasets, volume_ids, benchmark, split, download)
205    return get_mitoverse_data(path, volume_ids, download)
206
207
208def get_mitoverse_dataset(
209    path: Union[os.PathLike, str],
210    patch_shape: Tuple[int, int, int],
211    datasets: Optional[Union[str, Sequence[str]]] = None,
212    volume_ids: Optional[Union[str, Sequence[str]]] = None,
213    benchmark: Optional[str] = None,
214    split: Optional[Literal["train", "val", "test"]] = None,
215    download: bool = False,
216    offsets: Optional[List[List[int]]] = None,
217    boundaries: bool = False,
218    binary: bool = False,
219    **kwargs
220) -> Dataset:
221    """Get the MitoVerse dataset for mitochondria segmentation in EM.
222
223    The volumes have different voxel sizes, see 'voxel_nm' in `get_mitoverse_catalog`.
224    Without any selection all volumes are used, which downloads the whole collection (about 28 GB).
225
226    Args:
227        path: Filepath to a folder where the downloaded data will be saved.
228        patch_shape: The patch shape to use for training.
229        datasets: The source datasets to use, e.g. 'casser20' or 'han24'. By default all datasets are used.
230        volume_ids: The volumes to use, see `get_mitoverse_catalog`. By default all volumes are used.
231        benchmark: A benchmark that restricts the volumes to its own, one of `BENCHMARKS`.
232        split: The split of the benchmark to use, one of 'train', 'val' or 'test'. Requires a benchmark.
233        download: Whether to download the data if it is not present.
234        offsets: Offset values for affinity computation used as target.
235        boundaries: Whether to compute boundaries as the target.
236        binary: Whether to return a binary segmentation target.
237        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
238
239    Returns:
240        The segmentation dataset.
241    """
242    assert len(patch_shape) == 3
243
244    paths = get_mitoverse_paths(path, datasets, volume_ids, benchmark, split, download)
245
246    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
247    kwargs, _ = util.add_instance_label_transform(
248        kwargs, add_binary_target=True, binary=binary, boundaries=boundaries, offsets=offsets
249    )
250
251    return torch_em.default_segmentation_dataset(
252        raw_paths=paths, raw_key="img", label_paths=paths, label_key="mito", patch_shape=patch_shape, **kwargs
253    )
254
255
256def get_mitoverse_loader(
257    path: Union[os.PathLike, str],
258    batch_size: int,
259    patch_shape: Tuple[int, int, int],
260    datasets: Optional[Union[str, Sequence[str]]] = None,
261    volume_ids: Optional[Union[str, Sequence[str]]] = None,
262    benchmark: Optional[str] = None,
263    split: Optional[Literal["train", "val", "test"]] = None,
264    download: bool = False,
265    offsets: Optional[List[List[int]]] = None,
266    boundaries: bool = False,
267    binary: bool = False,
268    **kwargs
269) -> DataLoader:
270    """Get the MitoVerse dataloader for mitochondria segmentation in EM.
271
272    Args:
273        path: Filepath to a folder where the downloaded data will be saved.
274        batch_size: The batch size for training.
275        patch_shape: The patch shape to use for training.
276        datasets: The source datasets to use, e.g. 'casser20' or 'han24'. By default all datasets are used.
277        volume_ids: The volumes to use, see `get_mitoverse_catalog`. By default all volumes are used.
278        benchmark: A benchmark that restricts the volumes to its own, one of `BENCHMARKS`.
279        split: The split of the benchmark to use, one of 'train', 'val' or 'test'. Requires a benchmark.
280        download: Whether to download the data if it is not present.
281        offsets: Offset values for affinity computation used as target.
282        boundaries: Whether to compute boundaries as the target.
283        binary: Whether to return a binary segmentation target.
284        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
285
286    Returns:
287        The DataLoader.
288    """
289    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
290    dataset = get_mitoverse_dataset(
291        path=path, patch_shape=patch_shape, datasets=datasets, volume_ids=volume_ids, benchmark=benchmark, split=split,
292        download=download, offsets=offsets, boundaries=boundaries, binary=binary, **ds_kwargs,
293    )
294    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)