torch_em.data.datasets.medical.us_nerve

  1"""The US Nerve dataset contains annotations for segmentation of the Brachial Plexus (BP)
  2nerve structure in ultrasound images of the neck.
  3
  4NOTE: This dataset requires the Kaggle API. You need to install it via 'pip install kaggle'
  5and set up an API token, see https://www.kaggle.com/docs/api. You also need to accept the
  6competition rules on the Kaggle website (https://www.kaggle.com/c/ultrasound-nerve-segmentation/rules)
  7before the download will succeed.
  8
  9NOTE: Not all training images contain the nerve structure. For images where the Brachial
 10Plexus is not visible, the corresponding mask is empty (all-background).
 11
 12The dataset is located at https://www.kaggle.com/c/ultrasound-nerve-segmentation.
 13This dataset is from the "Ultrasound Nerve Segmentation" Kaggle competition, hosted by Kensho.
 14Please cite it if you use this dataset for your research.
 15"""
 16
 17import os
 18from glob import glob
 19from natsort import natsorted
 20from typing import Union, Tuple, List
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29def get_us_nerve_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the US Nerve dataset.
 31
 32    Args:
 33        path: Filepath to a folder where the data is downloaded for further processing.
 34        download: Whether to download the data if it is not present.
 35
 36    Returns:
 37        Filepath where the data is downloaded.
 38    """
 39    data_dir = os.path.join(path, "train")
 40    if os.path.exists(data_dir):
 41        return path
 42
 43    os.makedirs(path, exist_ok=True)
 44
 45    util.download_source_kaggle(
 46        path=path, dataset_name="ultrasound-nerve-segmentation", download=download, competition=True
 47    )
 48
 49    zip_path = os.path.join(path, "ultrasound-nerve-segmentation.zip")
 50    util.unzip(zip_path=zip_path, dst=path)
 51
 52    # The competition bundle ships 'train' and 'test' as nested zip archives.
 53    for name in ["train", "test"]:
 54        nested_zip = os.path.join(path, f"{name}.zip")
 55        if os.path.exists(nested_zip):
 56            util.unzip(zip_path=nested_zip, dst=path)
 57
 58    return path
 59
 60
 61def get_us_nerve_paths(
 62    path: Union[os.PathLike, str], download: bool = False
 63) -> Tuple[List[str], List[str]]:
 64    """Get paths to the US Nerve data.
 65
 66    Args:
 67        path: Filepath to a folder where the data is downloaded for further processing.
 68        download: Whether to download the data if it is not present.
 69
 70    Returns:
 71        List of filepaths for the image data.
 72        List of filepaths for the label data.
 73    """
 74    data_dir = get_us_nerve_data(path=path, download=download)
 75
 76    image_paths = natsorted([
 77        p for p in glob(os.path.join(data_dir, "train", "*.tif")) if not p.endswith("_mask.tif")
 78    ])
 79    gt_paths = natsorted(glob(os.path.join(data_dir, "train", "*_mask.tif")))
 80
 81    assert len(image_paths) == len(gt_paths), f"{len(image_paths)} != {len(gt_paths)}"
 82
 83    return image_paths, gt_paths
 84
 85
 86def get_us_nerve_dataset(
 87    path: Union[os.PathLike, str],
 88    patch_shape: Tuple[int, int],
 89    resize_inputs: bool = False,
 90    download: bool = False,
 91    **kwargs
 92) -> Dataset:
 93    """Get the US Nerve dataset for brachial plexus nerve segmentation.
 94
 95    Args:
 96        path: Filepath to a folder where the data is downloaded for further processing.
 97        patch_shape: The patch shape to use for training.
 98        resize_inputs: Whether to resize the inputs to the patch shape.
 99        download: Whether to download the data if it is not present.
100        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
101
102    Returns:
103        The segmentation dataset.
104    """
105    image_paths, gt_paths = get_us_nerve_paths(path, download)
106
107    if resize_inputs:
108        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
109        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
110            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
111        )
112
113    return torch_em.default_segmentation_dataset(
114        raw_paths=image_paths,
115        raw_key=None,
116        label_paths=gt_paths,
117        label_key=None,
118        patch_shape=patch_shape,
119        is_seg_dataset=False,
120        **kwargs
121    )
122
123
124def get_us_nerve_loader(
125    path: Union[os.PathLike, str],
126    patch_shape: Tuple[int, int],
127    batch_size: int,
128    resize_inputs: bool = False,
129    download: bool = False,
130    **kwargs
131) -> DataLoader:
132    """Get the US Nerve dataloader for brachial plexus nerve segmentation.
133
134    Args:
135        path: Filepath to a folder where the data is downloaded for further processing.
136        patch_shape: The patch shape to use for training.
137        batch_size: The batch size for training.
138        resize_inputs: Whether to resize the inputs to the patch shape.
139        download: Whether to download the data if it is not present.
140        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
141
142    Returns:
143        The DataLoader.
144    """
145    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
146    dataset = get_us_nerve_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
147    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)