torch_em.data.datasets.light_microscopy.cellapp

  1"""The Cell-APP dataset contains annotations for cell segmentation in
  2transmitted light microscopy images of cultured mammalian cells.
  3
  4The dataset provides per cell line subsets for HeLa, RPE1 and U2OS, and eight multi cell line
  5subsets of increasing size under 'general'. Every cell carries a mitotic or nonmitotic class, at a
  6ratio of about one to twenty. The loader stores the instance labels and the semantic labels, so you
  7can train on either target.
  8
  9NOTE: An automatic pipeline created the annotations, so an image does not show every cell as an
 10instance. A cell without an annotation becomes background in the label image. Take this into account
 11when you train on this data, and when you measure the recall of a model on it.
 12
 13NOTE: The archive does not contain the images of the HeLa train split, although its annotation file
 14declares 270 images. Only the HeLa test split is available.
 15
 16The dataset is located at https://doi.org/10.5281/zenodo.16738843 under the CC BY 4.0 license.
 17This dataset is from the publication https://doi.org/10.1091/mbc.E25-02-0076.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import json
 23from glob import glob
 24from pathlib import Path
 25from natsort import natsorted
 26from collections import defaultdict
 27from typing import List, Literal, Optional, Tuple, Union
 28
 29import numpy as np
 30import imageio.v3 as imageio
 31
 32from torch.utils.data import DataLoader, Dataset
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39URL = "https://zenodo.org/api/records/16738843/files/for_zenodo_d.zip/content"
 40CHECKSUM = "0e875ef3ec5e2937155e21ed0e905da6780108edace64d9c6c94dcfe52e28c18"
 41
 42# The image folder and the COCO file of each cell line split.
 43CELL_LINES = {
 44    ("hela", "test"): ("HeLa/test/images", "HeLa/test/hela_0.1_test.json"),
 45    ("rpe1", "train"): ("RPE1/train/images", "RPE1/train/rpe1_0.4_train.json"),
 46    ("rpe1", "test"): ("RPE1/test/images", "RPE1/test/rpe1_0.4_test.json"),
 47    ("u2os", "train"): ("U2OS/train/images", "U2OS/train/instances_u2os_0.8_train.json"),
 48    ("u2os", "test"): ("U2OS/test/images", "U2OS/test/instances_u2os_0.8_test.json"),
 49}
 50
 51SOURCES = ("hela", "rpe1", "u2os", "general")
 52GENERAL_SIZES = tuple(range(1, 9))
 53
 54
 55def _get_subset_layout(source: str, split: Optional[str], size: int) -> Tuple[str, str, str]:
 56    """Get the image folder, the COCO file and the name of the requested subset."""
 57    if source not in SOURCES:
 58        raise ValueError(f"'{source}' is not a valid source. Choose from {list(SOURCES)}.")
 59
 60    if source == "general":
 61        if size not in GENERAL_SIZES:
 62            raise ValueError(f"'{size}' is not a valid size. Choose an integer from {list(GENERAL_SIZES)}.")
 63        return f"general/dataset_{size}/data", f"general/dataset_{size}/labels.json", f"general_{size}"
 64
 65    if (source, split) not in CELL_LINES:
 66        if source == "hela" and split == "train":
 67            raise ValueError(
 68                "The Cell-APP archive does not contain the images of the HeLa train split. "
 69                "Use split='test' for HeLa, or use another source."
 70            )
 71        valid = [s for (line, s) in CELL_LINES if line == source]
 72        raise ValueError(f"'{split}' is not a valid split for '{source}'. Choose from {valid}.")
 73
 74    image_dir, coco_file = CELL_LINES[(source, split)]
 75    return image_dir, coco_file, f"{source}_{split}"
 76
 77
 78def _rasterize(annotation, shape: Tuple[int, int]) -> np.ndarray:
 79    """Draw all polygons of one annotation into a single binary mask."""
 80    from skimage.draw import polygon as draw_polygon
 81
 82    mask = np.zeros(shape, dtype=bool)
 83    for part in annotation["segmentation"]:
 84        polygon = np.array(part, dtype=float).reshape(-1, 2)
 85        rows, columns = draw_polygon(polygon[:, 1], polygon[:, 0], shape=shape)
 86        mask[rows, columns] = True
 87    return mask
 88
 89
 90def _create_labels(data_dir: str, source: str, split: Optional[str], size: int) -> str:
 91    """Rasterize the COCO polygons into instance labels and semantic labels."""
 92    import h5py
 93    from tqdm import tqdm
 94
 95    image_dir, coco_file, subset = _get_subset_layout(source, split, size)
 96
 97    preprocessed_dir = os.path.join(data_dir, "preprocessed", subset)
 98    os.makedirs(preprocessed_dir, exist_ok=True)
 99
100    with open(os.path.join(data_dir, coco_file)) as f:
101        coco = json.load(f)
102
103    annotations = defaultdict(list)
104    for annotation in coco["annotations"]:
105        annotations[annotation["image_id"]].append(annotation)
106
107    for image in tqdm(coco["images"], desc=f"Preprocess the '{subset}' subset"):
108        stem = Path(image["file_name"]).stem
109        output_path = os.path.join(preprocessed_dir, f"{stem}.h5")
110        if os.path.exists(output_path):
111            continue
112
113        image_path = os.path.join(data_dir, image_dir, image["file_name"])
114        if not os.path.exists(image_path):
115            raise RuntimeError(f"Could not find the Cell-APP image {image_path}.")
116
117        raw = imageio.imread(image_path)
118        if raw.ndim == 3:
119            raw = raw[..., 0]  # The channels of the transmitted light images are identical.
120
121        shape = (image["height"], image["width"])
122        masks = [(_rasterize(a, shape), a["category_id"]) for a in annotations[image["id"]]]
123
124        # Paint the large cells first, so that a small cell on top of a large one keeps its label.
125        labels = np.zeros(shape, dtype="uint16")
126        semantic = np.zeros(shape, dtype="uint8")
127        for instance_id, (mask, category_id) in enumerate(sorted(masks, key=lambda m: -m[0].sum()), start=1):
128            labels[mask] = instance_id
129            semantic[mask] = category_id + 1  # 0 is the background, 1 is nonmitotic, 2 is mitotic.
130
131        with h5py.File(output_path, "w") as f:
132            f.create_dataset("raw", data=raw, compression="gzip")
133            f.create_dataset("labels", data=labels, compression="gzip")
134            f.create_dataset("semantic", data=semantic, compression="gzip")
135
136    return preprocessed_dir
137
138
139def get_cellapp_data(path: Union[os.PathLike, str], download: bool = False) -> str:
140    """Download the Cell-APP dataset.
141
142    Args:
143        path: Filepath to a folder where the downloaded data will be saved.
144        download: Whether to download the data if it is not present.
145
146    Returns:
147        The filepath to the extracted data.
148    """
149    data_dir = os.path.join(path, "for_zenodo_d")
150    if os.path.exists(data_dir):
151        return data_dir
152
153    os.makedirs(path, exist_ok=True)
154    zip_path = os.path.join(path, "for_zenodo_d.zip")
155    util.download_source(zip_path, URL, download, CHECKSUM)
156    util.unzip(zip_path=zip_path, dst=path)
157
158    return data_dir
159
160
161def get_cellapp_paths(
162    path: Union[os.PathLike, str],
163    source: Literal["hela", "rpe1", "u2os", "general"] = "general",
164    split: Optional[Literal["train", "test"]] = None,
165    size: int = 8,
166    download: bool = False,
167) -> List[str]:
168    """Get paths to the Cell-APP data.
169
170    Args:
171        path: Filepath to a folder where the downloaded data will be saved.
172        source: The data source. Either a cell line, 'hela', 'rpe1' or 'u2os', or 'general'.
173        split: The data split, 'train' or 'test'. The source 'general' has no split.
174            The source 'hela' has no train split, because the archive misses its images.
175        size: The size of the 'general' subset, an integer from 1 to 8. Other sources ignore it.
176        download: Whether to download the data if it is not present.
177
178    Returns:
179        List of filepaths for the preprocessed h5 data.
180    """
181    data_dir = get_cellapp_data(path, download)
182    preprocessed_dir = _create_labels(data_dir, source, split, size)
183
184    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
185    if not volume_paths:
186        raise RuntimeError(f"Could not find any preprocessed Cell-APP data in {preprocessed_dir}.")
187
188    return volume_paths
189
190
191def get_cellapp_dataset(
192    path: Union[os.PathLike, str],
193    patch_shape: Tuple[int, int],
194    source: Literal["hela", "rpe1", "u2os", "general"] = "general",
195    split: Optional[Literal["train", "test"]] = None,
196    size: int = 8,
197    label_choice: Literal["instances", "semantic"] = "instances",
198    offsets: Optional[List[List[int]]] = None,
199    boundaries: bool = False,
200    binary: bool = False,
201    download: bool = False,
202    **kwargs,
203) -> Dataset:
204    """Get the Cell-APP dataset for cell segmentation.
205
206    Args:
207        path: Filepath to a folder where the downloaded data will be saved.
208        patch_shape: The 2D patch shape to use for training.
209        source: The data source. Either a cell line, 'hela', 'rpe1' or 'u2os', or 'general'.
210        split: The data split, 'train' or 'test'. The source 'general' has no split.
211            The source 'hela' has no train split, because the archive misses its images.
212        size: The size of the 'general' subset, an integer from 1 to 8. Other sources ignore it.
213        label_choice: The target. Either 'instances' for the cell instances, or 'semantic' for the
214            mitotic and nonmitotic classes.
215        offsets: Offset values for affinity computation used as target.
216        boundaries: Whether to compute boundaries as the target.
217        binary: Whether to use a binary segmentation target.
218        download: Whether to download the data if it is not present.
219        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
220
221    Returns:
222        The segmentation dataset.
223    """
224    if len(patch_shape) != 2:
225        raise ValueError(f"The Cell-APP patch shape must be two-dimensional, got {patch_shape}.")
226
227    if label_choice not in ("instances", "semantic"):
228        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose 'instances' or 'semantic'.")
229
230    volume_paths = get_cellapp_paths(path, source, split, size, download)
231    label_key = "labels" if label_choice == "instances" else "semantic"
232
233    if label_choice == "instances":
234        kwargs, _ = util.add_instance_label_transform(
235            kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
236        )
237    kwargs = util.ensure_transforms(ndim=2, **kwargs)
238
239    return torch_em.default_segmentation_dataset(
240        raw_paths=volume_paths,
241        raw_key="raw",
242        label_paths=volume_paths,
243        label_key=label_key,
244        patch_shape=patch_shape,
245        ndim=2,
246        **kwargs,
247    )
248
249
250def get_cellapp_loader(
251    path: Union[os.PathLike, str],
252    batch_size: int,
253    patch_shape: Tuple[int, int],
254    source: Literal["hela", "rpe1", "u2os", "general"] = "general",
255    split: Optional[Literal["train", "test"]] = None,
256    size: int = 8,
257    label_choice: Literal["instances", "semantic"] = "instances",
258    offsets: Optional[List[List[int]]] = None,
259    boundaries: bool = False,
260    binary: bool = False,
261    download: bool = False,
262    **kwargs,
263) -> DataLoader:
264    """Get the Cell-APP dataloader for cell segmentation.
265
266    Args:
267        path: Filepath to a folder where the downloaded data will be saved.
268        batch_size: The batch size for training.
269        patch_shape: The 2D patch shape to use for training.
270        source: The data source. Either a cell line, 'hela', 'rpe1' or 'u2os', or 'general'.
271        split: The data split, 'train' or 'test'. The source 'general' has no split.
272            The source 'hela' has no train split, because the archive misses its images.
273        size: The size of the 'general' subset, an integer from 1 to 8. Other sources ignore it.
274        label_choice: The target. Either 'instances' for the cell instances, or 'semantic' for the
275            mitotic and nonmitotic classes.
276        offsets: Offset values for affinity computation used as target.
277        boundaries: Whether to compute boundaries as the target.
278        binary: Whether to use a binary segmentation target.
279        download: Whether to download the data if it is not present.
280        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
281
282    Returns:
283        The DataLoader.
284    """
285    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
286    dataset = get_cellapp_dataset(
287        path=path,
288        patch_shape=patch_shape,
289        source=source,
290        split=split,
291        size=size,
292        label_choice=label_choice,
293        offsets=offsets,
294        boundaries=boundaries,
295        binary=binary,
296        download=download,
297        **ds_kwargs,
298    )
299    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)