torch_em.data.datasets.medical.qata_cov19

  1"""The QaTa-COV19 dataset contains annotations for COVID-19 infection region segmentation in chest X-rays.
  2
  3The dataset (QaTa-COV19-v1 subset) consists of 4,603 COVID-19 chest X-rays, 2,951 of which have
  4corresponding ground-truth segmentation masks for the infected lung regions. The masks are binary,
  5where the foreground marks the COVID-19 infection region.
  6
  7This dataset is located at https://www.kaggle.com/datasets/aysendegerli/qatacov19-dataset.
  8NOTE: There is another mirror of this dataset on Kaggle (`prashant268/chest-xray-covid19-pneumonia`)
  9that only provides classification labels, without any segmentation masks. Please make sure to use the
 10dataset mentioned above for the segmentation task.
 11
 12This dataset is from the publication https://doi.org/10.1007/s13755-021-00146-8.
 13Please cite it if you use this dataset for your research.
 14"""
 15
 16import os
 17import zipfile
 18from glob import glob
 19from natsort import natsorted
 20from typing import Union, Tuple, List
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29KAGGLE_DATASET_NAME = "aysendegerli/qatacov19-dataset"
 30
 31
 32def _extract_v1_subset(zip_path, path):
 33    with zipfile.ZipFile(zip_path) as f:
 34        members = [m for m in f.namelist() if "/QaTa-COV19-v1/" in m]
 35        f.extractall(path, members=members)
 36
 37
 38def get_qata_cov19_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 39    """Download the QaTa-COV19 dataset.
 40
 41    Args:
 42        path: Filepath to a folder where the data is downloaded for further processing.
 43        download: Whether to download the data if it is not present.
 44
 45    Returns:
 46        Filepath where the data is downloaded.
 47    """
 48    data_dir = os.path.join(path, "QaTa-COV19", "QaTa-COV19-v1")
 49    if os.path.exists(data_dir):
 50        return data_dir
 51
 52    os.makedirs(path, exist_ok=True)
 53
 54    zip_path = os.path.join(path, "qatacov19-dataset.zip")
 55    util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 56    _extract_v1_subset(zip_path, path)
 57    os.remove(zip_path)
 58
 59    if not os.path.exists(data_dir):
 60        raise RuntimeError(f"The dataset could not be found at '{data_dir}' after extraction.")
 61
 62    return data_dir
 63
 64
 65def get_qata_cov19_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 66    """Get paths to the QaTa-COV19 data.
 67
 68    Args:
 69        path: Filepath to a folder where the data is downloaded for further processing.
 70        download: Whether to download the data if it is not present.
 71
 72    Returns:
 73        List of filepaths for the image data.
 74        List of filepaths for the label data.
 75    """
 76    data_dir = get_qata_cov19_data(path=path, download=download)
 77
 78    gt_paths = natsorted(glob(os.path.join(data_dir, "Ground-truths", "mask_*.png")))
 79    image_paths = [
 80        os.path.join(data_dir, "Images", os.path.basename(p)[len("mask_"):]) for p in gt_paths
 81    ]
 82
 83    if len(image_paths) == 0 or len(image_paths) != len(gt_paths):
 84        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 85
 86    return image_paths, gt_paths
 87
 88
 89def get_qata_cov19_dataset(
 90    path: Union[os.PathLike, str],
 91    patch_shape: Tuple[int, int],
 92    resize_inputs: bool = False,
 93    download: bool = False,
 94    **kwargs
 95) -> Dataset:
 96    """Get the QaTa-COV19 dataset for COVID-19 infection region segmentation.
 97
 98    Args:
 99        path: Filepath to a folder where the data is downloaded for further processing.
100        patch_shape: The patch shape to use for training.
101        resize_inputs: Whether to resize the inputs.
102        download: Whether to download the data if it is not present.
103        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
104
105    Returns:
106        The segmentation dataset.
107    """
108    image_paths, gt_paths = get_qata_cov19_paths(path, download)
109
110    if resize_inputs:
111        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
112        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
113            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
114        )
115
116    return torch_em.default_segmentation_dataset(
117        raw_paths=image_paths,
118        raw_key=None,
119        label_paths=gt_paths,
120        label_key=None,
121        patch_shape=patch_shape,
122        **kwargs
123    )
124
125
126def get_qata_cov19_loader(
127    path: Union[os.PathLike, str],
128    batch_size: int,
129    patch_shape: Tuple[int, int],
130    resize_inputs: bool = False,
131    download: bool = False,
132    **kwargs
133) -> DataLoader:
134    """Get the QaTa-COV19 dataloader for COVID-19 infection region segmentation.
135
136    Args:
137        path: Filepath to a folder where the data is downloaded for further processing.
138        batch_size: The batch size for training.
139        patch_shape: The patch shape to use for training.
140        resize_inputs: Whether to resize the inputs.
141        download: Whether to download the data if it is not present.
142        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
143
144    Returns:
145        The DataLoader.
146    """
147    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
148    dataset = get_qata_cov19_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
149    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)