torch_em.data.datasets.light_microscopy.synthmt

  1"""The SynthMT dataset contains interference reflection microscopy (IRM) images of microtubules
  2with segmentation annotations. It is available from two sources, selected with the 'source' argument:
  3- 'synthetic_irm': 6,600 synthetically generated 512x512 RGB images with per-instance binary masks,
  4  designed to train foundation models (e.g. SAM) for automated in vitro microtubule analysis.
  5- 'real_irm': the IRM_InVitroMT dataset with 237 real images (206 train and 31 test) with manually annotated
  6  binary microtubule masks, acquired with four microscopes.
  7
  8The synthetic data is located at https://huggingface.co/datasets/HTW-KI-Werkstatt/SynthMT.
  9This dataset is from the publication https://doi.org/10.64898/2026.01.09.698597.
 10The real data is located at https://doi.org/10.5281/zenodo.20597920 and is distributed under a CC BY 4.0 license.
 11Please cite the publication if you use this dataset in your research.
 12"""
 13
 14import os
 15from glob import glob
 16from typing import Union, Tuple, List, Literal, Optional
 17
 18import numpy as np
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://huggingface.co/datasets/HTW-KI-Werkstatt/SynthMT/resolve/main/data/{FILENAME}"
 28NUM_PARQUET_FILES = 8
 29
 30REAL_URLS = {
 31    "train": "https://zenodo.org/records/20597920/files/Training_dataset.zip",
 32    "test": "https://zenodo.org/records/20597920/files/Test_dataset.zip",
 33}
 34REAL_CHECKSUMS = {
 35    "train": "99ce37c0ed886f905b8cb3dc26af060ec71277db50a1dbba37dbed9839e09e37",
 36    "test": "bd11f99abcdebaf3784e18c33dcfd75e6386416cfd33b1f1567ed8c4a771d711",
 37}
 38REAL_DIRS = {
 39    "train": ("Training_dataset", "Training_source", "Training_target"),
 40    "test": ("Test_dataset", "Prediction_source", "Prediction_target"),
 41}
 42
 43
 44def _download_parquets(path, download):
 45    """Download all parquet files for the dataset."""
 46    parquet_dir = os.path.join(path, "parquets")
 47    os.makedirs(parquet_dir, exist_ok=True)
 48
 49    for i in range(NUM_PARQUET_FILES):
 50        fname = f"train-{i:05d}-of-{NUM_PARQUET_FILES:05d}.parquet"
 51        fpath = os.path.join(parquet_dir, fname)
 52        if not os.path.exists(fpath):
 53            url = URL.format(FILENAME=fname)
 54            util.download_source(path=fpath, url=url, download=download, checksum=None)
 55
 56    return parquet_dir
 57
 58
 59def _create_images_from_parquets(path):
 60    """Extract images and instance labels from parquet files and save as TIF."""
 61    import imageio.v3 as imageio
 62    import pandas as pd
 63    from io import BytesIO
 64    from PIL import Image
 65    from tqdm import tqdm
 66
 67    image_dir = os.path.join(path, "images")
 68    label_dir = os.path.join(path, "labels")
 69    os.makedirs(image_dir, exist_ok=True)
 70    os.makedirs(label_dir, exist_ok=True)
 71
 72    parquet_dir = os.path.join(path, "parquets")
 73    parquet_files = sorted(glob(os.path.join(parquet_dir, "*.parquet")))
 74
 75    for pfile in tqdm(parquet_files, desc="Processing parquet files"):
 76        df = pd.read_parquet(pfile)
 77        for _, row in df.iterrows():
 78            sample_id = row["id"]
 79            img_path = os.path.join(image_dir, f"{sample_id}.tif")
 80            lbl_path = os.path.join(label_dir, f"{sample_id}.tif")
 81
 82            if os.path.exists(img_path) and os.path.exists(lbl_path):
 83                continue
 84
 85            # Decode the image.
 86            img = Image.open(BytesIO(row["image"]["bytes"])).convert("RGB")
 87            img_arr = np.array(img)
 88
 89            # Decode instance masks and merge into a single label map.
 90            masks = row["mask"]
 91            instances = np.zeros(img_arr.shape[:2], dtype="uint32")
 92            for i, mask_entry in enumerate(masks, start=1):
 93                mask = np.array(Image.open(BytesIO(mask_entry["bytes"])).convert("L"))
 94                instances[mask > 0] = i
 95
 96            imageio.imwrite(img_path, img_arr, compression="zlib")
 97            imageio.imwrite(lbl_path, instances, compression="zlib")
 98
 99
100def _get_real_data(path, split, download):
101    data_dir, _, _ = REAL_DIRS[split]
102    real_dir = os.path.join(path, "real_irm")
103    if os.path.exists(os.path.join(real_dir, data_dir)):
104        return real_dir
105
106    os.makedirs(real_dir, exist_ok=True)
107    zip_path = os.path.join(real_dir, f"{data_dir}.zip")
108    util.download_source(path=zip_path, url=REAL_URLS[split], download=download, checksum=REAL_CHECKSUMS[split])
109    util.unzip(zip_path=zip_path, dst=real_dir)
110    return real_dir
111
112
113def get_synthmt_data(
114    path: Union[os.PathLike, str],
115    source: Literal["synthetic_irm", "real_irm"],
116    split: Optional[Literal["train", "test"]] = None,
117    download: bool = False,
118) -> str:
119    """Download the SynthMT dataset.
120
121    Args:
122        path: Filepath to a folder where the downloaded data will be saved.
123        source: The data source. Either 'synthetic_irm' or 'real_irm'.
124        split: The data split of the real source. Either 'train' or 'test'. Must be None for the synthetic source.
125        download: Whether to download the data if it is not present.
126
127    Returns:
128        The filepath to the directory with the data.
129    """
130    assert source in ("synthetic_irm", "real_irm"), f"'{source}' is not a valid data source."
131
132    if source == "real_irm":
133        assert split in REAL_URLS, f"'{split}' is not a valid data split of the real source."
134        return _get_real_data(path, split, download)
135
136    assert split is None, "The synthetic source has no data splits."
137    _download_parquets(path, download)
138
139    image_dir = os.path.join(path, "images")
140    label_dir = os.path.join(path, "labels")
141    if not os.path.exists(image_dir) or not os.path.exists(label_dir):
142        _create_images_from_parquets(path)
143
144    return path
145
146
147def get_synthmt_paths(
148    path: Union[os.PathLike, str],
149    source: Literal["synthetic_irm", "real_irm"],
150    split: Optional[Literal["train", "test"]] = None,
151    download: bool = False,
152) -> Tuple[List[str], List[str]]:
153    """Get paths to the SynthMT data.
154
155    Args:
156        path: Filepath to a folder where the downloaded data will be saved.
157        source: The data source. Either 'synthetic_irm' or 'real_irm'.
158        split: The data split of the real source. Either 'train' or 'test'. Must be None for the synthetic source.
159        download: Whether to download the data if it is not present.
160
161    Returns:
162        List of filepaths for the image data.
163        List of filepaths for the label data.
164    """
165    from natsort import natsorted
166
167    data_path = get_synthmt_data(path, source, split, download)
168
169    if source == "real_irm":
170        data_dir, image_dir, label_dir = REAL_DIRS[split]
171        image_paths = natsorted(glob(os.path.join(data_path, data_dir, image_dir, "*.tif")))
172        label_paths = natsorted(glob(os.path.join(data_path, data_dir, label_dir, "*.tif")))
173        assert [os.path.basename(p) for p in image_paths] == [os.path.basename(p) for p in label_paths]
174    else:
175        image_paths = natsorted(glob(os.path.join(data_path, "images", "*.tif")))
176        label_paths = natsorted(glob(os.path.join(data_path, "labels", "*.tif")))
177
178    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
179
180    return image_paths, label_paths
181
182
183def get_synthmt_dataset(
184    path: Union[os.PathLike, str],
185    patch_shape: Tuple[int, int],
186    source: Literal["synthetic_irm", "real_irm"],
187    split: Optional[Literal["train", "test"]] = None,
188    download: bool = False,
189    **kwargs,
190) -> Dataset:
191    """Get the SynthMT dataset for microtubule segmentation.
192
193    The synthetic source provides instance labels, the real source provides binary labels.
194
195    Args:
196        path: Filepath to a folder where the downloaded data will be saved.
197        patch_shape: The patch shape to use for training.
198        source: The data source. Either 'synthetic_irm' or 'real_irm'.
199        split: The data split of the real source. Either 'train' or 'test'. Must be None for the synthetic source.
200        download: Whether to download the data if it is not present.
201        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
202
203    Returns:
204        The segmentation dataset.
205    """
206    image_paths, label_paths = get_synthmt_paths(path, source, split, download)
207
208    if source == "real_irm":
209        kwargs = util.update_kwargs(kwargs, "label_transform", torch_em.transform.label.labels_to_binary)
210    else:
211        kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
212    kwargs = util.update_kwargs(kwargs, "ndim", 2)
213
214    return torch_em.default_segmentation_dataset(
215        raw_paths=image_paths,
216        raw_key=None,
217        label_paths=label_paths,
218        label_key=None,
219        patch_shape=patch_shape,
220        is_seg_dataset=False,
221        **kwargs,
222    )
223
224
225def get_synthmt_loader(
226    path: Union[os.PathLike, str],
227    batch_size: int,
228    patch_shape: Tuple[int, int],
229    source: Literal["synthetic_irm", "real_irm"],
230    split: Optional[Literal["train", "test"]] = None,
231    download: bool = False,
232    **kwargs,
233) -> DataLoader:
234    """Get the SynthMT dataloader for microtubule segmentation.
235
236    Args:
237        path: Filepath to a folder where the downloaded data will be saved.
238        batch_size: The batch size for training.
239        patch_shape: The patch shape to use for training.
240        source: The data source. Either 'synthetic_irm' or 'real_irm'.
241        split: The data split of the real source. Either 'train' or 'test'. Must be None for the synthetic source.
242        download: Whether to download the data if it is not present.
243        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
244
245    Returns:
246        The DataLoader.
247    """
248    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
249    dataset = get_synthmt_dataset(path, patch_shape, source, split, download, **ds_kwargs)
250    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)