torch_em.data.datasets.medical.dentex

  1"""The DENTEX dataset contains annotations for semantic segmentation of abnormal teeth by their
  2diagnosis (caries, deep caries, periapical lesion or impacted tooth) in panoramic dental X-rays.
  3
  4The dataset was curated for the DENTEX 2023 MICCAI challenge (https://dentex.grand-challenge.org),
  5which comprises panoramic dental X-rays from three institutions, annotated for the quadrant,
  6tooth (FDI enumeration) and diagnosis of each abnormal tooth. The original annotations are per-tooth
  7polygons with a quadrant, tooth enumeration and diagnosis label each; this module rasterizes them
  8into a per-pixel label map of the diagnosis class (see `DIAGNOSIS_CLASSES`), which is the part of
  9the dataset that adds value over the other, purely structural, tooth segmentation datasets in this
 10package. It only uses the fully annotated 'quadrant-enumeration-diagnosis' subset: the training
 11split has ground truth for all classes, and the validation split ground truth was released after
 12the challenge. The test split annotations were never released and so are not covered here.
 13
 14The data is hosted on Hugging Face at https://huggingface.co/datasets/ibrahimhamamci/DENTEX
 15and is distributed under the CC BY-NC-SA 4.0 license.
 16
 17The dataset is from the publication https://doi.org/10.48550/arXiv.2305.19112.
 18Please cite it if you use this dataset for your research.
 19"""
 20
 21import os
 22import json
 23from tqdm import tqdm
 24from natsort import natsorted
 25from typing import Union, Tuple, List, Literal
 26
 27import numpy as np
 28from skimage.draw import polygon
 29
 30import imageio.v3 as imageio
 31
 32from torch.utils.data import Dataset, DataLoader
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39URLS = {
 40    "train": "https://huggingface.co/datasets/ibrahimhamamci/DENTEX/resolve/main/DENTEX/training_data.zip",
 41    "val_images": "https://huggingface.co/datasets/ibrahimhamamci/DENTEX/resolve/main/DENTEX/validation_data.zip",
 42    "val_labels": "https://huggingface.co/datasets/ibrahimhamamci/DENTEX/resolve/main/DENTEX/validation_triple.json",
 43}
 44
 45CHECKSUMS = {
 46    "train": "18b2a2dbc5a2b10b0cc6a7677c46a382f4709ab8c9c3bb94f57b74e38e11ffd3",
 47    "val_images": "6370bb4f1024bd610cde13242a465cb2eff195fc02f56ac22126555e7edc7bc3",
 48    "val_labels": "d058afd35d2849923c7c045e61fd3e05d231dcf74d55009993fff88bd9b6f5a2",
 49}
 50
 51DIAGNOSIS_CLASSES = {0: "background", 1: "impacted", 2: "caries", 3: "periapical_lesion", 4: "deep_caries"}
 52"""The label ids of the diagnosis classes. The 'category_id_3' in the annotations (0-indexed,
 53without background) is mapped to the label id (1-indexed, i.e. + 1) so that 0 marks background."""
 54
 55
 56def _extract_subset(zip_path, path, subdir):
 57    import zipfile
 58
 59    with zipfile.ZipFile(zip_path) as f:
 60        members = [m for m in f.namelist() if f"/{subdir}/" in m and "ipynb_checkpoints" not in m]
 61        f.extractall(path, members=members)
 62
 63
 64def get_dentex_data(
 65    path: Union[os.PathLike, str], split: Literal["train", "val"] = None, download: bool = False
 66) -> str:
 67    """Download the DENTEX dataset.
 68
 69    Args:
 70        path: Filepath to a folder where the data is downloaded for further processing.
 71        split: The data split to download. Either 'train' or 'val'. By default downloads both.
 72        download: Whether to download the data if it is not present.
 73
 74    Returns:
 75        Filepath where the data is downloaded.
 76    """
 77    os.makedirs(path, exist_ok=True)
 78
 79    if split in (None, "train"):
 80        # NOTE: The training archive names this folder with hyphens ('quadrant-enumeration-disease'),
 81        # unlike the validation archive, which uses underscores ('quadrant_enumeration_disease').
 82        train_dir = os.path.join(path, "training_data", "quadrant-enumeration-disease")
 83        if not os.path.exists(train_dir):
 84            zip_path = os.path.join(path, "training_data.zip")
 85            util.download_source(path=zip_path, url=URLS["train"], download=download, checksum=CHECKSUMS["train"])
 86            _extract_subset(zip_path, path, "quadrant-enumeration-disease")
 87            os.remove(zip_path)
 88
 89    if split in (None, "val"):
 90        val_dir = os.path.join(path, "validation_data", "quadrant_enumeration_disease")
 91        if not os.path.exists(val_dir):
 92            zip_path = os.path.join(path, "validation_data.zip")
 93            util.download_source(
 94                path=zip_path, url=URLS["val_images"], download=download, checksum=CHECKSUMS["val_images"]
 95            )
 96            _extract_subset(zip_path, path, "quadrant_enumeration_disease")
 97            os.remove(zip_path)
 98
 99        val_labels_path = os.path.join(path, "validation_triple.json")
100        util.download_source(
101            path=val_labels_path, url=URLS["val_labels"], download=download, checksum=CHECKSUMS["val_labels"]
102        )
103
104    return path
105
106
107def _rasterize_annotations(shape, annotations):
108    labels = np.zeros(shape, dtype="uint8")
109    for ann in annotations:
110        class_id = ann["category_id_3"] + 1
111        for seg in ann["segmentation"]:
112            c, r = np.asarray(seg[0::2]), np.asarray(seg[1::2])
113            rr, cc = polygon(r, c, shape=shape)
114            labels[rr, cc] = class_id
115    return labels
116
117
118def _preprocess_split(image_dir, annotation_path, preprocessed_dir):
119    os.makedirs(preprocessed_dir, exist_ok=True)
120
121    with open(annotation_path) as f:
122        annotations = json.load(f)
123
124    image_paths, gt_paths = [], []
125    for image_info in tqdm(annotations["images"], desc=f"Preprocessing labels for {image_dir}"):
126        image_path = os.path.join(image_dir, image_info["file_name"])
127        if not os.path.exists(image_path):
128            continue
129
130        gt_path = os.path.join(preprocessed_dir, f"{os.path.splitext(image_info['file_name'])[0]}.tif")
131        if not os.path.exists(gt_path):
132            image_annotations = [a for a in annotations["annotations"] if a["image_id"] == image_info["id"]]
133            shape = (image_info["height"], image_info["width"])
134            labels = _rasterize_annotations(shape, image_annotations)
135            imageio.imwrite(gt_path, labels)
136
137        image_paths.append(image_path)
138        gt_paths.append(gt_path)
139
140    return image_paths, gt_paths
141
142
143def get_dentex_paths(
144    path: Union[os.PathLike, str], split: Literal["train", "val"], download: bool = False
145) -> Tuple[List[str], List[str]]:
146    """Get paths to the DENTEX data.
147
148    Args:
149        path: Filepath to a folder where the data is downloaded for further processing.
150        split: The data split to use. Either 'train' or 'val'.
151        download: Whether to download the data if it is not present.
152
153    Returns:
154        List of filepaths for the image data.
155        List of filepaths for the label data.
156    """
157    if split not in ("train", "val"):
158        raise ValueError(f"'{split}' is not a valid split. Please choose either 'train' or 'val'.")
159
160    data_dir = get_dentex_data(path, split, download)
161
162    if split == "train":
163        image_dir = os.path.join(data_dir, "training_data", "quadrant-enumeration-disease", "xrays")
164        annotation_path = os.path.join(
165            data_dir, "training_data", "quadrant-enumeration-disease", "train_quadrant_enumeration_disease.json"
166        )
167    else:
168        image_dir = os.path.join(data_dir, "validation_data", "quadrant_enumeration_disease", "xrays")
169        annotation_path = os.path.join(data_dir, "validation_triple.json")
170
171    preprocessed_dir = os.path.join(data_dir, "preprocessed", split)
172    image_paths, gt_paths = _preprocess_split(image_dir, annotation_path, preprocessed_dir)
173
174    image_paths, gt_paths = natsorted(image_paths), natsorted(gt_paths)
175    return image_paths, gt_paths
176
177
178def get_dentex_dataset(
179    path: Union[os.PathLike, str],
180    patch_shape: Tuple[int, int],
181    split: Literal["train", "val"],
182    resize_inputs: bool = False,
183    download: bool = False,
184    **kwargs
185) -> Dataset:
186    """Get the DENTEX dataset for tooth diagnosis segmentation in panoramic dental radiographs.
187
188    Args:
189        path: Filepath to a folder where the data is downloaded for further processing.
190        patch_shape: The patch shape to use for training.
191        split: The data split to use. Either 'train' or 'val'.
192        resize_inputs: Whether to resize the inputs to the patch shape.
193        download: Whether to download the data if it is not present.
194        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
195
196    Returns:
197        The segmentation dataset.
198    """
199    image_paths, gt_paths = get_dentex_paths(path, split, download)
200
201    if resize_inputs:
202        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
203        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
204            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
205        )
206
207    return torch_em.default_segmentation_dataset(
208        raw_paths=image_paths,
209        raw_key=None,
210        label_paths=gt_paths,
211        label_key=None,
212        is_seg_dataset=False,
213        patch_shape=patch_shape,
214        **kwargs
215    )
216
217
218def get_dentex_loader(
219    path: Union[os.PathLike, str],
220    batch_size: int,
221    patch_shape: Tuple[int, int],
222    split: Literal["train", "val"],
223    resize_inputs: bool = False,
224    download: bool = False,
225    **kwargs
226) -> DataLoader:
227    """Get the DENTEX dataloader for tooth diagnosis segmentation in panoramic dental radiographs.
228
229    Args:
230        path: Filepath to a folder where the data is downloaded for further processing.
231        batch_size: The batch size for training.
232        patch_shape: The patch shape to use for training.
233        split: The data split to use. Either 'train' or 'val'.
234        resize_inputs: Whether to resize the inputs to the patch shape.
235        download: Whether to download the data if it is not present.
236        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
237
238    Returns:
239        The DataLoader.
240    """
241    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
242    dataset = get_dentex_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
243    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)