torch_em.data.datasets.medical.mbh_seg

  1"""The MBH Seg dataset contains annotations for intracranial hemorrhages
  2in non-contrast CT scans.
  3
  4This dataset is from the MBH-Seg challenge: https://mbh-seg.com
  5- original scans: https://kaggle.com/competitions/rsna-intracranial-hemorrhage-detection
  6Please cite these if you use this dataset for your publication.
  7"""
  8
  9import os
 10from glob import glob
 11from natsort import natsorted
 12from typing import Union, Tuple, List
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URL = "https://huggingface.co/datasets/WuBiao/BHSD/resolve/main/label_192.zip"
 22CHECKSUM = "582bf184af993541a4958a4d209a6a44e3bbe702a5daefaf9fb1733a4e7a6e39"
 23
 24
 25def get_mbh_seg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 26    """Download the MBH Seg dataset.
 27
 28    Args:
 29        path: Filepath to a folder where the data is downloaded for further processing.
 30        download: Whether to download the data if it is not present.
 31
 32    Returns:
 33        Filepath where the data is download.
 34    """
 35    data_dir = os.path.join(path, "label_192")
 36    if os.path.exists(data_dir):
 37        return data_dir
 38
 39    os.makedirs(path, exist_ok=True)
 40
 41    zip_path = os.path.join(path, "label_192.zip")
 42    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 43    util.unzip(zip_path=zip_path, dst=path)
 44
 45    return data_dir
 46
 47
 48def get_mbh_seg_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 49    """Get paths to the MBH Seg data.
 50
 51    Args:
 52        path: Filepath to a folder where the data is downloaded for further processing.
 53        download: Whether to download the data if it is not present.
 54
 55    Returns:
 56        List of filepaths for the image data.
 57        List of filepaths for the label data.
 58    """
 59    data_dir = get_mbh_seg_data(path=path, download=download)
 60    image_paths = natsorted(glob(os.path.join(data_dir, "images", "*.nii.gz")))
 61    gt_paths = natsorted(glob(os.path.join(data_dir, r"ground truths", "*.nii.gz")))
 62    return image_paths, gt_paths
 63
 64
 65def get_mbh_seg_dataset(
 66    path: Union[os.PathLike, str],
 67    patch_shape: Tuple[int, ...],
 68    resize_inputs: bool = False,
 69    download: bool = False,
 70    **kwargs
 71) -> Dataset:
 72    """Get the MBH Seg dataset for intracranial hemorrhage segmentation.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        patch_shape: The patch shape to use for training.
 77        resize_inputs: Whether to resize inputs to the desired patch shape.
 78        download: Whether to download the data if it is not present.
 79        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 80
 81    Returns:
 82        The segmentation dataset.
 83    """
 84    image_paths, gt_paths = get_mbh_seg_paths(path=path, download=download)
 85
 86    if resize_inputs:
 87        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
 88        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
 89            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
 90        )
 91
 92    return torch_em.default_segmentation_dataset(
 93        raw_paths=image_paths,
 94        raw_key="data",
 95        label_paths=gt_paths,
 96        label_key="data",
 97        patch_shape=patch_shape,
 98        **kwargs
 99    )
100
101
102def get_mbh_seg_loader(
103    path: Union[os.PathLike, str],
104    batch_size: int,
105    patch_shape: Tuple[int, ...],
106    resize_inputs: bool = False,
107    download: bool = False,
108    **kwargs
109) -> DataLoader:
110    """Get the MBH Seg dataloader for intracranial hemorrhage segmentation.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded for further processing.
114        batch_size: The batch size for training.
115        patch_shape: The patch shape to use for training.
116        resize_inputs: Whether to resize inputs to the desired patch shape.
117        download: Whether to download the data if it is not present.
118        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
119
120    Returns:
121        The DataLoader.
122    """
123    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
124    dataset = get_mbh_seg_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
125    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)