torch_em.data.datasets.medical.busbra

  1"""The BUS-BRA dataset contains annotations for breast tumor segmentation in ultrasound images.
  2
  3The dataset contains 1,875 breast ultrasound images from 1,064 patients, with biopsy-confirmed
  4tumor region masks (benign or malignant). The images were acquired at the National Institute of
  5Cancer (INCA, Brazil) with different ultrasound scanners.
  6
  7This dataset is located at https://zenodo.org/records/8231412 (CC BY 4.0). See also
  8https://github.com/wgomezf/BUS-BRA for further details on the dataset.
  9
 10This dataset is from the publication https://doi.org/10.1002/mp.16812.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from typing import Union, Tuple, List
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://zenodo.org/records/8231412/files/BUSBRA.zip?download=1"
 26CHECKSUM = "ba3e6ed19cc37c682d8d39e25435bbf8a555a12cb7e641b5f2117685c95580ff"
 27
 28
 29def get_busbra_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the BUS-BRA 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, "BUSBRA")
 40    if os.path.exists(data_dir):
 41        return data_dir
 42
 43    os.makedirs(path, exist_ok=True)
 44
 45    zip_path = os.path.join(path, "BUSBRA.zip")
 46    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 47    util.unzip(zip_path=zip_path, dst=path)
 48
 49    return data_dir
 50
 51
 52def get_busbra_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 53    """Get paths to the BUS-BRA data.
 54
 55    Args:
 56        path: Filepath to a folder where the data is downloaded for further processing.
 57        download: Whether to download the data if it is not present.
 58
 59    Returns:
 60        List of filepaths for the image data.
 61        List of filepaths for the label data.
 62    """
 63    data_dir = get_busbra_data(path=path, download=download)
 64
 65    image_paths = sorted(glob(os.path.join(data_dir, "Images", "*.png")))
 66    gt_paths = sorted(glob(os.path.join(data_dir, "Masks", "*.png")))
 67
 68    if len(image_paths) == 0 or len(image_paths) != len(gt_paths):
 69        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 70
 71    return image_paths, gt_paths
 72
 73
 74def get_busbra_dataset(
 75    path: Union[os.PathLike, str],
 76    patch_shape: Tuple[int, int],
 77    resize_inputs: bool = False,
 78    download: bool = False,
 79    **kwargs
 80) -> Dataset:
 81    """Get the BUS-BRA dataset for breast tumor segmentation.
 82
 83    Args:
 84        path: Filepath to a folder where the data is downloaded for further processing.
 85        patch_shape: The patch shape to use for training.
 86        resize_inputs: Whether to resize the inputs.
 87        download: Whether to download the data if it is not present.
 88        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 89
 90    Returns:
 91        The segmentation dataset.
 92    """
 93    image_paths, gt_paths = get_busbra_paths(path, download)
 94
 95    if resize_inputs:
 96        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
 97        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
 98            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
 99        )
100
101    return torch_em.default_segmentation_dataset(
102        raw_paths=image_paths,
103        raw_key=None,
104        label_paths=gt_paths,
105        label_key=None,
106        patch_shape=patch_shape,
107        **kwargs
108    )
109
110
111def get_busbra_loader(
112    path: Union[os.PathLike, str],
113    batch_size: int,
114    patch_shape: Tuple[int, int],
115    resize_inputs: bool = False,
116    download: bool = False,
117    **kwargs
118) -> DataLoader:
119    """Get the BUS-BRA dataloader for breast tumor segmentation.
120
121    Args:
122        path: Filepath to a folder where the data is downloaded for further processing.
123        batch_size: The batch size for training.
124        patch_shape: The patch shape to use for training.
125        resize_inputs: Whether to resize the inputs.
126        download: Whether to download the data if it is not present.
127        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
128
129    Returns:
130        The DataLoader.
131    """
132    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
133    dataset = get_busbra_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
134    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)