torch_em.data.datasets.medical.kipa

  1"""The KiPA dataset contains annotations for kidney, renal tumor, renal artery and renal vein segmentation
  2in abdominal CT angiography (CTA) scans.
  3
  4It comprises the training set of the KiPA22 challenge (https://kipa22.grand-challenge.org): 70 CTA volumes
  5that are cropped around the kidney, with a dense annotation of the four renal structures. The remaining
  630 open testing volumes are distributed without annotations and are therefore not included here.
  7
  8NOTE: The label legend is as follows:
  9- background: 0, renal vein: 1, kidney: 2, renal artery: 3, renal tumor: 4
 10The ids were verified on the data by comparing the per-label volumes with the official dataset statistics.
 11
 12The data is a redistribution of the official challenge data at https://huggingface.co/datasets/YongchengYAO/KiPA22
 13(CC BY-NC 4.0). See https://kipa22.grand-challenge.org for the official release.
 14
 15This dataset is from the publication https://doi.org/10.1016/j.media.2021.102055.
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20from glob import glob
 21from natsort import natsorted
 22from typing import Union, Tuple, List
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URL = "https://huggingface.co/datasets/YongchengYAO/KiPA22/resolve/main/train.zip"
 32CHECKSUM = "43ea3667cb7bb5e0e465f203ce2ea9008034eecc5ddede73c69983763b3a1a65"
 33
 34LABEL_IDS = {"background": 0, "renal_vein": 1, "kidney": 2, "renal_artery": 3, "renal_tumor": 4}
 35
 36
 37def get_kipa_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 38    """Download the KiPA dataset.
 39
 40    Args:
 41        path: Filepath to a folder where the data is downloaded for further processing.
 42        download: Whether to download the data if it is not present.
 43
 44    Returns:
 45        Filepath where the data is stored.
 46    """
 47    data_dir = os.path.join(path, "train")
 48    if os.path.exists(data_dir):
 49        return data_dir
 50
 51    os.makedirs(path, exist_ok=True)
 52
 53    zip_path = os.path.join(path, "train.zip")
 54    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 55    util.unzip(zip_path=zip_path, dst=path)
 56
 57    return data_dir
 58
 59
 60def get_kipa_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 61    """Get paths to the KiPA data.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        List of filepaths for the image data.
 69        List of filepaths for the label data.
 70    """
 71    data_dir = get_kipa_data(path, download)
 72
 73    raw_paths = natsorted(glob(os.path.join(data_dir, "image", "*.nii.gz")))
 74    label_paths = [p.replace(os.sep + "image" + os.sep, os.sep + "label" + os.sep) for p in raw_paths]
 75    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 76
 77    return raw_paths, label_paths
 78
 79
 80def get_kipa_dataset(
 81    path: Union[os.PathLike, str],
 82    patch_shape: Tuple[int, ...],
 83    resize_inputs: bool = False,
 84    download: bool = False,
 85    **kwargs
 86) -> Dataset:
 87    """Get the KiPA dataset for kidney, renal tumor, renal artery and renal vein 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 inputs to the desired patch shape.
 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    raw_paths, label_paths = get_kipa_paths(path, download)
100
101    if resize_inputs:
102        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
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=raw_paths,
109        raw_key="data",
110        label_paths=label_paths,
111        label_key="data",
112        patch_shape=patch_shape,
113        is_seg_dataset=True,
114        **kwargs
115    )
116
117
118def get_kipa_loader(
119    path: Union[os.PathLike, str],
120    batch_size: int,
121    patch_shape: Tuple[int, ...],
122    resize_inputs: bool = False,
123    download: bool = False,
124    **kwargs
125) -> DataLoader:
126    """Get the KiPA dataloader for kidney, renal tumor, renal artery and renal vein 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 inputs to the desired 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_kipa_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
141    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)