Repository: nipreps/niworkflows
Issue: RobustMNINormalizationRPT should mask images before plotting

If moving and or fixed masks were set as inputs the reportlet should present masked version of images (so the voxels not included in optimization would not distract the user like in this example: https://35-66973489-gh.circle-artifacts.com/0/home/ubuntu/ds005/scratch/work/workflow_enumerator/01/t1w_preprocessing/T1_2_MNI_Registration/report.svg)

Changed files:
niworkflows/tests/test_reports.py, niworkflows/common/report.py, niworkflows/viz/utils.py, niworkflows/interfaces/registration.py

diff --git a/niworkflows/common/report.py b/niworkflows/common/report.py
index 47e402db..11c60ef4 100644
--- a/niworkflows/common/report.py
+++ b/niworkflows/common/report.py
@@ -7,9 +7,11 @@ import os
 from sys import version_info
 from abc import abstractmethod
 from io import open
+import nibabel as nb
 
 from nipype.interfaces.base import File, traits, BaseInterface, BaseInterfaceInputSpec, TraitedSpec
 from niworkflows import NIWORKFLOWS_LOG
+from nilearn.masking import apply_mask, unmask
 
 PY3 = version_info[0] > 2
 
@@ -110,6 +112,8 @@ class RegistrationRC(ReportCapableInterface):
     def __init__(self, **inputs):
         self._fixed_image = None
         self._moving_image = None
+        self._fixed_image_mask = None
+        self._moving_image_mask = None
         super(RegistrationRC, self).__init__(**inputs)
 
     DEFAULT_MNI_CUTS = {
@@ -123,12 +127,26 @@ class RegistrationRC(ReportCapableInterface):
         from niworkflows.viz.utils import compose_view, plot_registration
         NIWORKFLOWS_LOG.info('Generating visual report')
 
+        fixed_image_nii = nb.load(self._fixed_image)
+        moving_image_nii = nb.load(self._moving_image)
+
+        if self._fixed_image_mask:
+            fixed_image_nii = unmask(apply_mask(fixed_image_nii,
+                                                self._fixed_image_mask),
+                                     self._fixed_image_mask)
+
+        if self._moving_image_mask:
+            moving_image_nii = unmask(apply_mask(moving_image_nii,
+                                                 self._moving_image_mask),
+                                      self._moving_image_mask)
+
+
         # Call composer
         compose_view(
-            plot_registration(self._fixed_image, 'fixed-image',
+            plot_registration(fixed_image_nii, 'fixed-image',
                               estimate_brightness=True,
                               cuts=self.DEFAULT_MNI_CUTS),
-            plot_registration(self._moving_image, 'moving-image',
+            plot_registration(moving_image_nii, 'moving-image',
                               estimate_brightness=True,
                               cuts=self.DEFAULT_MNI_CUTS),
             out_file=self._out_report)
diff --git a/niworkflows/interfaces/registration.py b/niworkflows/interfaces/registration.py
index ea69e440..0b7faf87 100644
--- a/niworkflows/interfaces/registration.py
+++ b/niworkflows/interfaces/registration.py
@@ -9,6 +9,7 @@ from nipype.interfaces import fsl, ants
 from niworkflows.anat import mni
 import niworkflows.common.report as nrc
 from niworkflows import NIWORKFLOWS_LOG
+from nipype.interfaces.base import isdefined
 
 class RobustMNINormalizationInputSpecRPT(
     nrc.RegistrationRCInputSpec, mni.RobustMNINormalizationInputSpec):
@@ -26,7 +27,11 @@ class RobustMNINormalizationRPT(
     def _post_run_hook(self, runtime):
         # We need to dig into the internal ants.Registration interface
         self._fixed_image = self.norm.inputs.fixed_image[0]  # and get first item
+        if isdefined(self.norm.inputs.fixed_image_mask):
+            self._fixed_image_mask = self.norm.inputs.fixed_image_mask
         self._moving_image = self.aggregate_outputs().warped_image
+        if isdefined(self.norm.inputs.moving_image_mask):
+            self._moving_image_mask = self.norm.inputs.moving_image_mask
         NIWORKFLOWS_LOG.info('Report - setting fixed (%s) and moving (%s) images',
                              self._fixed_image, self._moving_image)
 
diff --git a/niworkflows/tests/test_reports.py b/niworkflows/tests/test_reports.py
index aac26f3b..37ef8c8e 100644
--- a/niworkflows/tests/test_reports.py
+++ b/niworkflows/tests/test_reports.py
@@ -35,6 +35,7 @@ class TestRegistrationInterfaces(unittest.TestCase):
     """ tests the interfaces derived from RegistrationRC """
 
     reference = os.path.join(MNI_DIR, 'MNI152_T1_1mm.nii.gz')
+    reference_mask = os.path.join(MNI_DIR, 'MNI152_T1_1mm_brain_mask.nii.gz')
     moving = os.path.join(DS003_DIR, 'sub-01/anat/sub-01_T1w.nii.gz')
 
     def test_FLIRTRPT(self):
@@ -51,6 +52,14 @@ class TestRegistrationInterfaces(unittest.TestCase):
                 generate_report=True, moving_image=self.moving, testing=True)
             _smoke_test_report(ants_rpt, 'testRobustMNINormalizationRPT.svg')
 
+    def test_RobustMNINormalizationRPT_masked(self):
+        """ the RobustMNINormalizationRPT report capable test with masking """
+        with InTemporaryDirectory():
+            ants_rpt = RobustMNINormalizationRPT(
+                generate_report=True, moving_image=self.moving,
+                reference_mask=self.reference_mask, testing=True)
+            _smoke_test_report(ants_rpt, 'testRobustMNINormalizationRPT_masked.svg')
+
     def test_ANTSRegistrationRPT(self):
         """ the RobustMNINormalizationRPT report capable test """
         with InTemporaryDirectory():
diff --git a/niworkflows/viz/utils.py b/niworkflows/viz/utils.py
index f5e3501d..4ec0f013 100644
--- a/niworkflows/viz/utils.py
+++ b/niworkflows/viz/utils.py
@@ -173,7 +173,7 @@ def plot_xyz(image, plot_func, cuts, plot_params=None, dimensions=('z', 'x', 'y'
         plot_func(image, **kwargs)
 
 
-def plot_registration(anat_img, div_id, plot_params=None,
+def plot_registration(anat_nii, div_id, plot_params=None,
                       order=('z', 'x', 'y'), cuts=None,
                       estimate_brightness=False):
     """
@@ -189,8 +189,7 @@ def plot_registration(anat_img, div_id, plot_params=None,
 
     out_files = []
     if estimate_brightness:
-        from nibabel import load as loadnii
-        data = loadnii(anat_img).get_data().reshape(-1)
+        data = anat_nii.get_data().reshape(-1)
         vmin = np.percentile(data, 15)
         if plot_params.get('vmin', None) is None:
             plot_params['vmin'] = vmin
@@ -205,7 +204,7 @@ def plot_registration(anat_img, div_id, plot_params=None,
         plot_params['output_file'] = out_file
 
         # Generate nilearn figure
-        plot_anat(anat_img, **plot_params)
+        plot_anat(anat_nii, **plot_params)
         out_files.append(out_file)
 
         # Open generated svg file and fix id