torch_em.data.datasets.medical.soft_tissue_sarcoma

  1"""The Soft-tissue-Sarcoma dataset contains annotations for tumor segmentation in MRI and FDG-PET/CT
  2of patients with soft-tissue sarcomas of the extremities.
  3
  4It consists of 51 patients with a T1-weighted MRI, a T2-weighted fat-suppressed MRI (T2FS, or STIR if T2FS
  5was not available), a CT and a FDG-PET scan each. The tumor was manually delineated on the T2FS scan by a
  6radiation oncologist and the contours were propagated to the other scans by rigid registration. The contours
  7are distributed as DICOM RTSTRUCT and rasterized onto the image grid by this module
  8(see `torch_em.data.datasets.util.rasterize_rtstruct`). Images and labels are stored in hdf5 files.
  9The semantic label ids are: 1: tumor ('GTV_Mass'), 2: peritumoral edema ('GTV_Edema' outside of the tumor,
 10annotated for 32 of the 51 patients).
 11
 12NOTE: This requires the pydicom python package.
 13
 14The dataset is located at https://www.cancerimagingarchive.net/collection/soft-tissue-sarcoma/.
 15
 16This dataset is from the publication https://doi.org/10.1088/0031-9155/60/14/5471.
 17The data was released at https://doi.org/10.7937/K9/TCIA.2015.7GO2GSKS.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import csv
 23from glob import glob
 24from tqdm import tqdm
 25from natsort import natsorted
 26from typing import Union, Tuple, List, Literal
 27
 28import numpy as np
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URL = "https://www.cancerimagingarchive.net/wp-content/uploads/doiJNLP-zgVcrK7I.tcia"
 38
 39# The DICOM series are downloaded individually from TCIA.
 40CHECKSUM = None
 41
 42# The ROI names of the contours. 'GTV_Edema' incorporates the tumor, so the edema id is only assigned outside
 43# of the tumor (`rasterize_rtstruct` gives precedence to the lower label id). One patient (STS_030) uses the
 44# names 'GTV_Research' and 'GTV_Res+edema' for the same two structures.
 45LABEL_IDS = {"GTV_Mass": 1, "GTV_Research": 1, "GTV_Edema": 2, "GTV_Res+edema": 2}
 46
 47# The RTSTRUCT series descriptions per modality, lower-cased. The T2FS category consists of the T2-weighted
 48# fat-saturated scans (26 patients) and the STIR scans used where they were not available (25 patients).
 49# The dataset additionally contains the T1 and T2FS scans registered and resampled to the PET scan
 50# ('RTstructAlignedT1toPET' etc.), which are not used here.
 51MODALITIES = {
 52    "T1": ["rtstructt1"],
 53    "T2FS": ["rtstructt2fs", "rtstructstir"],
 54    "CT": ["rtstructct"],
 55    "PET": ["rtstructpet"],
 56}
 57
 58
 59def _get_referenced_series(rtstruct_path):
 60    import pydicom
 61
 62    rtstruct = pydicom.dcmread(rtstruct_path, stop_before_pixels=True)
 63    return str(
 64        rtstruct.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence[0]
 65        .RTReferencedSeriesSequence[0].SeriesInstanceUID
 66    )
 67
 68
 69def _preprocess_soft_tissue_sarcoma(dicom_dir, csv_path, preprocessed_dir, modality):
 70    import h5py
 71
 72    with open(csv_path, "r") as f:
 73        rows = list(csv.DictReader(f))
 74    # The series descriptions are matched case-insensitively, as the collection is not consistent
 75    # (e.g. 'RTstructCT' and 'RTStructCT').
 76    rtstruct_series = {
 77        row["Series UID"]: row["Subject ID"] for row in rows
 78        if row["Modality"] == "RTSTRUCT" and row["Series Description"].lower() in MODALITIES[modality]
 79    }
 80
 81    os.makedirs(preprocessed_dir, exist_ok=True)
 82    for series_uid, subject_id in tqdm(sorted(rtstruct_series.items()), desc=f"Preprocess STS {modality}"):
 83        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
 84        if os.path.exists(out_path):
 85            continue
 86
 87        rtstruct_path = glob(os.path.join(dicom_dir, series_uid, "*.dcm"))[0]
 88        image_dir = os.path.join(dicom_dir, _get_referenced_series(rtstruct_path))
 89        volume, geometry = util.load_dicom_series(image_dir)
 90        if modality == "CT":
 91            volume = np.round(volume).astype("int16")
 92        labels = util.rasterize_rtstruct(rtstruct_path, geometry, volume.shape, LABEL_IDS)
 93
 94        with h5py.File(out_path, "w") as f:
 95            f.create_dataset("raw", data=volume, compression="gzip")
 96            f.create_dataset("labels", data=labels, compression="gzip")
 97
 98
 99def get_soft_tissue_sarcoma_data(
100    path: Union[os.PathLike, str], modality: Literal["T1", "T2FS", "CT", "PET"], download: bool = False
101) -> str:
102    """Download the Soft-tissue-Sarcoma dataset.
103
104    Args:
105        path: Filepath to a folder where the data is downloaded for further processing.
106        modality: The imaging modality. One of 'T1', 'T2FS', 'CT' or 'PET'.
107        download: Whether to download the data if it is not present.
108
109    Returns:
110        Filepath where the preprocessed data is stored.
111    """
112    assert modality in MODALITIES, f"'{modality}' is not a valid modality. Choose one of {list(MODALITIES)}."
113    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
114    preprocessed_dir = os.path.join(path, "preprocessed", modality)
115    os.makedirs(path, exist_ok=True)
116
117    # Download the DICOM series (MR, CT, PT and RTSTRUCT) from the TCIA manifest. The series metadata are written
118    # after all series are downloaded, so their presence means the download is complete.
119    dicom_dir = os.path.join(path, "dicom")
120    csv_path = os.path.join(path, "soft_tissue_sarcoma_series")
121    if not os.path.exists(f"{csv_path}.csv"):
122        util.download_source_tcia(
123            path=os.path.join(path, os.path.basename(URL)), url=URL, dst=dicom_dir, csv_filename=csv_path,
124            download=download,
125        )
126
127    _preprocess_soft_tissue_sarcoma(dicom_dir, f"{csv_path}.csv", preprocessed_dir, modality)
128    return preprocessed_dir
129
130
131def get_soft_tissue_sarcoma_paths(
132    path: Union[os.PathLike, str], modality: Literal["T1", "T2FS", "CT", "PET"], download: bool = False
133) -> List[str]:
134    """Get paths to the Soft-tissue-Sarcoma data.
135
136    Args:
137        path: Filepath to a folder where the data is downloaded for further processing.
138        modality: The imaging modality. One of 'T1', 'T2FS', 'CT' or 'PET'.
139        download: Whether to download the data if it is not present.
140
141    Returns:
142        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
143    """
144    data_dir = get_soft_tissue_sarcoma_data(path, modality, download)
145    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
146    return volume_paths
147
148
149def get_soft_tissue_sarcoma_dataset(
150    path: Union[os.PathLike, str],
151    patch_shape: Tuple[int, ...],
152    modality: Literal["T1", "T2FS", "CT", "PET"],
153    resize_inputs: bool = False,
154    download: bool = False,
155    **kwargs
156) -> Dataset:
157    """Get the Soft-tissue-Sarcoma dataset for tumor segmentation.
158
159    Args:
160        path: Filepath to a folder where the data is downloaded for further processing.
161        patch_shape: The patch shape to use for training.
162        modality: The imaging modality. One of 'T1', 'T2FS', 'CT' or 'PET'.
163        resize_inputs: Whether to resize inputs to the desired patch shape.
164        download: Whether to download the data if it is not present.
165        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
166
167    Returns:
168        The segmentation dataset.
169    """
170    volume_paths = get_soft_tissue_sarcoma_paths(path, modality, download)
171
172    if resize_inputs:
173        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
174        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
175            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
176        )
177
178    return torch_em.default_segmentation_dataset(
179        raw_paths=volume_paths,
180        raw_key="raw",
181        label_paths=volume_paths,
182        label_key="labels",
183        patch_shape=patch_shape,
184        is_seg_dataset=True,
185        **kwargs
186    )
187
188
189def get_soft_tissue_sarcoma_loader(
190    path: Union[os.PathLike, str],
191    batch_size: int,
192    patch_shape: Tuple[int, ...],
193    modality: Literal["T1", "T2FS", "CT", "PET"],
194    resize_inputs: bool = False,
195    download: bool = False,
196    **kwargs
197) -> DataLoader:
198    """Get the Soft-tissue-Sarcoma dataloader for tumor segmentation.
199
200    Args:
201        path: Filepath to a folder where the data is downloaded for further processing.
202        batch_size: The batch size for training.
203        patch_shape: The patch shape to use for training.
204        modality: The imaging modality. One of 'T1', 'T2FS', 'CT' or 'PET'.
205        resize_inputs: Whether to resize inputs to the desired patch shape.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
208
209    Returns:
210        The DataLoader.
211    """
212    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
213    dataset = get_soft_tissue_sarcoma_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
214    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)