torch_em.data.datasets.medical.bus_uclm

  1"""The BUS-UCLM dataset contains annotations for breast lesion segmentation in ultrasound images.
  2
  3The dataset comprises breast ultrasound images from 38 patients, acquired at the Ciudad Real
  4General University Hospital with a Siemens ACUSON S2000 Ultrasound System between 2022 and 2023.
  5The ground-truth is provided as RGB masks where green denotes benign lesions, red denotes
  6malignant lesions and black denotes background (including normal images without any lesion).
  7The label ids are: 0 = background, 1 = benign, 2 = malignant.
  8
  9This dataset is located at https://data.mendeley.com/datasets/7fvgj4jsp7/3 (CC BY 4.0). We use the
 10mirror at https://www.kaggle.com/datasets/orvile/bus-uclm-breast-ultrasound-dataset for the download,
 11as the Mendeley download links are unreliable.
 12
 13This dataset is from the publication https://doi.org/10.1038/s41597-025-04562-3.
 14Please cite it if you use this dataset for your research.
 15"""
 16
 17import os
 18from glob import glob
 19from tqdm import tqdm
 20from pathlib import Path
 21from natsort import natsorted
 22from typing import Union, Tuple, List
 23
 24import numpy as np
 25import imageio.v3 as imageio
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34KAGGLE_DATASET_NAME = "orvile/bus-uclm-breast-ultrasound-dataset"
 35
 36LABEL_COLORS = {
 37    0: (0, 0, 0),  # background / normal
 38    1: (0, 255, 0),  # benign
 39    2: (255, 0, 0),  # malignant
 40}
 41
 42
 43def get_bus_uclm_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 44    """Download the BUS-UCLM dataset.
 45
 46    Args:
 47        path: Filepath to a folder where the data is downloaded for further processing.
 48        download: Whether to download the data if it is not present.
 49
 50    Returns:
 51        Filepath where the data is downloaded.
 52    """
 53    data_dirs = glob(os.path.join(path, "**", "BUS-UCLM"), recursive=True)
 54    if data_dirs:
 55        return data_dirs[0]
 56
 57    os.makedirs(path, exist_ok=True)
 58
 59    util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 60
 61    zip_path = os.path.join(path, "bus-uclm-breast-ultrasound-dataset.zip")
 62    util.unzip(zip_path=zip_path, dst=path)
 63
 64    data_dirs = glob(os.path.join(path, "**", "BUS-UCLM"), recursive=True)
 65    if not data_dirs:
 66        raise RuntimeError(f"The dataset could not be found at '{path}' after extraction.")
 67
 68    return data_dirs[0]
 69
 70
 71def get_bus_uclm_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 72    """Get paths to the BUS-UCLM data.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image data.
 80        List of filepaths for the label data.
 81    """
 82    data_dir = get_bus_uclm_data(path=path, download=download)
 83
 84    image_paths = natsorted(glob(os.path.join(data_dir, "images", "*.png")))
 85    mask_paths = natsorted(glob(os.path.join(data_dir, "masks", "*.png")))
 86
 87    if len(image_paths) == 0 or len(image_paths) != len(mask_paths):
 88        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 89
 90    neu_gt_dir = os.path.join(data_dir, "masks", "preprocessed")
 91    os.makedirs(neu_gt_dir, exist_ok=True)
 92
 93    reference_colors = np.array(list(LABEL_COLORS.values()))
 94
 95    gt_paths = []
 96    for mask_path in tqdm(mask_paths, desc="Preprocessing labels"):
 97        gt_path = os.path.join(neu_gt_dir, f"{Path(mask_path).stem}.tif")
 98        gt_paths.append(gt_path)
 99        if os.path.exists(gt_path):
100            continue
101
102        mask = imageio.imread(mask_path)[..., :3].astype("float32")
103        distances = np.linalg.norm(mask[..., None, :] - reference_colors[None, None, :, :], axis=-1)
104        semantic_gt = np.argmin(distances, axis=-1).astype("uint8")
105        imageio.imwrite(gt_path, semantic_gt, compression="zlib")
106
107    return image_paths, gt_paths
108
109
110def get_bus_uclm_dataset(
111    path: Union[os.PathLike, str],
112    patch_shape: Tuple[int, int],
113    resize_inputs: bool = False,
114    download: bool = False,
115    **kwargs
116) -> Dataset:
117    """Get the BUS-UCLM dataset for breast lesion segmentation.
118
119    Args:
120        path: Filepath to a folder where the data is downloaded for further processing.
121        patch_shape: The patch shape to use for training.
122        resize_inputs: Whether to resize the inputs.
123        download: Whether to download the data if it is not present.
124        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
125
126    Returns:
127        The segmentation dataset.
128    """
129    image_paths, gt_paths = get_bus_uclm_paths(path, download)
130
131    if resize_inputs:
132        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
133        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
134            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
135        )
136
137    return torch_em.default_segmentation_dataset(
138        raw_paths=image_paths,
139        raw_key=None,
140        label_paths=gt_paths,
141        label_key=None,
142        patch_shape=patch_shape,
143        is_seg_dataset=False,
144        **kwargs
145    )
146
147
148def get_bus_uclm_loader(
149    path: Union[os.PathLike, str],
150    batch_size: int,
151    patch_shape: Tuple[int, int],
152    resize_inputs: bool = False,
153    download: bool = False,
154    **kwargs
155) -> DataLoader:
156    """Get the BUS-UCLM dataloader for breast lesion segmentation.
157
158    Args:
159        path: Filepath to a folder where the data is downloaded for further processing.
160        batch_size: The batch size for training.
161        patch_shape: The patch shape to use for training.
162        resize_inputs: Whether to resize the inputs.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
165
166    Returns:
167        The DataLoader.
168    """
169    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
170    dataset = get_bus_uclm_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
171    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)