torch_em.data.datasets.electron_microscopy.probtem

  1"""ProbTEM dataset for mitochondria segmentation in 2D TEM images.
  2
  3The dataset contains TEM images of skeletal muscle with binary semantic segmentation
  4masks for mitochondria (0=background, 1=mitochondria). Images are 2560 x 2560 pixels
  5at 65 nm sample thickness acquired with a JEM-1011 microscope at 80 kV.
  6
  7The dataset has 21 training and 6 test images. There is no validation split.
  8
  9Masks are stored as grayscale PNGs (0=background, 255=mitochondria) with slight
 10anti-aliased edges. They are thresholded to binary during preprocessing.
 11One test mask (y_5) is stored at half the image resolution and is upsampled to it.
 12
 13This dataset is from the publication https://doi.org/10.1038/s41598-025-03311-1.
 14Please cite it if you use this dataset in your research.
 15
 16The data is available at https://yoonlab.unist.ac.kr/index.php/research/mitochondria-tem-dataset/
 17and requires a Google Drive download via gdown: pip install gdown.
 18Google Drive may deny the automated download of the image files. In that case, download the folder
 19https://drive.google.com/drive/folders/1n2ZqbJEHPyMB_6a6OTBBACt5Jct2PZJc in a browser and extract it in `path`,
 20so that `path/Deeppi-EM/mitoseg_deploy/datasets/Skeletal_muscle` exists.
 21"""
 22
 23import os
 24from glob import glob
 25from typing import List, Literal, Tuple, Union
 26
 27import h5py
 28import imageio.v3 as imageio
 29import numpy as np
 30
 31import torch_em
 32from torch.utils.data import Dataset, DataLoader
 33from .. import util
 34
 35
 36PROBTEM_GDRIVE_FOLDER = "1n2ZqbJEHPyMB_6a6OTBBACt5Jct2PZJc"
 37PROBTEM_DATA_ROOT = "Deeppi-EM/mitoseg_deploy/datasets/Skeletal_muscle"
 38
 39
 40def _preprocess_probtem(raw_dir, label_dir, out_dir):
 41    os.makedirs(out_dir, exist_ok=True)
 42    raw_paths = sorted(glob(os.path.join(raw_dir, "*.tif")) + glob(os.path.join(raw_dir, "*.tiff")))
 43    for rp in raw_paths:
 44        name = os.path.splitext(os.path.basename(rp))[0]
 45        out_path = os.path.join(out_dir, f"{name}.h5")
 46        if os.path.exists(out_path):
 47            continue
 48
 49        raw = imageio.imread(rp)
 50        if raw.ndim == 3:
 51            raw = raw[..., 0]
 52
 53        label_name = name.replace("x_", "y_")
 54        lp = os.path.join(label_dir, f"{label_name}.png")
 55        if not os.path.exists(lp):
 56            raise FileNotFoundError(f"The mask {lp} for {rp} is missing, the download is likely incomplete.")
 57
 58        labels = imageio.imread(lp)
 59        if labels.ndim == 3:
 60            labels = labels[..., 0]
 61        labels = (labels >= 127).astype(np.uint8)
 62
 63        # Some masks are stored at a lower resolution than their image, by an integer factor.
 64        if labels.shape != raw.shape:
 65            factors = [r / s for r, s in zip(raw.shape, labels.shape)]
 66            if factors[0] != factors[1] or not factors[0].is_integer():
 67                raise RuntimeError(f"The mask {lp} has shape {labels.shape}, which does not fit the image {raw.shape}.")
 68            labels = np.kron(labels, np.ones((int(factors[0]),) * 2, dtype=labels.dtype))
 69
 70        # The file is written under a temporary name, so that an interrupted conversion is redone.
 71        tmp_path = out_path + ".tmp"
 72        with h5py.File(tmp_path, "w") as f:
 73            f.create_dataset("raw", data=raw, compression="gzip")
 74            f.create_dataset("labels", data=labels, compression="gzip")
 75        os.replace(tmp_path, out_path)
 76
 77
 78def get_probtem_data(
 79    path: Union[os.PathLike, str],
 80    split: Literal["train", "test"] = "train",
 81    download: bool = False,
 82) -> str:
 83    """Download and preprocess the ProbTEM dataset.
 84
 85    Args:
 86        path: Filepath to a folder where the data will be saved.
 87        split: The data split to use, either "train" or "test".
 88        download: Whether to download the data if not present.
 89
 90    Returns:
 91        Path to the folder containing preprocessed HDF5 files.
 92    """
 93    processed_dir = os.path.join(str(path), "processed", split)
 94    raw_dir = os.path.join(str(path), PROBTEM_DATA_ROOT, split, "input")
 95    label_dir = os.path.join(str(path), PROBTEM_DATA_ROOT, split, "target")
 96
 97    # Without the original data there is nothing left to process, so the processed data is used as it is.
 98    if not os.path.isdir(raw_dir) and len(glob(os.path.join(processed_dir, "*.h5"))) > 0:
 99        return processed_dir
100
101    if not os.path.isdir(raw_dir):
102        if not download:
103            raise RuntimeError(
104                f"ProbTEM data not found at '{path}'. Set download=True or download manually from "
105                "https://yoonlab.unist.ac.kr/index.php/research/mitochondria-tem-dataset/ "
106                "and place in the given path."
107            )
108        try:
109            import gdown
110        except ImportError:
111            raise ImportError("gdown is required to download ProbTEM: pip install gdown")
112        gdown.download_folder(id=PROBTEM_GDRIVE_FOLDER, output=str(path), quiet=False)
113
114    _preprocess_probtem(raw_dir, label_dir, processed_dir)
115    return processed_dir
116
117
118def get_probtem_paths(
119    path: Union[os.PathLike, str],
120    split: Literal["train", "test"] = "train",
121    download: bool = False,
122) -> List[str]:
123    """Get paths to ProbTEM HDF5 files.
124
125    Args:
126        path: Filepath to a folder where the data will be saved.
127        split: The data split to use, either "train" or "test".
128        download: Whether to download the data if not present.
129
130    Returns:
131        List of paths to HDF5 files.
132    """
133    processed_dir = get_probtem_data(path, split, download)
134    return sorted(glob(os.path.join(processed_dir, "*.h5")))
135
136
137def get_probtem_dataset(
138    path: Union[os.PathLike, str],
139    patch_shape: Tuple[int, int],
140    split: Literal["train", "test"] = "train",
141    download: bool = False,
142    **kwargs,
143) -> Dataset:
144    """Get the ProbTEM dataset for mitochondria segmentation in 2D TEM images.
145
146    Args:
147        path: Filepath to a folder where the data will be saved.
148        patch_shape: The patch shape (H, W) for training.
149        split: The data split to use, either "train" or "test".
150        download: Whether to download the data if not present.
151        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
152
153    Returns:
154        The segmentation dataset.
155    """
156    assert len(patch_shape) == 2
157    data_paths = get_probtem_paths(path, split, download)
158
159    return torch_em.default_segmentation_dataset(
160        raw_paths=data_paths,
161        raw_key="raw",
162        label_paths=data_paths,
163        label_key="labels",
164        patch_shape=patch_shape,
165        is_seg_dataset=True,
166        **kwargs,
167    )
168
169
170def get_probtem_loader(
171    path: Union[os.PathLike, str],
172    batch_size: int,
173    patch_shape: Tuple[int, int],
174    split: Literal["train", "test"] = "train",
175    download: bool = False,
176    **kwargs,
177) -> DataLoader:
178    """Get the DataLoader for mitochondria segmentation in the ProbTEM dataset.
179
180    Args:
181        path: Filepath to a folder where the data will be saved.
182        batch_size: The batch size for training.
183        patch_shape: The patch shape (H, W) for training.
184        split: The data split to use, either "train" or "test".
185        download: Whether to download the data if not present.
186        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
187            or for the PyTorch DataLoader.
188
189    Returns:
190        The DataLoader.
191    """
192    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
193    ds = get_probtem_dataset(path=path, patch_shape=patch_shape, split=split, download=download, **ds_kwargs)
194    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)