torch_em.data.datasets.light_microscopy.fluo_neuronal_cells

  1"""The Fluorescent Neuronal Cells v2 dataset contains annotations for cell segmentation
  2in fluorescence microscopy images of rodent brain slices.
  3
  4The dataset holds three image collections, named after the marker color: green, yellow and red.
  5Each collection has a trainval and a test split. The archive ships a binary mask per image, and a
  6COCO file that stores one polygon per cell. This loader rasterizes the polygons, so the target is an
  7instance segmentation.
  8
  9NOTE: The collections differ a lot. The green images show small punctate nuclei, the yellow images
 10show sparse bright cells, and the red images show dense filamentous neurons at a larger size.
 11
 12NOTE: The COCO files of this dataset do not follow the usual layout. They store one annotation per
 13image, and its 'segmentation', 'bbox', 'dots' and 'area' fields are parallel lists over the cells.
 14
 15NOTE: The green trainval split ships a mask for '128.png' but not the image, so this loader skips
 16that sample. It yields 749 labelled images instead of the 750 that the publication reports.
 17
 18The dataset is located at https://amsacta.unibo.it/id/eprint/7347 under the CC BY 4.0 license.
 19This dataset is from the publication https://doi.org/10.1038/s41597-024-03005-9.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24import json
 25from glob import glob
 26from pathlib import Path
 27from natsort import natsorted
 28from typing import List, Literal, Optional, Sequence, Tuple, Union
 29
 30import numpy as np
 31import imageio.v3 as imageio
 32
 33from torch.utils.data import DataLoader, Dataset
 34
 35import torch_em
 36
 37from .. import util
 38
 39
 40URLS = {
 41    "green": "https://amsacta.unibo.it/id/eprint/7347/28/green.zip",
 42    "yellow": "https://amsacta.unibo.it/id/eprint/7347/27/yellow.zip",
 43    "red": "https://amsacta.unibo.it/id/eprint/7347/29/red.zip",
 44}
 45
 46CHECKSUMS = {
 47    "green": "7760c3c55d236c17f2225a0a63e8e61a86c93d0645c7188bbe9652f22f520517",
 48    "yellow": "08048f5c71d4afa726cc54fd887a06cc6894aff62f6704ff5a8c15aa2f9946ca",
 49    "red": "0a9a2ac9dbdd7222cadf44a7846f61f9d0adb08cf09b801efb917a76e8704533",
 50}
 51
 52COLLECTIONS = tuple(URLS)
 53SPLITS = ("trainval", "test")
 54
 55
 56def _rasterize(annotation, shape: Tuple[int, int]) -> np.ndarray:
 57    """Draw one label per cell polygon of an image."""
 58    from skimage.draw import polygon as draw_polygon
 59
 60    labels = np.zeros(shape, dtype="uint16")
 61    for instance_id, part in enumerate(annotation["segmentation"], start=1):
 62        polygon = np.array(part, dtype=float).reshape(-1, 2)
 63        rows, columns = draw_polygon(polygon[:, 1], polygon[:, 0], shape=shape)
 64        labels[rows, columns] = instance_id
 65    return labels
 66
 67
 68def _create_instance_labels(data_dir: str, collection: str, split: str) -> str:
 69    """Rasterize the COCO polygons of one collection split into instance label images."""
 70    from tqdm import tqdm
 71
 72    label_dir = os.path.join(data_dir, collection, split, "instance_labels")
 73    os.makedirs(label_dir, exist_ok=True)
 74
 75    coco_paths = glob(os.path.join(data_dir, collection, split, "ground_truths", "COCO", "*.json"))
 76    if not coco_paths:
 77        raise RuntimeError(f"Could not find the COCO file for '{collection}/{split}' in {data_dir}.")
 78
 79    with open(coco_paths[0]) as f:
 80        coco = json.load(f)
 81
 82    images = {image["id"]: image["file_name"] for image in coco["images"]}
 83    image_dir = os.path.join(data_dir, collection, split, "images")
 84
 85    for annotation in tqdm(coco["annotations"], desc=f"Preprocess '{collection}/{split}'"):
 86        file_name = images.get(annotation["image_id"])
 87        if file_name is None:
 88            continue
 89
 90        # The green trainval split lists an image that the archive does not contain.
 91        image_path = os.path.join(image_dir, file_name)
 92        if not os.path.exists(image_path):
 93            continue
 94
 95        output_path = os.path.join(label_dir, f"{Path(file_name).stem}.tif")
 96        if os.path.exists(output_path):
 97            continue
 98
 99        shape = imageio.imread(image_path).shape[:2]
100        imageio.imwrite(output_path, _rasterize(annotation, shape), compression="zlib")
101
102    return label_dir
103
104
105def get_fluo_neuronal_cells_data(
106    path: Union[os.PathLike, str],
107    collection: Literal["green", "yellow", "red"] = "green",
108    download: bool = False,
109) -> str:
110    """Download the Fluorescent Neuronal Cells dataset.
111
112    Args:
113        path: Filepath to a folder where the downloaded data will be saved.
114        collection: The image collection. Either 'green', 'yellow' or 'red'.
115        download: Whether to download the data if it is not present.
116
117    Returns:
118        The filepath to the folder that holds the collections.
119    """
120    if collection not in COLLECTIONS:
121        raise ValueError(f"'{collection}' is not a valid collection. Choose from {list(COLLECTIONS)}.")
122
123    collection_dir = os.path.join(path, collection)
124    if os.path.exists(collection_dir):
125        return path
126
127    os.makedirs(path, exist_ok=True)
128    zip_path = os.path.join(path, f"{collection}.zip")
129    util.download_source(zip_path, URLS[collection], download, CHECKSUMS[collection])
130    util.unzip(zip_path=zip_path, dst=path)
131
132    return path
133
134
135def get_fluo_neuronal_cells_paths(
136    path: Union[os.PathLike, str],
137    split: Literal["trainval", "test"] = "trainval",
138    collection: Optional[Union[str, Sequence[str]]] = None,
139    download: bool = False,
140) -> Tuple[List[str], List[str]]:
141    """Get paths to the Fluorescent Neuronal Cells data.
142
143    Args:
144        path: Filepath to a folder where the downloaded data will be saved.
145        split: The data split. Either 'trainval' or 'test'.
146        collection: The image collection or collections. Defaults to all of them.
147        download: Whether to download the data if it is not present.
148
149    Returns:
150        List of filepaths for the image data.
151        List of filepaths for the instance label data.
152    """
153    if split not in SPLITS:
154        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
155
156    if collection is None:
157        collections = list(COLLECTIONS)
158    else:
159        collections = [collection] if isinstance(collection, str) else list(collection)
160
161    image_paths, label_paths = [], []
162    for name in collections:
163        data_dir = get_fluo_neuronal_cells_data(path, name, download)
164        label_dir = _create_instance_labels(data_dir, name, split)
165
166        for label_path in natsorted(glob(os.path.join(label_dir, "*.tif"))):
167            image_path = os.path.join(data_dir, name, split, "images", f"{Path(label_path).stem}.png")
168            if not os.path.exists(image_path):
169                continue
170            image_paths.append(image_path)
171            label_paths.append(label_path)
172
173    if not image_paths:
174        raise RuntimeError(f"Could not find any Fluorescent Neuronal Cells data in {path}.")
175
176    return image_paths, label_paths
177
178
179def get_fluo_neuronal_cells_dataset(
180    path: Union[os.PathLike, str],
181    patch_shape: Tuple[int, int],
182    split: Literal["trainval", "test"] = "trainval",
183    collection: Optional[Union[str, Sequence[str]]] = None,
184    offsets: Optional[List[List[int]]] = None,
185    boundaries: bool = False,
186    binary: bool = False,
187    download: bool = False,
188    **kwargs,
189) -> Dataset:
190    """Get the Fluorescent Neuronal Cells dataset for cell segmentation.
191
192    Args:
193        path: Filepath to a folder where the downloaded data will be saved.
194        patch_shape: The 2D patch shape to use for training.
195        split: The data split. Either 'trainval' or 'test'.
196        collection: The image collection or collections. Defaults to all of them.
197        offsets: Offset values for affinity computation used as target.
198        boundaries: Whether to compute boundaries as the target.
199        binary: Whether to use a binary segmentation target.
200        download: Whether to download the data if it is not present.
201        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
202
203    Returns:
204        The segmentation dataset.
205    """
206    if len(patch_shape) != 2:
207        raise ValueError(f"The Fluorescent Neuronal Cells patch shape must be two-dimensional, got {patch_shape}.")
208
209    image_paths, label_paths = get_fluo_neuronal_cells_paths(path, split, collection, download)
210
211    kwargs, _ = util.add_instance_label_transform(
212        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
213    )
214    kwargs = util.ensure_transforms(ndim=2, **kwargs)
215
216    return torch_em.default_segmentation_dataset(
217        raw_paths=image_paths,
218        raw_key=None,
219        label_paths=label_paths,
220        label_key=None,
221        patch_shape=patch_shape,
222        is_seg_dataset=False,
223        ndim=2,
224        **kwargs,
225    )
226
227
228def get_fluo_neuronal_cells_loader(
229    path: Union[os.PathLike, str],
230    batch_size: int,
231    patch_shape: Tuple[int, int],
232    split: Literal["trainval", "test"] = "trainval",
233    collection: Optional[Union[str, Sequence[str]]] = None,
234    offsets: Optional[List[List[int]]] = None,
235    boundaries: bool = False,
236    binary: bool = False,
237    download: bool = False,
238    **kwargs,
239) -> DataLoader:
240    """Get the Fluorescent Neuronal Cells dataloader for cell segmentation.
241
242    Args:
243        path: Filepath to a folder where the downloaded data will be saved.
244        batch_size: The batch size for training.
245        patch_shape: The 2D patch shape to use for training.
246        split: The data split. Either 'trainval' or 'test'.
247        collection: The image collection or collections. Defaults to all of them.
248        offsets: Offset values for affinity computation used as target.
249        boundaries: Whether to compute boundaries as the target.
250        binary: Whether to use a binary segmentation target.
251        download: Whether to download the data if it is not present.
252        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
253
254    Returns:
255        The DataLoader.
256    """
257    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
258    dataset = get_fluo_neuronal_cells_dataset(
259        path=path,
260        patch_shape=patch_shape,
261        split=split,
262        collection=collection,
263        offsets=offsets,
264        boundaries=boundaries,
265        binary=binary,
266        download=download,
267        **ds_kwargs,
268    )
269    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)