torch_em.data.datasets.light_microscopy.oocyteseg

  1"""The OocyteSeg dataset contains annotations for binary membrane segmentation
  2in transmitted light microscopy images of oocytes from multiple species.
  3
  4NOTE: The dataset only has semantic (binary) segmentation.
  5
  6The dataset is from the publication https://doi.org/10.1242/jcs.260281.
  7Please cite it if you use this dataset in your research.
  8"""
  9
 10import os
 11from glob import glob
 12from typing import Union, Literal, Optional, Tuple, List
 13
 14import numpy as np
 15import imageio.v3 as imageio
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URL = "https://zenodo.org/records/6502830/files/SegmentationCortex.tar.gz"
 25# NOTE: This is the sha256 of the archive, which is what `util.download_source` verifies.
 26# Zenodo publishes an md5 (1da5d4fd102d8e903744db424f6114c6), which was stored here previously
 27# and always failed the check, leaving the completed download as a rejected .incomplete file.
 28CHECKSUM = "46852468b41c4e19c28f079de9e7f6ae565e45c8dc89273adb6f3bd7e9b50613"
 29
 30SPECIES = ["mouse", "human", "sea_urchin"]
 31
 32_SUBDIRS = {
 33    "mouse": {
 34        "train": ["exp1", "exp2"],
 35        "test": ["exp1_test", "exp2_test"],
 36    },
 37    "human": {
 38        "train": ["clin1", "clin2"],
 39        "test": ["clin1_test", "clin2_test"],
 40    },
 41    "sea_urchin": {
 42        "train": ["train"],
 43        "test": ["test"],
 44    },
 45}
 46
 47
 48def _preprocess_data(data_dir, processed_dir, species, split):
 49    """Preprocess images and masks to ensure consistent format.
 50
 51    Some sea urchin images are stored as RGB instead of grayscale.
 52    Masks are stored as 0/255 and need to be normalized to 0/1.
 53    This function converts all data to a consistent single-channel uint8 format.
 54    """
 55    img_out_dir = os.path.join(processed_dir, "images")
 56    mask_out_dir = os.path.join(processed_dir, "masks")
 57    os.makedirs(img_out_dir, exist_ok=True)
 58    os.makedirs(mask_out_dir, exist_ok=True)
 59
 60    subdirs = _SUBDIRS[species][split]
 61
 62    for subdir in subdirs:
 63        input_dir = os.path.join(data_dir, species, subdir, "input")
 64        mask_dir = os.path.join(data_dir, species, subdir, "mask")
 65
 66        input_names = {os.path.splitext(f)[0] for f in os.listdir(input_dir) if f.endswith(".png")}
 67        mask_names = {os.path.splitext(f)[0] for f in os.listdir(mask_dir) if f.endswith(".png")}
 68        matched = sorted(input_names & mask_names)
 69
 70        for name in matched:
 71            img_out = os.path.join(img_out_dir, f"{subdir}_{name}.tif")
 72            mask_out = os.path.join(mask_out_dir, f"{subdir}_{name}.tif")
 73
 74            if os.path.exists(img_out) and os.path.exists(mask_out):
 75                continue
 76
 77            img = imageio.imread(os.path.join(input_dir, f"{name}.png"))
 78            if img.ndim == 3:
 79                img = np.mean(img[..., :3], axis=-1).astype("uint8")
 80            imageio.imwrite(img_out, img, compression="zlib")
 81
 82            mask = imageio.imread(os.path.join(mask_dir, f"{name}.png"))
 83            if mask.ndim == 3:
 84                mask = mask[..., 0]
 85            mask = (mask > 0).astype("uint8")
 86            imageio.imwrite(mask_out, mask, compression="zlib")
 87
 88
 89def get_oocyteseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 90    """Download the OocyteSeg dataset.
 91
 92    Args:
 93        path: Filepath to a folder where the downloaded data will be saved.
 94        download: Whether to download the data if it is not present.
 95
 96    Returns:
 97        The filepath to the extracted data directory.
 98    """
 99    data_dir = os.path.join(path, "SegmentationCortex")
100    if os.path.exists(data_dir):
101        return data_dir
102
103    os.makedirs(path, exist_ok=True)
104    tar_path = os.path.join(path, "SegmentationCortex.tar.gz")
105    util.download_source(path=tar_path, url=URL, download=download, checksum=CHECKSUM)
106    util.unzip_tarfile(tar_path=tar_path, dst=path)
107
108    return data_dir
109
110
111def get_oocyteseg_paths(
112    path: Union[os.PathLike, str],
113    split: Literal["train", "test"] = "train",
114    species: Optional[str] = None,
115    download: bool = False,
116) -> Tuple[List[str], List[str]]:
117    """Get paths to the OocyteSeg data.
118
119    Args:
120        path: Filepath to a folder where the downloaded data will be saved.
121        split: The data split to use. One of 'train' or 'test'.
122        species: The species to select. One of 'mouse', 'human' or 'sea_urchin'.
123            If None, data from all species is returned.
124        download: Whether to download the data if it is not present.
125
126    Returns:
127        List of filepaths for the image data.
128        List of filepaths for the label data.
129    """
130    assert split in ("train", "test"), f"'{split}' is not a valid split. Choose from 'train' or 'test'."
131
132    if species is None:
133        species_list = SPECIES
134    else:
135        assert species in SPECIES, f"'{species}' is not a valid species. Choose from {SPECIES}."
136        species_list = [species]
137
138    data_dir = get_oocyteseg_data(path, download)
139
140    all_image_paths = []
141    all_seg_paths = []
142
143    from natsort import natsorted
144
145    for sp in species_list:
146        processed_dir = os.path.join(path, "processed", sp, split)
147        img_out_dir = os.path.join(processed_dir, "images")
148        mask_out_dir = os.path.join(processed_dir, "masks")
149
150        if not os.path.exists(img_out_dir) or len(glob(os.path.join(img_out_dir, "*.tif"))) == 0:
151            _preprocess_data(data_dir, processed_dir, sp, split)
152
153        image_paths = natsorted(glob(os.path.join(img_out_dir, "*.tif")))
154        seg_paths = natsorted(glob(os.path.join(mask_out_dir, "*.tif")))
155
156        assert len(image_paths) == len(seg_paths), \
157            f"Mismatch: {len(image_paths)} images vs {len(seg_paths)} masks for {sp}/{split}"
158        assert len(image_paths) > 0, f"No images found for {sp}/{split}"
159
160        all_image_paths.extend(image_paths)
161        all_seg_paths.extend(seg_paths)
162
163    return all_image_paths, all_seg_paths
164
165
166def get_oocyteseg_dataset(
167    path: Union[os.PathLike, str],
168    patch_shape: Tuple[int, int],
169    split: Literal["train", "test"] = "train",
170    species: Optional[str] = None,
171    download: bool = False,
172    **kwargs
173) -> Dataset:
174    """Get the OocyteSeg dataset for binary membrane segmentation.
175
176    Args:
177        path: Filepath to a folder where the downloaded data will be saved.
178        patch_shape: The patch shape to use for training.
179        split: The data split to use. One of 'train' or 'test'.
180        species: The species to select. One of 'mouse', 'human' or 'sea_urchin'.
181            If None, data from all species is returned.
182        download: Whether to download the data if it is not present.
183        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
184
185    Returns:
186        The segmentation dataset.
187    """
188    image_paths, seg_paths = get_oocyteseg_paths(path, split, species, download)
189
190    kwargs = util.ensure_transforms(ndim=2, **kwargs)
191
192    return torch_em.default_segmentation_dataset(
193        raw_paths=image_paths,
194        raw_key=None,
195        label_paths=seg_paths,
196        label_key=None,
197        patch_shape=patch_shape,
198        is_seg_dataset=False,
199        ndim=2,
200        **kwargs
201    )
202
203
204def get_oocyteseg_loader(
205    path: Union[os.PathLike, str],
206    batch_size: int,
207    patch_shape: Tuple[int, int],
208    split: Literal["train", "test"] = "train",
209    species: Optional[str] = None,
210    download: bool = False,
211    **kwargs
212) -> DataLoader:
213    """Get the OocyteSeg dataloader for binary 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        species: The species to select. One of 'mouse', 'human' or 'sea_urchin'.
221            If None, data from all species is returned.
222        download: Whether to download the data if it is not present.
223        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
224
225    Returns:
226        The DataLoader.
227    """
228    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
229    dataset = get_oocyteseg_dataset(
230        path=path,
231        patch_shape=patch_shape,
232        split=split,
233        species=species,
234        download=download,
235        **ds_kwargs,
236    )
237    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)