torch_em.data.datasets.medical.plethora

  1"""The PLETHORA dataset contains annotations for thoracic organs and pleural effusion in CT.
  2
  3This dataset is from the publication https://doi.org/10.1002/mp.14424/.
  4Please cite it if you use this dataset for your research.
  5"""
  6
  7import os
  8from glob import glob
  9from tqdm import tqdm
 10from pathlib import Path
 11from natsort import natsorted
 12from urllib.parse import urljoin
 13from typing import Union, Tuple, Literal, List
 14
 15import numpy as np
 16import pandas as pd
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25BASE_URL = "https://wiki.cancerimagingarchive.net/download/attachments/68551327/"
 26
 27
 28URL = {
 29    "image": urljoin(BASE_URL, "NSCLC-Radiomics-OriginalCTs.tcia"),
 30    "gt": {
 31        "thoracic": urljoin(
 32            BASE_URL, "PleThora%20Thoracic_Cavities%20June%202020.zip?version=1&modificationDate=1593202695428&api=v2"
 33        ),
 34        "pleural_effusion": urljoin(
 35            BASE_URL, "PleThora%20Effusions%20June%202020.zip?version=1&modificationDate=1593202778373&api=v2"
 36        )
 37    }
 38}
 39
 40
 41CHECKSUMS = {
 42    "image": None,
 43    "gt": {
 44        "thoracic": "6dfcb60e46c7b0ccf240bc5d13acb1c45c8d2f4922223f7b2fbd5e37acff2be0",
 45        "pleural_effusion": "5dd07c327fb5723c5bbb48f2a02d7f365513d3ad136811fbe4def330ef2d7f6a"
 46    }
 47}
 48
 49
 50ZIPFILES = {
 51    "thoracic": "thoracic.zip",
 52    "pleural_effusion": "pleural_effusion.zip"
 53}
 54
 55
 56def get_plethora_data(
 57    path: Union[os.PathLike, str], task: Literal["thoracic", "pleural_effusion"], download: bool = False
 58) -> Tuple[str, str, str]:
 59    """Get the PLETHORA dataset.
 60
 61    Args:
 62        path: Filepath to a folder where the data is downloaded for further processing.
 63        task: The choice of task.
 64        download: Whether to download the data if it is not present.
 65
 66    Returns:
 67        Filepath where the data is downloaded.
 68    """
 69    image_dir = os.path.join(path, "data", "images")
 70    gt_dir = os.path.join(path, "data", "gt", "Thoracic_Cavities" if task == "thoracic" else "Effusions")
 71    csv_path = os.path.join(path, "plethora_images")
 72    if os.path.exists(image_dir) and os.path.exists(gt_dir):
 73        return image_dir, gt_dir, Path(csv_path).with_suffix(".csv")
 74
 75    os.makedirs(path, exist_ok=True)
 76
 77    # let's download dicom files from the tcia manifest
 78    tcia_path = os.path.join(path, "NSCLC-Radiomics-OriginalCTs.tcia")
 79    util.download_source_tcia(path=tcia_path, url=URL["image"], dst=image_dir, csv_filename=csv_path, download=download)
 80
 81    # let's download the segmentations from zipfiles
 82    zip_path = os.path.join(path, ZIPFILES[task])
 83    util.download_source(
 84        path=zip_path, url=URL["gt"][task], download=download, checksum=CHECKSUMS["gt"][task]
 85    )
 86    util.unzip(zip_path=zip_path, dst=os.path.join(path, "data", "gt"))
 87
 88    return image_dir, gt_dir, Path(csv_path).with_suffix(".csv")
 89
 90
 91def _assort_plethora_inputs(image_dir, gt_dir, task, csv_path):
 92    import nibabel as nib
 93    import pydicom as dicom
 94
 95    df = pd.read_csv(csv_path)
 96
 97    task_gt_dir = os.path.join(gt_dir, )
 98
 99    os.makedirs(os.path.join(image_dir, "preprocessed"), exist_ok=True)
100    os.makedirs(os.path.join(task_gt_dir, "preprocessed"), exist_ok=True)
101
102    # let's get all the series uid of the volumes downloaded and spot their allocated subject id
103    all_series_uid_dirs = glob(os.path.join(image_dir, "1.3*"))
104    image_paths, gt_paths = [], []
105    for series_uid_dir in tqdm(all_series_uid_dirs):
106        series_uid = os.path.split(series_uid_dir)[-1]
107        subject_id = pd.Series.to_string(df.loc[df["Series UID"] == series_uid]["Subject ID"])[-9:]
108
109        try:
110            gt_path = glob(os.path.join(task_gt_dir, subject_id, "*.nii.gz"))[0]
111        except IndexError:
112            # - some patients do not have "Thoracic_Cavities" segmentation
113            print(f"The ground truth is missing for subject '{subject_id}'")
114            continue
115
116        assert os.path.exists(gt_path)
117
118        vol_path = os.path.join(image_dir, "preprocessed", f"{subject_id}.nii.gz")
119        neu_gt_path = os.path.join(task_gt_dir, "preprocessed", os.path.split(gt_path)[-1])
120
121        image_paths.append(vol_path)
122        gt_paths.append(neu_gt_path)
123        if os.path.exists(vol_path) and os.path.exists(neu_gt_path):
124            continue
125
126        # the individual slices for the inputs need to be merged into one volume.
127        if not os.path.exists(vol_path):
128            all_dcm_slices = natsorted(glob(os.path.join(series_uid_dir, "*.dcm")))
129            all_slices = []
130            for dcm_path in all_dcm_slices:
131                dcmfile = dicom.dcmread(dcm_path)
132                img = dcmfile.pixel_array
133                all_slices.append(img)
134
135            volume = np.stack(all_slices)
136            volume = volume.transpose(1, 2, 0)
137            nii_vol = nib.Nifti1Image(volume, np.eye(4))
138            nii_vol.header.get_xyzt_units()
139            nii_vol.to_filename(vol_path)
140
141        # the ground truth needs to be aligned as the inputs, let's take care of that.
142        gt = nib.load(gt_path)
143        gt = gt.get_fdata()
144        gt = gt.transpose(2, 1, 0)  # aligning w.r.t the inputs
145        gt = np.flip(gt, axis=(0, 1))
146
147        gt = gt.transpose(1, 2, 0)
148        gt_nii_vol = nib.Nifti1Image(gt, np.eye(4))
149        gt_nii_vol.header.get_xyzt_units()
150        gt_nii_vol.to_filename(neu_gt_path)
151
152    return image_paths, gt_paths
153
154
155def get_plethora_paths(
156    path: Union[os.PathLike, str], task: Literal["thoracic", "pleural_effusion"], download: bool = False
157) -> Tuple[List[str], List[str]]:
158    """Get paths to the PLETHORA data.
159
160    Args:
161        path: Filepath to a folder where the data is downloaded for further processing.
162        task: The choice of task.
163        download: Whether to download the data if it is not present.
164
165    Returns:
166        List of filepaths for the image data.
167        List of filepaths for the label data.
168    """
169    image_dir, gt_dir, csv_path = get_plethora_data(path, task, download)
170    image_paths, gt_paths = _assort_plethora_inputs(image_dir=image_dir, gt_dir=gt_dir, task=task, csv_path=csv_path)
171    return image_paths, gt_paths
172
173
174def get_plethora_dataset(
175    path: Union[os.PathLike, str],
176    patch_shape: Tuple[int, ...],
177    task: Literal["thoracic", "pleural_effusion"],
178    resize_inputs: bool = False,
179    download: bool = False,
180    **kwargs
181) -> Dataset:
182    """Get the PLETHORA dataset.
183
184    Args:
185        path: Filepath to a folder where the data is downloaded for further processing.
186        patch_shape: The patch shape to use for training.
187        task: The choice of task.
188        resize_inputs: Whether to resize inputs to the desired patch shape.
189        download: Whether to download the data if it is not present.
190        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
191
192    Returns:
193        The segmentation dataset.
194    """
195    image_paths, gt_paths = get_plethora_paths(path, task, download)
196
197    if resize_inputs:
198        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
199        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
200            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
201        )
202
203    return torch_em.default_segmentation_dataset(
204        raw_paths=image_paths,
205        raw_key="data",
206        label_paths=gt_paths,
207        label_key="data",
208        patch_shape=patch_shape,
209        **kwargs
210    )
211
212
213def get_plethora_loader(
214    path: Union[os.PathLike, str],
215    batch_size: int,
216    patch_shape: Tuple[int, ...],
217    task: Literal["thoracic", "pleural_effusion"],
218    resize_inputs: bool = False,
219    download: bool = False,
220    **kwargs
221) -> DataLoader:
222    """Get the PLETHORA dataloader.
223
224    Args:
225        path: Filepath to a folder where the data is downloaded for further processing.
226        batch_size: The batch size for training.
227        patch_shape: The patch shape to use for training.
228        task: The choice of task.
229        resize_inputs: Whether to resize inputs to the desired patch shape.
230        download: Whether to download the data if it is not present.
231        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
232
233    Returns:
234        The DataLoader.
235    """
236    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
237    dataset = get_plethora_dataset(path, patch_shape, task, resize_inputs, download, **ds_kwargs)
238    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)