torch_em.data.datasets.medical.imagetbad

  1"""The ImageTBAD dataset contains annotations for type-B aortic dissection segmentation in
  2computed tomography angiography (CTA).
  3
  4The dataset consists of 100 3D CTA scans collected at the Guangdong Provincial People's Hospital between
  5January 2013 and April 2019. The segmentation was performed by one cardiovascular radiologist and checked
  6by a second one. It labels three substructures of the aortic dissection: the true lumen (TL), the false
  7lumen (FL) and the false lumen thrombus (FLT), see `LABEL_IDS`. 68 of the 100 cases contain an FLT, the
  8remaining 32 cases are free of it. The data is distributed as one 'image.nii.gz' / 'label.nii.gz' pair per
  9case and hosted on Kaggle as a single, split zip archive (https://www.kaggle.com/datasets/xiaoweixumedicalai/
 10imagetbad), because the official GitHub repository (https://github.com/XiaoweiXu/Dataset_Type-B-Aortic-
 11Dissection) does not host the data itself.
 12
 13The Kaggle archive is itself split into 19 parts ('imageTBAD.change2zip', 'imageTBAD.z01' to 'imageTBAD.z18'):
 14this module downloads all parts, joins them into a single zip with the 'zip' CLI (Info-ZIP) and extracts it.
 15
 16NOTE: This requires a Kaggle account and API credentials (see https://www.kaggle.com/docs/api), as well as
 17the 'zip' CLI (Info-ZIP) to join the split archives.
 18
 19This dataset is from the publication https://doi.org/10.3389/fphys.2021.732711.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24from glob import glob
 25from shutil import which
 26from subprocess import run
 27from natsort import natsorted
 28from typing import Union, Tuple, List
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37KAGGLE_DATASET = "xiaoweixumedicalai/imagetbad"
 38
 39N_PARTS = 18
 40"""The number of split zip parts ('imageTBAD.z01' to 'imageTBAD.z18') on top of 'imageTBAD.change2zip'."""
 41
 42LABEL_IDS = {"background": 0, "true_lumen": 1, "false_lumen": 2, "false_lumen_thrombus": 3}
 43
 44
 45def _download_kaggle_file(filename: str, dst_dir: str, download: bool) -> str:
 46    """Download a single file from the ImageTBAD Kaggle dataset.
 47
 48    Kaggle wraps every single-file download in an outer zip container (even if the file is itself
 49    already an archive), which is unpacked here to recover the original file.
 50    """
 51    out_path = os.path.join(dst_dir, filename)
 52    if os.path.exists(out_path):
 53        return out_path
 54    if not download:
 55        raise RuntimeError(f"Cannot find the data at {out_path}, but download was set to False.")
 56
 57    try:
 58        from kaggle.api.kaggle_api_extended import KaggleApi
 59    except ModuleNotFoundError:
 60        msg = "Please install the Kaggle API. You can do this using 'pip install kaggle'. "
 61        msg += "After you have installed kaggle, you would need an API token. "
 62        msg += "Follow the instructions at https://www.kaggle.com/docs/api."
 63        raise ModuleNotFoundError(msg)
 64
 65    os.makedirs(dst_dir, exist_ok=True)
 66    api = KaggleApi()
 67    api.authenticate()
 68    api.dataset_download_file(KAGGLE_DATASET, filename, path=dst_dir)
 69
 70    wrapper_path = os.path.join(dst_dir, f"{filename}.zip")
 71    util.unzip(zip_path=wrapper_path, dst=dst_dir)
 72    return out_path
 73
 74
 75def get_imagetbad_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 76    """Download the ImageTBAD dataset.
 77
 78    Args:
 79        path: Filepath to a folder where the data is downloaded for further processing.
 80        download: Whether to download the data if it is not present.
 81
 82    Returns:
 83        Filepath where the data is stored.
 84    """
 85    raw_dir = os.path.join(path, "data")
 86    if len(glob(os.path.join(raw_dir, "**", "*_image.nii.gz"), recursive=True)) >= 100:
 87        return raw_dir
 88
 89    os.makedirs(raw_dir, exist_ok=True)
 90
 91    zip_dir = os.path.join(path, "zips")
 92    parts = ["imageTBAD.change2zip"] + [f"imageTBAD.z{i:02d}" for i in range(1, N_PARTS + 1)]
 93    for part in parts:
 94        _download_kaggle_file(part, zip_dir, download)
 95
 96    base_zip = os.path.join(zip_dir, "imageTBAD.zip")
 97    if not os.path.exists(base_zip):
 98        os.rename(os.path.join(zip_dir, "imageTBAD.change2zip"), base_zip)
 99
100    merged_zip = os.path.join(zip_dir, "imageTBAD.merged.zip")
101    if not os.path.exists(merged_zip):
102        if which("zip") is None:
103            raise RuntimeError(
104                "Need the 'zip' CLI (Info-ZIP) to join the split zip archive of the ImageTBAD dataset. "
105                "You can install it via 'conda install -c conda-forge zip'."
106            )
107        run(["zip", "-s", "0", base_zip, "--out", merged_zip], check=True, cwd=zip_dir)
108
109    util.unzip(zip_path=merged_zip, dst=raw_dir, remove=False)
110
111    return raw_dir
112
113
114def get_imagetbad_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
115    """Get paths to the ImageTBAD data.
116
117    Args:
118        path: Filepath to a folder where the data is downloaded for further processing.
119        download: Whether to download the data if it is not present.
120
121    Returns:
122        List of filepaths for the image data.
123        List of filepaths for the label data.
124    """
125    raw_dir = get_imagetbad_data(path, download)
126
127    image_paths = natsorted(glob(os.path.join(raw_dir, "**", "*_image.nii.gz"), recursive=True))
128    label_paths = natsorted(glob(os.path.join(raw_dir, "**", "*_label.nii.gz"), recursive=True))
129    assert len(image_paths) > 0 and len(image_paths) == len(label_paths), \
130        f"Could not find a matching number of images and labels in '{raw_dir}'."
131
132    return image_paths, label_paths
133
134
135def get_imagetbad_dataset(
136    path: Union[os.PathLike, str],
137    patch_shape: Tuple[int, ...],
138    resize_inputs: bool = False,
139    download: bool = False,
140    **kwargs
141) -> Dataset:
142    """Get the ImageTBAD dataset for aortic dissection segmentation.
143
144    Args:
145        path: Filepath to a folder where the data is downloaded for further processing.
146        patch_shape: The patch shape to use for training.
147        resize_inputs: Whether to resize inputs to the desired patch shape.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
150
151    Returns:
152        The segmentation dataset.
153    """
154    image_paths, label_paths = get_imagetbad_paths(path, download)
155
156    if resize_inputs:
157        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
158        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
159            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
160        )
161
162    return torch_em.default_segmentation_dataset(
163        raw_paths=image_paths,
164        raw_key="data",
165        label_paths=label_paths,
166        label_key="data",
167        patch_shape=patch_shape,
168        is_seg_dataset=True,
169        **kwargs
170    )
171
172
173def get_imagetbad_loader(
174    path: Union[os.PathLike, str],
175    batch_size: int,
176    patch_shape: Tuple[int, ...],
177    resize_inputs: bool = False,
178    download: bool = False,
179    **kwargs
180) -> DataLoader:
181    """Get the ImageTBAD dataloader for aortic dissection segmentation.
182
183    Args:
184        path: Filepath to a folder where the data is downloaded for further processing.
185        batch_size: The batch size for training.
186        patch_shape: The patch shape to use for training.
187        resize_inputs: Whether to resize inputs to the desired patch shape.
188        download: Whether to download the data if it is not present.
189        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
190
191    Returns:
192        The DataLoader.
193    """
194    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
195    dataset = get_imagetbad_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
196    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)