torch_em.data.datasets.light_microscopy.fusionx

  1"""FusionX contains light microscopy images of mono- and multi-nucleated cells with
  2instance segmentation annotations, curated for quantifying cell-to-cell fusion.
  3
  4The dataset is located at https://zenodo.org/records/18596465 and is licensed under
  5CC BY 4.0. A handful of annotated images are missing from the released archives; these
  6are skipped automatically when building the paths.
  7"""
  8
  9import os
 10from glob import glob
 11from natsort import natsorted
 12from typing import List, Literal, Optional, Tuple, Union
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URL = "https://zenodo.org/records/18596465/files/{}?download=1"
 22
 23CHECKSUMS = {
 24    "train_png.7z": "214988031e094b0aa9a181b52404a2eaef95bbe3f4921a6c92ac0c912faf7d70",
 25    "train_png.json": "216e4a33860b46f6712ddf1317db7d89cd77337477630af46f425cfb562c97d6",
 26    "test_png.7z": "f9769dc8a6587214d0f8919311fc13771def75830f04ba95b15289edfb8ed8d9",
 27    "test_png.json": "244063f5a5599af6983e8c601fea476b26e1b31defdd2acfb634d16aa86f5cc7",
 28}
 29
 30SPLITS = {"train": "train_png.json", "test": "test_png.json"}
 31
 32
 33def _create_segmentations_from_coco_annotations(path, split):
 34    """Convert COCO mask annotations to instance segmentation masks."""
 35    import numpy as np
 36    import imageio.v3 as imageio
 37    from tqdm import tqdm
 38
 39    try:
 40        from pycocotools.coco import COCO
 41    except ImportError:
 42        raise ImportError(
 43            "'pycocotools' is required for processing the FusionX ground-truth. "
 44            "Install it with 'conda install -c conda-forge pycocotools'."
 45        )
 46
 47    image_dir = os.path.join(path, f"{split}_png")
 48    label_dir = os.path.join(path, "labels", split)
 49    if os.path.exists(label_dir):
 50        label_paths = natsorted(glob(os.path.join(label_dir, "*.tif")))
 51        if len(label_paths) > 0:
 52            image_paths = [
 53                os.path.join(image_dir, f"{os.path.splitext(os.path.basename(p))[0]}.png") for p in label_paths
 54            ]
 55            return natsorted(image_paths), label_paths
 56
 57    os.makedirs(label_dir, exist_ok=True)
 58
 59    ann_file = os.path.join(path, SPLITS[split])
 60    coco = COCO(ann_file)
 61
 62    image_paths, label_paths = [], []
 63    for image_id in tqdm(coco.getImgIds(), desc=f"Creating FusionX segmentations ({split})"):
 64        image_metadata = coco.loadImgs(image_id)[0]
 65        file_name = image_metadata["file_name"]
 66
 67        image_path = os.path.join(image_dir, file_name)
 68        if not os.path.exists(image_path):
 69            continue  # A few annotated images are missing from the released archive.
 70        image_paths.append(image_path)
 71
 72        label_path = os.path.join(label_dir, f"{os.path.splitext(file_name)[0]}.tif")
 73        label_paths.append(label_path)
 74        if os.path.exists(label_path):
 75            continue
 76
 77        annotations = coco.loadAnns(coco.getAnnIds(imgIds=image_id))
 78        shape = (image_metadata["height"], image_metadata["width"])
 79        seg = np.zeros(shape, dtype="uint32")
 80
 81        # Paint the largest cells first, so smaller overlapping cells stay visible on top.
 82        masks = [coco.annToMask(a).astype(bool) for a in annotations]
 83        sorting = np.argsort([m.sum() for m in masks])[::-1]
 84        for seg_id, idx in enumerate(sorting, 1):
 85            seg[masks[idx]] = seg_id
 86
 87        imageio.imwrite(label_path, seg.astype("uint16"), compression="zlib")
 88
 89    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
 90    return natsorted(image_paths), natsorted(label_paths)
 91
 92
 93def get_fusionx_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 94    """Download the FusionX dataset.
 95
 96    Args:
 97        path: Filepath to a folder where the downloaded data will be saved.
 98        download: Whether to download the data if it is not present.
 99
100    Returns:
101        The filepath to the folder where the data is stored.
102    """
103    if os.path.exists(os.path.join(path, "train_png")) and os.path.exists(os.path.join(path, "test_png")):
104        return path
105
106    if not download:
107        raise RuntimeError(f"Cannot find the data at {path}, but 'download' is set to False.")
108
109    os.makedirs(path, exist_ok=True)
110    for fname, checksum in CHECKSUMS.items():
111        fpath = os.path.join(path, fname)
112        util.download_source(fpath, URL.format(fname), download, checksum=checksum)
113        if fname.endswith(".7z"):
114            util.unzip_7z(fpath, path, remove=True)
115
116    return path
117
118
119def get_fusionx_paths(
120    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False,
121) -> Tuple[List[str], List[str]]:
122    """Get paths to the FusionX data.
123
124    Args:
125        path: Filepath to a folder where the downloaded data will be saved.
126        split: The data split to use. Either 'train' or 'test'.
127        download: Whether to download the data if it is not present.
128
129    Returns:
130        List of filepaths for the image data.
131        List of filepaths for the label data.
132    """
133    assert split in SPLITS, f"'{split}' is not a valid split. Choose from {list(SPLITS.keys())}."
134    data_dir = get_fusionx_data(path, download)
135    return _create_segmentations_from_coco_annotations(data_dir, split)
136
137
138def get_fusionx_dataset(
139    path: Union[os.PathLike, str],
140    patch_shape: Tuple[int, int],
141    split: Literal["train", "test"],
142    offsets: Optional[List[List[int]]] = None,
143    boundaries: bool = False,
144    binary: bool = False,
145    download: bool = False,
146    **kwargs,
147) -> Dataset:
148    """Get the FusionX dataset for cell instance segmentation.
149
150    Args:
151        path: Filepath to a folder where the downloaded data will be saved.
152        patch_shape: The patch shape to use for training.
153        split: The data split to use. Either 'train' or 'test'.
154        offsets: Offset values for affinity computation used as target.
155        boundaries: Whether to compute boundaries as the target.
156        binary: Whether to use a binary segmentation target.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
159
160    Returns:
161        The segmentation dataset.
162    """
163    image_paths, label_paths = get_fusionx_paths(path, split, download)
164
165    kwargs, _ = util.add_instance_label_transform(
166        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
167    )
168    kwargs = util.update_kwargs(kwargs, "ndim", 2)
169
170    return torch_em.default_segmentation_dataset(
171        raw_paths=image_paths,
172        raw_key=None,
173        label_paths=label_paths,
174        label_key=None,
175        patch_shape=patch_shape,
176        is_seg_dataset=False,
177        with_channels=True,
178        **kwargs,
179    )
180
181
182def get_fusionx_loader(
183    path: Union[os.PathLike, str],
184    batch_size: int,
185    patch_shape: Tuple[int, int],
186    split: Literal["train", "test"],
187    offsets: Optional[List[List[int]]] = None,
188    boundaries: bool = False,
189    binary: bool = False,
190    download: bool = False,
191    **kwargs,
192) -> DataLoader:
193    """Get the FusionX dataloader for cell instance segmentation.
194
195    Args:
196        path: Filepath to a folder where the downloaded data will be saved.
197        batch_size: The batch size for training.
198        patch_shape: The patch shape to use for training.
199        split: The data split to use. Either 'train' or 'test'.
200        offsets: Offset values for affinity computation used as target.
201        boundaries: Whether to compute boundaries as the target.
202        binary: Whether to use a binary segmentation target.
203        download: Whether to download the data if it is not present.
204        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
205
206    Returns:
207        The DataLoader.
208    """
209    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
210    dataset = get_fusionx_dataset(
211        path=path,
212        patch_shape=patch_shape,
213        split=split,
214        offsets=offsets,
215        boundaries=boundaries,
216        binary=binary,
217        download=download,
218        **ds_kwargs,
219    )
220    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)