torch_em.data.datasets.medical.tg3k

  1"""The TG3K dataset contains annotations for thyroid gland segmentation in ultrasound images.
  2
  3The dataset consists of 3585 thyroid ultrasound frames with pixel-level thyroid gland masks. The
  4masks are binary: 0 for background and 1 for thyroid gland. Unlike TN3K, which is bundled in the same
  5archive, TG3K labels the gland region instead of individual nodules and has no predefined train/test split.
  6
  7The dataset is hosted at https://github.com/haifangong/TRFE-Net-for-thyroid-nodule-segmentation, which
  8bundles it together with the TN3K (thyroid nodule) 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, 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_tg3k_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the TG3K 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", "tg3k")
 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, "tg3k.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_tg3k_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 59    """Get paths to the TG3K data.
 60
 61    Args:
 62        path: Filepath to a folder where the data is downloaded for further processing.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    data_dir = get_tg3k_data(path=path, download=download)
 70
 71    image_paths = natsorted(glob(os.path.join(data_dir, "thyroid-image", "*.jpg")))
 72    mask_paths = natsorted(glob(os.path.join(data_dir, "thyroid-mask", "*.jpg")))
 73
 74    if len(image_paths) == 0 or len(image_paths) != len(mask_paths):
 75        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 76
 77    neu_gt_dir = os.path.join(data_dir, "thyroid-mask", "preprocessed")
 78    os.makedirs(neu_gt_dir, exist_ok=True)
 79
 80    gt_paths = []
 81    for mask_path in mask_paths:
 82        gt_path = os.path.join(neu_gt_dir, os.path.basename(mask_path).replace(".jpg", ".tif"))
 83        gt_paths.append(gt_path)
 84        if os.path.exists(gt_path):
 85            continue
 86
 87        # The masks are stored as jpegs, which introduces compression artifacts around otherwise binary values.
 88        mask = (imageio.imread(mask_path) > 127).astype("uint8")
 89        imageio.imwrite(gt_path, mask, compression="zlib")
 90
 91    return image_paths, gt_paths
 92
 93
 94def get_tg3k_dataset(
 95    path: Union[os.PathLike, str],
 96    patch_shape: Tuple[int, int],
 97    resize_inputs: bool = False,
 98    download: bool = False,
 99    **kwargs
100) -> Dataset:
101    """Get the TG3K dataset for thyroid gland segmentation.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        patch_shape: The patch shape to use for training.
106        resize_inputs: Whether to resize the inputs.
107        download: Whether to download the data if it is not present.
108        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
109
110    Returns:
111        The segmentation dataset.
112    """
113    image_paths, gt_paths = get_tg3k_paths(path, download)
114
115    if resize_inputs:
116        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
117        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
118            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
119        )
120
121    return torch_em.default_segmentation_dataset(
122        raw_paths=image_paths,
123        raw_key=None,
124        label_paths=gt_paths,
125        label_key=None,
126        patch_shape=patch_shape,
127        is_seg_dataset=False,
128        **kwargs
129    )
130
131
132def get_tg3k_loader(
133    path: Union[os.PathLike, str],
134    batch_size: int,
135    patch_shape: Tuple[int, int],
136    resize_inputs: bool = False,
137    download: bool = False,
138    **kwargs
139) -> DataLoader:
140    """Get the TG3K dataloader for thyroid gland segmentation.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        batch_size: The batch size for training.
145        patch_shape: The patch shape to use for training.
146        resize_inputs: Whether to resize the inputs.
147        download: Whether to download the data if it is not present.
148        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
149
150    Returns:
151        The DataLoader.
152    """
153    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
154    dataset = get_tg3k_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
155    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)