diff --git a/models/deepcac2/config/default.yml b/models/deepcac2/config/default.yml new file mode 100644 index 00000000..b5e4ae56 --- /dev/null +++ b/models/deepcac2/config/default.yml @@ -0,0 +1,53 @@ +general: + data_base_dir: /app/data + version: 1.0 + description: custom DeepCAC2 starting from dicom and exporting AGS and RISK only for non-mapped cac segmentations + +execute: +- DicomImporter +- NiftiConverter +- NNUnetRunner +- DeepCACRunner +- AGSCalculator +- ReportExporter +- module: ReportExporter + globalreport: true + meta: + scope: global +- DataOrganizer + +modules: + + DicomImporter: + meta: + mod: '%Modality' + pid: '%PatientID' + + NNUnetRunner: + folds: all + nnunet_task: Task400_OPEN_HEART_1FOLD + nnunet_model: 3d_lowres + roi: HEART + + ReportExporter: + meta: + mod: report + scope: instance + includes: + - attr: sid + label: SeriesInstanceUID + - attr: pid + label: PatientID + - data: ags + label: Predicted Agatston Score + value: value + - data: rc + label: Predicted Risk Group + value: value + + DataOrganizer: + targets: + - nrrd:mod=seg:roi=CAC AND NOT any:variant=mapped-->[i:sid]/nrrd/cac.seg.nrrd + - nrrd:mod=seg:roi=CAC:variant=mapped-->[i:sid]/nrrd/mcac.seg.nrrd + - json:mod=report:scope=instance-->[i:sid]/DeepCAC.report.json + - json:mod=report:scope=global-->DeepCAC.report.json \ No newline at end of file diff --git a/models/deepcac2/config/pid_nifti.yml b/models/deepcac2/config/pid_nifti.yml new file mode 100644 index 00000000..88c1afa2 --- /dev/null +++ b/models/deepcac2/config/pid_nifti.yml @@ -0,0 +1,65 @@ + +general: + data_base_dir: /app/data + version: 1.0 + description: DeepCAC2 starting from nifti file(s) named by PatientID and exporting AGS and RISK for mapped and non-mapped cac segmentations. + +execute: +- FileStructureImporter +- NNUnetRunner +- DeepCACRunner +#- CACMapping +- AGSCalculator +- ReportExporter +- module: ReportExporter + globalreport: true + meta: + scope: global +- DataOrganizer + +modules: + + FileStructureImporter: + input_dir: 'input_data' + import_id: _instance + structures: + - re:(\d+).nii.gz::$pid@instance@nifti:mod=ct + + NNUnetRunner: + folds: all + nnunet_task: Task400_OPEN_HEART_1FOLD + nnunet_model: 3d_lowres + roi: HEART + + ReportExporter: + meta: + mod: report + scope: instance + includes: + - attr: pid + label: PatientID + - data: ags + label: Agatston Score Description + value: description + - data: ags + label: Predicted Agatston Score + value: value + - data: rc + label: Risk Group Description + value: description + - data: rc + label: Predicted Risk Group + value: value + # - data: mags + # label: Mapped Predicted Agatston Score + # value: value + # - data: mrc + # label: Mapped Predicted Risk Group + # value: value + + DataOrganizer: + targets: + #- nrrd:mod=seg:roi=CAC AND NOT any:variant=mapped-->[i:pid]/cac.seg.nrrd + #- nrrd:mod=seg:roi=CAC:variant=mapped-->[i:pid]/mcac.seg.nrrd + #- json:mod=report:scope=instance-->[i:pid]/DeepCAC.report.json + - json:mod=report:scope=global-->DeepCAC.report.json diff --git a/models/deepcac2/dockerfiles/Dockerfile b/models/deepcac2/dockerfiles/Dockerfile new file mode 100644 index 00000000..44e1e412 --- /dev/null +++ b/models/deepcac2/dockerfiles/Dockerfile @@ -0,0 +1,46 @@ +FROM mhubai/base:latest + +# FIXME: set this environment variable as a shortcut to avoid nnunet crashing the build +# by pulling sklearn instead of scikit-learn +# N.B. this is a known issue: +# https://github.com/MIC-DKFZ/nnUNet/issues/1281 +# https://github.com/MIC-DKFZ/nnUNet/pull/1209 +ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True + +# Install dependencies +RUN pip3 install --no-cache-dir \ + nnunet==1.7.1 \ + torch==2.0.1 \ + torchvision==0.15.2 \ + torchio==0.19.1 \ + && pip3 install --no-cache-dir --force-reinstall \ + "dicom2nifti<2.6" \ + numpy==1.26.4 \ + pydicom==2.4.4 + +# pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run +ENV WEIGHTS_DIR="/root/.platipy/nnUNet_models/nnUNet/" +ENV WEIGHTS_URL="https://zenodo.org/record/6585664/files/Task400_OPEN_HEART_3d_lowres.zip" +ENV WEIGHTS_FN="Task400_OPEN_HEART_3d_lowres.zip" + +RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL} +RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR} +RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN} + +# specify nnunet specific environment variables +ENV WEIGHTS_FOLDER=$WEIGHTS_DIR + +# Import the MHub model definition +ARG MHUB_MODELS_REPO +RUN buildutils/import_mhub_model.sh deepcac2 ${MHUB_MODELS_REPO} + +# download the DeepCAC2 model weights +ENV DEEPCAC2_WEIGHTS_DIR="/app/models/deepcac2/src/weights/" +ENV DEEPCAC2_WEIGHTS_FN="vivid-haze-6_model.pt" +ENV DEEPCAC2_WEIGHTS_URL="https://www.dropbox.com/scl/fi/4rodcp9y2vula8loh1v6r/vivid-haze-6_model.pt?rlkey=7j3cyaltvdpthukchn5xbh01d&dl=1" +RUN mkdir -p ${DEEPCAC2_WEIGHTS_DIR} \ + && wget -o /dev/stdout -O ${DEEPCAC2_WEIGHTS_DIR}${DEEPCAC2_WEIGHTS_FN} ${DEEPCAC2_WEIGHTS_URL} + +# Default run script +ENTRYPOINT ["mhub.run"] +CMD ["--workflow", "default"] diff --git a/models/deepcac2/meta.json b/models/deepcac2/meta.json new file mode 100644 index 00000000..968f7bc8 --- /dev/null +++ b/models/deepcac2/meta.json @@ -0,0 +1,174 @@ +{ + "id": "7c7c7e7b-3c4d-4a84-9a45-5a6c2b46c1a1", + "name": "deepcac2", + "title": "DeepCAC2", + "summary": { + "description": "DeepCAC2 is a two-stage deep learning pipeline for coronary artery calcification assessment in chest CT. It first uses an nnU-Net-based cardiac segmentation model to localize the heart, then applies a patch-based 3D U-Net to segment coronary artery calcifications and derive Agatston calcium scores and risk categories.", + "inputs": [ + { + "label": "Chest CT", + "description": "Non-contrast, low-dose chest CT scan in DICOM format.", + "format": "DICOM", + "modality": "CT", + "bodypartexamined": "CHEST", + "slicethickness": "1-5 mm", + "non-contrast": true, + "contrast": false + } + ], + "outputs": [ + { + "type": "Segmentation", + "classes": [ + "heart", + "coronary artery calcification" + ] + }, + { + "type": "Prediction", + "valueType": "float", + "label": "Predicted Agatston Score", + "description": "Automated coronary artery calcium score derived from the predicted CAC segmentation." + }, + { + "type": "Classification", + "classes": [ + "very low", + "low", + "moderate", + "high" + ] + } + ], + "model": { + "architecture": "Two-stage pipeline with nnU-Net-based cardiac segmentation followed by patch-based 3D U-Net CAC segmentation.", + "training": "supervised", + "cmpapproach": "3D" + }, + "data": { + "training": { + "vol_samples": 622 + }, + "evaluation": { + "vol_samples": 390 + }, + "public": true, + "external": true + } + }, + "details": { + "name": "DeepCAC2", + "version": "1.0.0", + "devteam": "Artificial Intelligence in Medicine Program, Mass General Brigham / Harvard Medical School", + "type": "Segmentation and prediction", + "date": { + "weights": "2026-01-01", + "code": "2026-01-01", + "pub": "2026-03-25" + }, + "cite": "Nürnberg L, Bernatz S, Foldyna B, Lu MT, Fedorov A, Aerts HJWL. Coronary artery calcification assessment in National Lung Screening Trial CT images (DeepCAC2).", + "license": { + "code": "MIT", + "weights": "MIT" + }, + "publications": [ + { + "title": "Coronary artery calcification assessment in National Lung Screening Trial CT images (DeepCAC2)", + "uri": "https://arxiv.org/" + } + ], + "github": "https://github.com/MHubAI/models/tree/main/models/deepcac2", + "slicer": false + }, + "info": { + "use": { + "title": "Intended use", + "text": "DeepCAC2 is designed for automated coronary artery calcification (CAC) segmentation and cardiovascular risk assessment from non-contrast chest CT scans, enabling opportunistic screening without dedicated cardiac CT protocols. The generated outputs and dataset can be explored via an interactive public dashboard [1] and are derived from the National Lung Screening Trial (NLST) cohort available through the Imaging Data Commons (IDC) [2].", + "references": [ + { + "label": "DeepCAC2 interactive dashboard", + "uri": "https://lookerstudio.google.com/reporting/2c656fab-89ce-4ccd-a6a6-6c76ba7ecc51/page/P3qpF" + }, + { + "label": "Imaging Data Commons (NLST collection)", + "uri": "https://imaging.datacommons.cancer.gov/" + } + ] + }, + "analyses": { + "title": "Analyses", + "text": "The pipeline produces volumetric segmentations of the heart and coronary artery calcifications, from which Agatston calcium scores (CACS) are computed and mapped to categorical risk groups (0, 1–100, 101–300, >300). These standardized outputs enable cohort stratification, longitudinal analyses, survival modeling, and reproducible imaging biomarker research at scale [1].", + "references": [ + { + "label": "DeepCAC2 preprint", + "uri": "https://arxiv.org/abs/2601.10154" + } + ], + "tables": [ + { + "label": "CAC risk groups", + "entries": { + "0": "CACS = 0 (very low)", + "1": "CACS 1–100 (low)", + "2": "CACS 101–300 (moderate)", + "3": "CACS >300 (high)" + } + } + ] + }, + "evaluation": { + "title": "Evaluation", + "text": "DeepCAC2 was evaluated for both technical validity and clinical relevance. Technical validation compared automated CAC scores and risk categories against expert annotations on a subset of NLST scans, demonstrating high agreement. Clinical validation using survival analysis showed that model-derived CAC risk groups are strong independent predictors of all-cause mortality, with clear stratification across risk categories [1].", + "references": [ + { + "label": "DeepCAC2 preprint", + "uri": "https://arxiv.org/abs/2601.10154" + } + ], + "tables": [ + { + "label": "Technical validation metrics", + "entries": { + "Evaluation dataset": "NLST subset", + "Sample size": "390 CT scans", + "Spearman correlation (CACS)": "0.879", + "Weighted Cohen's kappa (risk groups)": "0.844", + "Agreement interpretation": "High correlation and near-perfect agreement" + } + }, + { + "label": "Survival analysis summary", + "entries": { + "Cohort size": "23,011 participants", + "Events": "1,624 deaths", + "Median follow-up": "6.7 years", + "Model type": "Cox proportional hazards", + "Adjusted covariates": "Age, sex, smoking status, heart disease", + "Highest risk group HR": "2.05 (95% CI 1.77–2.39)", + "Concordance index": "0.685" + } + } + ] + }, + "training": { + "title": "Training", + "text": "The CAC segmentation model was trained in a supervised manner using 778 ECG-gated chest CT scans from the Framingham Heart Study with expert CAC annotations. The pipeline uses a two-stage approach: nnU-Net-based cardiac segmentation followed by patch-based 3D U-Net training on balanced overlapping 3D patches extracted from heart-centered volumes [1].", + "references": [ + { + "label": "DeepCAC2 preprint", + "uri": "https://arxiv.org/abs/2601.10154" + } + ] + }, + "limitations": { + "title": "Limitations", + "text": "The model is optimized for non-contrast, low-dose chest CT scans and may be sensitive to variations in acquisition protocols, reconstruction parameters, and image quality. CAC quantification is derived from non-ECG-gated scans, which may introduce motion-related inaccuracies compared to dedicated cardiac CT. Generalization to other populations, scanners, and clinical workflows requires further validation [1].", + "references": [ + { + "label": "DeepCAC2 preprint", + "uri": "https://arxiv.org/abs/2601.10154" + } + ] + } + } +} \ No newline at end of file diff --git a/models/deepcac2/mhub.toml b/models/deepcac2/mhub.toml new file mode 100644 index 00000000..5b59fc25 --- /dev/null +++ b/models/deepcac2/mhub.toml @@ -0,0 +1,2 @@ +[model.deployment] +test = "https://zenodo.org/records/19921552/files/mhub_test_deepcac2.zip" diff --git a/models/deepcac2/src/.gitignore b/models/deepcac2/src/.gitignore new file mode 100644 index 00000000..26ded4a1 --- /dev/null +++ b/models/deepcac2/src/.gitignore @@ -0,0 +1 @@ +weights/ \ No newline at end of file diff --git a/models/deepcac2/src/__init__.py b/models/deepcac2/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/models/deepcac2/src/agatston.py b/models/deepcac2/src/agatston.py new file mode 100644 index 00000000..55454744 --- /dev/null +++ b/models/deepcac2/src/agatston.py @@ -0,0 +1,114 @@ +# https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5487233/ +# Agatston method - The Agatston method uses the weighted sum of lesions with a density above 130 HU, +# multiplying the area of calcium by a factor related to maximum plaque attenuation: +# 130-199 HU, factor 1; 200-299 HU, factor 2; 300-399 HU, factor 3; and ≥ 400 HU, factor 4. + +# https://www-sciencedirect-com.mu.idm.oclc.org/science/article/pii/S1939865421001569 +# Agatston score is a summed value of all calcified coronary lesions, based on both the total area and the maximal density of coronary calcification. + +# ---------------------------------------------------- +# Calculation +# detect continuous voxels (connected shapes) with HU > 130 and minimal size of 1 mm^3 (number of voxels depend on teh spacing) +# for each individual calcified leason in all coronary arteries +# dwf = { 1 if max HU in leasion in [130, 199] +# 2 if max HU in leasion in [200, 299] +# 3 if max HU in leasion in [300, 399] +# 4 if max HU in leasion > 400 } +# s_i = leasion area * dwf +# overall Agatston score is the sum of all individual leasion scores +# s = \sum_i s_i + +# Interpretation +# AS = 0: indicates no identifiable atherosclerotic plaque and very low cardiovascular disease (CVD) risk (Fig. 1). +# AS \in [1, 10]: indicates minimal plaque burden and low CVD risk. +# AS \in ]10,100]: indicates mild plaque burden and moderate CVD risk (Fig. 2). +# AS \in ]100, 400]: indicates moderate plaque burden and high CVD risk. +# AS > 400: indicates extensive plaque burden and very high CVD risk (Fig. 3). + +# Questions +# "[..] for the detection of calcium in contiguous voxels of 1 sq mm in area to be counted as individual lesions." +# --> why area and why square not cubic? + +import numpy as np +from scipy.ndimage import measurements + +def agatston_score_slices2(cac_np, img_np, nrConPx, spacing, allow_diagonal_connections=True, verbose=False): + AG_DIV = 3 + pxArea = round(spacing[0] * spacing[1] * spacing[2] / AG_DIV, 3) + + AG = 0 + + for z_slice in range(cac_np.shape[2]): + cac_slice_np = cac_np[:, :, z_slice] + img_slice_np = img_np[:, :, z_slice] + + # NOTE: allegedly (he didn't verify if intentionally) Roman used np.ones((3, 3)) thus allowed for diagonal connections + if allow_diagonal_connections: + structure = np.array([ + [1, 1, 1], + [1, 1, 1], + [1, 1, 1] + ]) + else: + structure = np.array([ + [0, 1, 0], + [1, 1, 1], + [0, 1, 0] + ]) + + # extract connected shapes (objects) + # objects_lblmask, n_objects + labeledMask, numLabels = measurements.label(cac_slice_np, structure=structure) + + for labelNr in range(1, numLabels + 1): + label = np.zeros(cac_slice_np.shape) + label[labeledMask == labelNr] = 1 + + clcObject = img_slice_np * label + + # 1) Remove small objects + if np.sum(label) <= nrConPx: # FIXME: this is copied from Roman's code but should be < instead. + continue + + # 3) Calculate volume + objectArea = np.sum(label) * pxArea + + # Get object max HU and DFW + objectMaxHU = clcObject.max() + objectDFW = dfw(objectMaxHU, disable_hu_check=True) + + # 4) Calculate AG for object + objectAG = round(objectArea * objectDFW, 3) + + # 5) Sum up scores + AG += objectAG + + #if verbose: print(f"z {z_slice}, plaque {plaque_label:<4}| s:{plaque_slice_score:<8.3f} v:{plaque_slice_volume:<7.2f} hu:{plaque_slice_max_hu:<8.2f} dfw:{dfw(plaque_slice_max_hu, disable_hu_check=True)}") + #if verbose: print(f"∑ p{pid:<8}> {patient_score:<8.3f}") + + # return cummulated sum as agatston score + return AG + + +# agatston score object density factor +def dfw(max_hu, disable_hu_check=False): + assert disable_hu_check or max_hu >= 130, "HU value below TH." + if max_hu < 200: # [130, 200[ + return 1 + elif max_hu < 300: # [200, 300[ + return 2 + elif max_hu < 400: # [300, 400[ + return 3 + else: # [400, + return 4 + +# as implemented by Roman (getAGclass(AG) method) +def agatston_score_interpretation(AG): + AG = round(AG, 3) + classAG = None + if AG == 0: classAG = 0 + if AG > 0 and AG <= 100: classAG = 1 + if AG > 100 and AG <= 300: classAG = 2 + if AG > 300: classAG = 3 + return classAG + diff --git a/models/deepcac2/src/cacmapping.py b/models/deepcac2/src/cacmapping.py new file mode 100644 index 00000000..8b75d315 --- /dev/null +++ b/models/deepcac2/src/cacmapping.py @@ -0,0 +1,59 @@ +import numpy as np +from scipy.ndimage import measurements + +def register_cac_mask(img_np, cac_np, allow_diagonal_connections=True): + """Generates a thresholded mask from the input image and removes all objects which do not overlap with the provided mask. + The idea is, that due to resampling operations, the resulting mask might not exactly match to the thresholded objects which + this method aims to correct for. Both, the input image and the cac mask must already be in the same spacing and dimensions. + NOTE: Does not remove small objects, which has to be done in the agatston score calculation. + + Args: + img_np (3d numpy array): the original input image + cac_np (3d binary numpy array): the (auto-) generated binary cac mask + """ + assert img_np.shape == cac_np.shape, "shape missmatch" + TH = 130 + + # re-assemble a cac mask based on the hu threshold and the provided cac mask + registered_mask = np.zeros(cac_np.shape) + + # generate a mask from the thresholded input image + thmask_np = np.zeros(img_np.shape) + thmask_np[img_np >= TH] = 1 # >= 130, verified + + # iterate through the slices + for z_slice in range(cac_np.shape[2]): + cac_slice_np = cac_np[:, :, z_slice] + thmask_slice_np = thmask_np[:, :, z_slice] + + # NOTE: allegedly (he didn't verify if intentionally) Roman used np.ones((3, 3)) thus allowed for diagonal connections + if allow_diagonal_connections: + structure = np.array([ + [1, 1, 1], + [1, 1, 1], + [1, 1, 1] + ]) + else: + structure = np.array([ + [0, 1, 0], + [1, 1, 1], + [0, 1, 0] + ]) + + # extract connected shapes (objects) + objects_lblmask, n_objects = measurements.label(thmask_slice_np, structure=structure) + + # iterate over all objects, identified on the thresholded image mask + for object_nr in range(1, n_objects + 1): + + # label mask + #object_mask = np.zeros(cac_slice_np.shape) + #object_mask[objects_lblmask == object_nr] = 1 + object_mask = (objects_lblmask == object_nr).astype(int) + + # overlap / registration match + if object_mask[cac_slice_np > 0].sum() > 0: # is present on both + registered_mask[:, :, z_slice] += object_mask + + assert registered_mask.max() <= 1, f"overlappings detected ({registered_mask.max()})" + return registered_mask \ No newline at end of file diff --git a/models/deepcac2/src/inference.py b/models/deepcac2/src/inference.py new file mode 100644 index 00000000..6ea46871 --- /dev/null +++ b/models/deepcac2/src/inference.py @@ -0,0 +1,157 @@ +""" +------------------------------------------------- +DeepCAC2 - 3D Unet Pipeline + +This script runs the entrire processing pieline based on nrrd input chest ct scans: +(loading data, preprocessing, patching, model execution, reassembling, ags computation) +on all patients from the validation split from start to finish based on a +pre-trained 3D Unet (see model.py). +Therefore, this script is INDEPENDENT from any preparations (see preparedata.py) +which is used soley for training speedup. + +Metrics can then be calculated in the next step based on the predicted and assembled cac segmentation. + +WORK IN PROGRESS. +------------------------------------------------- + +------------------------------------------------- +Author: Leonard Nürnberg +Email: leonard.nuernberg@maastrichtuniversity.nl +------------------------------------------------- +""" + +# imports +from typing import Tuple + +import os +import numpy as np +import torch +import SimpleITK as sitk +import torchio as tio + +from .model import UNET3D +from .subject import Sample, get_subject, get_tfx + +# alpha +cWEIGHTS = 'vivid-haze-6_model.pt' +cCOMMON_SPACING = 0.7, 0.7, 2.5 + +# beta +#cWEIGHTS = 'glad-fire-13_model.pt' +#cCOMMON_SPACING = 0.68, 0.68, 2.5 + +cDEVICE = 'cuda:0' +cBOUNDING_BOX = 256, 256, 58 +cPATCHSIZE = 64, 64, 16 +cSTRIDE = 32, 32, 8 +cTH = 0.5 +cAPPLY_ZNORM = True + +# instantiate model & load trained weights +model = UNET3D(1, 1) +model.to(cDEVICE) +model.load_state_dict(torch.load(os.path.join(os.path.dirname(__file__), 'weights', cWEIGHTS))) +model.eval() + +# +def predict_sample( + sample: Sample, + tfx: tio.Transform, + patchsize: Tuple[int, int, int], + stride: Tuple[int, int, int], + threshold: float, + predcac_seg_file: str + ): + + # load subject + subject = get_subject(sample) + + # preprocess subject + subject_tfx = tfx(subject) + assert hasattr(subject_tfx, 'image') and hasattr(subject_tfx, 'heart') + + # extract data as numpy arrays + image_np = subject_tfx.image.numpy().squeeze() # type: ignore + heart_np = subject_tfx.heart.numpy().squeeze() # type: ignore + + # + img_vol = sitk.ReadImage(sample['img_path']) + + # get shape dim + w, h, d = image_np.shape + pw, ph, pd = patchsize + + # calculate / estimate the number of patches + num_patches = 0 + for wi in range(0, w-pw, stride[0]): + for hi in range(0, h-ph, stride[1]): + for di in range(0, d-pd, stride[2]): + heart_patch = heart_np[wi:wi+pw,hi:hi+ph,di:di+pd] + assert heart_patch.shape == patchsize + + # ignore patches with no heart present + if heart_patch.sum() > 0: + num_patches += 1 + + out = np.zeros((w, h, d, num_patches)) + div = np.zeros((w, h, d)) + + patch_i = 0 + for wi in range(0, w-pw, stride[0]): + for hi in range(0, h-ph, stride[1]): + for di in range(0, d-pd, stride[2]): + image_patch = image_np[wi:wi+pw,hi:hi+ph,di:di+pd] + heart_patch = heart_np[wi:wi+pw,hi:hi+ph,di:di+pd] + assert image_patch.shape == patchsize + + # ignore patches with no heart present + if heart_patch.sum() == 0: + continue + + ins = torch.tensor(np.array(image_patch)).to(cDEVICE).unsqueeze(0).unsqueeze(1) + pred = model(ins).sigmoid().squeeze(1).squeeze(0).detach().cpu().numpy() + + out[wi:wi+pw,hi:hi+ph,di:di+pd,patch_i] = pred + div[wi:wi+pw,hi:hi+ph,di:di+pd] += 1 + + patch_i += 1 + + div[div == 0] = 1 + + out = np.sum(out, axis=3) / div + + out_final = (out > threshold).astype(int) + + # add the prediction to the transformed subject + subject_tfx['cacpred'] = tio.LabelMap(tensor=torch.Tensor(out_final).unsqueeze(0), affine=subject_tfx.image.affine) # type: ignore + + # inverse (for some reason affien is ignored by apply_inverse_transform although it shoul dbe invertible?!) + resample_itfx = tio.Resample(subject.image) # type: ignore + + # inverse pre-processing transformations on the subject + subject_tt = resample_itfx(subject_tfx.apply_inverse_transform(image_interpolation='linear')) # type: ignore + + # again etract an numpy array from the subject and pass it to sitk + # NOTE sitk expects data in z, y, x orientation (thus transpose)! + sitk_vol = sitk.GetImageFromArray(subject_tt['cacpred'].numpy().squeeze(0).transpose(2, 1, 0)) # type: ignore + + # now load the original image using sitk and apply all meta-data (origin, dimension, spacing) etc. to the predicted mask sitk volume + # NOTE: this is 'just' meta data. The actual data is already in the matching shape since we applied the inverse transformations using tio + sitk_vol.CopyInformation(img_vol) + + # save the prediction to the file system + sitk.WriteImage(sitk_vol, predcac_seg_file) + + +def run_inference(sample: Sample, predcac_seg_file: str): + # static parameters + kwargs = { + 'tfx': get_tfx(cCOMMON_SPACING, cBOUNDING_BOX, cAPPLY_ZNORM), + 'patchsize': cPATCHSIZE, + 'stride': cSTRIDE, + 'threshold': cTH, + 'predcac_seg_file': predcac_seg_file + } + + # predict sample + predict_sample(sample, **kwargs) \ No newline at end of file diff --git a/models/deepcac2/src/model.py b/models/deepcac2/src/model.py new file mode 100644 index 00000000..8ccf69e3 --- /dev/null +++ b/models/deepcac2/src/model.py @@ -0,0 +1,66 @@ +import torch +import torch.nn as nn + +class DoubleConv3D(nn.Module): + def __init__(self, in_channels, out_channels): + super(DoubleConv3D, self).__init__() + + self.conv = nn.Sequential( + nn.Conv3d(in_channels, out_channels, 3, 1, 1, bias=False), # 3 1 1 + nn.BatchNorm3d(out_channels), + nn.ReLU(inplace=True), + nn.Conv3d(out_channels, out_channels, 3, 1, 1, bias=False), # 3 1 1 + nn.BatchNorm3d(out_channels), + nn.ReLU(inplace=True), + ) + + def forward(self, x): + return self.conv(x) + + +class UNET3D(nn.Module): + def __init__(self, in_channels=1, out_channels=1, features=[64, 128, 256, 512], up_stop=0): + super(UNET3D, self).__init__() + + self.ups = nn.ModuleList() + self.downs = nn.ModuleList() + self.pool = nn.MaxPool3d(kernel_size=2, stride=2) + + # down part + for feature in features: + self.downs.append(DoubleConv3D(in_channels, feature)) + in_channels = feature + + # up part + for feature in reversed(features[up_stop:]): + self.ups.append( + nn.ConvTranspose3d(feature * 2, feature, kernel_size=2, stride=2) + ) + + self.ups.append(DoubleConv3D(feature * 2, feature)) + + # bottleneck & final + self.bottleneck = DoubleConv3D(features[-1], features[-1] * 2) + self.final_conv = nn.Conv3d(features[up_stop], out_channels, kernel_size=1) + + def forward(self, x): + skip_connections = [] + + for down in self.downs: + x = down(x) + skip_connections.append(x) + x = self.pool(x) + + x = self.bottleneck(x) + skip_connections = skip_connections[::-1] + + for idx in range(0, len(self.ups), 2): + x = self.ups[idx](x) #ConvTranspose2d + skip_connection = skip_connections[idx//2] + + assert x.shape == skip_connection.shape + + concat_skip = torch.cat((skip_connection, x), dim=1) # dim 1 = channel dimension (batch, channel, height, with) + x = self.ups[idx+1](concat_skip) + + return self.final_conv(x) diff --git a/models/deepcac2/src/subject.py b/models/deepcac2/src/subject.py new file mode 100644 index 00000000..db3a6dfc --- /dev/null +++ b/models/deepcac2/src/subject.py @@ -0,0 +1,54 @@ +from typing import Tuple, TypedDict, Dict, Any, Optional +import torchio as tio + +class Sample(TypedDict): + id: str + img_path: str + hrt_path: Optional[str] + meta: Dict[str, Any] + +def get_subject(sample: Sample) -> tio.Subject: + return tio.Subject( + image=tio.ScalarImage(sample['img_path']), + heart=tio.LabelMap(sample['hrt_path']) + ) + +def get_tfx(spacing: Tuple[float, float, float], hbb: Tuple[int, int, int], apply_znorm: bool) -> tio.Transform: + + # resample + resample_tfx = tio.Resample( + target = spacing # x, y, z + ) + + # cropping (105, 184, 212) -> (184, 184, 212) + crop_tfx = tio.CropOrPad( + target_shape = hbb, + mask_name = 'heart' + ) + + # windowing (apply abdomen soft tissue window: W:350, L:40) + clamp_tfx = tio.Clamp( + out_min = -135, + out_max = 500, + keep = {'image': 'image_original_hu'} + ) + + # per patient z-norm of HU values + znorm_tfx = tio.ZNormalization( + exclude = ['image_original_hu'] + ) + + # map values to [-1, 1] + intmap_tfx = tio.RescaleIntensity( + out_min_max = (0, 1), + exclude = ['image_original_hu'] + ) + + # composed transformatiuon chain + return tio.Compose([ + resample_tfx, + crop_tfx, + clamp_tfx, + *([znorm_tfx] if apply_znorm else []), + intmap_tfx + ]) diff --git a/models/deepcac2/utils/AGSCalculator.py b/models/deepcac2/utils/AGSCalculator.py new file mode 100644 index 00000000..08e05dd0 --- /dev/null +++ b/models/deepcac2/utils/AGSCalculator.py @@ -0,0 +1,50 @@ +# import mhub fw +from mhubio.core import Module, IO, Instance, InstanceData, ValueOutput, ClassOutput + +# import pipeline +import os +import SimpleITK as sitk +from ..src.agatston import agatston_score_slices2, agatston_score_interpretation + + +@ValueOutput.Name('ags') +@ValueOutput.Label('AgatstonScore') +@ValueOutput.Type(int) +@ValueOutput.Description('Prediction of the agatson score.') +class AgatstonScore(ValueOutput): + pass + +@ClassOutput.Name('rc') +@ClassOutput.Label('RiskCategory') +@ClassOutput.Description('Prediction of the risk category.') +@ClassOutput.Class(0, 'No', the='The zero risk group for AGS equal zero.') +@ClassOutput.Class(1, 'Low', the='Class describing the lowest risk group.') +@ClassOutput.Class(2, 'Moderate', the='Moderate risk.') +@ClassOutput.Class(3, 'High', 'High risk.') +class RiskCategory(ClassOutput): + pass + + +@IO.ConfigInput('cac', 'nrrd:mod=seg:roi=CAC AND NOT any:variant=mapped', the='input ct scan') +class AGSCalculator(Module): + + @IO.Instance() + @IO.Input('image', 'nifti:mod=ct', the='input ct scan') + @IO.Input('cac', the='detected coronary artery calcification') + @IO.OutputData('ags', AgatstonScore, data='cac', the='agatston score') + @IO.OutputData('risk', RiskCategory, data='cac', the='risk classification') + def task(self, instance: Instance, image: InstanceData, cac: InstanceData, ags: AgatstonScore, risk: RiskCategory) -> InstanceData: + + # load img using sitk as x, y, z numpy array + img_vol = sitk.ReadImage(image.abspath) + img_np = sitk.GetArrayFromImage(img_vol).transpose(2, 1, 0) + + # load prediction using sitk + cacpred_file = cac.abspath + assert os.path.exists(cacpred_file), f"no prediction found, expected {cacpred_file}" + cacpred_vol = sitk.ReadImage(cacpred_file) + cacpred_np = sitk.GetArrayFromImage(cacpred_vol).transpose(2, 1, 0) + + # calculate ags and risk + ags.value = round(agatston_score_slices2(cacpred_np, img_np, nrConPx=3, spacing=img_vol.GetSpacing())) + risk.value = agatston_score_interpretation(ags.value) diff --git a/models/deepcac2/utils/CACMapping.py b/models/deepcac2/utils/CACMapping.py new file mode 100644 index 00000000..d8114585 --- /dev/null +++ b/models/deepcac2/utils/CACMapping.py @@ -0,0 +1,36 @@ +# import mhub fw +from mhubio.core import Module, IO, Instance, InstanceData + +# import pipeline +import os +import SimpleITK as sitk +from ..src.cacmapping import register_cac_mask + +class CACMapping(Module): + + @IO.Instance() + @IO.Input('image', 'nifti:mod=ct', the='input ct scan') + @IO.Input('cac', 'nrrd:mod=seg:roi=CAC', the='detected coronary artery calcification') + @IO.Output('mapped_cac', 'maped_pcac.nrrd', 'nrrd:mod=seg:variant=mapped', data='cac', the='delineation of combined structures with HU > 130 where overlapping with detected cac') + def task(self, instance: Instance, image: InstanceData, cac: InstanceData, mapped_cac: InstanceData) -> InstanceData: + + # load img using sitk as x, y, z numpy array + img_vol = sitk.ReadImage(image.abspath) + img_np = sitk.GetArrayFromImage(img_vol).transpose(2, 1, 0) + + # load prediction using sitk + cacpred_file = os.path.join(cac.abspath) + assert os.path.exists(cacpred_file), f"no prediction found, expected {cacpred_file}" + cacpred_vol = sitk.ReadImage(cacpred_file) + cacpred_np = sitk.GetArrayFromImage(cacpred_vol).transpose(2, 1, 0) + + # mapping + rcacpred_np = register_cac_mask(img_np, cacpred_np) + + # store mapped + rcacpred_vol = sitk.GetImageFromArray(rcacpred_np.transpose(2, 1, 0)) # type: ignore + rcacpred_vol.CopyInformation(img_vol) + + # save the prediction to the file system + sitk.WriteImage(rcacpred_vol, mapped_cac.abspath) + diff --git a/models/deepcac2/utils/DeepCACPostProcessor.py b/models/deepcac2/utils/DeepCACPostProcessor.py new file mode 100644 index 00000000..be4ece17 --- /dev/null +++ b/models/deepcac2/utils/DeepCACPostProcessor.py @@ -0,0 +1,83 @@ +# import mhub fw +from mhubio.core import Module, IO, Instance, InstanceData, ValueOutput, ClassOutput + +# import pipeline +import os +import SimpleITK as sitk +from ..src.cacmapping import register_cac_mask +from ..src.agatston import agatston_score_slices2, agatston_score_interpretation + +@ValueOutput.Name('ags') +@ValueOutput.Label('AgatstonScore') +@ValueOutput.Type(int) +@ValueOutput.Description('Prediction of the agatson score.') +class AgatstonScore(ValueOutput): + pass + +@ValueOutput.Name('mags') +@ValueOutput.Label('MappedCAC AgatstonScore') +@ValueOutput.Type(int) +@ValueOutput.Description('Prediction of the agatson score using the mapped cac segmentation.') +class MappedAgatstonScore(ValueOutput): + pass + +@ClassOutput.Name('rc') +@ClassOutput.Label('RiskCategory') +@ClassOutput.Description('Prediction of the risk category.') +@ClassOutput.Class(0, 'No', the='The zero risk group for AGS equal zero.') +@ClassOutput.Class(1, 'Low', the='Class describing the lowest risk group.') +@ClassOutput.Class(2, 'Moderate', the='Moderate risk.') +@ClassOutput.Class(3, 'High', 'High risk.') +class RiskCategory(ClassOutput): + pass + +@ClassOutput.Name('mrc') +@ClassOutput.Label('MappedCAC RiskCategory') +@ClassOutput.Description('Prediction of the risk category based on the mapped cac segmentation.') +@ClassOutput.Class(0, 'No', the='The zero risk group for AGS equal zero.') +@ClassOutput.Class(1, 'Low', the='Class describing the lowest risk group.') +@ClassOutput.Class(2, 'Moderate', the='Moderate risk.') +@ClassOutput.Class(3, 'High', 'High risk.') +class MappedRiskCategory(ClassOutput): + pass + +class DeepCACPostProcessor(Module): + + @IO.Instance() + @IO.Input('image', 'nifti:mod=ct', the='input ct scan') + @IO.Input('cac', 'nrrd:mod=seg:roi=CAC', the='detected coronary artery calcification') + @IO.Output('mapped_cac', 'maped_pcac.nrrd', 'nrrd:mod=seg:variant=mapped', data='cac', the='delineation of combined structures with HU > 130 where overlapping with detected cac') + @IO.OutputData('ags', AgatstonScore, data='cac', the='agatston score') + @IO.OutputData('mags', MappedAgatstonScore, data='cac', the='agatston score') + @IO.OutputData('risk', RiskCategory, data='cac', the='risk classification') + @IO.OutputData('mrisk', MappedRiskCategory, data='cac', the='risk classification') + def task(self, instance: Instance, image: InstanceData, cac: InstanceData, mapped_cac: InstanceData, ags: AgatstonScore, risk: RiskCategory, mags: MappedAgatstonScore, mrisk: MappedRiskCategory) -> InstanceData: + + # load img using sitk as x, y, z numpy array + img_vol = sitk.ReadImage(image.abspath) + img_np = sitk.GetArrayFromImage(img_vol).transpose(2, 1, 0) + + # load prediction using sitk + cacpred_file = os.path.join(cac.abspath) + assert os.path.exists(cacpred_file), f"no prediction found, expected {cacpred_file}" + cacpred_vol = sitk.ReadImage(cacpred_file) + cacpred_np = sitk.GetArrayFromImage(cacpred_vol).transpose(2, 1, 0) + + # mapping + mcacpred_np = register_cac_mask(img_np, cacpred_np) + + # store mapped + mcacpred_vol = sitk.GetImageFromArray(mcacpred_np.transpose(2, 1, 0)) # type: ignore + mcacpred_vol.CopyInformation(img_vol) + + # save the prediction to the file system + sitk.WriteImage(mcacpred_vol, mapped_cac.abspath) + + # calculate ags and risk for original cac prediction + ags.value = round(agatston_score_slices2(cacpred_np, img_np, nrConPx=3, spacing=img_vol.GetSpacing())) + risk.value = agatston_score_interpretation(ags.value) + + # calculate ags and risk for mapped cac prediction + mags.value = round(agatston_score_slices2(mcacpred_np, img_np, nrConPx=3, spacing=img_vol.GetSpacing())) + mrisk.value = agatston_score_interpretation(mags.value) + diff --git a/models/deepcac2/utils/DeepCACRunner.py b/models/deepcac2/utils/DeepCACRunner.py new file mode 100644 index 00000000..07716e8b --- /dev/null +++ b/models/deepcac2/utils/DeepCACRunner.py @@ -0,0 +1,27 @@ +# import mhub fw +from mhubio.core import Module, IO, Instance, InstanceData + +# import pipeline +from ..src.inference import run_inference +from ..src.subject import Sample + +class DeepCACRunner(Module): + + @IO.Instance() + @IO.Input('image', 'nifti:mod=ct', the='input ct scan') + @IO.Input('heart', 'nifti:mod=seg', the='input heart segmentation') + @IO.Output('cac', 'pcac.nrrd', 'nrrd:mod=seg:model=DeepCAC2:roi=CAC', the='detected coronary artery calcification') + def task(self, instance: Instance, image: InstanceData, heart: InstanceData, cac: InstanceData) -> InstanceData: + + + # create sample + sample: Sample = { + 'id': instance.attr['id'], + 'img_path': image.abspath, + 'hrt_path': heart.abspath, + 'meta': {} + } + + # run pipeline + run_inference(sample, cac.abspath) +