torch_em.data.datasets.light_microscopy.ccagt

  1"""The CCAgT dataset contains annotations for nucleus segmentation in
  2bright-field microscopy images of cervical cells with the AgNOR stain.
  3
  4CCAgT stands for Cervical Cells with the AgNOR Technique. The stain marks the nucleolar organizer
  5regions (NORs) of a nucleus as dark dots, so the annotations cover both the nuclei and the NORs
  6inside them. The dataset holds 9339 image tiles of 15 slides with 63190 annotated objects.
  7
  8NOTE: This dataset uses the CC BY-NC 3.0 license, which forbids commercial use. The other datasets
  9of this module use a permissive license, so check the terms before you use this one in your project.
 10
 11NOTE: A NOR sits inside a nucleus, and in the released masks the NOR label overwrites the nucleus.
 12A nucleus therefore looks like a ring with holes. Use the label choice 'nuclei' to merge the nucleus
 13classes and the NOR classes back into whole nuclei.
 14
 15NOTE: The class 'Satellite' does not hold a real outline. The annotators clicked a point, and the
 16export drew a disc of a fixed radius around it, so 6457 of its 6478 objects have the same area. The
 17label choice 'instances' drops this class.
 18
 19NOTE: The publication defines no split. This loader splits by slide, so that no slide appears in two
 20splits. A split over single tiles would put tiles of one patient into several splits.
 21
 22The dataset is located at https://doi.org/10.17632/wg4bpm33hj.2 under the CC BY-NC 3.0 license.
 23This dataset is from the publication https://doi.org/10.1109/CBMS49503.2020.00110.
 24Please cite it if you use this dataset in your research.
 25"""
 26
 27import os
 28import json
 29from glob import glob
 30from natsort import natsorted
 31from collections import defaultdict
 32from typing import List, Literal, Optional, Sequence, Tuple, Union
 33
 34import numpy as np
 35import imageio.v3 as imageio
 36
 37from torch.utils.data import DataLoader, Dataset
 38
 39import torch_em
 40
 41from .. import util
 42
 43
 44URL = "https://data.mendeley.com/public-api/zip/wg4bpm33hj/download/2"
 45CHECKSUM = "eb9f01e8feae0029056dac2dfa6f5e9780692e1a54c0df1bbd33cd279b9eb557"
 46
 47ARCHIVE_ROOT = "wg4bpm33hj-2"
 48
 49CLASS_NAMES = {
 50    1: "nucleus",
 51    2: "cluster",
 52    3: "satellite",
 53    4: "nucleus_out_of_focus",
 54    5: "overlapped_nuclei",
 55    6: "non_viable_nucleus",
 56    7: "leukocyte_nucleus",
 57}
 58
 59# The nucleus classes and the NOR classes together form a whole nucleus.
 60NUCLEUS_CLASSES = (1, 2, 3, 4, 5, 6)
 61
 62# The annotators clicked a point for this class, so its outline is a disc of a fixed radius.
 63SYNTHETIC_CLASS = 3
 64
 65# The slides of one split never appear in another split, so a patient stays in one split.
 66SPLITS = {
 67    "train": ("A", "D", "F", "I", "J", "K", "M", "N", "O"),
 68    "val": ("B", "G", "H"),
 69    "test": ("C", "E", "L"),
 70}
 71
 72LABEL_CHOICES = ("semantic", "nuclei", "instances")
 73
 74
 75def _extract_archive(zip_path: str, path: str) -> None:
 76    """Extract the archive and the per slide archives that it holds."""
 77    import zipfile
 78
 79    with zipfile.ZipFile(zip_path) as archive:
 80        archive.extractall(path)
 81
 82    data_dir = os.path.join(path, ARCHIVE_ROOT)
 83    for kind in ("images", "masks"):
 84        for slide_path in natsorted(glob(os.path.join(data_dir, kind, "*.zip"))):
 85            with zipfile.ZipFile(slide_path) as archive:
 86                archive.extractall(os.path.join(data_dir, kind))
 87            os.remove(slide_path)
 88
 89
 90def _rasterize(annotations, shape: Tuple[int, int]) -> np.ndarray:
 91    """Draw one label per object, and let a small object win over a large one."""
 92    from skimage.draw import polygon as draw_polygon
 93
 94    labels = np.zeros(shape, dtype="uint16")
 95    # Draw the large objects first, so that a small object on top keeps its label.
 96    ordered = sorted(annotations, key=lambda a: -a.get("area", 0))
 97    for instance_id, annotation in enumerate(ordered, start=1):
 98        for part in annotation["segmentation"]:
 99            polygon = np.array(part, dtype=float).reshape(-1, 2)
100            rows, columns = draw_polygon(polygon[:, 1], polygon[:, 0], shape=shape)
101            labels[rows, columns] = instance_id
102    return labels
103
104
105def _create_instance_labels(data_dir: str) -> str:
106    """Rasterize the COCO polygons of every tile into an instance label image."""
107    from tqdm import tqdm
108
109    label_dir = os.path.join(data_dir, "instance_labels")
110    if os.path.exists(label_dir) and len(glob(os.path.join(label_dir, "*", "*.tif"))) > 0:
111        return label_dir
112
113    with open(os.path.join(data_dir, "CCAgT_COCO_OD.json")) as f:
114        coco = json.load(f)
115
116    per_image = defaultdict(list)
117    for annotation in coco["annotations"]:
118        # The synthetic class holds a disc around a click, so it is not a real outline.
119        if annotation["category_id"] == SYNTHETIC_CLASS:
120            continue
121        per_image[annotation["image_id"]].append(annotation)
122
123    for image in tqdm(coco["images"], desc="Preprocess the CCAgT annotations"):
124        name = image["file_name"]
125        slide = name.split("_")[0]
126        output_dir = os.path.join(label_dir, slide)
127        os.makedirs(output_dir, exist_ok=True)
128        output_path = os.path.join(output_dir, f"{os.path.splitext(name)[0]}.tif")
129        if os.path.exists(output_path):
130            continue
131
132        shape = (image["height"], image["width"])
133        labels = _rasterize(per_image.get(image["id"], []), shape)
134        imageio.imwrite(output_path, labels, compression="zlib")
135
136    return label_dir
137
138
139def _create_nuclei_labels(data_dir: str) -> str:
140    """Merge the nucleus classes and the NOR classes of the released masks into whole nuclei."""
141    from tqdm import tqdm
142
143    label_dir = os.path.join(data_dir, "nuclei_labels")
144    mask_paths = natsorted(glob(os.path.join(data_dir, "masks", "*", "*.png")))
145    if os.path.exists(label_dir) and len(glob(os.path.join(label_dir, "*", "*.tif"))) == len(mask_paths):
146        return label_dir
147
148    for mask_path in tqdm(mask_paths, desc="Merge the CCAgT nucleus classes"):
149        slide = os.path.basename(os.path.dirname(mask_path))
150        output_dir = os.path.join(label_dir, slide)
151        os.makedirs(output_dir, exist_ok=True)
152        output_path = os.path.join(output_dir, os.path.basename(mask_path).replace(".png", ".tif"))
153        if os.path.exists(output_path):
154            continue
155
156        mask = imageio.imread(mask_path)
157        nuclei = np.isin(mask, NUCLEUS_CLASSES).astype("uint8")
158        imageio.imwrite(output_path, nuclei, compression="zlib")
159
160    return label_dir
161
162
163def get_ccagt_data(path: Union[os.PathLike, str], download: bool = False) -> str:
164    """Download the CCAgT dataset.
165
166    Args:
167        path: Filepath to a folder where the downloaded data will be saved.
168        download: Whether to download the data if it is not present.
169
170    Returns:
171        The filepath to the extracted data.
172    """
173    data_dir = os.path.join(path, ARCHIVE_ROOT)
174    if os.path.exists(data_dir):
175        return data_dir
176
177    os.makedirs(path, exist_ok=True)
178    zip_path = os.path.join(path, "ccagt.zip")
179    util.download_source(zip_path, URL, download, CHECKSUM)
180    _extract_archive(zip_path, path)
181
182    return data_dir
183
184
185def get_ccagt_paths(
186    path: Union[os.PathLike, str],
187    split: Optional[Literal["train", "val", "test"]] = "train",
188    slides: Optional[Sequence[str]] = None,
189    label_choice: Literal["semantic", "nuclei", "instances"] = "instances",
190    download: bool = False,
191) -> Tuple[List[str], List[str]]:
192    """Get paths to the CCAgT data.
193
194    Args:
195        path: Filepath to a folder where the downloaded data will be saved.
196        split: The data split, which groups the slides. Ignored when you pass `slides`.
197        slides: The slides to use, for example ('A', 'D'). Overrides `split`.
198        label_choice: The target. Either 'instances' for the single objects, 'nuclei' for the whole
199            nuclei, or 'semantic' for the seven classes as the authors released them.
200        download: Whether to download the data if it is not present.
201
202    Returns:
203        List of filepaths for the image data.
204        List of filepaths for the label data.
205    """
206    if label_choice not in LABEL_CHOICES:
207        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose from {list(LABEL_CHOICES)}.")
208
209    if slides is None:
210        if split not in SPLITS:
211            raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}, or pass slides.")
212        slides = SPLITS[split]
213    else:
214        known = {slide for group in SPLITS.values() for slide in group}
215        for slide in slides:
216            if slide not in known:
217                raise ValueError(f"'{slide}' is not a valid slide. Choose from {sorted(known)}.")
218
219    data_dir = get_ccagt_data(path, download)
220
221    if label_choice == "semantic":
222        label_dir, extension = os.path.join(data_dir, "masks"), ".png"
223    elif label_choice == "nuclei":
224        label_dir, extension = _create_nuclei_labels(data_dir), ".tif"
225    else:
226        label_dir, extension = _create_instance_labels(data_dir), ".tif"
227
228    image_paths, label_paths = [], []
229    for slide in slides:
230        for image_path in natsorted(glob(os.path.join(data_dir, "images", slide, "*.jpg"))):
231            name = os.path.splitext(os.path.basename(image_path))[0]
232            label_path = os.path.join(label_dir, slide, f"{name}{extension}")
233            if not os.path.exists(label_path):
234                continue
235            image_paths.append(image_path)
236            label_paths.append(label_path)
237
238    if not image_paths:
239        raise RuntimeError(f"Could not find any CCAgT data in {data_dir}.")
240
241    return image_paths, label_paths
242
243
244def get_ccagt_dataset(
245    path: Union[os.PathLike, str],
246    patch_shape: Tuple[int, int],
247    split: Optional[Literal["train", "val", "test"]] = "train",
248    slides: Optional[Sequence[str]] = None,
249    label_choice: Literal["semantic", "nuclei", "instances"] = "instances",
250    offsets: Optional[List[List[int]]] = None,
251    boundaries: bool = False,
252    binary: bool = False,
253    download: bool = False,
254    **kwargs,
255) -> Dataset:
256    """Get the CCAgT dataset for nucleus segmentation.
257
258    Args:
259        path: Filepath to a folder where the downloaded data will be saved.
260        patch_shape: The 2D patch shape to use for training.
261        split: The data split, which groups the slides. Ignored when you pass `slides`.
262        slides: The slides to use, for example ('A', 'D'). Overrides `split`.
263        label_choice: The target. Either 'instances' for the single objects, 'nuclei' for the whole
264            nuclei, or 'semantic' for the seven classes as the authors released them.
265        offsets: Offset values for affinity computation used as target.
266        boundaries: Whether to compute boundaries as the target.
267        binary: Whether to use a binary segmentation target.
268        download: Whether to download the data if it is not present.
269        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
270
271    Returns:
272        The segmentation dataset.
273    """
274    if len(patch_shape) != 2:
275        raise ValueError(f"The CCAgT patch shape must be two-dimensional, got {patch_shape}.")
276
277    image_paths, label_paths = get_ccagt_paths(path, split, slides, label_choice, download)
278
279    if label_choice == "instances":
280        kwargs, _ = util.add_instance_label_transform(
281            kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
282        )
283    kwargs = util.ensure_transforms(ndim=2, **kwargs)
284
285    return torch_em.default_segmentation_dataset(
286        raw_paths=image_paths,
287        raw_key=None,
288        label_paths=label_paths,
289        label_key=None,
290        patch_shape=patch_shape,
291        is_seg_dataset=False,
292        ndim=2,
293        **kwargs,
294    )
295
296
297def get_ccagt_loader(
298    path: Union[os.PathLike, str],
299    batch_size: int,
300    patch_shape: Tuple[int, int],
301    split: Optional[Literal["train", "val", "test"]] = "train",
302    slides: Optional[Sequence[str]] = None,
303    label_choice: Literal["semantic", "nuclei", "instances"] = "instances",
304    offsets: Optional[List[List[int]]] = None,
305    boundaries: bool = False,
306    binary: bool = False,
307    download: bool = False,
308    **kwargs,
309) -> DataLoader:
310    """Get the CCAgT dataloader for nucleus segmentation.
311
312    Args:
313        path: Filepath to a folder where the downloaded data will be saved.
314        batch_size: The batch size for training.
315        patch_shape: The 2D patch shape to use for training.
316        split: The data split, which groups the slides. Ignored when you pass `slides`.
317        slides: The slides to use, for example ('A', 'D'). Overrides `split`.
318        label_choice: The target. Either 'instances' for the single objects, 'nuclei' for the whole
319            nuclei, or 'semantic' for the seven classes as the authors released them.
320        offsets: Offset values for affinity computation used as target.
321        boundaries: Whether to compute boundaries as the target.
322        binary: Whether to use a binary segmentation target.
323        download: Whether to download the data if it is not present.
324        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
325
326    Returns:
327        The DataLoader.
328    """
329    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
330    dataset = get_ccagt_dataset(
331        path=path,
332        patch_shape=patch_shape,
333        split=split,
334        slides=slides,
335        label_choice=label_choice,
336        offsets=offsets,
337        boundaries=boundaries,
338        binary=binary,
339        download=download,
340        **ds_kwargs,
341    )
342    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)