synapse_net.sample_data

 1import os
 2import tempfile
 3import pooch
 4
 5from .file_utils import read_mrc, get_cache_dir
 6
 7
 8def get_sample_data(name: str) -> str:
 9    """Get the filepath to SynapseNet sample data, stored as mrc file.
10
11    Args:
12        name: The name of the sample data. Currently, we only provide 'tem_2d' and 'tem_tomo'.
13
14    Returns:
15        The filepath to the downloaded sample data.
16    """
17    registry = {
18        "tem_2d.mrc": "3c6f9ff6d7673d9bf2fd46c09750c3c7dbb8fa1aa59dcdb3363b65cc774dcf28",
19        "tem_tomo.mrc": "fe862ce7c22000d4440e3aa717ca9920b42260f691e5b2ab64cd61c928693c99",
20        "tomo_small.mrc": "057b214777157682e220958e7ca5c90104eada67210a5589572134ac0d8d177f",
21        "mito_small.mrc": "643534ac080f13a4ba8c9f12f5ea59442cb6d1e7ae1a5aec5ea98b7ae9ab45cc",
22    }
23    urls = {
24        "tem_2d.mrc": "https://owncloud.gwdg.de/index.php/s/5sAQ0U4puAspcHg/download",
25        "tem_tomo.mrc": "https://owncloud.gwdg.de/index.php/s/FJDhDfbT4UxhtOn/download",
26        "tomo_small.mrc": "https://owncloud.gwdg.de/index.php/s/03NvEuIye08kkkL/download",
27        "mito_small.mrc": "https://owncloud.gwdg.de/index.php/s/TUhvVMrA7GE8sJv/download",
28    }
29    key = f"{name}.mrc"
30
31    if key not in registry:
32        valid_names = [k[:-4] for k in registry.keys()]
33        raise ValueError(f"Invalid sample name {name}, please choose one of {valid_names}.")
34
35    cache_dir = get_cache_dir()
36    data_registry = pooch.create(
37        path=os.path.join(cache_dir, "sample_data"),
38        base_url="",
39        registry=registry,
40        urls=urls,
41    )
42    file_path = data_registry.fetch(key)
43    return file_path
44
45
46def _sample_data(name):
47    file_path = get_sample_data(name)
48    data, voxel_size = read_mrc(file_path)
49    metadata = {"file_path": file_path, "voxel_size": voxel_size}
50    add_image_kwargs = {"name": name, "metadata": metadata, "colormap": "gray"}
51    return [(data, add_image_kwargs)]
52
53
54def sample_data_tem_2d():
55    return _sample_data("tem_2d")
56
57
58def sample_data_tem_tomo():
59    return _sample_data("tem_tomo")
60
61
62def sample_data_tomo_small():
63    return _sample_data("tomo_small")
64
65
66def sample_data_mito_small():
67    return _sample_data("mito_small")
68
69
70def download_data_from_zenodo(path: str, name: str):
71    """Download data uploaded for the SynapseNet manuscript from zenodo.
72
73    Args:
74        path: The path where the downloaded data will be saved.
75        name: The name of the zenodi dataset.
76    """
77    from torch_em.data.datasets.util import download_source, unzip
78
79    urls = {
80        "2d_tem": "https://zenodo.org/records`/14236382/files/tem_2d.zip?download=1",
81        "inner_ear_ribbon_synapse": "https://zenodo.org/records/14232607/files/inner-ear-ribbon-synapse-tomgrams.zip?download=1",  # noqa
82        "training_data": "https://zenodo.org/records/14330011/files/synapse-net.zip?download=1"
83    }
84    assert name in urls
85    url = urls[name]
86
87    # May need to adapt this for other datasets.
88    # Check if the download already exists.
89    dl_path = path
90    if os.path.exists(dl_path):
91        return
92
93    with tempfile.TemporaryDirectory() as tmp:
94        tmp_path = os.path.join(tmp, f"{name}.zip")
95        download_source(tmp_path, url, download=True, checksum=None)
96        unzip(tmp_path, path, remove=False)