torch_em.data.datasets.light_microscopy.micronucml

  1"""The MicroNucML dataset contains annotations for micronucleus segmentation
  2in live-cell fluorescence microscopy images of MCF10A and RPE-1 cells.
  3
  4The dataset is located at https://data.mendeley.com/datasets/hrjn4dy6z9/1.
  5This dataset is from the publication https://doi.org/10.1016/j.crmeth.2026.101573.
  6Please cite it if you use this dataset for your research.
  7"""
  8
  9import os
 10import shutil
 11from glob import glob
 12from tqdm import tqdm
 13from natsort import natsorted
 14from typing import Union, Tuple, Literal, List, Optional
 15
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23URLS = {
 24    "h2b_gfp": "https://data.mendeley.com/public-files/datasets/hrjn4dy6z9/files/622ca122-7e2f-4203-81d9-ce01ad7bf590/file_downloaded",  # noqa
 25    "h2b_mcherry": "https://data.mendeley.com/public-files/datasets/hrjn4dy6z9/files/2b8206d4-aa79-4879-b638-895d4ededf3c/file_downloaded",  # noqa
 26}
 27CHECKSUMS = {
 28    "h2b_gfp": "13028b103f40c70fe8868f9c85d84bc2e3fa333dc01418914653993cad5122bd",
 29    "h2b_mcherry": "ca89f6b9fa8d9912da721e8dda6397e6675c6e353e5480f32b6658a18c36c1ad",
 30}
 31
 32# The archives are named after the fluorescent reporter, "H2B-GFP" and "H2B-mCherry" in the data record.
 33# 'train' and 'test' are the official split of the H2B-GFP images, the mCherry images are evaluation sets
 34# that capture a different colour and image quality.
 35SPLITS = {
 36    "train": ("h2b_gfp", "H2B-GFP", "train_image", "train_mask"),
 37    "test": ("h2b_gfp", "H2B-GFP", "test_image", "test_mask"),
 38    "mcherry_red": ("h2b_mcherry", "H2B-mCherry", "red_image", "red_mask"),
 39    "mcherry_grey": ("h2b_mcherry", "H2B-mCherry", "grey_image", "grey_mask"),
 40}
 41
 42
 43def _preprocess_labels(data_dir, mask_dir):
 44    import numpy as np
 45    import imageio.v3 as imageio
 46
 47    label_dir = os.path.join(data_dir, f"{mask_dir}_instances")
 48    mask_paths = natsorted(glob(os.path.join(data_dir, mask_dir, "*.npy")))
 49    if os.path.exists(label_dir) and len(glob(os.path.join(label_dir, "*.tif"))) == len(mask_paths):
 50        return label_dir
 51
 52    os.makedirs(label_dir, exist_ok=True)
 53    for mask_path in tqdm(mask_paths, desc=f"Preprocessing '{mask_dir}'"):
 54        masks = np.load(mask_path)
 55        # The masks are stored as one binary plane per micronucleus, in some files with an extra singleton axis.
 56        masks = masks.reshape(masks.shape[0], *masks.shape[-2:])
 57
 58        labels = np.zeros(masks.shape[-2:], dtype="uint16")
 59        for instance_id, mask in enumerate(masks, start=1):
 60            labels[mask > 0] = instance_id
 61
 62        fname = os.path.splitext(os.path.basename(mask_path))[0]
 63        imageio.imwrite(os.path.join(label_dir, f"{fname}.tif"), labels, compression="zlib")
 64
 65    return label_dir
 66
 67
 68def get_micronucml_data(
 69    path: Union[os.PathLike, str],
 70    split: Literal["train", "test", "mcherry_red", "mcherry_grey"],
 71    download: bool = False,
 72) -> str:
 73    """Download the MicroNucML dataset.
 74
 75    Args:
 76        path: Filepath to a folder where the downloaded data will be saved.
 77        split: The choice of data split.
 78        download: Whether to download the data if it is not present.
 79
 80    Returns:
 81        Filepath where the data is stored.
 82    """
 83    if split not in SPLITS:
 84        raise ValueError(f"'{split}' is not a valid split choice.")
 85
 86    source, dname = SPLITS[split][:2]
 87    data_dir = os.path.join(path, dname)
 88    if os.path.exists(data_dir):
 89        return data_dir
 90
 91    os.makedirs(path, exist_ok=True)
 92
 93    zip_path = os.path.join(path, f"{source}.zip")
 94    util.download_source(path=zip_path, url=URLS[source], download=download, checksum=CHECKSUMS[source])
 95    util.unzip(zip_path=zip_path, dst=path)
 96    shutil.rmtree(os.path.join(path, "__MACOSX"), ignore_errors=True)
 97
 98    return data_dir
 99
100
101def get_micronucml_paths(
102    path: Union[os.PathLike, str],
103    split: Literal["train", "test", "mcherry_red", "mcherry_grey"],
104    download: bool = False,
105) -> Tuple[List[str], List[str]]:
106    """Get paths to the MicroNucML data.
107
108    Args:
109        path: Filepath to a folder where the downloaded data will be saved.
110        split: The choice of data split.
111        download: Whether to download the data if it is not present.
112
113    Returns:
114        List of filepaths for the image data.
115        List of filepaths for the label data.
116    """
117    data_dir = get_micronucml_data(path, split, download)
118    image_dir, mask_dir = SPLITS[split][2:]
119    label_dir = _preprocess_labels(data_dir, mask_dir)
120
121    raw_paths = natsorted(glob(os.path.join(data_dir, image_dir, "*.png")))
122    label_paths = [
123        os.path.join(label_dir, f"{os.path.splitext(os.path.basename(p))[0]}.tif") for p in raw_paths
124    ]
125
126    assert raw_paths and len(raw_paths) == len(label_paths)
127    assert all(os.path.exists(p) for p in label_paths)
128
129    return raw_paths, label_paths
130
131
132def get_micronucml_dataset(
133    path: Union[os.PathLike, str],
134    patch_shape: Tuple[int, int],
135    split: Literal["train", "test", "mcherry_red", "mcherry_grey"] = "train",
136    offsets: Optional[List[List[int]]] = None,
137    boundaries: bool = False,
138    binary: bool = False,
139    download: bool = False,
140    **kwargs
141) -> Dataset:
142    """Get the MicroNucML dataset for micronucleus segmentation.
143
144    Args:
145        path: Filepath to a folder where the downloaded data will be saved.
146        patch_shape: The patch shape to use for training.
147        split: The choice of data split.
148        offsets: Offset values for affinity computation used as target.
149        boundaries: Whether to compute boundaries as the target.
150        binary: Whether to use a binary segmentation target.
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
153
154    Returns:
155        The segmentation dataset.
156    """
157    raw_paths, label_paths = get_micronucml_paths(path, split, download)
158
159    kwargs = util.ensure_transforms(ndim=2, **kwargs)
160    kwargs, _ = util.add_instance_label_transform(
161        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
162    )
163
164    return torch_em.default_segmentation_dataset(
165        raw_paths=raw_paths,
166        raw_key=None,
167        label_paths=label_paths,
168        label_key=None,
169        patch_shape=patch_shape,
170        is_seg_dataset=False,
171        with_channels=split != "mcherry_grey",
172        ndim=2,
173        **kwargs
174    )
175
176
177def get_micronucml_loader(
178    path: Union[os.PathLike, str],
179    batch_size: int,
180    patch_shape: Tuple[int, int],
181    split: Literal["train", "test", "mcherry_red", "mcherry_grey"] = "train",
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 MicroNucML dataloader for micronucleus 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 patch shape to use for training.
194        split: The choice of data split.
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 for 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_micronucml_dataset(path, patch_shape, split, offsets, boundaries, binary, download, **ds_kwargs)
206    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)