torch_em.data.datasets.medical.mosmed_plus

  1"""The MosMedData+ dataset contains annotations for COVID-19 lung lesion segmentation in chest CT scans.
  2
  3This dataset is a subset of the MosMedData collection (https://mosmed.ai/datasets/covid19_1110/):
  450 out of the 1110 chest CT studies come with expert-annotated binary pixel masks for ground-glass
  5opacifications and consolidations (label ids: background 0, lesion 1). The official release requires
  6registering with an email address, so this loader relies on the public Kaggle mirror at
  7https://www.kaggle.com/datasets/mathurinache/mosmeddata-chest-ct-scans-with-covid19, which redistributes
  8the same volumes and masks as uncompressed NIfTI files (*.nii).
  9
 10The dataset is from the publication https://doi.org/10.1101/2020.05.20.20100362.
 11Please cite it if you use this dataset in your research.
 12
 13The dataset is distributed under the CC BY-NC-ND 3.0 license.
 14"""
 15
 16import os
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28KAGGLE_DATASET_NAME = "mathurinache/mosmeddata-chest-ct-scans-with-covid19"
 29
 30
 31def get_mosmed_plus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the MosMedData+ dataset.
 33
 34    Args:
 35        path: Filepath to a folder where the data is downloaded for further processing.
 36        download: Whether to download the data if it is not present.
 37
 38    Returns:
 39        Filepath where the data is downloaded.
 40    """
 41    data_dir = os.path.join(path, "data")
 42    if os.path.exists(data_dir):
 43        return data_dir
 44
 45    os.makedirs(path, exist_ok=True)
 46
 47    util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 48    zip_path = os.path.join(path, "mosmeddata-chest-ct-scans-with-covid19.zip")
 49    util.unzip(zip_path=zip_path, dst=data_dir)
 50
 51    return data_dir
 52
 53
 54def get_mosmed_plus_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 55    """Get paths to the MosMedData+ data.
 56
 57    Args:
 58        path: Filepath to a folder where the data is downloaded for further processing.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        List of filepaths for the image data.
 63        List of filepaths for the label data.
 64    """
 65    data_dir = get_mosmed_plus_data(path, download)
 66
 67    mask_paths = natsorted(glob(os.path.join(data_dir, "**", "masks", "study_*_mask.nii"), recursive=True))
 68    assert len(mask_paths) == 50, f"Expected 50 annotated studies, found {len(mask_paths)} in '{data_dir}'."
 69
 70    image_paths = []
 71    for mask_path in mask_paths:
 72        fname = os.path.basename(mask_path).replace("_mask.nii", ".nii")
 73        matches = glob(os.path.join(data_dir, "**", "studies", "*", fname), recursive=True)
 74        assert len(matches) == 1, f"Could not find a unique image volume for '{mask_path}'."
 75        image_paths.append(matches[0])
 76
 77    return image_paths, mask_paths
 78
 79
 80def get_mosmed_plus_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 MosMedData+ dataset for COVID-19 lung lesion 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 the inputs to the 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    image_paths, label_paths = get_mosmed_plus_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=image_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_mosmed_plus_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 MosMedData+ dataloader for COVID-19 lung lesion 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 the inputs to the 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_mosmed_plus_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
141    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)