torch_em.data.datasets.medical.fumpe

  1"""The FUMPE dataset contains annotations for pulmonary embolism in computed tomography angiography.
  2
  3The dataset consists of 35 CT angiography scans of different patients, with a binary mask marking the
  4pulmonary emboli. The scans are distributed as DICOM series and the masks as MATLAB files, which this
  5module converts into hdf5 files, one per patient.
  6
  7NOTE: This requires the pydicom and scipy python packages.
  8
  9NOTE: The slices of a scan are ordered by their instance number, which is the order the masks were
 10created in. Two of the scans are acquired from head to feet, so sorting the slices along the slice
 11normal instead would flip their masks.
 12
 13The dataset is located at https://www.kaggle.com/datasets/andrewmvd/pulmonary-embolism-in-ct-images.
 14This dataset is from the publication https://doi.org/10.1038/sdata.2018.180.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from tqdm import tqdm
 21from natsort import natsorted
 22from typing import Union, Tuple, List
 23
 24import numpy as np
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33KAGGLE_DATASET_NAME = "andrewmvd/pulmonary-embolism-in-ct-images"
 34
 35
 36def _load_dicom_volume(series_dir):
 37    """Stack a DICOM series into a volume with axes (z, y, x), ordered by the instance number.
 38
 39    The masks of this dataset were created in this order, so the slices must not be sorted along the
 40    slice normal, which would flip the scans that are acquired from head to feet.
 41    """
 42    import pydicom
 43
 44    slices = [pydicom.dcmread(p) for p in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 45    slices.sort(key=lambda dcm: int(dcm.InstanceNumber))
 46
 47    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 48    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 49    return np.round(volume).astype("int16")
 50
 51
 52def _preprocess_fumpe(data_dir, preprocessed_dir):
 53    import h5py
 54    import scipy.io as sio
 55
 56    os.makedirs(preprocessed_dir, exist_ok=True)
 57    patient_dirs = natsorted(glob(os.path.join(data_dir, "CT_scans", "*")))
 58    for patient_dir in tqdm(patient_dirs, desc="Preprocess FUMPE"):
 59        patient_id = os.path.basename(patient_dir)
 60        out_path = os.path.join(preprocessed_dir, f"{patient_id}.h5")
 61        if os.path.exists(out_path):
 62            continue
 63
 64        volume = _load_dicom_volume(patient_dir)
 65        # The masks are stored with axes (y, x, z) and have to be transposed to match the volume.
 66        labels = sio.loadmat(os.path.join(data_dir, "GroundTruth", f"{patient_id}.mat"))["Mask"]
 67        labels = np.transpose(labels, (2, 0, 1)).astype("uint8")
 68        assert labels.shape == volume.shape, f"The mask of '{patient_id}' does not match its scan."
 69
 70        with h5py.File(out_path, "w") as f:
 71            f.create_dataset("raw", data=volume, compression="gzip")
 72            f.create_dataset("labels", data=labels, compression="gzip")
 73
 74
 75def get_fumpe_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 76    """Download the FUMPE 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 preprocessed data is stored.
 84    """
 85    preprocessed_dir = os.path.join(path, "preprocessed")
 86    if os.path.exists(preprocessed_dir):
 87        return preprocessed_dir
 88
 89    os.makedirs(path, exist_ok=True)
 90    data_dir = os.path.join(path, "FUMPE")
 91    if not os.path.exists(data_dir):
 92        util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 93        util.unzip(zip_path=os.path.join(path, "pulmonary-embolism-in-ct-images.zip"), dst=path, remove=False)
 94
 95    _preprocess_fumpe(data_dir, preprocessed_dir)
 96    return preprocessed_dir
 97
 98
 99def get_fumpe_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
100    """Get paths to the FUMPE data.
101
102    Args:
103        path: Filepath to a folder where the data is downloaded for further processing.
104        download: Whether to download the data if it is not present.
105
106    Returns:
107        List of filepaths for the stored data.
108    """
109    preprocessed_dir = get_fumpe_data(path, download)
110    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
111    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
112    return volume_paths
113
114
115def get_fumpe_dataset(
116    path: Union[os.PathLike, str],
117    patch_shape: Tuple[int, ...],
118    resize_inputs: bool = False,
119    download: bool = False,
120    **kwargs
121) -> Dataset:
122    """Get the FUMPE dataset for pulmonary embolism segmentation.
123
124    Args:
125        path: Filepath to a folder where the data is downloaded for further processing.
126        patch_shape: The patch shape to use for training.
127        resize_inputs: Whether to resize inputs to the desired patch shape.
128        download: Whether to download the data if it is not present.
129        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
130
131    Returns:
132        The segmentation dataset.
133    """
134    volume_paths = get_fumpe_paths(path, download)
135
136    if resize_inputs:
137        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
138        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
139            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
140        )
141
142    return torch_em.default_segmentation_dataset(
143        raw_paths=volume_paths,
144        raw_key="raw",
145        label_paths=volume_paths,
146        label_key="labels",
147        patch_shape=patch_shape,
148        is_seg_dataset=True,
149        **kwargs
150    )
151
152
153def get_fumpe_loader(
154    path: Union[os.PathLike, str],
155    batch_size: int,
156    patch_shape: Tuple[int, ...],
157    resize_inputs: bool = False,
158    download: bool = False,
159    **kwargs
160) -> DataLoader:
161    """Get the FUMPE dataloader for pulmonary embolism segmentation.
162
163    Args:
164        path: Filepath to a folder where the data is downloaded for further processing.
165        batch_size: The batch size for training.
166        patch_shape: The patch shape to use for training.
167        resize_inputs: Whether to resize inputs to the desired patch shape.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
170
171    Returns:
172        The DataLoader.
173    """
174    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
175    dataset = get_fumpe_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
176    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)