torch_em.data.datasets.medical.tn3k

  1"""The TN3K dataset contains annotations for thyroid nodule segmentation in ultrasound images.
  2
  3The dataset consists of 3493 thyroid ultrasound images (2879 in the 'trainval' split and 614 in the
  4'test' split) collected from 2421 patients across multiple devices and views, with pixel-level nodule
  5masks. The masks are binary: 0 for background and 1 for thyroid nodule.
  6
  7The dataset is hosted at https://github.com/haifangong/TRFE-Net-for-thyroid-nodule-segmentation, which
  8bundles it together with the TG3K (thyroid gland) dataset and a copy of DDTI inside a single archive.
  9We download it from the Google Drive mirror linked in that repository's README.
 10
 11This dataset is from the publications https://doi.org/10.1016/j.compbiomed.2022.106389 and
 12https://doi.org/10.1109/ISBI48211.2021.9434087. Please cite them if you use this dataset for your research.
 13"""
 14
 15import os
 16from glob import glob
 17from natsort import natsorted
 18from typing import Union, Tuple, Literal, List
 19
 20import imageio.v3 as imageio
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29GDRIVE_ID = "1reHyY5eTZ5uePXMVMzFOq5j3eFOSp50F"
 30CHECKSUM = "0ba1770076dab01b1f8fd661a227d80982168bc433de61c80fd266427b20cf60"
 31
 32
 33def get_tn3k_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the TN3K dataset.
 35
 36    Args:
 37        path: Filepath to a folder where the data is downloaded for further processing.
 38        download: Whether to download the data if it is not present.
 39
 40    Returns:
 41        Filepath where the data is downloaded.
 42    """
 43    data_dir = os.path.join(path, "Thyroid Dataset", "tn3k")
 44    if os.path.exists(data_dir):
 45        return data_dir
 46
 47    os.makedirs(path, exist_ok=True)
 48
 49    zip_path = os.path.join(path, "tn3k.zip")
 50    util.download_source_gdrive(
 51        path=zip_path, url=f"https://drive.google.com/uc?id={GDRIVE_ID}", download=download, checksum=CHECKSUM,
 52    )
 53    util.unzip(zip_path=zip_path, dst=path)
 54
 55    return data_dir
 56
 57
 58def get_tn3k_paths(
 59    path: Union[os.PathLike, str], split: Literal["trainval", "test"], download: bool = False
 60) -> Tuple[List[str], List[str]]:
 61    """Get paths to the TN3K data.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        split: The choice of data split. Either 'trainval' or 'test'.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        List of filepaths for the image data.
 70        List of filepaths for the label data.
 71    """
 72    data_dir = get_tn3k_data(path=path, download=download)
 73
 74    if split not in ["trainval", "test"]:
 75        raise ValueError(f"'{split}' is not a valid split.")
 76
 77    image_paths = natsorted(glob(os.path.join(data_dir, f"{split}-image", "*.jpg")))
 78    mask_paths = natsorted(glob(os.path.join(data_dir, f"{split}-mask", "*.jpg")))
 79
 80    if len(image_paths) == 0 or len(image_paths) != len(mask_paths):
 81        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 82
 83    neu_gt_dir = os.path.join(data_dir, f"{split}-mask", "preprocessed")
 84    os.makedirs(neu_gt_dir, exist_ok=True)
 85
 86    gt_paths = []
 87    for mask_path in mask_paths:
 88        gt_path = os.path.join(neu_gt_dir, os.path.basename(mask_path).replace(".jpg", ".tif"))
 89        gt_paths.append(gt_path)
 90        if os.path.exists(gt_path):
 91            continue
 92
 93        # The masks are stored as jpegs, which introduces compression artifacts around otherwise binary values.
 94        mask = (imageio.imread(mask_path) > 127).astype("uint8")
 95        imageio.imwrite(gt_path, mask, compression="zlib")
 96
 97    return image_paths, gt_paths
 98
 99
100def get_tn3k_dataset(
101    path: Union[os.PathLike, str],
102    patch_shape: Tuple[int, int],
103    split: Literal["trainval", "test"],
104    resize_inputs: bool = False,
105    download: bool = False,
106    **kwargs
107) -> Dataset:
108    """Get the TN3K dataset for thyroid nodule segmentation.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        patch_shape: The patch shape to use for training.
113        split: The choice of data split. Either 'trainval' or 'test'.
114        resize_inputs: Whether to resize the inputs.
115        download: Whether to download the data if it is not present.
116        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
117
118    Returns:
119        The segmentation dataset.
120    """
121    image_paths, gt_paths = get_tn3k_paths(path, split, download)
122
123    if resize_inputs:
124        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
125        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
126            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
127        )
128
129    return torch_em.default_segmentation_dataset(
130        raw_paths=image_paths,
131        raw_key=None,
132        label_paths=gt_paths,
133        label_key=None,
134        patch_shape=patch_shape,
135        is_seg_dataset=False,
136        **kwargs
137    )
138
139
140def get_tn3k_loader(
141    path: Union[os.PathLike, str],
142    batch_size: int,
143    patch_shape: Tuple[int, int],
144    split: Literal["trainval", "test"],
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs
148) -> DataLoader:
149    """Get the TN3K dataloader for thyroid nodule segmentation.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        batch_size: The batch size for training.
154        patch_shape: The patch shape to use for training.
155        split: The choice of data split. Either 'trainval' or 'test'.
156        resize_inputs: Whether to resize the inputs.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
159
160    Returns:
161        The DataLoader.
162    """
163    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
164    dataset = get_tn3k_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
165    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)