torch_em.data.datasets.medical.optic_nerve_sheaths

  1"""The OpticNerveSheaths dataset contains annotations for optic nerve and optic nerve
  2sheath segmentation in transorbital ultrasound images.
  3
  4The dataset contains 464 B-mode transorbital ultrasound images collected on a multidevice,
  5multicenter cohort, with pixel-level annotations of the optic nerve and its sheath, used to
  6estimate the optic nerve sheath diameter (ONSD), a marker of increased intracranial pressure.
  7
  8This dataset is located at https://doi.org/10.17632/kw8gvp8m8x.1 (CC BY 4.0).
  9This dataset is from the publication https://doi.org/10.1016/j.ultrasmedbio.2023.05.011.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14from glob import glob
 15from typing import Union, Tuple, List
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URL = "https://data.mendeley.com/public-files/datasets/kw8gvp8m8x/files/2cc2372d-9edb-4b61-aa84-9513a0156cf0/file_downloaded"  # noqa
 25CHECKSUM = "f8fc47b345462aa585e511bd73cf66eba05d317c2bbc86025c8b875f9b4fdcff"
 26
 27
 28def get_optic_nerve_sheaths_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 29    """Download the OpticNerveSheaths dataset.
 30
 31    Args:
 32        path: Filepath to a folder where the data is downloaded for further processing.
 33        download: Whether to download the data if it is not present.
 34
 35    Returns:
 36        Filepath where the data is downloaded.
 37    """
 38    data_dir = os.path.join(path, "Ultrasound-OpticNerveSheaths")
 39    if os.path.exists(data_dir):
 40        return data_dir
 41
 42    os.makedirs(path, exist_ok=True)
 43
 44    zip_path = os.path.join(path, "Ultrasound-OpticNerveSheaths.zip")
 45    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 46    util.unzip(zip_path=zip_path, dst=path)
 47
 48    return data_dir
 49
 50
 51def get_optic_nerve_sheaths_paths(
 52    path: Union[os.PathLike, str], download: bool = False
 53) -> Tuple[List[str], List[str]]:
 54    """Get paths to the OpticNerveSheaths data.
 55
 56    Args:
 57        path: Filepath to a folder where the data is downloaded for further processing.
 58        download: Whether to download the data if it is not present.
 59
 60    Returns:
 61        List of filepaths for the image data.
 62        List of filepaths for the label data.
 63    """
 64    data_dir = get_optic_nerve_sheaths_data(path=path, download=download)
 65
 66    image_paths = sorted(glob(os.path.join(data_dir, "DATA", "IMAGES_256", "*.png")))
 67    gt_paths = sorted(glob(os.path.join(data_dir, "DATA", "LABELS_256", "*.png")))
 68
 69    if len(image_paths) == 0 or len(image_paths) != len(gt_paths):
 70        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 71
 72    return image_paths, gt_paths
 73
 74
 75def get_optic_nerve_sheaths_dataset(
 76    path: Union[os.PathLike, str],
 77    patch_shape: Tuple[int, int],
 78    resize_inputs: bool = False,
 79    download: bool = False,
 80    **kwargs
 81) -> Dataset:
 82    """Get the OpticNerveSheaths dataset for optic nerve and sheath segmentation.
 83
 84    Args:
 85        path: Filepath to a folder where the data is downloaded for further processing.
 86        patch_shape: The patch shape to use for training.
 87        resize_inputs: Whether to resize the inputs.
 88        download: Whether to download the data if it is not present.
 89        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 90
 91    Returns:
 92        The segmentation dataset.
 93    """
 94    image_paths, gt_paths = get_optic_nerve_sheaths_paths(path, download)
 95
 96    if resize_inputs:
 97        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
 98        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
 99            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
100        )
101
102    return torch_em.default_segmentation_dataset(
103        raw_paths=image_paths,
104        raw_key=None,
105        label_paths=gt_paths,
106        label_key=None,
107        patch_shape=patch_shape,
108        is_seg_dataset=False,
109        **kwargs
110    )
111
112
113def get_optic_nerve_sheaths_loader(
114    path: Union[os.PathLike, str],
115    batch_size: int,
116    patch_shape: Tuple[int, int],
117    resize_inputs: bool = False,
118    download: bool = False,
119    **kwargs
120) -> DataLoader:
121    """Get the OpticNerveSheaths dataloader for optic nerve and sheath segmentation.
122
123    Args:
124        path: Filepath to a folder where the data is downloaded for further processing.
125        batch_size: The batch size for training.
126        patch_shape: The patch shape to use for training.
127        resize_inputs: Whether to resize the inputs.
128        download: Whether to download the data if it is not present.
129        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
130
131    Returns:
132        The DataLoader.
133    """
134    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
135    dataset = get_optic_nerve_sheaths_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
136    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)