torch_em.data.datasets.medical.covid_qu_ex

  1"""The COVID QU EX dataset contains annotations for segmentations of
  2lung and infection in X-Ray images.
  3
  4The dataset is located at https://www.kaggle.com/datasets/anasmohammedtahir/covidqu.
  5This dataset is from the publication https://doi.org/10.1016/j.compbiomed.2021.104319.
  6Please cite them if you use this dataset for your research.
  7"""
  8
  9import os
 10from glob import glob
 11from natsort import natsorted
 12from typing import Tuple, Union, Optional, Literal, List
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21def get_covid_qu_ex_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 22    """Download the COVID QU EX dataset.
 23
 24    Args:
 25        path: Filepath to a folder where the data is downloaded for further processing.
 26        download: Whether to download the data if it is not present.
 27
 28    Returns:
 29        Filepath where the data is downlaoded.
 30    """
 31    data_dir = os.path.join(path, "data")
 32    if os.path.exists(data_dir):
 33        return data_dir
 34
 35    os.makedirs(path, exist_ok=True)
 36
 37    util.download_source_kaggle(path=path, dataset_name="anasmohammedtahir/covidqu", download=download)
 38    zip_path = os.path.join(path, "covidqu.zip")
 39    util.unzip(zip_path=zip_path, dst=data_dir)
 40
 41    return data_dir
 42
 43
 44def get_covid_qu_ex_paths(
 45    path: Union[os.PathLike, str],
 46    split: Literal['train', 'val', 'test'],
 47    task: Literal['lung', 'infection'],
 48    patient_type: Optional[Literal['covid19', 'non-covid', 'normal']] = None,
 49    segmentation_mask: Literal['lung', 'infection'] = "lung",
 50    download: bool = False
 51) -> Tuple[List[str], List[str]]:
 52    """Get paths to the COVID QU EX data.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        split: The data split to use. Either 'train', 'val' or 'test'.
 57        task: The choice for the subset of dataset. Either 'lung' or 'infection'.
 58        patient_type: The choice of subset of patients. Either 'covid19', 'non-covid' or 'normal'.
 59            By default is None, i.e. all the patient data will be chosen.
 60        segmentation_mask: The choice of segmentation labels. Either 'lung' or 'infection'.
 61        download: Whether to download the data if it is not present.
 62
 63    Returns:
 64        List of filepaths for the image data.
 65        List of filepaths for the label data.
 66    """
 67    data_dir = get_covid_qu_ex_data(path=path, download=download)
 68
 69    assert split.lower() in ["train", "val", "test"], f"'{split}' is not a valid split."
 70
 71    if task == "lung":
 72        _task = r"Lung Segmentation Data/Lung Segmentation Data"
 73    elif task == "infection":
 74        _task = r"Infection Segmentation Data/Infection Segmentation Data"
 75    else:
 76        raise ValueError(f"'{task}' is not a valid task.")
 77
 78    if patient_type == "covid19":
 79        patient_type = "COVID-19"
 80    elif patient_type == "non-covid":
 81        patient_type = "Non-COVID"
 82    elif patient_type == "normal":
 83        patient_type = "Normal"
 84    else:
 85        if patient_type is None:
 86            patient_type = "*"
 87        else:
 88            raise ValueError(f"'{patient_type}' is not a valid patient type.")
 89
 90    base_dir = os.path.join(data_dir, _task, split.title(), patient_type)
 91
 92    if segmentation_mask == "lung":
 93        segmentation_mask = r"lung masks"
 94    elif segmentation_mask == "infection":
 95        if task == "lung":
 96            raise AssertionError("The 'lung' data subset does not have infection masks.")
 97        segmentation_mask = r"infection masks"
 98    else:
 99        if segmentation_mask is None:
100            segmentation_mask = "*"
101        else:
102            raise ValueError(f"'{segmentation_mask}' is not a valid segmentation task.")
103
104    image_paths = natsorted(glob(os.path.join(base_dir, "images", "*")))
105    gt_paths = natsorted(glob(os.path.join(base_dir, segmentation_mask, "*")))
106
107    return image_paths, gt_paths
108
109
110def get_covid_qu_ex_dataset(
111    path: Union[os.PathLike, str],
112    patch_shape: Tuple[int, int],
113    split: Literal['train', 'val', 'test'],
114    task: Literal['lung', 'infection'],
115    patient_type: Optional[Literal['covid19', 'non-covid', 'normal']] = None,
116    segmentation_mask: Literal['lung', 'infection'] = "lung",
117    resize_inputs: bool = False,
118    download: bool = False,
119    **kwargs
120) -> Dataset:
121    """Get the COVID QU EX dataset for lung and infection segmentation.
122
123    Args:
124        path: Filepath to a folder where the data is downloaded for further processing.
125        patch_shape: The patch shape to use for training.
126        split: The data split to use. Either 'train', 'val' or 'test'.
127        task: The choice for the subset of dataset. Either 'lung' or 'infection'.
128        patient_type: The choice of subset of patients. Either 'covid19', 'non-covid' or 'normal'.
129            By default is None, i.e. all the patient data will be chosen.
130        segmentation_mask: The choice of segmentation labels. Either 'lung' or 'infection'.
131        resize_inputs: Whether to resize the inputs to the 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`.
134
135    Returns:
136        The segmentation dataset.
137    """
138    image_paths, gt_paths = get_covid_qu_ex_paths(path, split, task, patient_type, segmentation_mask, download)
139
140    if resize_inputs:
141        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
142        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
143            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
144        )
145
146    return torch_em.default_segmentation_dataset(
147        raw_paths=image_paths,
148        raw_key=None,
149        label_paths=gt_paths,
150        label_key=None,
151        patch_shape=patch_shape,
152        **kwargs
153    )
154
155
156def get_covid_qu_ex_loader(
157    path: Union[os.PathLike, str],
158    batch_size: int,
159    patch_shape: Tuple[int, int],
160    split: Literal['train', 'val', 'test'],
161    task: Literal['lung', 'infection'],
162    patient_type: Optional[Literal['covid19', 'non-covid', 'normal']] = None,
163    segmentation_mask: Literal['lung', 'infection'] = "lung",
164    resize_inputs: bool = False,
165    download: bool = False,
166    **kwargs
167) -> DataLoader:
168    """Get the COVID QU EX dataloader for lung and infection segmentation.
169
170    Args:
171        path: Filepath to a folder where the data is downloaded for further processing.
172        batch_size: The batch size for training.
173        patch_shape: The patch shape to use for training.
174        split: The data split to use. Either 'train', 'val' or 'test'.
175        task: The choice for the subset of dataset. Either 'lung' or 'infection'.
176        patient_type: The choice of subset of patients. Either 'covid19', 'non-covid' or 'normal'.
177            By default is None, i.e. all the patient data will be chosen.
178        segmentation_mask: The choice of segmentation labels. Either 'lung' or 'infection'.
179        resize_inputs: Whether to resize the inputs to the patch shape.
180        download: Whether to download the data if it is not present.
181        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
182
183    Returns:
184        The DataLoader.
185    """
186    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
187    dataset = get_covid_qu_ex_dataset(
188        path, patch_shape, split, task, patient_type, segmentation_mask, resize_inputs, download, **ds_kwargs
189    )
190    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)