torch_em.data.datasets.medical.covid19_20

  1"""The COVID-19-20 dataset contains annotations for COVID-19 lung lesion segmentation in chest CT scans.
  2
  3It comprises the training set of the COVID-19-20 Lung CT Lesion Segmentation Challenge
  4(https://covid-segmentation.grand-challenge.org/COVID-19-20/): 199 non-contrast chest CT volumes of
  5COVID-19 positive patients with a radiologist-verified binary lesion mask. The 50 validation volumes
  6and the 46 test volumes of the challenge are distributed without annotations and are not included here.
  7
  8NOTE: The label legend is as follows:
  9- background: 0, covid-19 lesion: 1
 10Verified on the data: the label volumes only contain the ids 0 and 1.
 11
 12NOTE: This is a different dataset than the one in `torch_em.data.datasets.medical.covid19_seg`, which
 13contains the 20 annotated volumes from https://doi.org/10.5281/zenodo.3757476.
 14
 15The data is a redistribution of the official challenge data at
 16https://huggingface.co/datasets/MedOtter/COVID-19-20 (CC BY 4.0). The official release at
 17https://covid-segmentation.grand-challenge.org/Data/ requires registration, so please make sure that
 18you are allowed to use the data for your purpose.
 19
 20This dataset is from the publication https://doi.org/10.1016/j.media.2022.102605.
 21Please cite it if you use this dataset in your research.
 22"""
 23
 24import os
 25from glob import glob
 26from natsort import natsorted
 27from typing import Union, Tuple, List
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36HF_REPO = "MedOtter/COVID-19-20"
 37
 38LABEL_IDS = {"background": 0, "covid19_lesion": 1}
 39
 40
 41def get_covid19_20_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 42    """Download the COVID-19-20 dataset.
 43
 44    Args:
 45        path: Filepath to a folder where the data is downloaded for further processing.
 46        download: Whether to download the data if it is not present.
 47
 48    Returns:
 49        Filepath where the data is stored.
 50    """
 51    data_dir = os.path.join(path, "Train")
 52    if os.path.exists(data_dir):
 53        return data_dir
 54
 55    if not download:
 56        raise RuntimeError(f"Cannot find the data at '{path}', but download was set to False.")
 57
 58    try:
 59        from huggingface_hub import snapshot_download
 60    except ImportError:
 61        raise ImportError("'huggingface_hub' is required to download COVID-19-20. Install it via conda/pip.")
 62
 63    os.makedirs(path, exist_ok=True)
 64    snapshot_download(repo_id=HF_REPO, repo_type="dataset", local_dir=path, allow_patterns=["Train/*"])
 65
 66    return data_dir
 67
 68
 69def get_covid19_20_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 70    """Get paths to the COVID-19-20 data.
 71
 72    Args:
 73        path: Filepath to a folder where the data is downloaded for further processing.
 74        download: Whether to download the data if it is not present.
 75
 76    Returns:
 77        List of filepaths for the image data.
 78        List of filepaths for the label data.
 79    """
 80    data_dir = get_covid19_20_data(path, download)
 81
 82    raw_paths = natsorted(glob(os.path.join(data_dir, "*_ct.nii.gz")))
 83    label_paths = [p.replace("_ct.nii.gz", "_seg.nii.gz") for p in raw_paths]
 84    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 85
 86    return raw_paths, label_paths
 87
 88
 89def get_covid19_20_dataset(
 90    path: Union[os.PathLike, str],
 91    patch_shape: Tuple[int, ...],
 92    resize_inputs: bool = False,
 93    download: bool = False,
 94    **kwargs
 95) -> Dataset:
 96    """Get the COVID-19-20 dataset for COVID-19 lung lesion 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 inputs to the desired patch shape.
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    raw_paths, label_paths = get_covid19_20_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=raw_paths,
118        raw_key="data",
119        label_paths=label_paths,
120        label_key="data",
121        patch_shape=patch_shape,
122        is_seg_dataset=True,
123        **kwargs
124    )
125
126
127def get_covid19_20_loader(
128    path: Union[os.PathLike, str],
129    batch_size: int,
130    patch_shape: Tuple[int, ...],
131    resize_inputs: bool = False,
132    download: bool = False,
133    **kwargs
134) -> DataLoader:
135    """Get the COVID-19-20 dataloader for COVID-19 lung lesion segmentation.
136
137    Args:
138        path: Filepath to a folder where the data is downloaded for further processing.
139        batch_size: The batch size for training.
140        patch_shape: The patch shape to use for training.
141        resize_inputs: Whether to resize inputs to the desired patch shape.
142        download: Whether to download the data if it is not present.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
144
145    Returns:
146        The DataLoader.
147    """
148    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
149    dataset = get_covid19_20_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
150    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)