torch_em.data.datasets.medical.ctrus

  1"""C-TRUS is the Colon Wall Segmentation in Transabdominal Ultrasound dataset, with
  2annotations for colon wall segmentation in transabdominal ultrasound images of patients
  3with ulcerative colitis.
  4
  5The dataset is located at https://github.com/wwu-mmll/c-trus (no explicit license is
  6stated in the repository).
  7This dataset is from the publication https://doi.org/10.1007/978-3-031-73647-6_10.
  8Please cite it if you use this dataset for your research.
  9
 10NOTE: The labels are stored as JPEG images, so the (originally binary) colon wall masks
 11have lossy compression artifacts near the mask boundaries. This module binarizes them
 12with a fixed intensity threshold when caching the labels to disk.
 13"""
 14
 15import os
 16from glob import glob
 17from tqdm import tqdm
 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
 29URL = "https://github.com/wwu-mmll/c-trus/archive/ad3ce4d4bbc4b89792f757c7aeb83f31f9a229bd.zip"
 30CHECKSUM = "e584471c8a1340de1ca4180e50469aba1a3ecb554099c5df9aa56f57a9110245"
 31
 32
 33def get_ctrus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the C-TRUS 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 to the folder with the downloaded images and colon wall annotations.
 42    """
 43    data_dir = os.path.join(path, "c-trus-ad3ce4d4bbc4b89792f757c7aeb83f31f9a229bd")
 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, "c-trus.zip")
 50    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 51    util.unzip(zip_path=zip_path, dst=path)
 52
 53    return data_dir
 54
 55
 56def get_ctrus_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 57    """Get paths to the C-TRUS data.
 58
 59    Args:
 60        path: Filepath to a folder where the data is downloaded for further processing.
 61        download: Whether to download the data if it is not present.
 62
 63    Returns:
 64        List of filepaths for the image data.
 65        List of filepaths for the label data.
 66    """
 67    data_dir = get_ctrus_data(path=path, download=download)
 68
 69    image_paths = sorted(glob(os.path.join(data_dir, "original", "*.jpg")))
 70
 71    label_dir = os.path.join(data_dir, "labels_binary")
 72    os.makedirs(label_dir, exist_ok=True)
 73
 74    gt_paths = []
 75    for image_path in tqdm(image_paths, desc="Preprocessing C-TRUS labels"):
 76        fname = os.path.splitext(os.path.basename(image_path))[0]
 77        gt_path = os.path.join(label_dir, f"{fname}.tif")
 78        gt_paths.append(gt_path)
 79        if os.path.exists(gt_path):
 80            continue
 81
 82        label_path = os.path.join(data_dir, "labels", f"{fname}.jpg")
 83        label = imageio.imread(label_path)
 84        label = (label > 127).astype("uint8")
 85        imageio.imwrite(gt_path, label, compression="zlib")
 86
 87    return image_paths, gt_paths
 88
 89
 90def get_ctrus_dataset(
 91    path: Union[os.PathLike, str],
 92    patch_shape: Tuple[int, int],
 93    resize_inputs: bool = False,
 94    download: bool = False,
 95    **kwargs
 96) -> Dataset:
 97    """Get the C-TRUS dataset for colon wall segmentation.
 98
 99    Args:
100        path: Filepath to a folder where the data is downloaded for further processing.
101        patch_shape: The patch shape to use for training.
102        resize_inputs: Whether to resize the inputs to the patch shape.
103        download: Whether to download the data if it is not present.
104        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
105
106    Returns:
107        The segmentation dataset.
108    """
109    image_paths, gt_paths = get_ctrus_paths(path, download)
110
111    if resize_inputs:
112        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
113        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
114            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
115        )
116
117    return torch_em.default_segmentation_dataset(
118        raw_paths=image_paths,
119        raw_key=None,
120        label_paths=gt_paths,
121        label_key=None,
122        patch_shape=patch_shape,
123        is_seg_dataset=False,
124        **kwargs
125    )
126
127
128def get_ctrus_loader(
129    path: Union[os.PathLike, str],
130    patch_shape: Tuple[int, int],
131    batch_size: int,
132    resize_inputs: bool = False,
133    download: bool = False,
134    **kwargs
135) -> DataLoader:
136    """Get the C-TRUS dataloader for colon wall segmentation.
137
138    Args:
139        path: Filepath to a folder where the data is downloaded for further processing.
140        patch_shape: The patch shape to use for training.
141        batch_size: The batch size for training.
142        resize_inputs: Whether to resize the inputs to the patch shape.
143        download: Whether to download the data if it is not present.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
145
146    Returns:
147        The DataLoader.
148    """
149    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
150    dataset = get_ctrus_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
151    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)