torch_em.data.datasets.medical.parse22

  1"""The PARSE dataset contains annotations for pulmonary artery segmentation in
  2contrast-enhanced CT pulmonary angiography (CTPA) scans.
  3
  4It comprises the training set of the PARSE2022 challenge (https://parse2022.grand-challenge.org):
  5100 CTPA volumes with a refined multi-level pulmonary artery annotation. The 30 validation volumes and
  670 test volumes of the challenge are distributed without annotations and are therefore not included here.
  7
  8NOTE: The label legend is as follows:
  9- background: 0, pulmonary artery: 1
 10
 11The dataset is located at https://parse2022.grand-challenge.org/Dataset/, from where the organizers share
 12the training set via a public Google Drive link (mirrored on Baidu Netdisk).
 13
 14This dataset is from the publication https://doi.org/10.48550/arXiv.2304.03708.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, List
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://drive.google.com/uc?id=1_-w8kNc2k4ttHTrVnRWaEaLbjGmNWRZD"
 31CHECKSUM = "4d51cae38b4e9ca530d9f577dea0e77e3e08d5610bfac0bd632f173a07554468"
 32
 33LABEL_IDS = {"background": 0, "pulmonary_artery": 1}
 34
 35
 36def get_parse22_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 37    """Download the PARSE dataset.
 38
 39    Args:
 40        path: Filepath to a folder where the data is downloaded for further processing.
 41        download: Whether to download the data if it is not present.
 42
 43    Returns:
 44        Filepath where the data is stored.
 45    """
 46    data_dir = os.path.join(path, "train")
 47    if os.path.exists(data_dir):
 48        return data_dir
 49
 50    os.makedirs(path, exist_ok=True)
 51
 52    rar_path = os.path.join(path, "train.rar")
 53    util.download_source_gdrive(path=rar_path, url=URL, download=download, checksum=CHECKSUM)
 54    util.unzip_rarfile(rar_path=rar_path, dst=path, remove=True)
 55
 56    return data_dir
 57
 58
 59def get_parse22_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 60    """Get paths to the PARSE data.
 61
 62    Args:
 63        path: Filepath to a folder where the data is downloaded for further processing.
 64        download: Whether to download the data if it is not present.
 65
 66    Returns:
 67        List of filepaths for the image data.
 68        List of filepaths for the label data.
 69    """
 70    data_dir = get_parse22_data(path, download)
 71
 72    raw_paths = natsorted(glob(os.path.join(data_dir, "*", "image", "*.nii.gz")))
 73    label_paths = [p.replace(os.sep + "image" + os.sep, os.sep + "label" + os.sep) for p in raw_paths]
 74    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 75
 76    return raw_paths, label_paths
 77
 78
 79def get_parse22_dataset(
 80    path: Union[os.PathLike, str],
 81    patch_shape: Tuple[int, ...],
 82    resize_inputs: bool = False,
 83    download: bool = False,
 84    **kwargs
 85) -> Dataset:
 86    """Get the PARSE dataset for pulmonary artery segmentation.
 87
 88    Args:
 89        path: Filepath to a folder where the data is downloaded for further processing.
 90        patch_shape: The patch shape to use for training.
 91        resize_inputs: Whether to resize inputs to the desired patch shape.
 92        download: Whether to download the data if it is not present.
 93        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 94
 95    Returns:
 96        The segmentation dataset.
 97    """
 98    raw_paths, label_paths = get_parse22_paths(path, download)
 99
100    if resize_inputs:
101        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
102        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
103            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
104        )
105
106    return torch_em.default_segmentation_dataset(
107        raw_paths=raw_paths,
108        raw_key="data",
109        label_paths=label_paths,
110        label_key="data",
111        patch_shape=patch_shape,
112        is_seg_dataset=True,
113        **kwargs
114    )
115
116
117def get_parse22_loader(
118    path: Union[os.PathLike, str],
119    batch_size: int,
120    patch_shape: Tuple[int, ...],
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> DataLoader:
125    """Get the PARSE dataloader for pulmonary artery segmentation.
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        resize_inputs: Whether to resize inputs to the desired patch shape.
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
134
135    Returns:
136        The DataLoader.
137    """
138    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
139    dataset = get_parse22_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
140    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)