torch_em.data.datasets.medical.hil_toothseg

  1"""The HIL ToothSeg dataset contains annotations for teeth segmentation
  2in panoramic dental radiographs.
  3
  4This dataset is from the publication https://www.mdpi.com/1424-8220/21/9/3110.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9from glob import glob
 10from tqdm import tqdm
 11from pathlib import Path
 12from natsort import natsorted
 13from typing import Union, Literal, Tuple, List
 14
 15import numpy as np
 16import imageio.v3 as imageio
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://hitl-public-datasets.s3.eu-central-1.amazonaws.com/Teeth+Segmentation.zip"
 26CHECKSUM = "3b628165a218a5e8d446d1313e6ecbe7cfc599a3d6418cd60b4fb78745becc2e"
 27
 28
 29def get_hil_toothseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the HIL ToothSeg dataset.
 31
 32    Args:
 33        path: Filepath to a folder where the data is downloaded for further processing.
 34        download: Whether to download the data if it is not present.
 35    """
 36    data_dir = os.path.join(path, r"Teeth Segmentation PNG")
 37    if os.path.exists(data_dir):
 38        return data_dir
 39
 40    os.makedirs(path, exist_ok=True)
 41
 42    zip_path = os.path.join(path, "Teeth_Segmentation.zip")
 43    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 44    util.unzip(zip_path=zip_path, dst=path)
 45
 46    return data_dir
 47
 48
 49def get_hil_toothseg_paths(
 50    path: Union[os.PathLike, str], split: Literal['train', 'val', 'test'], download: bool = False
 51) -> Tuple[List[str], List[str]]:
 52    """Get paths to the HIL ToothSeg data.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        split: The data split to use. Either 'train', 'val' or 'test'.
 57        download: Whether to download the data if it is not present.
 58
 59    Returns:
 60        List of filepaths for the image data.
 61        List of filepaths for the label data.
 62    """
 63    import cv2 as cv
 64
 65    data_dir = get_hil_toothseg_data(path=path, download=download)
 66
 67    image_paths = natsorted(glob(os.path.join(data_dir, "d2", "img", "*")))
 68    gt_paths = natsorted(glob(os.path.join(data_dir, "d2", "masks_machine", "*")))
 69
 70    neu_gt_dir = os.path.join(data_dir, "preprocessed", "gt")
 71    os.makedirs(neu_gt_dir, exist_ok=True)
 72
 73    neu_gt_paths = []
 74    for gt_path in tqdm(gt_paths, desc="Preprocessing inputs"):
 75        neu_gt_path = os.path.join(neu_gt_dir, f"{Path(gt_path).stem}.tif")
 76        neu_gt_paths.append(neu_gt_path)
 77        if os.path.exists(neu_gt_path):
 78            continue
 79
 80        rgb_gt = cv.imread(gt_path)
 81        rgb_gt = cv.cvtColor(rgb_gt, cv.COLOR_BGR2RGB)
 82        incolors = np.unique(rgb_gt.reshape(-1, rgb_gt.shape[2]), axis=0)
 83
 84        # the first id is always background, let's remove it
 85        if np.array_equal(incolors[0], np.array([0, 0, 0])):
 86            incolors = incolors[1:]
 87
 88        instances = np.zeros(rgb_gt.shape[:2])
 89
 90        color_to_id = {tuple(cvalue): i for i, cvalue in enumerate(incolors, start=1)}
 91        for cvalue, idx in color_to_id.items():
 92            binary_map = (rgb_gt == cvalue).all(axis=2)
 93            instances[binary_map] = idx
 94
 95        imageio.imwrite(neu_gt_path, instances)
 96
 97    if split == "train":
 98        image_paths, neu_gt_paths = image_paths[:450], neu_gt_paths[:450]
 99    elif split == "val":
100        image_paths, neu_gt_paths = image_paths[425:475], neu_gt_paths[425:475]
101    elif split == "test":
102        image_paths, neu_gt_paths = image_paths[475:], neu_gt_paths[475:]
103    else:
104        raise ValueError(f"{split} is not a valid split.")
105
106    return image_paths, neu_gt_paths
107
108
109def get_hil_toothseg_dataset(
110    path: Union[os.PathLike, str],
111    patch_shape: Tuple[int, int],
112    split: Literal["train", "val", "test"],
113    resize_inputs: bool = False,
114    download: bool = False,
115    **kwargs
116) -> Dataset:
117    """Get the HIL ToothSeg dataset for teeth segmentation.
118
119    Args:
120        path: Filepath to a folder where the data is downloaded for further processing.
121        patch_shape: The patch shape to use for training.
122        split: The data split to use. Either 'train', 'val' or 'test'.
123        resize_inputs: Whether to resize the inputs to the patch shape.
124        download: Whether to download the data if it is not present.
125        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
126
127    Returns:
128        The segmentation dataset.
129    """
130    image_paths, gt_paths = get_hil_toothseg_paths(path, split, download)
131
132    if resize_inputs:
133        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
134        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
135            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
136        )
137
138    return torch_em.default_segmentation_dataset(
139        raw_paths=image_paths,
140        raw_key=None,
141        label_paths=gt_paths,
142        label_key=None,
143        is_seg_dataset=False,
144        patch_shape=patch_shape,
145        **kwargs
146    )
147
148
149def get_hil_toothseg_loader(
150    path: Union[os.PathLike, str],
151    batch_size: int,
152    patch_shape: Tuple[int, int],
153    split: Literal["train", "val", "test"],
154    resize_inputs: bool = False,
155    download: bool = False,
156    **kwargs
157) -> DataLoader:
158    """Get the HIL ToothSeg dataloader for teeth segmentation.
159
160    Args:
161        path: Filepath to a folder where the data is downloaded for further processing.
162        batch_size: The batch size for training.
163        patch_shape: The patch shape to use for training.
164        split: The data split to use. Either 'train', 'val' or 'test'.
165        resize_inputs: Whether to resize the inputs to the patch shape.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
168
169    Returns:
170        The DataLoader.
171    """
172    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
173    dataset = get_hil_toothseg_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
174    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)