torch_em.data.datasets.medical.acrin_hnscc

  1"""The ACRIN-HNSCC-FDG-PET-CT dataset contains annotations for head and neck tumor and lymph node segmentation
  2in FDG-PET and CT of head and neck squamous cell carcinoma patients from the ACRIN 6685 trial.
  3
  4The lesion segmentations were created by radiologists in the 'ACRIN 6685-Tumor-Annotations' analysis result
  5of the collection and are distributed as DICOM RTSTRUCT (one file per lesion, plus a seed point file per lesion
  6and 'negative' assessments for scans without findings). They reference 257 PET, 378 CT and 31 MR series.
  7The RTSTRUCT annotations are public, but the images are only available under the NIH Controlled Data Access
  8Policy and have to be downloaded manually, see `get_acrin_hnscc_data`. This module rasterizes the lesion contours
  9onto the image grid (see `torch_em.data.datasets.util.rasterize_rtstruct`) and stores images and labels in hdf5
 10files, one per annotated PET or CT series. The labels are instance labels: each annotated lesion (primary tumor
 11or lymph node) gets its own id, starting from 1.
 12
 13NOTE: This requires the pydicom python package.
 14
 15The dataset is located at https://www.cancerimagingarchive.net/collection/acrin-hnscc-fdg-pet-ct/ and the annotations
 16at https://www.cancerimagingarchive.net/analysis-result/acrin-6685-tumor-annotations/.
 17
 18This dataset is from the publication https://doi.org/10.1200/JCO.18.01182.
 19The data was released at https://doi.org/10.7937/K9/TCIA.2016.JQEJZZNG (images) and https://doi.org/10.7937/jvgc-aq36
 20(annotations). Please cite them if you use this dataset in your research.
 21"""
 22
 23import os
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from collections import defaultdict
 28from typing import Union, Tuple, List, Literal
 29
 30import numpy as np
 31import pandas as pd
 32
 33from torch.utils.data import Dataset, DataLoader
 34
 35import torch_em
 36
 37from .. import util
 38
 39
 40URLS = {
 41    "annotations": "https://www.cancerimagingarchive.net/wp-content/uploads/ACRIN6685_Tumor-Annotations-manifest_2026-01-08.tcia",  # noqa
 42    "metadata": "https://www.cancerimagingarchive.net/wp-content/uploads/Metadata_Report_ACRIN6685_2026-01-08.csv",
 43}
 44
 45CHECKSUMS = {
 46    "annotations": None,  # The DICOM series are downloaded individually from TCIA.
 47    "metadata": "2462ec93293e8a4d9fdf1c1ae7c6badb0766a383507ea2947c4abb5b6223a4e0",
 48}
 49
 50MODALITIES = {"PET": "PT", "CT": "CT"}
 51
 52
 53def _find_image_series(image_dir):
 54    """Map the series instance UIDs to the folders with the DICOM files, independent of the folder layout."""
 55    import pydicom
 56
 57    series_dirs = {}
 58    for root, _, files in os.walk(image_dir):
 59        dcm_files = [fname for fname in files if fname.endswith(".dcm")]
 60        if dcm_files:
 61            header = pydicom.dcmread(os.path.join(root, dcm_files[0]), stop_before_pixels=True)
 62            series_dirs[str(header.SeriesInstanceUID)] = root
 63    return series_dirs
 64
 65
 66def _preprocess_acrin_hnscc(annotation_dir, image_dir, metadata_path, preprocessed_dir, modality):
 67    import h5py
 68
 69    metadata = pd.read_csv(metadata_path)
 70    metadata = metadata[
 71        (metadata["AnnotationType"] == "Segmentation") &
 72        (metadata["ReferencedSeriesModality"] == MODALITIES[modality])
 73    ]
 74    lesions_per_series = defaultdict(list)
 75    for _, row in metadata.iterrows():
 76        lesions_per_series[row["ReferencedSeriesInstanceUID"]].append(row)
 77
 78    image_series = _find_image_series(image_dir)
 79    missing = [uid for uid in lesions_per_series if uid not in image_series]
 80    if missing:
 81        print(f"{len(missing)} of {len(lesions_per_series)} annotated {modality} series were not found in {image_dir}.")
 82
 83    os.makedirs(preprocessed_dir, exist_ok=True)
 84    for series_uid, lesions in tqdm(sorted(lesions_per_series.items()), desc=f"Preprocess ACRIN-HNSCC {modality}"):
 85        if series_uid not in image_series:
 86            continue
 87        patient_id = lesions[0]["PatientID"]
 88        time_point = lesions[0]["ClinicalTrialTimePointID"].lower().replace(" ", "_").replace("#", "")
 89        out_path = os.path.join(preprocessed_dir, f"{patient_id}_{time_point}_{series_uid[-8:]}.h5")
 90        if os.path.exists(out_path):
 91            continue
 92
 93        volume, geometry = util.load_dicom_series(image_series[series_uid])
 94        if modality == "CT":
 95            volume = np.round(volume).astype("int16")
 96
 97        labels = np.zeros(volume.shape, dtype="uint8")
 98        for lesion_id, lesion in enumerate(sorted(lesions, key=lambda row: row["TrackingID"]), start=1):
 99            rtstruct_path = glob(os.path.join(annotation_dir, lesion["SeriesInstanceUID"], "*.dcm"))[0]
100            lesion_mask = util.rasterize_rtstruct(
101                rtstruct_path, geometry, volume.shape, lambda roi_number, roi_name: 1
102            )
103            labels[lesion_mask > 0] = lesion_id
104
105        with h5py.File(out_path, "w") as f:
106            f.create_dataset("raw", data=volume, compression="gzip")
107            f.create_dataset("labels", data=labels, compression="gzip")
108
109
110def get_acrin_hnscc_data(
111    path: Union[os.PathLike, str], modality: Literal["PET", "CT"], download: bool = False
112) -> str:
113    """Download the ACRIN-HNSCC-FDG-PET-CT dataset.
114
115    The RTSTRUCT annotations and their metadata are downloaded automatically. The images are only available
116    under the NIH Controlled Data Access Policy and cannot be downloaded automatically. To obtain them:
117    1. Request access following https://www.cancerimagingarchive.net/access-data/ (NIH Controlled Data Access).
118    2. On https://www.cancerimagingarchive.net/analysis-result/acrin-6685-tumor-annotations/ download the manifest
119       'Original ACRIN 6685 Images used to create Segmentations and Seed Points' (CT, PT, MR, 39.9 GB) and open it
120       with the NBIA Data Retriever, logged in with your account.
121    3. Store the downloaded DICOM series in the folder '<path>/images'. Any folder layout works, as the series are
122       identified by the 'SeriesInstanceUID' in the DICOM headers.
123
124    Args:
125        path: Filepath to a folder where the data is downloaded for further processing.
126        modality: The imaging modality. One of 'PET' or 'CT'.
127        download: Whether to download the data if it is not present.
128
129    Returns:
130        Filepath where the preprocessed data is stored.
131    """
132    assert modality in MODALITIES, f"'{modality}' is not a valid modality. Choose one of {list(MODALITIES)}."
133    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
134    preprocessed_dir = os.path.join(path, "preprocessed", modality)
135    os.makedirs(path, exist_ok=True)
136
137    # Download the public RTSTRUCT annotations and the annotation metadata.
138    annotation_dir = os.path.join(path, "annotations")
139    csv_path = os.path.join(path, "acrin_hnscc_annotations")
140    if not os.path.exists(f"{csv_path}.csv"):
141        util.download_source_tcia(
142            path=os.path.join(path, os.path.basename(URLS["annotations"])), url=URLS["annotations"],
143            dst=annotation_dir, csv_filename=csv_path, download=download,
144        )
145    metadata_path = os.path.join(path, os.path.basename(URLS["metadata"]))
146    util.download_source(path=metadata_path, url=URLS["metadata"], download=download, checksum=CHECKSUMS["metadata"])
147
148    image_dir = os.path.join(path, "images")
149    if not os.path.exists(image_dir):
150        raise RuntimeError(
151            f"The ACRIN-HNSCC-FDG-PET-CT images were not found at {image_dir}. They are only available under the "
152            "NIH Controlled Data Access Policy and have to be downloaded manually: request access at "
153            "https://www.cancerimagingarchive.net/access-data/, then download the manifest 'Original ACRIN 6685 "
154            "Images used to create Segmentations and Seed Points' from "
155            "https://www.cancerimagingarchive.net/analysis-result/acrin-6685-tumor-annotations/ with the NBIA Data "
156            f"Retriever and store the DICOM series in {image_dir}."
157        )
158
159    _preprocess_acrin_hnscc(annotation_dir, image_dir, metadata_path, preprocessed_dir, modality)
160    return preprocessed_dir
161
162
163def get_acrin_hnscc_paths(
164    path: Union[os.PathLike, str], modality: Literal["PET", "CT"], download: bool = False
165) -> List[str]:
166    """Get paths to the ACRIN-HNSCC-FDG-PET-CT data.
167
168    Args:
169        path: Filepath to a folder where the data is downloaded for further processing.
170        modality: The imaging modality. One of 'PET' or 'CT'.
171        download: Whether to download the data if it is not present.
172
173    Returns:
174        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
175    """
176    data_dir = get_acrin_hnscc_data(path, modality, download)
177    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
178    return volume_paths
179
180
181def get_acrin_hnscc_dataset(
182    path: Union[os.PathLike, str],
183    patch_shape: Tuple[int, ...],
184    modality: Literal["PET", "CT"],
185    resize_inputs: bool = False,
186    download: bool = False,
187    **kwargs
188) -> Dataset:
189    """Get the ACRIN-HNSCC-FDG-PET-CT dataset for head and neck lesion segmentation.
190
191    Args:
192        path: Filepath to a folder where the data is downloaded for further processing.
193        patch_shape: The patch shape to use for training.
194        modality: The imaging modality. One of 'PET' or 'CT'.
195        resize_inputs: Whether to resize inputs to the desired patch shape.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
198
199    Returns:
200        The segmentation dataset.
201    """
202    volume_paths = get_acrin_hnscc_paths(path, modality, download)
203
204    if resize_inputs:
205        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
206        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
207            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
208        )
209
210    return torch_em.default_segmentation_dataset(
211        raw_paths=volume_paths,
212        raw_key="raw",
213        label_paths=volume_paths,
214        label_key="labels",
215        patch_shape=patch_shape,
216        is_seg_dataset=True,
217        **kwargs
218    )
219
220
221def get_acrin_hnscc_loader(
222    path: Union[os.PathLike, str],
223    batch_size: int,
224    patch_shape: Tuple[int, ...],
225    modality: Literal["PET", "CT"],
226    resize_inputs: bool = False,
227    download: bool = False,
228    **kwargs
229) -> DataLoader:
230    """Get the ACRIN-HNSCC-FDG-PET-CT dataloader for head and neck lesion segmentation.
231
232    Args:
233        path: Filepath to a folder where the data is downloaded for further processing.
234        batch_size: The batch size for training.
235        patch_shape: The patch shape to use for training.
236        modality: The imaging modality. One of 'PET' or 'CT'.
237        resize_inputs: Whether to resize inputs to the desired patch shape.
238        download: Whether to download the data if it is not present.
239        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
240
241    Returns:
242        The DataLoader.
243    """
244    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
245    dataset = get_acrin_hnscc_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
246    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)