torch_em.data.datasets.medical.psma_pet_ct

  1"""The PSMA-PET-CT-Lesions dataset contains annotations for tumor lesions in whole-body PSMA PET/CT.
  2
  3The dataset consists of whole-body PSMA (prostate-specific membrane antigen) PET scans of prostate
  4cancer patients, with an expert DICOM-SEG object marking the detected tumor lesions on each PET series.
  5Lesions are frequently small metastatic foci, so most of a scan is background.
  6
  7NOTE: This requires the pydicom python package.
  8
  9The dataset is located at https://doi.org/10.7937/r7ep-3x37 and is distributed under the
 10CC BY 4.0 license.
 11Please cite it if you use this dataset in your research.
 12"""
 13
 14import os
 15import json
 16from glob import glob
 17from tqdm import tqdm
 18from natsort import natsorted
 19from typing import Union, Tuple, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .adrenal_acc import _load_dicom_volume, _load_dicom_seg, _resample_labels
 26from .. import util
 27
 28
 29COLLECTION = "PSMA-PET-CT-Lesions"
 30
 31
 32def _get_series_metadata(path, download):
 33    """Get the metadata of all series in the collection from the NBIA REST API."""
 34    import requests
 35
 36    metadata_path = os.path.join(path, "psma_pet_ct_series.json")
 37    if not os.path.exists(metadata_path):
 38        if not download:
 39            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
 40        response = requests.get(f"{util.NBIA_API_URL}getSeries", params={"Collection": COLLECTION})
 41        response.raise_for_status()
 42        with open(metadata_path, "w") as f:
 43            json.dump(response.json(), f, indent=2)
 44
 45    with open(metadata_path, "r") as f:
 46        return json.load(f)
 47
 48
 49def _referenced_series_uid(seg):
 50    return str(seg.ReferencedSeriesSequence[0].SeriesInstanceUID)
 51
 52
 53def _preprocess_psma_pet_ct(dicom_dir, series_metadata, preprocessed_dir):
 54    import h5py
 55    import pydicom
 56
 57    seg_series = [series for series in series_metadata if series.get("Modality") == "SEG"]
 58
 59    os.makedirs(preprocessed_dir, exist_ok=True)
 60    for series in tqdm(seg_series, desc="Preprocess PSMA-PET-CT-Lesions"):
 61        seg_paths = glob(os.path.join(dicom_dir, series["SeriesInstanceUID"], "*.dcm"))
 62        if not seg_paths:
 63            continue
 64
 65        out_path = os.path.join(preprocessed_dir, f"{series['SeriesInstanceUID']}.h5")
 66        if os.path.exists(out_path):
 67            continue
 68
 69        seg = pydicom.dcmread(seg_paths[0])
 70        pet_dir = os.path.join(dicom_dir, _referenced_series_uid(seg))
 71        if not glob(os.path.join(pet_dir, "*.dcm")):
 72            continue
 73
 74        volume, pet_affine = _load_dicom_volume(pet_dir)
 75        seg_labels, seg_affine = _load_dicom_seg(seg_paths[0])
 76        labels = _resample_labels(seg_labels, seg_affine, volume.shape, pet_affine)
 77        if labels.max() == 0:
 78            continue
 79
 80        with h5py.File(out_path, "w") as f:
 81            f.create_dataset("raw", data=volume, compression="gzip")
 82            f.create_dataset("labels", data=labels, compression="gzip")
 83
 84
 85def get_psma_pet_ct_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 86    """Download the PSMA-PET-CT-Lesions dataset.
 87
 88    Args:
 89        path: Filepath to a folder where the data is downloaded for further processing.
 90        download: Whether to download the data if it is not present.
 91
 92    Returns:
 93        Filepath where the preprocessed data is stored.
 94    """
 95    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
 96    preprocessed_dir = os.path.join(path, "preprocessed")
 97
 98    os.makedirs(path, exist_ok=True)
 99    series_metadata = _get_series_metadata(path, download)
100
101    seg_uids = [series["SeriesInstanceUID"] for series in series_metadata if series.get("Modality") == "SEG"]
102
103    dicom_dir = os.path.join(path, "dicom")
104    if download:  # The SEG series are downloaded first, so the PET series they reference can be found.
105        util.download_tcia_series(seg_uids, dst=dicom_dir, csv_filename=os.path.join(path, "psma_pet_ct_seg"))
106
107        import pydicom
108        pet_uids = set()
109        for uid in seg_uids:
110            seg_paths = glob(os.path.join(dicom_dir, uid, "*.dcm"))
111            if seg_paths:
112                seg = pydicom.dcmread(seg_paths[0])
113                pet_uids.add(_referenced_series_uid(seg))
114        util.download_tcia_series(sorted(pet_uids), dst=dicom_dir, csv_filename=os.path.join(path, "psma_pet_ct_pet"))
115    elif not glob(os.path.join(dicom_dir, "*", "*.dcm")):
116        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
117
118    _preprocess_psma_pet_ct(dicom_dir, series_metadata, preprocessed_dir)
119    return preprocessed_dir
120
121
122def get_psma_pet_ct_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
123    """Get paths to the PSMA-PET-CT-Lesions data.
124
125    Args:
126        path: Filepath to a folder where the data is downloaded for further processing.
127        download: Whether to download the data if it is not present.
128
129    Returns:
130        List of filepaths for the stored data.
131    """
132    preprocessed_dir = get_psma_pet_ct_data(path, download)
133    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
134    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
135    return volume_paths
136
137
138def get_psma_pet_ct_dataset(
139    path: Union[os.PathLike, str],
140    patch_shape: Tuple[int, ...],
141    resize_inputs: bool = False,
142    download: bool = False,
143    **kwargs
144) -> Dataset:
145    """Get the PSMA-PET-CT-Lesions dataset for tumor lesion segmentation.
146
147    Args:
148        path: Filepath to a folder where the data is downloaded for further processing.
149        patch_shape: The patch shape to use for training.
150        resize_inputs: Whether to resize inputs to the desired patch shape.
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
153
154    Returns:
155        The segmentation dataset.
156    """
157    volume_paths = get_psma_pet_ct_paths(path, download)
158
159    if resize_inputs:
160        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
161        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
162            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
163        )
164
165    return torch_em.default_segmentation_dataset(
166        raw_paths=volume_paths,
167        raw_key="raw",
168        label_paths=volume_paths,
169        label_key="labels",
170        patch_shape=patch_shape,
171        is_seg_dataset=True,
172        **kwargs
173    )
174
175
176def get_psma_pet_ct_loader(
177    path: Union[os.PathLike, str],
178    batch_size: int,
179    patch_shape: Tuple[int, ...],
180    resize_inputs: bool = False,
181    download: bool = False,
182    **kwargs
183) -> DataLoader:
184    """Get the PSMA-PET-CT-Lesions dataloader for tumor lesion segmentation.
185
186    Args:
187        path: Filepath to a folder where the data is downloaded for further processing.
188        batch_size: The batch size for training.
189        patch_shape: The patch shape to use for training.
190        resize_inputs: Whether to resize inputs to the desired patch shape.
191        download: Whether to download the data if it is not present.
192        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
193
194    Returns:
195        The DataLoader.
196    """
197    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
198    dataset = get_psma_pet_ct_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
199    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)