torch_em.data.datasets.medical.dprd

  1"""DPRD is a dataset for caries segmentation in children's dental panoramic radiographs.
  2
  3This dataset is part of the "Children's Dental Panoramic Radiographs Dataset", which is hosted on
  4figshare at https://doi.org/10.6084/m9.figshare.21621705.v1 (part of the collection
  5https://doi.org/10.6084/m9.figshare.c.6317013.v1) and distributed under the CC0 license. This module
  6only makes use of the "Children's dental caries segmentation dataset" subset, which is the part of
  7the archive with pixel-level segmentation masks for dental caries. The raw masks are RGB images
  8with black background and a fixed color marking the caries region; this module collapses them to
  9a single-channel binary label map, where 0 is background and 1 marks caries.
 10
 11The dataset is from the publication https://doi.org/10.1038/s41597-023-02237-5.
 12Please cite it if you use this dataset for your research.
 13"""
 14
 15import os
 16from glob import glob
 17from tqdm import tqdm
 18from pathlib import Path
 19from natsort import natsorted
 20from typing import Union, Tuple, List, Literal
 21
 22import numpy as np
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URL = "https://ndownloader.figshare.com/files/38322366"
 33CHECKSUM = "2e41d862a0787828d659cfd13035e0fbaeb8687995dbfe29a84fdeac09b83945"
 34
 35ZIP_SUBDIR = "Children's dental caries segmentation dataset"
 36
 37
 38def get_dprd_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 39    """Download the DPRD dataset.
 40
 41    Args:
 42        path: Filepath to a folder where the data is downloaded for further processing.
 43        download: Whether to download the data if it is not present.
 44
 45    Returns:
 46        Filepath where the data is downloaded.
 47    """
 48    data_dir = os.path.join(path, ZIP_SUBDIR)
 49    if os.path.exists(data_dir):
 50        return data_dir
 51
 52    os.makedirs(path, exist_ok=True)
 53
 54    zip_path = os.path.join(path, "Dental_dataset.zip")
 55    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 56
 57    import zipfile
 58    with zipfile.ZipFile(zip_path) as f:
 59        members = [m for m in f.namelist() if m.startswith(f"{ZIP_SUBDIR}/")]
 60        f.extractall(path, members=members)
 61    os.remove(zip_path)
 62
 63    return data_dir
 64
 65
 66def get_dprd_paths(
 67    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 68) -> Tuple[List[str], List[str]]:
 69    """Get paths to the DPRD data.
 70
 71    Args:
 72        path: Filepath to a folder where the data is downloaded for further processing.
 73        split: The data split to use. Either 'train' or 'test'.
 74        download: Whether to download the data if it is not present.
 75
 76    Returns:
 77        List of filepaths for the image data.
 78        List of filepaths for the label data.
 79    """
 80    if split not in ("train", "test"):
 81        raise ValueError(f"'{split}' is not a valid split. Please choose either 'train' or 'test'.")
 82
 83    data_dir = get_dprd_data(path, download)
 84
 85    split_dir = "Train" if split == "train" else "Test"
 86    image_paths = natsorted(glob(os.path.join(data_dir, split_dir, "images", "*.png")))
 87    raw_gt_paths = natsorted(glob(os.path.join(data_dir, split_dir, "mask", "*.png")))
 88
 89    assert len(image_paths) == len(raw_gt_paths) and len(image_paths) > 0
 90
 91    neu_gt_dir = os.path.join(data_dir, "preprocessed", split)
 92    os.makedirs(neu_gt_dir, exist_ok=True)
 93
 94    gt_paths = []
 95    for raw_gt_path in tqdm(raw_gt_paths, desc="Preprocessing labels"):
 96        gt_path = os.path.join(neu_gt_dir, f"{Path(raw_gt_path).stem}.tif")
 97        gt_paths.append(gt_path)
 98        if os.path.exists(gt_path):
 99            continue
100
101        # The raw masks are RGB images with black background and a fixed color (53, 119, 181)
102        # marking the caries region. We collapse this to a single-channel binary label map,
103        # where 0 is background and 1 marks caries.
104        raw_gt = imageio.imread(raw_gt_path)
105        binary_gt = (raw_gt.sum(axis=-1) > 0).astype(np.uint8)
106        imageio.imwrite(gt_path, binary_gt)
107
108    return image_paths, gt_paths
109
110
111def get_dprd_dataset(
112    path: Union[os.PathLike, str],
113    patch_shape: Tuple[int, int],
114    split: Literal["train", "test"],
115    resize_inputs: bool = False,
116    download: bool = False,
117    **kwargs
118) -> Dataset:
119    """Get the DPRD dataset for caries segmentation in panoramic dental radiographs.
120
121    Args:
122        path: Filepath to a folder where the data is downloaded for further processing.
123        patch_shape: The patch shape to use for training.
124        split: The data split to use. Either 'train' or 'test'.
125        resize_inputs: Whether to resize the inputs to the patch shape.
126        download: Whether to download the data if it is not present.
127        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
128
129    Returns:
130        The segmentation dataset.
131    """
132    image_paths, gt_paths = get_dprd_paths(path, split, download)
133
134    if resize_inputs:
135        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
136        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
137            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
138        )
139
140    return torch_em.default_segmentation_dataset(
141        raw_paths=image_paths,
142        raw_key=None,
143        label_paths=gt_paths,
144        label_key=None,
145        is_seg_dataset=False,
146        patch_shape=patch_shape,
147        **kwargs
148    )
149
150
151def get_dprd_loader(
152    path: Union[os.PathLike, str],
153    batch_size: int,
154    patch_shape: Tuple[int, int],
155    split: Literal["train", "test"],
156    resize_inputs: bool = False,
157    download: bool = False,
158    **kwargs
159) -> DataLoader:
160    """Get the DPRD dataloader for caries segmentation in panoramic dental radiographs.
161
162    Args:
163        path: Filepath to a folder where the data is downloaded for further processing.
164        batch_size: The batch size for training.
165        patch_shape: The patch shape to use for training.
166        split: The data split to use. Either 'train' or 'test'.
167        resize_inputs: Whether to resize the inputs to the patch shape.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
170
171    Returns:
172        The DataLoader.
173    """
174    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
175    dataset = get_dprd_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
176    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)