torch_em.data.datasets.medical.abus

  1"""The ABUS dataset contains annotations for breast cancer segmentation in ultrasound images.
  2
  3This dataset is located at https://www.kaggle.com/datasets/mohammedtgadallah/mt-small-dataset.
  4The dataset is from the publication https://doi.org/10.1371/journal.pone.0251899.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9from glob import glob
 10from natsort import natsorted
 11from typing import Tuple, List, Union, Literal
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20def get_abus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 21    """Download the ABUS dataset.
 22
 23    Args:
 24        path: Filepath to a folder where the data is downloaded for further processing.
 25        download: Whether to download the data if it is not present.
 26
 27    Returns:
 28        Filepath where the data is downloaded.
 29    """
 30    data_dir = os.path.join(path, "MT_Small_Dataset")
 31    if os.path.exists(data_dir):
 32        return data_dir
 33
 34    os.makedirs(path, exist_ok=True)
 35
 36    util.download_source_kaggle(path=path, dataset_name="mohammedtgadallah/mt-small-dataset", download=download)
 37    util.unzip(zip_path=os.path.join(path, "mt-small-dataset.zip"), dst=path)
 38
 39    return data_dir
 40
 41
 42def get_abus_paths(
 43    path: Union[os.PathLike, str],
 44    split: Literal["train", "val", "test"],
 45    category: Literal["benign", "malign"],
 46    image_choice: Literal["raw", "fuzzy"] = "raw",
 47    download: bool = False,
 48) -> Tuple[List[str], List[str]]:
 49    """Get paths to the ABUS data.
 50
 51    Args:
 52        path: Filepath to a folder where the data is downloaded for further processing.
 53        split: The choice of data split.
 54        category: The choice of tumor category.
 55        image_choice: The choice of input data.
 56        download: Whether to download the data if it is not present.
 57
 58    Returns:
 59        List of filepaths for the image data.
 60        List of filepaths for the label data.
 61    """
 62    data_dir = get_abus_data(path, download)
 63
 64    if image_choice not in ["raw", "fuzzy"]:
 65        raise ValueError("Invalid input choice provided.", image_choice)
 66
 67    if split not in ["train", "val", "test"]:
 68        raise ValueError("Invalid split choice provided.")
 69
 70    if category not in ["benign", "malign"]:
 71        raise ValueError("Invalid tumor category provided.")
 72
 73    cname = "Benign" if category == "benign" else "Malignant"
 74    raw_iname = f"Original_{cname}" if image_choice == "raw" else f"Fuzzy_{cname}"
 75    gt_iname = f"Ground_Truth_{cname}"
 76
 77    image_paths = natsorted(glob(os.path.join(data_dir, cname, raw_iname, "*.png")))
 78    gt_paths = natsorted(glob(os.path.join(data_dir, cname, gt_iname, "*.png")))
 79
 80    assert len(image_paths) and len(image_paths) == len(gt_paths)
 81
 82    if split == "train":
 83        image_paths, gt_paths = image_paths[:125], gt_paths[:125]
 84    elif split == "val":
 85        image_paths, gt_paths = image_paths[125:150], gt_paths[125:150]
 86    else:
 87        image_paths, gt_paths = image_paths[150:], gt_paths[150:]
 88
 89    return image_paths, gt_paths
 90
 91
 92def get_abus_dataset(
 93    path: Union[os.PathLike, str],
 94    patch_shape: Tuple[int, int],
 95    category: Literal["benign", "malign"],
 96    split: Literal["train", "val", "test"],
 97    image_choice: Literal["raw", "fuzzy"] = "raw",
 98    resize_inputs: bool = False,
 99    download: bool = False,
100    **kwargs
101) -> Dataset:
102    """Get the ABUS dataset for breast cancer segmentation.
103
104    Args:
105        path: Filepath to a folder where the data is downloaded for further processing.
106        patch_shape: The patch shape to use for training.
107        category: The choice of tumor category.
108        split: The choice of data split.
109        image_choice: The choice of input data.
110        resize_inputs: Whether to resize the inputs.
111        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
112
113    Returns:
114        The segmentation dataset.
115    """
116    image_paths, gt_paths = get_abus_paths(path, split, category, image_choice, download)
117
118    if resize_inputs:
119        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
120        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
121            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
122        )
123
124    return torch_em.default_segmentation_dataset(
125        raw_paths=image_paths,
126        raw_key=None,
127        label_paths=gt_paths,
128        label_key=None,
129        patch_shape=patch_shape,
130        ndim=2,
131        with_channels=True,
132        is_seg_dataset=False,
133        **kwargs
134    )
135
136
137def get_abus_loader(
138    path: Union[os.PathLike, str],
139    batch_size: int,
140    patch_shape: Tuple[int, int],
141    category: Literal["benign", "malign"],
142    split: Literal["train", "val", "test"],
143    image_choice: Literal["raw", "fuzzy"] = "raw",
144    resize_inputs: bool = False,
145    download: bool = False,
146    **kwargs
147) -> DataLoader:
148    """Get the ABUS dataloader for breast cancer segmentation.
149
150    Args:
151        path: Filepath to a folder where the data is downloaded for further processing.
152        batch_size: The batch size for training.
153        patch_shape: The patch shape to use for training.
154        category: The choice of tumor category.
155        split: The choice of data split.
156        image_choice: The choice of input data.
157        resize_inputs: Whether to resize the inputs.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
159
160    Returns:
161        The DataLoader.
162    """
163    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
164    dataset = get_abus_dataset(path, patch_shape, category, split, image_choice, resize_inputs, download, **ds_kwargs)
165    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)