torch_em.data.datasets.electron_microscopy.ai4mat

  1"""The AI4Mat dataset contains microscopy images of materials with semantic mask annotations.
  2
  3It is a curated collection of seven datasets with predefined train, val and test splits:
  4- Aachen-Heerlen: bainitic steel microstructures (1,705 images).
  5- EBC: environmental barrier coatings (24 images).
  6- EMPS: electron microscopy nanoparticles (465 images).
  7- Grain: grain structures in stainless steel 316L (816 images).
  8- MetalDAM: multi-phase metallographic images of advanced high-strength steel (42 images).
  9- Super: Ni-based superalloy microstructures (23 images).
 10- UHCS: ultrahigh carbon steel micrographs (550 images).
 11
 12The dataset is located at https://doi.org/10.5281/zenodo.22039777 (version 1.0) and is distributed under the
 13CC BY 4.0 license. The datasets were curated from previously published sources, see the Zenodo record for the
 14references of the respective source datasets. Please cite the Zenodo record and the source datasets
 15if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, List, Literal
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30BASE_URL = "https://zenodo.org/records/22039777/files"
 31CHECKSUMS = {
 32    "Aachen-Heerlen": "24ab2a412e818492a6b9b0fafd86d6bd5dc59e98e422ca12d122558c4f700f8c",
 33    "EBC": "b870428a07bcd242f6db221b9635ec2fd9f70a4ddb8f566f97b67f2d2773dcd0",
 34    "EMPS": "1e9c4dca11ae5eda59ac63f65e60be12ffbe08837f746b0b49633d16c12e795d",
 35    "Grain": "30ca8e80392519bb790d6093b5deb86df4d4aedf8cbe9beedbf329106d54d9d0",
 36    "MetalDAM": "df0da3dc5690764b6f45feacb3862dfd9a0433e82476c96fbfff6f7725957fa4",
 37    "Super": "a77726051e1f6c4dc68d0e44193f3ff268a08a767206d0b1eaae066cf101a993",
 38    "UHCS": "48790928559dea9238f543f6c34912925f6753ffcbb93bfa0cbbd07219001dcc",
 39}
 40URLS = {name: f"{BASE_URL}/{name}.zip?download=1" for name in CHECKSUMS}
 41SPLITS = ("train", "val", "test")
 42
 43
 44def get_ai4mat_data(
 45    path: Union[os.PathLike, str],
 46    name: Literal[tuple(CHECKSUMS)],
 47    download: bool = False,
 48) -> str:
 49    """Download the AI4Mat dataset.
 50
 51    NOTE: The 'Aachen-Heerlen' archive is 1.2 GB, the other archives are between 7 and 210 MB.
 52
 53    Args:
 54        path: Filepath to a folder where the downloaded data will be saved.
 55        name: The name of the dataset. One of 'Aachen-Heerlen', 'EBC', 'EMPS', 'Grain', 'MetalDAM', 'Super', 'UHCS'.
 56        download: Whether to download the data if it is not present.
 57
 58    Returns:
 59        The filepath to the folder with the data of the selected dataset.
 60    """
 61    assert name in CHECKSUMS, f"'{name}' is not a valid dataset name. Choose from {list(CHECKSUMS)}."
 62
 63    data_dir = os.path.join(path, name)
 64    if os.path.exists(data_dir):
 65        return data_dir
 66
 67    os.makedirs(path, exist_ok=True)
 68    zip_path = os.path.join(path, f"{name}.zip")
 69    util.download_source(path=zip_path, url=URLS[name], download=download, checksum=CHECKSUMS[name])
 70    util.unzip(zip_path=zip_path, dst=path)
 71    return data_dir
 72
 73
 74def get_ai4mat_paths(
 75    path: Union[os.PathLike, str],
 76    name: Literal[tuple(CHECKSUMS)],
 77    split: Literal["train", "val", "test"],
 78    download: bool = False,
 79) -> Tuple[List[str], List[str]]:
 80    """Get paths to the AI4Mat data.
 81
 82    Args:
 83        path: Filepath to a folder where the downloaded data will be saved.
 84        name: The name of the dataset. One of 'Aachen-Heerlen', 'EBC', 'EMPS', 'Grain', 'MetalDAM', 'Super', 'UHCS'.
 85        split: The data split to use. Either 'train', 'val' or 'test'.
 86        download: Whether to download the data if it is not present.
 87
 88    Returns:
 89        List of filepaths for the image data.
 90        List of filepaths for the label data.
 91    """
 92    assert split in SPLITS, f"'{split}' is not a valid data split. Choose from {SPLITS}."
 93
 94    data_dir = get_ai4mat_data(path, name, download)
 95    image_paths = natsorted(glob(os.path.join(data_dir, split, "images", "*")))
 96    label_paths = natsorted(glob(os.path.join(data_dir, split, "masks", "*")))
 97    assert len(image_paths) > 0 and len(image_paths) == len(label_paths)
 98    return image_paths, label_paths
 99
100
101def get_ai4mat_dataset(
102    path: Union[os.PathLike, str],
103    patch_shape: Tuple[int, int],
104    name: Literal[tuple(CHECKSUMS)],
105    split: Literal["train", "val", "test"],
106    download: bool = False,
107    **kwargs
108) -> Dataset:
109    """Get the AI4Mat dataset for semantic segmentation of materials microscopy images.
110
111    The labels contain the semantic class ids, where 0 is the background.
112    The number of classes is 1 for Aachen-Heerlen, EMPS, Grain; 2 for EBC, Super; 4 for MetalDAM; 6 for UHCS.
113
114    Args:
115        path: Filepath to a folder where the downloaded data will be saved.
116        patch_shape: The patch shape to use for training.
117        name: The name of the dataset. One of 'Aachen-Heerlen', 'EBC', 'EMPS', 'Grain', 'MetalDAM', 'Super', 'UHCS'.
118        split: The data split to use. Either 'train', 'val' or 'test'.
119        download: Whether to download the data if it is not present.
120        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
121
122    Returns:
123        The segmentation dataset.
124    """
125    image_paths, label_paths = get_ai4mat_paths(path, name, split, download)
126
127    kwargs = util.update_kwargs(kwargs, "ndim", 2)
128
129    return torch_em.default_segmentation_dataset(
130        raw_paths=image_paths,
131        raw_key=None,
132        label_paths=label_paths,
133        label_key=None,
134        patch_shape=patch_shape,
135        is_seg_dataset=False,
136        **kwargs
137    )
138
139
140def get_ai4mat_loader(
141    path: Union[os.PathLike, str],
142    patch_shape: Tuple[int, int],
143    batch_size: int,
144    name: Literal[tuple(CHECKSUMS)],
145    split: Literal["train", "val", "test"],
146    download: bool = False,
147    **kwargs
148) -> DataLoader:
149    """Get the DataLoader for semantic segmentation of materials microscopy images in the AI4Mat dataset.
150
151    Args:
152        path: Filepath to a folder where the downloaded data will be saved.
153        patch_shape: The patch shape to use for training.
154        batch_size: The batch size for training.
155        name: The name of the dataset. One of 'Aachen-Heerlen', 'EBC', 'EMPS', 'Grain', 'MetalDAM', 'Super', 'UHCS'.
156        split: The data split to use. Either 'train', 'val' or 'test'.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
159
160    Returns:
161        The DataLoader.
162    """
163    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
164    dataset = get_ai4mat_dataset(path, patch_shape, name, split, download, **ds_kwargs)
165    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)