torch_em.data.datasets.medical.arcade

  1"""The ARCADE dataset contains annotations for coronary artery segmentation and stenosis
  2detection in X-ray coronary angiography images.
  3
  4The dataset provides two tasks, each with 1000 training, 200 validation and 300 test images:
  5- 'syntax': multiclass segmentation of 25 coronary artery segments, following the SYNTAX score
  6  methodology (see `SYNTAX_LABELS`).
  7- 'stenosis': binary segmentation of atherosclerotic plaques (stenotic lesions).
  8
  9The dataset is located at https://zenodo.org/records/10390295 (DOI: 10.5281/zenodo.10390295)
 10and is distributed under the CC0 1.0 license.
 11
 12This dataset is from the publication https://doi.org/10.1038/s41597-023-02871-z.
 13Please cite it if you use this dataset for your research.
 14"""
 15
 16import os
 17import json
 18from glob import glob
 19from tqdm import tqdm
 20from natsort import natsorted
 21from typing import Union, Tuple, List, Literal
 22
 23import numpy as np
 24from skimage.draw import polygon
 25
 26import imageio.v3 as imageio
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://zenodo.org/records/10390295/files/arcade.zip"
 36CHECKSUM = "a396cdea7c92c55dc97bbf3dd8e3df517d76872b289a8bcb45513bdb3350837f"
 37
 38SYNTAX_LABELS = {
 39    0: "background", 1: "1", 2: "2", 3: "3", 4: "4", 5: "5", 6: "6", 7: "7", 8: "8", 9: "9", 10: "9a",
 40    11: "10", 12: "10a", 13: "11", 14: "12", 15: "12a", 16: "13", 17: "14", 18: "14a", 19: "15", 20: "16",
 41    21: "16a", 22: "16b", 23: "16c", 24: "12b", 25: "14b",
 42}
 43"""The label ids of the coronary artery segment classes for the 'syntax' task, following the SYNTAX
 44score segment definitions (https://syntaxscore.org/index.php/tutorial/definitions/14-appendix-i-segment-definitions).
 45"""
 46
 47STENOSIS_LABELS = {0: "background", 1: "stenosis"}
 48"""The label ids for the 'stenosis' task. The raw annotations use a single category id (26) for
 49all stenotic lesions, which is remapped to label id 1 here."""
 50
 51
 52def get_arcade_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 53    """Download the ARCADE dataset.
 54
 55    Args:
 56        path: Filepath to a folder where the data is downloaded for further processing.
 57        download: Whether to download the data if it is not present.
 58
 59    Returns:
 60        Filepath where the data is downloaded.
 61    """
 62    data_dir = os.path.join(path, "arcade")
 63    if os.path.exists(data_dir):
 64        return data_dir
 65
 66    os.makedirs(path, exist_ok=True)
 67
 68    zip_path = os.path.join(path, "arcade.zip")
 69    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 70    util.unzip(zip_path=zip_path, dst=path)
 71
 72    return data_dir
 73
 74
 75def _rasterize_annotations(shape, annotations, task):
 76    labels = np.zeros(shape, dtype="uint8")
 77    for ann in annotations:
 78        # The 'stenosis' task only has one category (id 26, 'stenosis'), which is remapped to label 1.
 79        # The 'syntax' task categories (ids 1-25) already match the desired label ids.
 80        class_id = 1 if task == "stenosis" else ann["category_id"]
 81        for seg in ann["segmentation"]:
 82            c, r = np.asarray(seg[0::2]), np.asarray(seg[1::2])
 83            rr, cc = polygon(r, c, shape=shape)
 84            labels[rr, cc] = class_id
 85    return labels
 86
 87
 88def _preprocess_split(image_dir, annotation_path, preprocessed_dir, task):
 89    os.makedirs(preprocessed_dir, exist_ok=True)
 90
 91    with open(annotation_path) as f:
 92        annotations = json.load(f)
 93
 94    annotations_by_image = {}
 95    for ann in annotations["annotations"]:
 96        annotations_by_image.setdefault(ann["image_id"], []).append(ann)
 97
 98    image_paths, gt_paths = [], []
 99    for image_info in tqdm(annotations["images"], desc=f"Preprocessing labels for {image_dir}"):
100        image_path = os.path.join(image_dir, image_info["file_name"])
101        if not os.path.exists(image_path):
102            continue
103
104        gt_path = os.path.join(preprocessed_dir, image_info["file_name"])
105        if not os.path.exists(gt_path):
106            image_annotations = annotations_by_image.get(image_info["id"], [])
107            shape = (image_info["height"], image_info["width"])
108            labels = _rasterize_annotations(shape, image_annotations, task)
109            imageio.imwrite(gt_path, labels)
110
111        image_paths.append(image_path)
112        gt_paths.append(gt_path)
113
114    return image_paths, gt_paths
115
116
117def get_arcade_paths(
118    path: Union[os.PathLike, str],
119    task: Literal["syntax", "stenosis"],
120    split: Literal["train", "val", "test"],
121    download: bool = False,
122) -> Tuple[List[str], List[str]]:
123    """Get paths to the ARCADE data.
124
125    Args:
126        path: Filepath to a folder where the data is downloaded for further processing.
127        task: The choice of task, either 'syntax' (multiclass vessel segment segmentation) or
128            'stenosis' (binary stenosis segmentation).
129        split: The choice of data split.
130        download: Whether to download the data if it is not present.
131
132    Returns:
133        List of filepaths for the image data.
134        List of filepaths for the label data.
135    """
136    if task not in ("syntax", "stenosis"):
137        raise ValueError(f"'{task}' is not a valid task. Please choose either 'syntax' or 'stenosis'.")
138
139    if split not in ("train", "val", "test"):
140        raise ValueError(f"'{split}' is not a valid split.")
141
142    data_dir = get_arcade_data(path, download)
143
144    image_dir = os.path.join(data_dir, task, split, "images")
145    annotation_path = os.path.join(data_dir, task, split, "annotations", f"{split}.json")
146    preprocessed_dir = os.path.join(data_dir, "preprocessed", task, split)
147
148    if os.path.exists(preprocessed_dir) and len(glob(os.path.join(preprocessed_dir, "*.png"))) > 0:
149        image_paths = natsorted(glob(os.path.join(image_dir, "*.png")))
150        gt_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.png")))
151        return image_paths, gt_paths
152
153    image_paths, gt_paths = _preprocess_split(image_dir, annotation_path, preprocessed_dir, task)
154    return natsorted(image_paths), natsorted(gt_paths)
155
156
157def get_arcade_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, int],
160    task: Literal["syntax", "stenosis"],
161    split: Literal["train", "val", "test"],
162    resize_inputs: bool = False,
163    download: bool = False,
164    **kwargs
165) -> Dataset:
166    """Get the ARCADE dataset for coronary artery and stenosis segmentation in X-ray angiography.
167
168    Args:
169        path: Filepath to a folder where the data is downloaded for further processing.
170        patch_shape: The patch shape to use for training.
171        task: The choice of task, either 'syntax' (multiclass vessel segment segmentation) or
172            'stenosis' (binary stenosis segmentation).
173        split: The choice of data split.
174        resize_inputs: Whether to resize the inputs to the expected patch shape.
175        download: Whether to download the data if it is not present.
176        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
177
178    Returns:
179        The segmentation dataset.
180    """
181    image_paths, gt_paths = get_arcade_paths(path, task, split, download)
182
183    if resize_inputs:
184        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
185        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
186            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
187        )
188
189    return torch_em.default_segmentation_dataset(
190        raw_paths=image_paths,
191        raw_key=None,
192        label_paths=gt_paths,
193        label_key=None,
194        patch_shape=patch_shape,
195        is_seg_dataset=False,
196        **kwargs
197    )
198
199
200def get_arcade_loader(
201    path: Union[os.PathLike, str],
202    batch_size: int,
203    patch_shape: Tuple[int, int],
204    task: Literal["syntax", "stenosis"],
205    split: Literal["train", "val", "test"],
206    resize_inputs: bool = False,
207    download: bool = False,
208    **kwargs
209) -> DataLoader:
210    """Get the ARCADE dataloader for coronary artery and stenosis segmentation in X-ray angiography.
211
212    Args:
213        path: Filepath to a folder where the data is downloaded for further processing.
214        batch_size: The batch size for training.
215        patch_shape: The patch shape to use for training.
216        task: The choice of task, either 'syntax' (multiclass vessel segment segmentation) or
217            'stenosis' (binary stenosis segmentation).
218        split: The choice of data split.
219        resize_inputs: Whether to resize the inputs to the expected patch shape.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
222
223    Returns:
224        The DataLoader.
225    """
226    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
227    dataset = get_arcade_dataset(path, patch_shape, task, split, resize_inputs, download, **ds_kwargs)
228    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)