torch_em.data.datasets.medical.mediastinal_ct

  1"""The Mediastinal CT datasets contain annotations for mediastinal lymph node and anatomical structure
  2segmentation in contrast-enhanced chest CT scans of lung cancer patients (St. Olavs University Hospital, Trondheim).
  3
  4Two datasets are provided, selected via the `task` argument:
  5- 'lymph_nodes': The benchmark subset of Bouget et al. (2022). 15 CT volumes with 3D lymph node annotations
  6  (proofread by an expert radiologist). The label volume contains one instance id per lymph node; the mediastinal
  7  station of each node is listed in 'Benchmark/stations_sto.csv'. The archive also provides binary masks for the
  8  esophagus, azygos vein, subclavian / carotid arteries and brachiocephalic veins next to the CT of each patient.
  9- 'structures': The mediastinal CT dataset of Bouget et al. (2019). 15 CT volumes with annotations for 14 mediastinal
 10  anatomical structures (see `STRUCTURE_IDS`) and, in a separate file, the lymph nodes. The data is distributed as
 11  MetaImage (.mhd / .raw) files and is converted to nifti once by `get_mediastinal_ct_data`.
 12  Label ids: 1: esophagus, 2: pulmonary trunk, 3: aortic arch, 4: ascending aorta, 5: descending aorta, 6: azygos vein,
 13  7: heart, 8: vena cava, 9: brachiocephalic veins, 11: spine, 12: pulmonary vein, 13: subclavian and carotid arteries,
 14  14: lungs, 16: airways (ids 10 and 15 are not used).
 15
 16The datasets are located at https://github.com/dbouget/ct_mediastinal_structures_segmentation.
 17
 18The datasets are from the publications https://doi.org/10.1080/21681163.2022.2043778 (lymph nodes)
 19and https://doi.org/10.1007/s11548-019-01948-8 (structures).
 20Please cite them if you use these datasets in your research.
 21"""
 22
 23import os
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from typing import Union, Tuple, Literal, List
 28
 29import numpy as np
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38URLS = {
 39    "lymph_nodes": "https://drive.google.com/uc?id=1ZsFq7PslqQ5ow_dXB01kDkaKPqYDXD5d",
 40    "structures": "https://drive.google.com/uc?id=1YqCRcBpsFoE4JsBq5NROqIpeijnITpe1",
 41}
 42
 43CHECKSUMS = {
 44    "lymph_nodes": "c6b4bab94e8e69d72a9e9707b8e1c261944784c49652772519d3155d4d052478",
 45    "structures": "b7e0fa1c8e242259fdc062e68f56e9e561d92e36264c556854d2d9b55e282531",
 46}
 47
 48STRUCTURE_IDS = {
 49    "esophagus": 1,
 50    "pulmonary_trunk": 2,
 51    "aortic_arch": 3,
 52    "ascending_aorta": 4,
 53    "descending_aorta": 5,
 54    "azygos": 6,
 55    "heart": 7,
 56    "vena_cava": 8,
 57    "brachiocephalic_veins": 9,
 58    "spine": 11,
 59    "pulmonary_vein": 12,
 60    "subclavian_and_carotid_arteries": 13,
 61    "lungs": 14,
 62    "airways": 16,
 63}
 64"""The label ids of the anatomical structures in the 'structures' dataset."""
 65
 66MHD_DTYPES = {
 67    "MET_CHAR": np.int8, "MET_UCHAR": np.uint8, "MET_SHORT": np.int16, "MET_USHORT": np.uint16,
 68    "MET_INT": np.int32, "MET_UINT": np.uint32, "MET_FLOAT": np.float32, "MET_DOUBLE": np.float64,
 69}
 70
 71
 72def read_mhd(path: str) -> Tuple[np.ndarray, Tuple[float, ...]]:
 73    """Read an uncompressed MetaImage (.mhd + .raw) volume.
 74
 75    Args:
 76        path: The filepath to the .mhd header.
 77
 78    Returns:
 79        The volume with axis order (x, y, z), matching the axis order nibabel uses for nifti files.
 80        The voxel spacing in (x, y, z) order.
 81    """
 82    header = {}
 83    with open(path) as f:
 84        for line in f:
 85            if "=" in line:
 86                key, value = line.split("=", 1)
 87                header[key.strip()] = value.strip()
 88
 89    if header.get("CompressedData", "False") == "True":
 90        raise NotImplementedError(f"Compressed MetaImage data is not supported: {path}")
 91
 92    shape = tuple(int(s) for s in header["DimSize"].split())
 93    spacing = tuple(float(s) for s in header["ElementSpacing"].split())
 94    raw_path = os.path.join(os.path.split(path)[0], header["ElementDataFile"])
 95    data = np.fromfile(raw_path, dtype=MHD_DTYPES[header["ElementType"]])
 96
 97    # The raw data is stored with x as the fastest axis, i.e. in (z, y, x) order.
 98    data = data.reshape(shape[::-1]).transpose(2, 1, 0)
 99    return data, spacing
100
101
102def _convert_structures_to_nifti(data_dir):
103    import nibabel as nib
104
105    for patient_dir in tqdm(natsorted(glob(os.path.join(data_dir, "pat*"))), desc="Converting MetaImage to nifti"):
106        patient_id = os.path.basename(patient_dir)
107        inputs = {
108            "data": os.path.join(patient_dir, "Data", f"{patient_id}.mhd"),
109            "structures": os.path.join(patient_dir, "Structures", f"{patient_id}_organs_gt.mhd"),
110            "lymph_nodes": os.path.join(patient_dir, "LN", f"{patient_id}_LN_gt.mhd"),
111        }
112        for name, mhd_path in inputs.items():
113            out_path = os.path.join(patient_dir, f"{patient_id}_{name}.nii.gz")
114            if os.path.exists(out_path):
115                continue
116            data, spacing = read_mhd(mhd_path)
117            if name != "data":
118                data = data.astype("uint16")
119            nib.save(nib.Nifti1Image(data, np.diag(list(spacing) + [1.0])), out_path)
120
121
122def get_mediastinal_ct_data(
123    path: Union[os.PathLike, str], task: Literal["lymph_nodes", "structures"], download: bool = False
124) -> str:
125    """Download the Mediastinal CT dataset.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        task: The dataset to download. Either 'lymph_nodes' (benchmark subset) or 'structures' (mediastinal CT dataset).
130        download: Whether to download the data if it is not present.
131
132    Returns:
133        Filepath where the data is downloaded.
134    """
135    if task not in URLS:
136        raise ValueError(f"'{task}' is not a valid task. Choose from {list(URLS.keys())}.")
137
138    data_dir = os.path.join(path, "Benchmark" if task == "lymph_nodes" else "mediastinal_2019")
139    if os.path.exists(data_dir):
140        return data_dir
141
142    os.makedirs(path, exist_ok=True)
143
144    if task == "lymph_nodes":
145        zip_path = os.path.join(path, "benchmark_subset.zip")
146        util.download_source_gdrive(path=zip_path, url=URLS[task], download=download, checksum=CHECKSUMS[task])
147        util.unzip(zip_path=zip_path, dst=path, remove=False)
148    else:
149        # NOTE: The file is shared as 'mediastinal_ct_dataset.zip', but it is a gzipped tarball.
150        tar_path = os.path.join(path, "mediastinal_ct_dataset.tar.gz")
151        util.download_source_gdrive(path=tar_path, url=URLS[task], download=download, checksum=CHECKSUMS[task])
152        util.unzip_tarfile(tar_path=tar_path, dst=data_dir, remove=False)
153        _convert_structures_to_nifti(data_dir)
154
155    return data_dir
156
157
158def get_mediastinal_ct_paths(
159    path: Union[os.PathLike, str], task: Literal["lymph_nodes", "structures"], download: bool = False
160) -> Tuple[List[str], List[str]]:
161    """Get paths to the Mediastinal CT data.
162
163    Args:
164        path: Filepath to a folder where the data is downloaded for further processing.
165        task: The dataset to use. Either 'lymph_nodes' (benchmark subset) or 'structures' (mediastinal CT dataset).
166        download: Whether to download the data if it is not present.
167
168    Returns:
169        List of filepaths for the image data.
170        List of filepaths for the label data.
171    """
172    data_dir = get_mediastinal_ct_data(path, task, download)
173
174    if task == "lymph_nodes":
175        raw_paths = natsorted(glob(os.path.join(data_dir, "Pat*", "*_data.nii.gz")))
176        label_paths = [p.replace("_data.nii.gz", "_labels_LymphNodes.nii.gz") for p in raw_paths]
177    else:
178        raw_paths = natsorted(glob(os.path.join(data_dir, "pat*", "*_data.nii.gz")))
179        label_paths = [p.replace("_data.nii.gz", "_structures.nii.gz") for p in raw_paths]
180
181    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
182    assert all(os.path.exists(p) for p in label_paths)
183
184    return raw_paths, label_paths
185
186
187def get_mediastinal_ct_dataset(
188    path: Union[os.PathLike, str],
189    patch_shape: Tuple[int, ...],
190    task: Literal["lymph_nodes", "structures"],
191    resize_inputs: bool = False,
192    download: bool = False,
193    **kwargs
194) -> Dataset:
195    """Get the Mediastinal CT dataset for lymph node or anatomical structure segmentation.
196
197    Args:
198        path: Filepath to a folder where the data is downloaded for further processing.
199        patch_shape: The patch shape to use for training.
200        task: The dataset to use. Either 'lymph_nodes' (benchmark subset) or 'structures' (mediastinal CT dataset).
201        resize_inputs: Whether to resize inputs to the desired patch shape.
202        download: Whether to download the data if it is not present.
203        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
204
205    Returns:
206        The segmentation dataset.
207    """
208    raw_paths, label_paths = get_mediastinal_ct_paths(path, task, download)
209
210    if resize_inputs:
211        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
212        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
213            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
214        )
215
216    return torch_em.default_segmentation_dataset(
217        raw_paths=raw_paths,
218        raw_key="data",
219        label_paths=label_paths,
220        label_key="data",
221        patch_shape=patch_shape,
222        is_seg_dataset=True,
223        **kwargs
224    )
225
226
227def get_mediastinal_ct_loader(
228    path: Union[os.PathLike, str],
229    batch_size: int,
230    patch_shape: Tuple[int, ...],
231    task: Literal["lymph_nodes", "structures"],
232    resize_inputs: bool = False,
233    download: bool = False,
234    **kwargs
235) -> DataLoader:
236    """Get the Mediastinal CT dataloader for lymph node or anatomical structure segmentation.
237
238    Args:
239        path: Filepath to a folder where the data is downloaded for further processing.
240        batch_size: The batch size for training.
241        patch_shape: The patch shape to use for training.
242        task: The dataset to use. Either 'lymph_nodes' (benchmark subset) or 'structures' (mediastinal CT dataset).
243        resize_inputs: Whether to resize inputs to the desired patch shape.
244        download: Whether to download the data if it is not present.
245        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
246
247    Returns:
248        The DataLoader.
249    """
250    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
251    dataset = get_mediastinal_ct_dataset(path, patch_shape, task, resize_inputs, download, **ds_kwargs)
252    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)