torch_em.data.datasets.medical.micro_usp

  1"""The MicroUSP dataset contains annotations for prostate segmentation
  2in micro-ultrasound scans.
  3
  4The dataset is from the publication https://doi.org/10.1016/j.compmedimag.2024.102326.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9from glob import glob
 10from pathlib import Path
 11from natsort import natsorted
 12from typing import Union, Tuple, List, Literal
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URL = "https://zenodo.org/records/10475293/files/Micro_Ultrasound_Prostate_Segmentation_Dataset.zip"
 22CHECKSUM = "031645dc30948314e379d0a0a7d54bad1cd4e1f3f918b77455d69810aa05dce3"
 23
 24
 25def get_micro_usp_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 26    """Download the MicroUSP 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 downloaded.
 34    """
 35    fname = Path(URL).stem
 36    data_dir = os.path.join(path, fname)
 37    if os.path.exists(data_dir):
 38        return data_dir
 39
 40    os.makedirs(path, exist_ok=True)
 41
 42    zip_path = os.path.join(path, f"{fname}.zip")
 43    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 44    util.unzip(zip_path=zip_path, dst=path)
 45
 46    return data_dir
 47
 48
 49def get_micro_usp_paths(
 50    path: Union[os.PathLike, str], split: Literal['train', 'val', 'test'], download: bool = False
 51) -> Tuple[List[str], List[str]]:
 52    """Get paths to the MicroUSP data.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        split: The choice of data split.
 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_micro_usp_data(path=path, download=download)
 64
 65    if split == "test":
 66        image_paths = natsorted(glob(os.path.join(data_dir, split, "micro_ultrasound_scans", "*.nii.gz")))
 67        gt_paths = natsorted(glob(os.path.join(data_dir, split, "expert_annotations", "*.nii.gz")))
 68    else:
 69        image_paths = natsorted(glob(os.path.join(data_dir, "train", "micro_ultrasound_scans", "*.nii.gz")))
 70        gt_paths = natsorted(glob(os.path.join(data_dir, "train", "expert_annotations", "*.nii.gz")))
 71
 72        if split == "train":
 73            image_paths, gt_paths = image_paths[:50], gt_paths[:50]
 74        elif split == "val":
 75            image_paths, gt_paths = image_paths[50:], gt_paths[50:]
 76        else:
 77            raise ValueError(f"'{split}' is not a valid split.")
 78
 79    return image_paths, gt_paths
 80
 81
 82def get_micro_usp_dataset(
 83    path: Union[os.PathLike, str],
 84    patch_shape: Tuple[int, ...],
 85    split: Literal['train', 'val', 'test'],
 86    resize_inputs: bool = False,
 87    download: bool = False,
 88    **kwargs
 89) -> Dataset:
 90    """Get the MicroUSP dataset for segmentation of prostate in micro-ultrasound scans.
 91
 92    Args:
 93        path: Filepath to a folder where the data is downloaded for further processing.
 94        patch_shape: The patch shape to use for training.
 95        split: The choice of data split.
 96        resize_inputs: Whether to resize the inputs to the patch shape.
 97        download: Whether to download the data if it is not present.
 98        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 99
100    Returns:
101        The segmentation dataset.
102    """
103    image_paths, gt_paths = get_micro_usp_paths(path, split, download)
104
105    if resize_inputs:
106        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
107        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
108            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
109        )
110
111    return torch_em.default_segmentation_dataset(
112        raw_paths=image_paths, raw_key="data", label_paths=gt_paths, label_key="data", patch_shape=patch_shape, **kwargs
113    )
114
115
116def get_micro_usp_loader(
117    path: Union[os.PathLike, str],
118    batch_size: int,
119    patch_shape: Tuple[int, ...],
120    split: Literal['train', 'val', 'test'],
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> DataLoader:
125    """Get the MicroUSP dataloader for segmentation of prostate in micro-ultrasound scans.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        batch_size: The batch size for training.
130        patch_shape: The patch shape to use for training.
131        split: The choice of data split.
132        resize_inputs: Whether to resize the inputs to the patch shape.
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_micro_usp_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
141    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)