torch_em.data.datasets.medical.sts2024

  1"""The STS2024 dataset contains annotations for instance-level tooth segmentation, labeled by
  2FDI tooth id, in panoramic dental X-rays (OPGs).
  3
  4The dataset was curated for the 2nd Semi-supervised Teeth Segmentation (STS 2024) MICCAI challenge
  5(https://sts-challenge.github.io/miccai2024/index.html), which extends the STS 2023 challenge to
  6multi-instance, multi-class (FDI enumeration) instance-level tooth segmentation in both panoramic
  7X-rays and CBCT scans. The full challenge data comprises 2,380 OPGs and 330 CBCT volumes, but only
  830 OPG cases and 30 CBCT cases carry full ground truth (the remaining images support the challenge's
  9semi-supervised learning setting and have no public labels). This module only covers the 30 fully
 10labeled OPG cases, released on Zenodo; the CBCT ground truth is not part of that release.
 11
 12The original annotations are per-tooth polygons (in labelme format) with the FDI tooth id as the
 13label; this module rasterizes them into a per-pixel label map, where the pixel value is the FDI id
 14of the tooth (0 marks background).
 15
 16The data is hosted on Zenodo at https://zenodo.org/records/17712688 and is distributed under the
 17CC BY 4.0 license.
 18
 19The dataset is from the publication https://doi.org/10.1016/j.media.2026.103986.
 20Please cite it if you use this dataset for your research.
 21"""
 22
 23import os
 24import json
 25from glob import glob
 26from tqdm import tqdm
 27from natsort import natsorted
 28from typing import Union, Tuple, List
 29
 30import numpy as np
 31from skimage.draw import polygon
 32
 33import imageio.v3 as imageio
 34
 35from torch.utils.data import Dataset, DataLoader
 36
 37import torch_em
 38
 39from .. import util
 40
 41
 42URL = "https://zenodo.org/records/17712688/files/Train-Labeled.zip?download=1"
 43CHECKSUM = "8d0e1491368f15c770f592d6c9284861629140e6bb550659257a3ff13b4b9761"
 44
 45
 46def get_sts2024_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 47    """Download the STS2024 dataset.
 48
 49    Args:
 50        path: Filepath to a folder where the data is downloaded for further processing.
 51        download: Whether to download the data if it is not present.
 52
 53    Returns:
 54        Filepath where the data is downloaded.
 55    """
 56    data_dir = os.path.join(path, "Train-Labeled")
 57    if os.path.exists(data_dir):
 58        return data_dir
 59
 60    os.makedirs(path, exist_ok=True)
 61    zip_path = os.path.join(path, "Train-Labeled.zip")
 62    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 63    util.unzip(zip_path=zip_path, dst=path)
 64
 65    return data_dir
 66
 67
 68def _rasterize_annotations(shape, annotation_path):
 69    with open(annotation_path) as f:
 70        annotations = json.load(f)
 71
 72    labels = np.zeros(shape, dtype="uint8")
 73    for shape_annotation in annotations["shapes"]:
 74        fdi_id = int(shape_annotation["label"])
 75        points = np.asarray(shape_annotation["points"])
 76        c, r = points[:, 0], points[:, 1]
 77        rr, cc = polygon(r, c, shape=shape)
 78        labels[rr, cc] = fdi_id
 79
 80    return labels
 81
 82
 83def get_sts2024_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 84    """Get paths to the STS2024 data.
 85
 86    Args:
 87        path: Filepath to a folder where the data is downloaded for further processing.
 88        download: Whether to download the data if it is not present.
 89
 90    Returns:
 91        List of filepaths for the image data.
 92        List of filepaths for the label data.
 93    """
 94    data_dir = get_sts2024_data(path, download)
 95
 96    image_dir = os.path.join(data_dir, "Images")
 97    annotation_dir = os.path.join(data_dir, "Masks")
 98    preprocessed_dir = os.path.join(data_dir, "preprocessed")
 99    os.makedirs(preprocessed_dir, exist_ok=True)
100
101    image_paths = natsorted(glob(os.path.join(image_dir, "*.jpg")))
102
103    gt_paths = []
104    for image_path in tqdm(image_paths, desc="Preprocessing labels"):
105        fname = os.path.splitext(os.path.basename(image_path))[0]
106        annotation_path = os.path.join(annotation_dir, f"{fname}_Mask.json")
107        gt_path = os.path.join(preprocessed_dir, f"{fname}.tif")
108
109        if not os.path.exists(gt_path):
110            image_shape = imageio.imread(image_path).shape[:2]
111            labels = _rasterize_annotations(image_shape, annotation_path)
112            imageio.imwrite(gt_path, labels)
113
114        gt_paths.append(gt_path)
115
116    return image_paths, gt_paths
117
118
119def get_sts2024_dataset(
120    path: Union[os.PathLike, str],
121    patch_shape: Tuple[int, int],
122    resize_inputs: bool = False,
123    download: bool = False,
124    **kwargs
125) -> Dataset:
126    """Get the STS2024 dataset for FDI tooth segmentation in panoramic dental radiographs.
127
128    Args:
129        path: Filepath to a folder where the data is downloaded for further processing.
130        patch_shape: The patch shape to use for training.
131        resize_inputs: Whether to resize the inputs to the patch shape.
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
134
135    Returns:
136        The segmentation dataset.
137    """
138    image_paths, gt_paths = get_sts2024_paths(path, download)
139
140    if resize_inputs:
141        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
142        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
143            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
144        )
145
146    return torch_em.default_segmentation_dataset(
147        raw_paths=image_paths,
148        raw_key=None,
149        label_paths=gt_paths,
150        label_key=None,
151        is_seg_dataset=False,
152        patch_shape=patch_shape,
153        **kwargs
154    )
155
156
157def get_sts2024_loader(
158    path: Union[os.PathLike, str],
159    batch_size: int,
160    patch_shape: Tuple[int, int],
161    resize_inputs: bool = False,
162    download: bool = False,
163    **kwargs
164) -> DataLoader:
165    """Get the STS2024 dataloader for FDI tooth segmentation in panoramic dental radiographs.
166
167    Args:
168        path: Filepath to a folder where the data is downloaded for further processing.
169        batch_size: The batch size for training.
170        patch_shape: The patch shape to use for training.
171        resize_inputs: Whether to resize the inputs to the patch shape.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
174
175    Returns:
176        The DataLoader.
177    """
178    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
179    dataset = get_sts2024_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
180    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)