torch_em.data.datasets.light_microscopy.revvity25

  1"""Revvity-25 contains brightfield microscopy images of cancer cells with instance
  2segmentation annotations for cell cytoplasm, including detailed overlap and border
  3annotations for adjacent and overlapping cells.
  4
  5The dataset is located at https://huggingface.co/datasets/YaroslavPrytula/Revvity-25
  6and is licensed under CC BY-NC 4.0. This dataset is from the publication
  7https://arxiv.org/abs/2508.01928.
  8Please cite it if you use this dataset in your research.
  9"""
 10
 11import os
 12from glob import glob
 13from natsort import natsorted
 14from typing import List, Literal, Optional, Tuple, Union
 15
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23HF_REPO = "YaroslavPrytula/Revvity-25"
 24
 25SPLITS = {"train": "train.json", "val": "valid.json"}
 26
 27
 28def _create_segmentations_from_coco_annotations(path, split):
 29    """Convert COCO polygon annotations to instance segmentation masks."""
 30    import numpy as np
 31    import imageio.v3 as imageio
 32    from tqdm import tqdm
 33
 34    try:
 35        from pycocotools.coco import COCO
 36    except ImportError:
 37        raise ImportError(
 38            "'pycocotools' is required for processing the Revvity-25 ground-truth. "
 39            "Install it with 'conda install -c conda-forge pycocotools'."
 40        )
 41
 42    label_dir = os.path.join(path, "labels", split)
 43    if os.path.exists(label_dir):
 44        label_paths = natsorted(glob(os.path.join(label_dir, "*.tif")))
 45        if len(label_paths) > 0:
 46            image_paths = [
 47                os.path.join(path, "images", f"{os.path.splitext(os.path.basename(p))[0]}.png")
 48                for p in label_paths
 49            ]
 50            return natsorted(image_paths), label_paths
 51
 52    os.makedirs(label_dir, exist_ok=True)
 53
 54    ann_file = os.path.join(path, "annotations", SPLITS[split])
 55    coco = COCO(ann_file)
 56
 57    image_paths, label_paths = [], []
 58    for image_id in tqdm(coco.getImgIds(), desc=f"Creating Revvity-25 segmentations ({split})"):
 59        image_metadata = coco.loadImgs(image_id)[0]
 60        file_name = image_metadata["file_name"]
 61
 62        image_path = os.path.join(path, "images", file_name)
 63        assert os.path.exists(image_path), image_path
 64        image_paths.append(image_path)
 65
 66        label_path = os.path.join(label_dir, f"{os.path.splitext(file_name)[0]}.tif")
 67        label_paths.append(label_path)
 68        if os.path.exists(label_path):
 69            continue
 70
 71        annotations = coco.loadAnns(coco.getAnnIds(imgIds=image_id))
 72        shape = (image_metadata["height"], image_metadata["width"])
 73        seg = np.zeros(shape, dtype="uint32")
 74
 75        # Paint the largest cells first, so smaller overlapping cells stay visible on top.
 76        masks = [coco.annToMask(a).astype(bool) for a in annotations]
 77        sorting = np.argsort([m.sum() for m in masks])[::-1]
 78        for seg_id, idx in enumerate(sorting, 1):
 79            seg[masks[idx]] = seg_id
 80
 81        imageio.imwrite(label_path, seg.astype("uint16"), compression="zlib")
 82
 83    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
 84    return natsorted(image_paths), natsorted(label_paths)
 85
 86
 87def get_revvity25_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 88    """Download the Revvity-25 dataset.
 89
 90    Args:
 91        path: Filepath to a folder where the downloaded data will be saved.
 92        download: Whether to download the data if it is not present.
 93
 94    Returns:
 95        The filepath to the folder where the data is stored.
 96    """
 97    if os.path.exists(os.path.join(path, "images")) and os.path.exists(os.path.join(path, "annotations")):
 98        return path
 99
100    if not download:
101        raise RuntimeError(f"Cannot find the data at {path}, but 'download' is set to False.")
102
103    try:
104        from huggingface_hub import snapshot_download
105    except ImportError:
106        raise ImportError("huggingface_hub is required. Install with: pip install huggingface_hub")
107
108    os.makedirs(path, exist_ok=True)
109    snapshot_download(repo_id=HF_REPO, repo_type="dataset", local_dir=path)
110
111    return path
112
113
114def get_revvity25_paths(
115    path: Union[os.PathLike, str], split: Literal["train", "val"], download: bool = False,
116) -> Tuple[List[str], List[str]]:
117    """Get paths to the Revvity-25 data.
118
119    Args:
120        path: Filepath to a folder where the downloaded data will be saved.
121        split: The data split to use. Either 'train' or 'val'.
122        download: Whether to download the data if it is not present.
123
124    Returns:
125        List of filepaths for the image data.
126        List of filepaths for the label data.
127    """
128    assert split in SPLITS, f"'{split}' is not a valid split. Choose from {list(SPLITS.keys())}."
129    data_dir = get_revvity25_data(path, download)
130    return _create_segmentations_from_coco_annotations(data_dir, split)
131
132
133def get_revvity25_dataset(
134    path: Union[os.PathLike, str],
135    patch_shape: Tuple[int, int],
136    split: Literal["train", "val"],
137    offsets: Optional[List[List[int]]] = None,
138    boundaries: bool = False,
139    binary: bool = False,
140    download: bool = False,
141    **kwargs,
142) -> Dataset:
143    """Get the Revvity-25 dataset for cell instance segmentation.
144
145    Args:
146        path: Filepath to a folder where the downloaded data will be saved.
147        patch_shape: The patch shape to use for training.
148        split: The data split to use. Either 'train' or 'val'.
149        offsets: Offset values for affinity computation used as target.
150        boundaries: Whether to compute boundaries as the target.
151        binary: Whether to use a binary segmentation target.
152        download: Whether to download the data if it is not present.
153        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
154
155    Returns:
156        The segmentation dataset.
157    """
158    image_paths, label_paths = get_revvity25_paths(path, split, download)
159
160    kwargs, _ = util.add_instance_label_transform(
161        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
162    )
163    kwargs = util.update_kwargs(kwargs, "ndim", 2)
164
165    return torch_em.default_segmentation_dataset(
166        raw_paths=image_paths,
167        raw_key=None,
168        label_paths=label_paths,
169        label_key=None,
170        patch_shape=patch_shape,
171        is_seg_dataset=False,
172        with_channels=True,
173        **kwargs,
174    )
175
176
177def get_revvity25_loader(
178    path: Union[os.PathLike, str],
179    batch_size: int,
180    patch_shape: Tuple[int, int],
181    split: Literal["train", "val"],
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 Revvity-25 dataloader for cell instance 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 data split to use. Either 'train' or 'val'.
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_revvity25_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=dataset, batch_size=batch_size, **loader_kwargs)