torch_em.data.datasets.medical.pu2756

  1"""The PU2756 dataset contains annotations for pulmonary tumor segmentation in B-mode
  2lung ultrasound images.
  3
  4The dataset contains 2,756 ultrasound images of peripheral pulmonary lesions from 2,756
  5unique patients, with expert sonographer pixel-level tumor masks and biopsy-confirmed
  6benign / malignant pathology labels.
  7
  8This dataset is located at https://doi.org/10.6084/m9.figshare.32672274.v1 (CC BY 4.0).
  9This dataset is from the publication https://doi.org/10.1038/s41597-026-07715-0.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14from glob import glob
 15from typing import Union, Tuple, List
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URLS = {
 25    "images": "https://ndownloader.figshare.com/files/65547621",
 26    "masks": "https://ndownloader.figshare.com/files/65547618",
 27}
 28CHECKSUMS = {
 29    "images": "f91821641147cbc4d2fd25a5e590e7028b8b1ec5a368987866f2c5551a011a2e",
 30    "masks": "3e88ec64da10555b70c170d31265a8474570df6592c23ec4906a57c89f78a7e3",
 31}
 32
 33
 34def get_pu2756_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the PU2756 dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the data is downloaded for further processing.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        Filepath where the data is downloaded.
 43    """
 44    data_dir = os.path.join(path, "masks")
 45    if os.path.exists(data_dir):
 46        return path
 47
 48    os.makedirs(path, exist_ok=True)
 49
 50    for name, url in URLS.items():
 51        zip_path = os.path.join(path, f"{name}.zip")
 52        util.download_source(path=zip_path, url=url, download=download, checksum=CHECKSUMS[name])
 53        util.unzip(zip_path=zip_path, dst=path)
 54
 55    return path
 56
 57
 58def get_pu2756_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 59    """Get paths to the PU2756 data.
 60
 61    Args:
 62        path: Filepath to a folder where the data is downloaded for further processing.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    data_dir = get_pu2756_data(path=path, download=download)
 70
 71    image_paths = sorted(glob(os.path.join(data_dir, "images", "*", "*.png")))
 72    gt_paths = sorted(glob(os.path.join(data_dir, "masks", "*", "*.png")))
 73
 74    if len(image_paths) == 0 or len(image_paths) != len(gt_paths):
 75        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 76
 77    return image_paths, gt_paths
 78
 79
 80def get_pu2756_dataset(
 81    path: Union[os.PathLike, str],
 82    patch_shape: Tuple[int, int],
 83    resize_inputs: bool = False,
 84    download: bool = False,
 85    **kwargs
 86) -> Dataset:
 87    """Get the PU2756 dataset for pulmonary tumor segmentation.
 88
 89    Args:
 90        path: Filepath to a folder where the data is downloaded for further processing.
 91        patch_shape: The patch shape to use for training.
 92        resize_inputs: Whether to resize the inputs.
 93        download: Whether to download the data if it is not present.
 94        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 95
 96    Returns:
 97        The segmentation dataset.
 98    """
 99    image_paths, gt_paths = get_pu2756_paths(path, download)
100
101    if resize_inputs:
102        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
103        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
104            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
105        )
106
107    return torch_em.default_segmentation_dataset(
108        raw_paths=image_paths,
109        raw_key=None,
110        label_paths=gt_paths,
111        label_key=None,
112        patch_shape=patch_shape,
113        is_seg_dataset=False,
114        **kwargs
115    )
116
117
118def get_pu2756_loader(
119    path: Union[os.PathLike, str],
120    batch_size: int,
121    patch_shape: Tuple[int, int],
122    resize_inputs: bool = False,
123    download: bool = False,
124    **kwargs
125) -> DataLoader:
126    """Get the PU2756 dataloader for pulmonary tumor segmentation.
127
128    Args:
129        path: Filepath to a folder where the data is downloaded for further processing.
130        batch_size: The batch size for training.
131        patch_shape: The patch shape to use for training.
132        resize_inputs: Whether to resize the inputs.
133        download: Whether to download the data if it is not present.
134        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
135
136    Returns:
137        The DataLoader.
138    """
139    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
140    dataset = get_pu2756_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
141    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)