torch_em.data.datasets.light_microscopy.dememseg

  1"""The DeMemSeg dataset contains annotations for prospore membrane segmentation
  2in fluorescence microscopy images of sporulating budding yeast.
  3
  4The dataset contains 2133 single cell crops of Saccharomyces cerevisiae with 7593 expert annotated
  5prospore membranes (PSMs). Each crop comes from a 2d maximum intensity projection of a 3d z-stack.
  6The original annotations overlap, because several prospore membranes can cover the same pixel. This
  7loader flattens them into one instance label image. A pixel that belongs to several membranes keeps
  8the label of the last membrane, so about 13 percent of the foreground pixels get an arbitrary label.
  9
 10The dataset is located at https://ssbd.riken.jp/repository/443/ under the CC BY 4.0 license.
 11This dataset is from the publication https://doi.org/10.1247/csf.25032.
 12Please cite it if you use this dataset in your research.
 13"""
 14
 15import os
 16from glob import glob
 17from pathlib import Path
 18from natsort import natsorted
 19from typing import List, Literal, Optional, Tuple, Union
 20
 21import numpy as np
 22import imageio.v3 as imageio
 23
 24from torch.utils.data import DataLoader, Dataset
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URL = "https://ssbd.riken.jp/data/ssbd-000443/zip/OriginalData_MMdetDataset.zip"
 32CHECKSUM = "b9772e343956358cf5e89459a28b02e6e3f0e05f403f640e8705957a0185a4c6"
 33
 34SPLITS = ("train", "val", "test")
 35
 36
 37def _create_instance_labels(data_dir: str, split: str) -> str:
 38    """Merge the per-instance masks of each crop into one instance label image."""
 39    import h5py
 40    from tqdm import tqdm
 41
 42    image_dir = os.path.join(data_dir, "images", split)
 43    mask_dir = os.path.join(data_dir, "masks", split)
 44
 45    preprocessed_dir = os.path.join(data_dir, "preprocessed", split)
 46    os.makedirs(preprocessed_dir, exist_ok=True)
 47
 48    image_paths = natsorted(glob(os.path.join(image_dir, "*.png")))
 49    for image_path in tqdm(image_paths, desc=f"Preprocess the '{split}' split"):
 50        stem = Path(image_path).stem
 51        output_path = os.path.join(preprocessed_dir, f"{stem}.h5")
 52        if os.path.exists(output_path):
 53            continue
 54
 55        mask_paths = natsorted(glob(os.path.join(mask_dir, f"{stem}_RoiRegion_*.png")))
 56        if not mask_paths:
 57            raise RuntimeError(f"Could not find any mask for the DeMemSeg image {image_path}.")
 58
 59        raw = imageio.imread(image_path)
 60        if raw.ndim == 3:
 61            raw = raw[..., 0]  # Only the first channel holds the membrane signal.
 62
 63        masks = [imageio.imread(p) > 0 for p in mask_paths]
 64        shape = masks[0].shape
 65
 66        # A crop at the border of the field is clipped, but its masks keep the full size.
 67        # The clipped crop sits in the center of the mask canvas, so pad it on both sides.
 68        if raw.shape != shape:
 69            raw = np.pad(raw, [((s - r) // 2, s - r - (s - r) // 2) for r, s in zip(raw.shape, shape)])
 70
 71        # Paint the large membranes first, so that a membrane inside a larger one keeps its label.
 72        labels = np.zeros(shape, dtype="uint16")
 73        for instance_id, mask in enumerate(sorted(masks, key=lambda mask: -mask.sum()), start=1):
 74            labels[mask] = instance_id
 75
 76        with h5py.File(output_path, "w") as f:
 77            f.create_dataset("raw", data=raw, compression="gzip")
 78            f.create_dataset("labels", data=labels, compression="gzip")
 79
 80    return preprocessed_dir
 81
 82
 83def get_dememseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 84    """Download the DeMemSeg dataset.
 85
 86    Args:
 87        path: Filepath to a folder where the downloaded data will be saved.
 88        download: Whether to download the data if it is not present.
 89
 90    Returns:
 91        The filepath to the extracted data.
 92    """
 93    data_dir = os.path.join(path, "OriginalData_MMdetDataset")
 94    if os.path.exists(data_dir):
 95        return data_dir
 96
 97    os.makedirs(path, exist_ok=True)
 98    zip_path = os.path.join(path, "OriginalData_MMdetDataset.zip")
 99    util.download_source(zip_path, URL, download, CHECKSUM)
100    util.unzip(zip_path=zip_path, dst=path)
101
102    return data_dir
103
104
105def get_dememseg_paths(
106    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False,
107) -> List[str]:
108    """Get paths to the DeMemSeg data.
109
110    Args:
111        path: Filepath to a folder where the downloaded data will be saved.
112        split: The data split. Either 'train', 'val' or 'test'.
113        download: Whether to download the data if it is not present.
114
115    Returns:
116        List of filepaths for the preprocessed h5 data.
117    """
118    if split not in SPLITS:
119        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
120
121    data_dir = get_dememseg_data(path, download)
122    preprocessed_dir = _create_instance_labels(data_dir, split)
123
124    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
125    if not volume_paths:
126        raise RuntimeError(f"Could not find any preprocessed DeMemSeg data in {preprocessed_dir}.")
127
128    return volume_paths
129
130
131def get_dememseg_dataset(
132    path: Union[os.PathLike, str],
133    patch_shape: Tuple[int, int],
134    split: Literal["train", "val", "test"],
135    offsets: Optional[List[List[int]]] = None,
136    boundaries: bool = False,
137    binary: bool = False,
138    download: bool = False,
139    **kwargs,
140) -> Dataset:
141    """Get the DeMemSeg dataset for prospore membrane segmentation.
142
143    Args:
144        path: Filepath to a folder where the downloaded data will be saved.
145        patch_shape: The 2D patch shape to use for training.
146        split: The data split. Either 'train', 'val' or 'test'.
147        offsets: Offset values for affinity computation used as target.
148        boundaries: Whether to compute boundaries as the target.
149        binary: Whether to use a binary segmentation target.
150        download: Whether to download the data if it is not present.
151        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
152
153    Returns:
154        The segmentation dataset.
155    """
156    if len(patch_shape) != 2:
157        raise ValueError(f"The DeMemSeg patch shape must be two-dimensional, got {patch_shape}.")
158
159    volume_paths = get_dememseg_paths(path, split, download)
160
161    kwargs, _ = util.add_instance_label_transform(
162        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
163    )
164    kwargs = util.ensure_transforms(ndim=2, **kwargs)
165
166    return torch_em.default_segmentation_dataset(
167        raw_paths=volume_paths,
168        raw_key="raw",
169        label_paths=volume_paths,
170        label_key="labels",
171        patch_shape=patch_shape,
172        ndim=2,
173        **kwargs,
174    )
175
176
177def get_dememseg_loader(
178    path: Union[os.PathLike, str],
179    batch_size: int,
180    patch_shape: Tuple[int, int],
181    split: Literal["train", "val", "test"],
182    offsets: Optional[List[List[int]]] = None,
183    boundaries: bool = False,
184    binary: bool = False,
185    download: bool = False,
186    **kwargs,
187) -> DataLoader:
188    """Get the DeMemSeg dataloader for prospore membrane segmentation.
189
190    Args:
191        path: Filepath to a folder where the downloaded data will be saved.
192        batch_size: The batch size for training.
193        patch_shape: The 2D patch shape to use for training.
194        split: The data split. Either 'train', 'val' or 'test'.
195        offsets: Offset values for affinity computation used as target.
196        boundaries: Whether to compute boundaries as the target.
197        binary: Whether to use a binary segmentation target.
198        download: Whether to download the data if it is not present.
199        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
200
201    Returns:
202        The DataLoader.
203    """
204    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
205    dataset = get_dememseg_dataset(
206        path=path,
207        patch_shape=patch_shape,
208        split=split,
209        offsets=offsets,
210        boundaries=boundaries,
211        binary=binary,
212        download=download,
213        **ds_kwargs,
214    )
215    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)