torch_em.data.datasets.histopathology.icc

  1"""The ICC dataset contains annotations for interstitial cells of Cajal (ICC) segmentation
  2in anti-CD117 immunohistochemistry stained anal canal wall images from haemorrhoidal disease
  3patients.
  4
  5The dataset is located at https://doi.org/10.5281/zenodo.14900511 under the
  6CC BY-NC-SA 4.0 license. This dataset is from the publication https://doi.org/10.3390/cells14070550.
  7Please cite it if you use this dataset in your research.
  8"""
  9
 10import os
 11import json
 12from glob import glob
 13from tqdm import tqdm
 14from natsort import natsorted
 15from collections import defaultdict
 16from typing import List, Tuple, Union
 17
 18import numpy as np
 19import imageio.v3 as imageio
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URLS = {
 29    "images": "https://zenodo.org/records/14900511/files/Images.zip",
 30    "annotations": "https://zenodo.org/records/14900511/files/Annotations.coco.json",
 31}
 32CHECKSUMS = {
 33    "images": "d51cb75091891437df80c4ab69412a622c2181ddb9b974ddbc84aeb326542380",
 34    "annotations": "e98facfdbe04226eb004b40ef06a4c25179ce952ab036d9f00738ceda49708be",
 35}
 36
 37
 38def _rasterize_labels(coco, data_dir):
 39    label_dir = os.path.join(data_dir, "labels")
 40    if os.path.exists(label_dir) and len(glob(os.path.join(label_dir, "*.tif"))) == len(coco["images"]):
 41        return label_dir
 42
 43    os.makedirs(label_dir, exist_ok=True)
 44
 45    from skimage.draw import polygon as draw_polygon
 46
 47    annotations_per_image = defaultdict(list)
 48    for annotation in coco["annotations"]:
 49        annotations_per_image[annotation["image_id"]].append(annotation)
 50
 51    for image in tqdm(coco["images"], desc="Rasterize the ICC annotations"):
 52        shape = (image["height"], image["width"])
 53        labels = np.zeros(shape, dtype="uint16")
 54        for instance_id, annotation in enumerate(annotations_per_image[image["id"]], start=1):
 55            polygon = np.array(annotation["segmentation"][0], dtype=float).reshape(-1, 2)
 56            rows, columns = draw_polygon(polygon[:, 1], polygon[:, 0], shape=shape)
 57            labels[rows, columns] = instance_id
 58
 59        name = os.path.splitext(image["file_name"])[0]
 60        imageio.imwrite(os.path.join(label_dir, f"{name}.tif"), labels, compression="zlib")
 61
 62    return label_dir
 63
 64
 65def get_icc_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 66    """Download the ICC dataset.
 67
 68    Args:
 69        path: Filepath to a folder where the downloaded data will be saved.
 70        download: Whether to download the data if it is not present.
 71
 72    Returns:
 73        The filepath to the folder where the images are stored.
 74    """
 75    image_dir = os.path.join(path, "images")
 76    if os.path.exists(image_dir):
 77        return path
 78
 79    os.makedirs(path, exist_ok=True)
 80
 81    zip_path = os.path.join(path, "Images.zip")
 82    util.download_source(path=zip_path, url=URLS["images"], download=download, checksum=CHECKSUMS["images"])
 83    util.unzip(zip_path=zip_path, dst=image_dir)
 84
 85    annotation_path = os.path.join(path, "Annotations.coco.json")
 86    util.download_source(
 87        path=annotation_path, url=URLS["annotations"], download=download, checksum=CHECKSUMS["annotations"]
 88    )
 89
 90    return path
 91
 92
 93def get_icc_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 94    """Get paths to the ICC data.
 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        List of filepaths for the image data.
102        List of filepaths for the label data.
103    """
104    data_dir = get_icc_data(path, download)
105
106    with open(os.path.join(data_dir, "Annotations.coco.json")) as f:
107        coco = json.load(f)
108
109    label_dir = _rasterize_labels(coco, data_dir)
110
111    raw_paths = natsorted(glob(os.path.join(data_dir, "images", "*.jpg")))
112    label_paths = natsorted(glob(os.path.join(label_dir, "*.tif")))
113
114    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
115    assert all(
116        os.path.splitext(os.path.basename(raw_path))[0] == os.path.splitext(os.path.basename(label_path))[0]
117        for raw_path, label_path in zip(raw_paths, label_paths)
118    )
119
120    return raw_paths, label_paths
121
122
123def get_icc_dataset(
124    path: Union[os.PathLike, str],
125    patch_shape: Tuple[int, int],
126    resize_inputs: bool = False,
127    download: bool = False,
128    **kwargs,
129) -> Dataset:
130    """Get the ICC dataset for interstitial cells of Cajal segmentation.
131
132    Args:
133        path: Filepath to a folder where the downloaded data will be saved.
134        patch_shape: The patch shape to use for training.
135        resize_inputs: Whether to resize the inputs.
136        download: Whether to download the data if it is not present.
137        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
138
139    Returns:
140        The segmentation dataset.
141    """
142    raw_paths, label_paths = get_icc_paths(path, download)
143
144    if resize_inputs:
145        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
146        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
147            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
148        )
149
150    return torch_em.default_segmentation_dataset(
151        raw_paths=raw_paths,
152        raw_key=None,
153        label_paths=label_paths,
154        label_key=None,
155        patch_shape=patch_shape,
156        is_seg_dataset=False,
157        ndim=2,
158        with_channels=True,
159        **kwargs,
160    )
161
162
163def get_icc_loader(
164    path: Union[os.PathLike, str],
165    batch_size: int,
166    patch_shape: Tuple[int, int],
167    resize_inputs: bool = False,
168    download: bool = False,
169    **kwargs,
170) -> DataLoader:
171    """Get the ICC dataloader for interstitial cells of Cajal segmentation.
172
173    Args:
174        path: Filepath to a folder where the downloaded data will be saved.
175        batch_size: The batch size for training.
176        patch_shape: The patch shape to use for training.
177        resize_inputs: Whether to resize the inputs.
178        download: Whether to download the data if it is not present.
179        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
180
181    Returns:
182        The DataLoader.
183    """
184    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
185    dataset = get_icc_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
186    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)