torch_em.data.datasets.electron_microscopy.emcopilot

  1"""The EMcopilot dataset contains electron micrographs of catalyst materials with manual semantic annotations.
  2
  3It contains two sub-datasets of (scanning transmission) electron micrographs with binary manual labels:
  4- PtSn_Al2O3_clusters: PtSn clusters on an Al2O3 support (23 labeled images).
  5- Au_ZSM5: Au on a ZSM-5 zeolite support (110 labeled images).
  6
  7The dataset is located at https://doi.org/10.5281/zenodo.18929384 (version 4) and is distributed under the
  8CC BY 4.0 license. It belongs to the publication https://doi.org/10.1038/s41524-025-01790-y.
  9Please cite the publication if you use this dataset in your research.
 10
 11NOTE: The record also contains a 'PtSn_Al2O3_nanoparticle' sub-dataset, which is not supported here because its labels
 12do not appear to be aligned with the images (some pairs differ in shape and the masks do not match the particles).
 13"""
 14
 15import os
 16import zipfile
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, List, Literal
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://zenodo.org/records/18929384/files/EMcopilot_dataset.zip?download=1"
 29CHECKSUM = "67629f015460ae59c01a24991a12f97f8853ab582cfd18010add5a1fbd8ccc12"
 30
 31FOLDERS = {
 32    "PtSn_Al2O3_clusters": ("01_02_PtSn_Al2O3_clusters", "02_02_PtSn_Al2O3_clusters"),
 33    "Au_ZSM5": ("01_03_Au_ZSM5", "02_03_Au_ZSM5"),
 34}
 35
 36
 37def get_emcopilot_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 38    """Download the EMcopilot dataset.
 39
 40    NOTE: The archive is 9.8 GB large, but only the experimental images and manual labels (~600 MB) are extracted.
 41
 42    Args:
 43        path: Filepath to a folder where the downloaded data will be saved.
 44        download: Whether to download the data if it is not present.
 45
 46    Returns:
 47        The filepath to the folder with the data.
 48    """
 49    data_dir = os.path.join(path, "EMcopilot_dataset")
 50    if os.path.exists(os.path.join(data_dir, "02_manual_labels")):
 51        return data_dir
 52
 53    os.makedirs(path, exist_ok=True)
 54    zip_path = os.path.join(path, "EMcopilot_dataset.zip")
 55    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 56
 57    with zipfile.ZipFile(zip_path, "r") as f:
 58        members = [m for m in f.namelist() if m.startswith(("EMcopilot_dataset/01_", "EMcopilot_dataset/02_"))]
 59        f.extractall(path, members=members)
 60    os.remove(zip_path)
 61    return data_dir
 62
 63
 64def get_emcopilot_paths(
 65    path: Union[os.PathLike, str],
 66    name: Literal["PtSn_Al2O3_clusters", "Au_ZSM5"],
 67    download: bool = False,
 68) -> Tuple[List[str], List[str]]:
 69    """Get paths to the EMcopilot data.
 70
 71    Args:
 72        path: Filepath to a folder where the downloaded data will be saved.
 73        name: The name of the sub-dataset. One of 'PtSn_Al2O3_clusters' or 'Au_ZSM5'.
 74        download: Whether to download the data if it is not present.
 75
 76    Returns:
 77        List of filepaths for the image data.
 78        List of filepaths for the label data.
 79    """
 80    assert name in FOLDERS, f"'{name}' is not a valid sub-dataset name. Choose from {list(FOLDERS)}."
 81
 82    data_dir = get_emcopilot_data(path, download)
 83    image_folder, label_folder = FOLDERS[name]
 84    image_paths = natsorted(glob(os.path.join(data_dir, "01_experimental_images", image_folder, "*.png")))
 85    label_paths = natsorted(glob(os.path.join(data_dir, "02_manual_labels", label_folder, "*.png")))
 86
 87    assert len(image_paths) > 0 and len(image_paths) == len(label_paths)
 88    return image_paths, label_paths
 89
 90
 91def get_emcopilot_dataset(
 92    path: Union[os.PathLike, str],
 93    patch_shape: Tuple[int, int],
 94    name: Literal["PtSn_Al2O3_clusters", "Au_ZSM5"],
 95    download: bool = False,
 96    **kwargs
 97) -> Dataset:
 98    """Get the EMcopilot dataset for semantic segmentation of catalyst particles in electron micrographs.
 99
100    The labels are binary, where 1 is the foreground (particles) and 0 is the background.
101
102    Args:
103        path: Filepath to a folder where the downloaded data will be saved.
104        patch_shape: The patch shape to use for training.
105        name: The name of the sub-dataset. One of 'PtSn_Al2O3_clusters' or 'Au_ZSM5'.
106        download: Whether to download the data if it is not present.
107        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
108
109    Returns:
110        The segmentation dataset.
111    """
112    image_paths, label_paths = get_emcopilot_paths(path, name, download)
113
114    kwargs = util.update_kwargs(kwargs, "ndim", 2)
115    kwargs = util.update_kwargs(kwargs, "label_transform", torch_em.transform.label.labels_to_binary)
116
117    return torch_em.default_segmentation_dataset(
118        raw_paths=image_paths,
119        raw_key=None,
120        label_paths=label_paths,
121        label_key=None,
122        patch_shape=patch_shape,
123        is_seg_dataset=False,
124        **kwargs
125    )
126
127
128def get_emcopilot_loader(
129    path: Union[os.PathLike, str],
130    patch_shape: Tuple[int, int],
131    batch_size: int,
132    name: Literal["PtSn_Al2O3_clusters", "Au_ZSM5"],
133    download: bool = False,
134    **kwargs
135) -> DataLoader:
136    """Get the DataLoader for semantic segmentation of catalyst particles in the EMcopilot dataset.
137
138    Args:
139        path: Filepath to a folder where the downloaded data will be saved.
140        patch_shape: The patch shape to use for training.
141        batch_size: The batch size for training.
142        name: The name of the sub-dataset. One of 'PtSn_Al2O3_clusters' or 'Au_ZSM5'.
143        download: Whether to download the data if it is not present.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
145
146    Returns:
147        The DataLoader.
148    """
149    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
150    dataset = get_emcopilot_dataset(path, patch_shape, name, download, **ds_kwargs)
151    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)