torch_em.data.datasets.medical.bhsd

  1"""The BHSD dataset contains annotations for multiclass intracranial hemorrhage segmentation in
  2brain CT.
  3
  4The Brain Hemorrhage Segmentation Dataset (BHSD) provides 192 non-contrast head CT volumes with
  5pixel-level annotations of five intracranial hemorrhage (ICH) subtypes, see `LABEL_IDS`. The
  6dataset also contains a much larger set of 1980 volumes with only slice-level (i.e. not
  7pixel-level) annotations, which is not exposed by this module.
  8
  9The data is hosted on Hugging Face at https://huggingface.co/datasets/Wendy-Fly/BHSD and is
 10distributed under the MIT license.
 11
 12This dataset is from the publication https://doi.org/10.1007/978-3-031-45673-2_15.
 13Please cite it if you use this dataset for your research.
 14"""
 15
 16import os
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://huggingface.co/datasets/Wendy-Fly/BHSD/resolve/main/label_192.zip"
 29CHECKSUM = "582bf184af993541a4958a4d209a6a44e3bbe702a5daefaf9fb1733a4e7a6e39"
 30
 31LABEL_IDS = {
 32    0: "background", 1: "epidural", 2: "intraparenchymal", 3: "intraventricular", 4: "subarachnoid", 5: "subdural",
 33}
 34"""The label ids of the intracranial hemorrhage subtypes, as defined by the dataset authors."""
 35
 36
 37def get_bhsd_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 38    """Download the BHSD dataset.
 39
 40    Args:
 41        path: Filepath to a folder where the data is downloaded for further processing.
 42        download: Whether to download the data if it is not present.
 43
 44    Returns:
 45        Filepath where the data is downloaded.
 46    """
 47    data_dir = os.path.join(path, "label_192")
 48    if os.path.exists(data_dir):
 49        return data_dir
 50
 51    os.makedirs(path, exist_ok=True)
 52
 53    zip_path = os.path.join(path, "label_192.zip")
 54    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 55    util.unzip(zip_path=zip_path, dst=path)
 56
 57    return data_dir
 58
 59
 60def get_bhsd_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 61    """Get paths to the BHSD data.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        List of filepaths for the image data.
 69        List of filepaths for the label data.
 70    """
 71    data_dir = get_bhsd_data(path, download)
 72
 73    image_paths = natsorted(glob(os.path.join(data_dir, "images", "*.nii.gz")))
 74    gt_paths = natsorted(glob(os.path.join(data_dir, "ground truths", "*.nii.gz")))
 75
 76    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0
 77
 78    return image_paths, gt_paths
 79
 80
 81def get_bhsd_dataset(
 82    path: Union[os.PathLike, str],
 83    patch_shape: Tuple[int, ...],
 84    resize_inputs: bool = False,
 85    download: bool = False,
 86    **kwargs
 87) -> Dataset:
 88    """Get the BHSD dataset for multiclass intracranial hemorrhage segmentation.
 89
 90    Args:
 91        path: Filepath to a folder where the data is downloaded for further processing.
 92        patch_shape: The patch shape to use for training.
 93        resize_inputs: Whether to resize the inputs to the patch shape.
 94        download: Whether to download the data if it is not present.
 95        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 96
 97    Returns:
 98        The segmentation dataset.
 99    """
100    image_paths, gt_paths = get_bhsd_paths(path, download)
101
102    if resize_inputs:
103        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
104        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
105            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
106        )
107
108    return torch_em.default_segmentation_dataset(
109        raw_paths=image_paths,
110        raw_key="data",
111        label_paths=gt_paths,
112        label_key="data",
113        patch_shape=patch_shape,
114        is_seg_dataset=True,
115        **kwargs
116    )
117
118
119def get_bhsd_loader(
120    path: Union[os.PathLike, str],
121    batch_size: int,
122    patch_shape: Tuple[int, ...],
123    resize_inputs: bool = False,
124    download: bool = False,
125    **kwargs
126) -> DataLoader:
127    """Get the BHSD dataloader for multiclass intracranial hemorrhage segmentation.
128
129    Args:
130        path: Filepath to a folder where the data is downloaded for further processing.
131        batch_size: The batch size for training.
132        patch_shape: The patch shape to use for training.
133        resize_inputs: Whether to resize the inputs to the patch shape.
134        download: Whether to download the data if it is not present.
135        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
136
137    Returns:
138        The DataLoader.
139    """
140    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
141    dataset = get_bhsd_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
142    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)