torch_em.data.datasets.medical.bpd_neo

  1"""BPD-Neo is a dataset for segmentation of the lungs and trachea in neonatal MRI.
  2
  3The dataset consists of 40 free-breathing 3D stack-of-stars radial gradient echo (StarVIBE) MRI
  4scans of neonates, most of whom are diagnosed with bronchopulmonary dysplasia (BPD), together with
  5semantic segmentation masks. All 40 scans have a lung mask, and 36 of them additionally have a
  6trachea mask (4 patients: BPD-Neo-01, BPD-Neo-03, BPD-Neo-11 and BPD-Neo-22 have no trachea mask).
  7
  8This dataset is located at https://doi.org/10.5281/zenodo.15768091 (Zenodo, CC BY 4.0).
  9The dataset is from the publication https://doi.org/10.1038/s41597-026-07006-8.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14from glob import glob
 15from natsort import natsorted
 16from typing import Union, Tuple, Literal, List
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://zenodo.org/records/15768091/files/BPD-Neo-data.zip"
 26CHECKSUM = "cae6eb7ad6cbd3b1bf488fe88af9b17147f02bf2d754a304ec3b94d0ec1c75a9"
 27
 28
 29def get_bpd_neo_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the BPD-Neo data.
 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, "Nifti-data")
 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, "BPD-Neo-data.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_bpd_neo_paths(
 53    path: Union[os.PathLike, str], structure: Literal['lung', 'trachea'] = "lung", download: bool = False
 54) -> Tuple[List[str], List[str]]:
 55    """Get paths to the BPD-Neo data.
 56
 57    Args:
 58        path: Filepath to a folder where the data is downloaded for further processing.
 59        structure: The anatomical structure to segment. Either 'lung' or 'trachea'.
 60        download: Whether to download the data if it is not present.
 61
 62    Returns:
 63        List of filepaths for the image data.
 64        List of filepaths for the label data.
 65    """
 66    data_dir = get_bpd_neo_data(path, download)
 67
 68    if structure not in ("lung", "trachea"):
 69        raise ValueError(f"'{structure}' is not a valid structure. Choose either 'lung' or 'trachea'.")
 70
 71    patient_dirs = natsorted(glob(os.path.join(data_dir, "BPD-Neo-*")))
 72
 73    raw_paths, label_paths = [], []
 74    for patient_dir in patient_dirs:
 75        raw_path = os.path.join(patient_dir, "image.nii.gz")
 76        label_path = os.path.join(patient_dir, f"{structure}_seg.nii.gz")
 77        if os.path.exists(raw_path) and os.path.exists(label_path):
 78            raw_paths.append(raw_path)
 79            label_paths.append(label_path)
 80
 81    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 82
 83    return raw_paths, label_paths
 84
 85
 86def get_bpd_neo_dataset(
 87    path: Union[os.PathLike, str],
 88    patch_shape: Tuple[int, ...],
 89    structure: Literal['lung', 'trachea'] = "lung",
 90    resize_inputs: bool = False,
 91    download: bool = False,
 92    **kwargs
 93) -> Dataset:
 94    """Get the BPD-Neo dataset for lung and trachea segmentation in neonatal MRI.
 95
 96    Args:
 97        path: Filepath to a folder where the data is downloaded for further processing.
 98        patch_shape: The patch shape to use for training.
 99        structure: The anatomical structure to segment. Either 'lung' or 'trachea'.
100        resize_inputs: Whether to resize inputs to the desired patch shape.
101        download: Whether to download the data if it is not present.
102        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
103
104    Returns:
105        The segmentation dataset.
106    """
107    raw_paths, label_paths = get_bpd_neo_paths(path, structure, download)
108
109    if resize_inputs:
110        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
111        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
112            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
113        )
114
115    return torch_em.default_segmentation_dataset(
116        raw_paths=raw_paths,
117        raw_key="data",
118        label_paths=label_paths,
119        label_key="data",
120        patch_shape=patch_shape,
121        is_seg_dataset=True,
122        **kwargs
123    )
124
125
126def get_bpd_neo_loader(
127    path: Union[os.PathLike, str],
128    batch_size: int,
129    patch_shape: Tuple[int, ...],
130    structure: Literal['lung', 'trachea'] = "lung",
131    resize_inputs: bool = False,
132    download: bool = False,
133    **kwargs
134) -> DataLoader:
135    """Get the BPD-Neo dataloader for lung and trachea segmentation in neonatal MRI.
136
137    Args:
138        path: Filepath to a folder where the data is downloaded for further processing.
139        batch_size: The batch size for training.
140        patch_shape: The patch shape to use for training.
141        structure: The anatomical structure to segment. Either 'lung' or 'trachea'.
142        resize_inputs: Whether to resize inputs to the desired patch shape.
143        download: Whether to download the data if it is not present.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
145
146    Returns:
147        The DataLoader.
148    """
149    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
150    dataset = get_bpd_neo_dataset(path, patch_shape, structure, resize_inputs, download, **ds_kwargs)
151    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)