torch_em.data.datasets.light_microscopy.medussa

  1"""The MeDuSSA dataset contains annotations for bacterial membrane
  2instance segmentation in fluorescence microscopy images stained with FM 4-64.
  3
  4The dataset provides 143 training images and 16 benchmarking images of
  5membrane-stained bacteria (primarily Bacillus subtilis PY79) with corresponding
  6instance segmentation masks annotated using JFilament in FIJI.
  7
  8The dataset is located at https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD2350.
  9This dataset is from the publication https://doi.org/10.1101/2025.10.26.684635.
 10Please cite it if you use this dataset in your research.
 11"""
 12
 13import os
 14import json
 15import warnings
 16from glob import glob
 17from typing import Union, Tuple, List, Literal
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26BASE_URL = "https://www.ebi.ac.uk/biostudies/files/S-BIAD2350"
 27
 28SPLIT_FILE_LISTS = {
 29    "train": {
 30        "images": "submission_segmentation_training_images_raw.json",
 31        "masks": "submission_segmentation_training_masks.json",
 32    },
 33    "test": {
 34        "images": "submission_segmentation_benchmarking_images_raw.json",
 35        "masks": "submission_segmentation_benchmarking_masks.json",
 36    },
 37}
 38
 39
 40def _download_file_lists(path, split):
 41    """Download and parse JSON file lists from BioStudies to get relative file paths."""
 42    file_list_dir = os.path.join(path, "file_lists")
 43    os.makedirs(file_list_dir, exist_ok=True)
 44
 45    result = {}
 46    for key in ("images", "masks"):
 47        json_fname = SPLIT_FILE_LISTS[split][key]
 48        json_path = os.path.join(file_list_dir, json_fname)
 49
 50        if not os.path.exists(json_path):
 51            url = f"{BASE_URL}/{json_fname}"
 52            util.download_source(path=json_path, url=url, download=True, checksum=None)
 53
 54        with open(json_path) as f:
 55            data = json.load(f)
 56
 57        result[key] = sorted([entry["path"] for entry in data])
 58
 59    return result["images"], result["masks"]
 60
 61
 62def _create_h5_data(path, split, image_paths_rel, mask_paths_rel):
 63    """Create h5 files with raw images and instance labels."""
 64    import h5py
 65    import imageio.v3 as imageio
 66    from tqdm import tqdm
 67
 68    h5_dir = os.path.join(path, "h5_data", split)
 69    os.makedirs(h5_dir, exist_ok=True)
 70
 71    assert len(image_paths_rel) == len(mask_paths_rel), \
 72        f"Mismatch: {len(image_paths_rel)} images vs {len(mask_paths_rel)} masks for split '{split}'"
 73
 74    for img_rel, mask_rel in tqdm(
 75        zip(image_paths_rel, mask_paths_rel),
 76        total=len(image_paths_rel),
 77        desc=f"Creating h5 files for '{split}'"
 78    ):
 79        fname = os.path.splitext(os.path.basename(img_rel))[0]
 80        h5_path = os.path.join(h5_dir, f"{fname}.h5")
 81
 82        if os.path.exists(h5_path):
 83            continue
 84
 85        raw = imageio.imread(os.path.join(path, img_rel))
 86        labels = imageio.imread(os.path.join(path, mask_rel))
 87
 88        # Handle potential multi-dimensional images (e.g. Z-stacks not fully max-projected).
 89        if raw.ndim > 2:
 90            raw = raw.max(axis=0)
 91
 92        if labels.ndim > 2:
 93            labels = labels.max(axis=0)
 94
 95        # Five benchmark masks belong to restored images that the study does not ship, so they cannot be paired.
 96        if raw.shape != labels.shape:
 97            warnings.warn(f"Skipping '{fname}': image shape {raw.shape} does not match mask shape {labels.shape}.")
 98            continue
 99
100        with h5py.File(h5_path, "w") as f:
101            f.create_dataset("raw", data=raw, compression="gzip")
102            f.create_dataset("labels", data=labels.astype("int64"), compression="gzip")
103
104    return h5_dir
105
106
107def get_medussa_data(
108    path: Union[os.PathLike, str],
109    split: Literal["train", "test"] = "train",
110    download: bool = False,
111) -> str:
112    """Download the MeDuSSA dataset.
113
114    Args:
115        path: Filepath to a folder where the downloaded data will be saved.
116        split: The data split to use. One of 'train' or 'test'.
117        download: Whether to download the data if it is not present.
118
119    Returns:
120        The filepath to the directory with the downloaded data.
121    """
122    assert split in ("train", "test"), f"'{split}' is not a valid split."
123
124    image_paths_rel, mask_paths_rel = _download_file_lists(path, split)
125
126    for rel_path in image_paths_rel + mask_paths_rel:
127        local_path = os.path.join(path, rel_path)
128        if os.path.exists(local_path):
129            continue
130
131        os.makedirs(os.path.dirname(local_path), exist_ok=True)
132        url = f"{BASE_URL}/{rel_path}"
133        util.download_source(path=local_path, url=url, download=download, checksum=None)
134
135    return path
136
137
138def get_medussa_paths(
139    path: Union[os.PathLike, str],
140    split: Literal["train", "test"] = "train",
141    download: bool = False,
142) -> List[str]:
143    """Get paths to the MeDuSSA data.
144
145    Args:
146        path: Filepath to a folder where the downloaded data will be saved.
147        split: The data split to use. One of 'train' or 'test'.
148        download: Whether to download the data if it is not present.
149
150    Returns:
151        List of filepaths for the h5 data.
152    """
153    from natsort import natsorted
154
155    get_medussa_data(path, split, download)
156
157    h5_dir = os.path.join(path, "h5_data", split)
158    if not os.path.exists(h5_dir) or len(glob(os.path.join(h5_dir, "*.h5"))) == 0:
159        image_paths_rel, mask_paths_rel = _download_file_lists(path, split)
160        _create_h5_data(path, split, image_paths_rel, mask_paths_rel)
161
162    h5_paths = natsorted(glob(os.path.join(h5_dir, "*.h5")))
163    assert len(h5_paths) > 0, f"No data found for split '{split}'"
164
165    return h5_paths
166
167
168def get_medussa_dataset(
169    path: Union[os.PathLike, str],
170    patch_shape: Tuple[int, int],
171    split: Literal["train", "test"] = "train",
172    download: bool = False,
173    **kwargs
174) -> Dataset:
175    """Get the MeDuSSA dataset for bacterial membrane segmentation.
176
177    Args:
178        path: Filepath to a folder where the downloaded data will be saved.
179        patch_shape: The patch shape to use for training.
180        split: The data split to use. One of 'train' or 'test'.
181        download: Whether to download the data if it is not present.
182        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
183
184    Returns:
185        The segmentation dataset.
186    """
187    h5_paths = get_medussa_paths(path, split, download)
188
189    kwargs, _ = util.add_instance_label_transform(
190        kwargs, add_binary_target=True,
191    )
192    kwargs = util.ensure_transforms(ndim=2, **kwargs)
193
194    return torch_em.default_segmentation_dataset(
195        raw_paths=h5_paths,
196        raw_key="raw",
197        label_paths=h5_paths,
198        label_key="labels",
199        patch_shape=patch_shape,
200        ndim=2,
201        **kwargs
202    )
203
204
205def get_medussa_loader(
206    path: Union[os.PathLike, str],
207    batch_size: int,
208    patch_shape: Tuple[int, int],
209    split: Literal["train", "test"] = "train",
210    download: bool = False,
211    **kwargs
212) -> DataLoader:
213    """Get the MeDuSSA dataloader for bacterial membrane segmentation.
214
215    Args:
216        path: Filepath to a folder where the downloaded data will be saved.
217        batch_size: The batch size for training.
218        patch_shape: The patch shape to use for training.
219        split: The data split to use. One of 'train' or 'test'.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
222
223    Returns:
224        The DataLoader.
225    """
226    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
227    dataset = get_medussa_dataset(
228        path=path,
229        patch_shape=patch_shape,
230        split=split,
231        download=download,
232        **ds_kwargs,
233    )
234    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)