torch_em.data.datasets.histopathology.tevg

  1"""The TEVG dataset contains annotations for microvascular segmentation in H&E stained
  2histology images of tissue-engineered vascular grafts (TEVGs) explanted from sheep carotid
  3arteries. Each patch is labeled with a semantic mask of 9 histological classes: arteriole
  4lumen, arteriole media, arteriole adventitia, venule lumen, venule wall, capillary lumen,
  5capillary wall, immune cells, and nerve trunks.
  6
  7The dataset is located at https://doi.org/10.5281/zenodo.10838384 under the CC-BY-4.0 license.
  8This dataset is from the publication https://doi.org/10.3389/fbioe.2024.1411680.
  9Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13from glob import glob
 14from natsort import natsorted
 15from typing import List, Literal, Tuple, Union
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URLS = {
 25    1: "https://zenodo.org/records/10838384/files/fold_1.zip",
 26    2: "https://zenodo.org/records/10838384/files/fold_2.zip",
 27    3: "https://zenodo.org/records/10838384/files/fold_3.zip",
 28    4: "https://zenodo.org/records/10838384/files/fold_4.zip",
 29    5: "https://zenodo.org/records/10838384/files/fold_5.zip",
 30}
 31CHECKSUMS = {
 32    1: "96a8be9ed361d2658670e5bde36e406a415ef5a9c30df39442265dbdb0e667a0",
 33    2: None,
 34    3: None,
 35    4: None,
 36    5: None,
 37}
 38
 39
 40def get_tevg_data(path: Union[os.PathLike, str], fold: Literal[1, 2, 3, 4, 5], download: bool = False) -> str:
 41    """Download the TEVG dataset for one cross-validation fold.
 42
 43    Args:
 44        path: Filepath to a folder where the downloaded data will be saved.
 45        fold: The choice of cross-validation fold.
 46        download: Whether to download the data if it is not present.
 47
 48    Returns:
 49        The filepath to the folder where the fold data is stored.
 50    """
 51    fold_dir = os.path.join(path, f"fold_{fold}")
 52    if os.path.exists(fold_dir):
 53        return fold_dir
 54
 55    os.makedirs(path, exist_ok=True)
 56
 57    zip_path = os.path.join(path, f"fold_{fold}.zip")
 58    util.download_source(path=zip_path, url=URLS[fold], download=download, checksum=CHECKSUMS[fold])
 59    util.unzip(zip_path=zip_path, dst=path)
 60
 61    return fold_dir
 62
 63
 64def get_tevg_paths(
 65    path: Union[os.PathLike, str],
 66    fold: Literal[1, 2, 3, 4, 5],
 67    split: Literal["train", "test"],
 68    download: bool = False,
 69) -> Tuple[List[str], List[str]]:
 70    """Get paths to the TEVG data.
 71
 72    Args:
 73        path: Filepath to a folder where the downloaded data will be saved.
 74        fold: The choice of cross-validation fold.
 75        split: The choice of data split.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image data.
 80        List of filepaths for the label data.
 81    """
 82    fold_dir = get_tevg_data(path, fold, download)
 83
 84    raw_paths = natsorted(glob(os.path.join(fold_dir, split, "img", "*.jpg")))
 85    label_paths = natsorted(glob(os.path.join(fold_dir, split, "mask", "*.png")))
 86
 87    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 88    assert all(
 89        os.path.splitext(os.path.basename(raw_path))[0] == os.path.splitext(os.path.basename(label_path))[0]
 90        for raw_path, label_path in zip(raw_paths, label_paths)
 91    )
 92
 93    return raw_paths, label_paths
 94
 95
 96def get_tevg_dataset(
 97    path: Union[os.PathLike, str],
 98    patch_shape: Tuple[int, int],
 99    fold: Literal[1, 2, 3, 4, 5] = 1,
100    split: Literal["train", "test"] = "train",
101    resize_inputs: bool = False,
102    download: bool = False,
103    **kwargs,
104) -> Dataset:
105    """Get the TEVG dataset for microvascular segmentation in tissue-engineered vascular grafts.
106
107    Args:
108        path: Filepath to a folder where the downloaded data will be saved.
109        patch_shape: The patch shape to use for training.
110        fold: The choice of cross-validation fold.
111        split: The choice of data split.
112        resize_inputs: Whether to resize the inputs.
113        download: Whether to download the data if it is not present.
114        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
115
116    Returns:
117        The segmentation dataset.
118    """
119    raw_paths, label_paths = get_tevg_paths(path, fold, split, download)
120
121    if resize_inputs:
122        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
123        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
124            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
125        )
126
127    return torch_em.default_segmentation_dataset(
128        raw_paths=raw_paths,
129        raw_key=None,
130        label_paths=label_paths,
131        label_key=None,
132        patch_shape=patch_shape,
133        is_seg_dataset=False,
134        ndim=2,
135        with_channels=True,
136        **kwargs,
137    )
138
139
140def get_tevg_loader(
141    path: Union[os.PathLike, str],
142    batch_size: int,
143    patch_shape: Tuple[int, int],
144    fold: Literal[1, 2, 3, 4, 5] = 1,
145    split: Literal["train", "test"] = "train",
146    resize_inputs: bool = False,
147    download: bool = False,
148    **kwargs,
149) -> DataLoader:
150    """Get the TEVG dataloader for microvascular segmentation in tissue-engineered vascular grafts.
151
152    Args:
153        path: Filepath to a folder where the downloaded data will be saved.
154        batch_size: The batch size for training.
155        patch_shape: The patch shape to use for training.
156        fold: The choice of cross-validation fold.
157        split: The choice of data split.
158        resize_inputs: Whether to resize the inputs.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_tevg_dataset(path, patch_shape, fold, split, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)