torch_em.data.datasets.medical.ct2us_kidney

  1"""CT2US for Kidney Segmentation is a dataset that contains synthetic ultrasound images of
  2kidneys with their corresponding kidney segmentation masks.
  3
  4The synthetic ultrasound images are generated from annotated CT slices via CycleGAN-based
  5cross-modal style transfer, in order to construct a large "transition" dataset for training
  6segmentation models on ultrasound data despite the scarcity of annotated real ultrasound
  7kidney images. NOTE: The Kaggle package for this dataset only ships the resulting synthetic
  8ultrasound slices (`slice`) paired with their kidney masks (`mask`); the source CT volumes
  9and the CycleGAN model are not part of this package. We therefore use the synthetic
 10ultrasound images with their masks for the segmentation loader below.
 11
 12The dataset is located at https://www.kaggle.com/datasets/siatsyx/ct2usforkidneyseg.
 13The corresponding code repository is at https://github.com/SIAT-SongYuxin/CT2USforKidneySeg.
 14This dataset is from the publication https://doi.org/10.1016/j.ultras.2022.106706.
 15Please cite it if you use this dataset for your research.
 16"""
 17
 18import os
 19from glob import glob
 20from tqdm import tqdm
 21from pathlib import Path
 22from natsort import natsorted
 23from typing import Union, Tuple, List
 24
 25import imageio.v3 as imageio
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34def get_ct2us_kidney_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the CT2US for Kidney Segmentation dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the data is downloaded for further processing.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        Filepath where the data is downloaded.
 43    """
 44    data_dir = os.path.join(path, "slice", "slice")
 45    if os.path.exists(data_dir):
 46        return path
 47
 48    os.makedirs(path, exist_ok=True)
 49
 50    util.download_source_kaggle(path=path, dataset_name="siatsyx/ct2usforkidneyseg", download=download)
 51    zip_path = os.path.join(path, "ct2usforkidneyseg.zip")
 52    util.unzip(zip_path=zip_path, dst=path)
 53
 54    return path
 55
 56
 57def get_ct2us_kidney_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 58    """Get paths to the CT2US for Kidney Segmentation data.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        download: Whether to download the data if it is not present.
 63
 64    Returns:
 65        List of filepaths for the image data.
 66        List of filepaths for the label data.
 67    """
 68    data_dir = get_ct2us_kidney_data(path=path, download=download)
 69
 70    image_paths = natsorted(glob(os.path.join(data_dir, "slice", "slice", "*.png")))
 71    gt_paths = natsorted(glob(os.path.join(data_dir, "mask", "mask", "*.png")))
 72    assert len(image_paths) > 0 and len(image_paths) == len(gt_paths)
 73
 74    neu_gt_dir = os.path.join(data_dir, "mask", "preprocessed")
 75    os.makedirs(neu_gt_dir, exist_ok=True)
 76
 77    neu_gt_paths = []
 78    for gt_path in tqdm(gt_paths, desc="Preprocessing labels"):
 79        neu_gt_path = os.path.join(neu_gt_dir, f"{Path(gt_path).stem}.tif")
 80        neu_gt_paths.append(neu_gt_path)
 81        if os.path.exists(neu_gt_path):
 82            continue
 83
 84        gt = imageio.imread(gt_path)
 85        gt = (gt >= 128).astype("uint8")
 86        imageio.imwrite(neu_gt_path, gt, compression="zlib")
 87
 88    return image_paths, neu_gt_paths
 89
 90
 91def get_ct2us_kidney_dataset(
 92    path: Union[os.PathLike, str],
 93    patch_shape: Tuple[int, int],
 94    resize_inputs: bool = False,
 95    download: bool = False,
 96    **kwargs
 97) -> Dataset:
 98    """Get the CT2US for Kidney Segmentation dataset for kidney segmentation in synthetic ultrasound images.
 99
100    Args:
101        path: Filepath to a folder where the data is downloaded for further processing.
102        patch_shape: The patch shape to use for training.
103        resize_inputs: Whether to resize the inputs to the patch shape.
104        download: Whether to download the data if it is not present.
105        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
106
107    Returns:
108        The segmentation dataset.
109    """
110    image_paths, gt_paths = get_ct2us_kidney_paths(path, download)
111
112    if resize_inputs:
113        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
114        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
115            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
116        )
117
118    return torch_em.default_segmentation_dataset(
119        raw_paths=image_paths,
120        raw_key=None,
121        label_paths=gt_paths,
122        label_key=None,
123        patch_shape=patch_shape,
124        is_seg_dataset=False,
125        **kwargs
126    )
127
128
129def get_ct2us_kidney_loader(
130    path: Union[os.PathLike, str],
131    batch_size: int,
132    patch_shape: Tuple[int, int],
133    resize_inputs: bool = False,
134    download: bool = False,
135    **kwargs
136) -> DataLoader:
137    """Get the CT2US for Kidney Segmentation dataloader for kidney segmentation in synthetic ultrasound images.
138
139    Args:
140        path: Filepath to a folder where the data is downloaded for further processing.
141        batch_size: The batch size for training.
142        patch_shape: The patch shape to use for training.
143        resize_inputs: Whether to resize the inputs to the patch shape.
144        download: Whether to download the data if it is not present.
145        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
146
147    Returns:
148        The DataLoader.
149    """
150    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
151    dataset = get_ct2us_kidney_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
152    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)