From cf3076931bce2b1a7d7f51608555b388b702fd94 Mon Sep 17 00:00:00 2001
From: Matt Cieslak <mattcieslak@gmail.com>
Date: Tue, 4 Apr 2023 11:48:52 -0400
Subject: [PATCH] [ENH] add CNR to the imageqc.csv (#541)

---
 qsiprep/interfaces/reports.py       | 30 +++++++++++++++++++++++++++--
 qsiprep/workflows/dwi/resampling.py |  1 +
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/qsiprep/interfaces/reports.py b/qsiprep/interfaces/reports.py
index fb2f1e12..5de22d3d 100644
--- a/qsiprep/interfaces/reports.py
+++ b/qsiprep/interfaces/reports.py
@@ -17,8 +17,10 @@
 from collections import defaultdict
 #from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
 import seaborn as sns
+import nibabel as nb
 import matplotlib
 import matplotlib.pyplot as plt
+from nilearn.maskers import NiftiLabelsMasker
 from matplotlib import animation
 from scipy.io.matlab import loadmat
 import pandas as pd
@@ -535,8 +537,9 @@ def _run_interface(self, runtime):
 
         if isdefined(self.inputs.t1_mask_file):
             if isdefined(self.inputs.t1_cnr_file):
-                # Add a function here to get CNR
-                pass
+                image_qc.update(
+                    get_cnr_values(self.inputs.t1_cnr_file,
+                                   self.inputs.t1_mask_file))
             if isdefined(self.inputs.t1_b0_series):
                 # Add a function to get b=0 TSNR
                 pass
@@ -559,6 +562,29 @@ def _load_qc_file(fname, prefix=""):
     return renamed
 
 
+def get_cnr_values(cnr_image, brain_mask):
+    cnr_img = nb.load(cnr_image)
+    mask_img = nb.load(brain_mask)
+
+    # Determine which CNRs we's getting
+    num_cnrs = 1 if cnr_img.ndim == 3 else cnr_img.shape[3]
+    if num_cnrs == 1:
+        cnr_labels = ["CNR"]
+    else:
+        cnr_labels = [ "CNR%d" % value for value in
+                      range(num_cnrs)]
+
+    cnrs = {}
+    strategies = ["mean", "median", "standard_deviation"]
+    for strategy in strategies:
+        masker = NiftiLabelsMasker(mask_img, strategy=strategy)
+        cnr_values = masker.fit_transform(cnr_img).flatten()
+        for cnr_name, cnr_value in zip(cnr_labels, cnr_values):
+            cnrs[cnr_name + "_" + strategy] = cnr_value
+
+    return cnrs
+
+
 def motion_derivatives(translations, rotations, framewise_disp,
                        original_files):
 
diff --git a/qsiprep/workflows/dwi/resampling.py b/qsiprep/workflows/dwi/resampling.py
index 5043d4c7..85267a62 100644
--- a/qsiprep/workflows/dwi/resampling.py
+++ b/qsiprep/workflows/dwi/resampling.py
@@ -274,6 +274,7 @@ def init_dwi_trans_wf(source_file,
     return workflow
 
 
+
 def _first(inlist):
     return inlist[0]