torch_em.data.datasets.medical.hcc_tace

  1"""The HCC-TACE-Seg dataset contains annotations for liver, tumor and vessel segmentation in multiphasic
  2contrast-enhanced CT of patients with hepatocellular carcinoma treated with transarterial chemoembolization (TACE).
  3
  4It consists of 105 CT volumes (the pre-treatment CT phase that was segmented) with semantic labels for:
  51: liver, 2: tumor, 3: tumor necrosis (only annotated for one patient), 4: portal vein,
  65: abdominal aorta. The CT scans are distributed as DICOM series and the labels as DICOM-SEG objects, which are
  7converted and stored in hdf5 files by this module. The segments are painted in the order of their ids, i.e. the
  8tumor overwrites the liver and the necrosis overwrites the tumor.
  9
 10The collection contains several CT phases per patient (28.6 GB in total), but only one of them is segmented.
 11Hence, this module downloads the DICOM-SEG objects and the CT series they reference via the TCIA REST API
 12instead of the full collection manifest.
 13
 14NOTE: This requires the pydicom python package.
 15
 16The dataset is located at https://www.cancerimagingarchive.net/collection/hcc-tace-seg/.
 17
 18This dataset is from the publication https://doi.org/10.1038/s41597-023-01928-3.
 19The data was released at https://doi.org/10.7937/TCIA.5FNA-0924.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24from glob import glob
 25from tqdm import tqdm
 26from warnings import warn
 27from natsort import natsorted
 28from collections import Counter, defaultdict
 29from typing import Union, Tuple, List
 30
 31import numpy as np
 32import requests
 33
 34from torch.utils.data import Dataset, DataLoader
 35
 36import torch_em
 37
 38from .. import util
 39
 40
 41NBIA_API_URL = "https://services.cancerimagingarchive.net/nbia-api/services/v1"
 42COLLECTION = "HCC-TACE-Seg"
 43
 44# The DICOM series are downloaded individually from TCIA.
 45CHECKSUM = None
 46
 47NUM_VOLUMES = 105
 48
 49LABEL_IDS = {"liver": 1, "tumor": 2, "necrosis": 3, "portal_vein": 4, "abdominal_aorta": 5}
 50
 51# The segment labels used in the DICOM-SEG objects.
 52SEGMENT_NAMES = {
 53    "Liver": "liver", "Mass": "tumor", "Necrosis": "necrosis", "Portal vein": "portal_vein",
 54    "Abdominal aorta": "abdominal_aorta",
 55}
 56
 57
 58def _download_series(series_uid, dicom_dir, download):
 59    """Download a DICOM series via the NBIA REST API and extract it to '<dicom_dir>/<series_uid>'."""
 60    series_dir = os.path.join(dicom_dir, series_uid)
 61    if os.path.exists(series_dir):
 62        return series_dir
 63
 64    zip_path = os.path.join(dicom_dir, f"{series_uid}.zip")
 65    url = f"{NBIA_API_URL}/getImage?SeriesInstanceUID={series_uid}"
 66    util.download_source(path=zip_path, url=url, download=download, checksum=None)
 67    util.unzip(zip_path=zip_path, dst=f"{series_dir}.tmp")
 68    os.rename(f"{series_dir}.tmp", series_dir)
 69    return series_dir
 70
 71
 72def _get_series_per_patient(download):
 73    """Get the metadata of all series in the collection from the NBIA REST API, grouped by patient."""
 74    if not download:
 75        raise RuntimeError("Cannot find the data, but download was set to False.")
 76    response = requests.get(f"{NBIA_API_URL}/getSeries", params={"Collection": COLLECTION})
 77    response.raise_for_status()
 78
 79    series_per_patient = defaultdict(list)
 80    for series in response.json():
 81        series_per_patient[series["PatientID"]].append(series)
 82    return series_per_patient
 83
 84
 85def _get_candidate_series(seg_path, ct_series):
 86    """Rank the CT series of a patient by how likely it is that the DICOM-SEG object was drawn on them.
 87
 88    The series referenced by the DICOM-SEG object is tried first. This reference is missing for one patient
 89    (HCC_048) and points to a series that does not overlap with the segmentation for another one (HCC_089),
 90    so the other CT series are ranked by the difference between their slice count and the number of segmented
 91    slices and are used as fall-backs.
 92
 93    Returns a list of (series UID, SOP instance UIDs of the referenced slices) tuples.
 94    """
 95    import pydicom
 96
 97    seg = pydicom.dcmread(seg_path, stop_before_pixels=True)
 98    z_positions = {
 99        round(float(group.PlanePositionSequence[0].ImagePositionPatient[2]), 2)
100        for group in seg.PerFrameFunctionalGroupsSequence
101    }
102
103    referenced_uid, referenced_sop_uids = None, set()
104    if "ReferencedSeriesSequence" in seg:
105        assert len(seg.ReferencedSeriesSequence) == 1, f"Expected a single referenced CT series in {seg_path}."
106        referenced_series = seg.ReferencedSeriesSequence[0]
107        referenced_uid = str(referenced_series.SeriesInstanceUID)
108        referenced_sop_uids = {
109            str(instance.ReferencedSOPInstanceUID) for instance in referenced_series.ReferencedInstanceSequence
110        }
111
112    ranked = sorted(ct_series, key=lambda series: abs(int(series["ImageCount"]) - len(z_positions)))
113    candidates = [(referenced_uid, referenced_sop_uids)] if referenced_uid is not None else []
114    candidates += [(str(series["SeriesInstanceUID"]), set()) for series in ranked
115                   if str(series["SeriesInstanceUID"]) != referenced_uid]
116    return candidates
117
118
119def _load_dicom_volume(series_dir, referenced_sop_uids):
120    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
121
122    Some series contain several acquisitions (CT phases) with overlapping slice positions. In this case, only the
123    acquisition referenced by the DICOM-SEG object is kept, so that each slice position occurs once.
124    Returns the volume in Hounsfield units and the geometry needed to align the DICOM-SEG frames with the volume.
125    """
126    import pydicom
127
128    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
129
130    acquisitions = Counter(dcm.get("AcquisitionNumber") for dcm in slices if dcm.SOPInstanceUID in referenced_sop_uids)
131    acquisition = acquisitions.most_common(1)[0][0] if acquisitions else None
132    slices_per_position = {}
133    for dcm in slices:
134        z = round(float(dcm.ImagePositionPatient[2]), 2)
135        priority = (dcm.SOPInstanceUID in referenced_sop_uids, dcm.get("AcquisitionNumber") == acquisition)
136        if z not in slices_per_position or priority > slices_per_position[z][0]:
137            slices_per_position[z] = (priority, dcm)
138    slices = [dcm for _, dcm in sorted(slices_per_position.values(), key=lambda item: item[1].ImagePositionPatient[2])]
139
140    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
141    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
142    volume = np.round(volume).astype("int16")
143
144    geometry = {
145        "sop_uids": {str(dcm.SOPInstanceUID): i for i, dcm in enumerate(slices)},
146        "z_positions": np.array([float(dcm.ImagePositionPatient[2]) for dcm in slices]),
147        "orientation": np.round([float(v) for v in slices[0].ImageOrientationPatient]).astype("int"),
148    }
149    return volume, geometry
150
151
152def _load_dicom_seg(seg_path, shape, geometry):
153    """Convert a DICOM-SEG object into binary masks (one per segment) aligned with the reference CT volume.
154
155    Each frame is mapped to its CT slice via the source image it was derived from (or its z position).
156    """
157    import pydicom
158
159    seg = pydicom.dcmread(seg_path)
160    frames = seg.pixel_array
161    if frames.ndim == 2:  # A segmentation with a single frame.
162        frames = frames[None]
163
164    segment_names = {int(segment.SegmentNumber): str(segment.SegmentLabel) for segment in seg.SegmentSequence}
165
166    # The segmentation frames may use a different in-plane orientation than the CT slices,
167    # in which case they have to be flipped to align them.
168    seg_orientation = seg.SharedFunctionalGroupsSequence[0].PlaneOrientationSequence[0].ImageOrientationPatient
169    seg_orientation = np.round([float(v) for v in seg_orientation]).astype("int")
170    orientation = geometry["orientation"]
171    assert np.all(np.abs(seg_orientation) == np.abs(orientation)), f"Unexpected orientation in {seg_path}."
172    flip_axes = []
173    if np.any(seg_orientation[3:] != orientation[3:]):  # The direction of the rows differs.
174        flip_axes.append(0)
175    if np.any(seg_orientation[:3] != orientation[:3]):  # The direction of the columns differs.
176        flip_axes.append(1)
177
178    # Frames are matched to CT slices via their z position, if they cannot be matched via the source image.
179    z_positions = geometry["z_positions"]
180    tolerance = np.diff(z_positions).min() / 2 if len(z_positions) > 1 else 1.0
181
182    masks = {name: np.zeros(shape, dtype="bool") for name in segment_names.values()}
183    for frame, frame_group in zip(frames, seg.PerFrameFunctionalGroupsSequence):
184        segment_number = int(frame_group.SegmentIdentificationSequence[0].ReferencedSegmentNumber)
185        mask = frame.astype("bool")
186
187        z = None
188        derivation = frame_group.get("DerivationImageSequence", [])
189        if derivation and derivation[0].get("SourceImageSequence"):
190            z = geometry["sop_uids"].get(str(derivation[0].SourceImageSequence[0].ReferencedSOPInstanceUID))
191        if z is None:
192            frame_z = float(frame_group.PlanePositionSequence[0].ImagePositionPatient[2])
193            z = int(np.argmin(np.abs(z_positions - frame_z)))
194            if abs(z_positions[z] - frame_z) > tolerance:
195                if mask.any():
196                    warn(f"Skipping a frame of '{segment_names[segment_number]}' at z={frame_z} in {seg_path}, "
197                         "which does not match a CT slice.")
198                continue
199
200        if flip_axes:
201            mask = np.flip(mask, axis=flip_axes)
202        masks[segment_names[segment_number]][z] |= mask
203
204    return masks
205
206
207def _preprocess_hcc_tace(seg_path, ct_dir, referenced_sop_uids, out_path):
208    """Convert a CT series and the segmentation drawn on it into an hdf5 file.
209
210    Returns whether the segmentation could be mapped onto the CT series.
211    """
212    import h5py
213
214    volume, geometry = _load_dicom_volume(ct_dir, referenced_sop_uids)
215    masks = _load_dicom_seg(seg_path, volume.shape, geometry)
216
217    unknown = [name for name in masks if name not in SEGMENT_NAMES]
218    if unknown:
219        raise ValueError(f"Unknown segment labels {unknown} in {seg_path}.")
220
221    labels = np.zeros(volume.shape, dtype="uint8")
222    for name in sorted(masks, key=lambda name: LABEL_IDS[SEGMENT_NAMES[name]]):
223        labels[masks[name]] = LABEL_IDS[SEGMENT_NAMES[name]]
224
225    if labels.max() == 0:  # The segmentation does not belong to this CT series.
226        return False
227
228    with h5py.File(out_path, "w") as f:
229        f.create_dataset("raw", data=volume, compression="gzip")
230        f.create_dataset("labels", data=labels, compression="gzip")
231    return True
232
233
234def get_hcc_tace_data(path: Union[os.PathLike, str], download: bool = False) -> str:
235    """Download the HCC-TACE-Seg dataset.
236
237    Args:
238        path: Filepath to a folder where the data is downloaded for further processing.
239        download: Whether to download the data if it is not present.
240
241    Returns:
242        Filepath where the preprocessed data is stored.
243    """
244    preprocessed_dir = os.path.join(path, "preprocessed")
245    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == NUM_VOLUMES:
246        return preprocessed_dir
247
248    dicom_dir = os.path.join(path, "dicom")
249    os.makedirs(dicom_dir, exist_ok=True)
250
251    # Download the DICOM-SEG object of each patient and the CT series it references, then convert them.
252    series_per_patient = _get_series_per_patient(download)
253    for patient_id, series in tqdm(sorted(series_per_patient.items()), desc="Download and preprocess HCC-TACE-Seg"):
254        out_path = os.path.join(preprocessed_dir, f"{patient_id}.h5")
255        if os.path.exists(out_path):
256            continue
257
258        seg_series = [s for s in series if s["Modality"] == "SEG"]
259        assert len(seg_series) == 1, f"Expected a single segmentation for the patient '{patient_id}'."
260        seg_dir = _download_series(str(seg_series[0]["SeriesInstanceUID"]), dicom_dir, download)
261        seg_path = glob(os.path.join(seg_dir, "*.dcm"))[0]
262
263        os.makedirs(preprocessed_dir, exist_ok=True)
264        candidates = _get_candidate_series(seg_path, [s for s in series if s["Modality"] == "CT"])
265        for ct_uid, referenced_sop_uids in candidates:
266            ct_dir = _download_series(ct_uid, dicom_dir, download)
267            if _preprocess_hcc_tace(seg_path, ct_dir, referenced_sop_uids, out_path):
268                break
269        else:
270            raise RuntimeError(f"Could not find the CT series that belongs to {seg_path}.")
271
272    return preprocessed_dir
273
274
275def get_hcc_tace_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
276    """Get paths to the HCC-TACE-Seg data.
277
278    Args:
279        path: Filepath to a folder where the data is downloaded for further processing.
280        download: Whether to download the data if it is not present.
281
282    Returns:
283        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
284    """
285    data_dir = get_hcc_tace_data(path, download)
286    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
287    return volume_paths
288
289
290def get_hcc_tace_dataset(
291    path: Union[os.PathLike, str],
292    patch_shape: Tuple[int, ...],
293    resize_inputs: bool = False,
294    download: bool = False,
295    **kwargs
296) -> Dataset:
297    """Get the HCC-TACE-Seg dataset for liver, tumor and vessel segmentation.
298
299    Args:
300        path: Filepath to a folder where the data is downloaded for further processing.
301        patch_shape: The patch shape to use for training.
302        resize_inputs: Whether to resize inputs to the desired patch shape.
303        download: Whether to download the data if it is not present.
304        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
305
306    Returns:
307        The segmentation dataset.
308    """
309    volume_paths = get_hcc_tace_paths(path, download)
310
311    if resize_inputs:
312        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
313        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
314            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
315        )
316
317    return torch_em.default_segmentation_dataset(
318        raw_paths=volume_paths,
319        raw_key="raw",
320        label_paths=volume_paths,
321        label_key="labels",
322        patch_shape=patch_shape,
323        is_seg_dataset=True,
324        **kwargs
325    )
326
327
328def get_hcc_tace_loader(
329    path: Union[os.PathLike, str],
330    batch_size: int,
331    patch_shape: Tuple[int, ...],
332    resize_inputs: bool = False,
333    download: bool = False,
334    **kwargs
335) -> DataLoader:
336    """Get the HCC-TACE-Seg dataloader for liver, tumor and vessel segmentation.
337
338    Args:
339        path: Filepath to a folder where the data is downloaded for further processing.
340        batch_size: The batch size for training.
341        patch_shape: The patch shape to use for training.
342        resize_inputs: Whether to resize inputs to the desired patch shape.
343        download: Whether to download the data if it is not present.
344        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
345
346    Returns:
347        The DataLoader.
348    """
349    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
350    dataset = get_hcc_tace_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
351    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)