ENH: add helmet view option in coreg (#10200)

* add helmet

* update test

* update changes

doc/changes/latest.inc

@@ -74,6 +74,8 @@ Enhancements
 
 - Annotations from a :class:`~mne.io.Raw` object are now preserved by the :class:`~mne.Epochs` constructor and are supported when saving Epochs (:gh:`9969` and :gh:`10019` by `Adam Li`_)
 
+- Add a button to display the MEG helmet in the coregistration GUI (:gh:`10200` but `Guillaume Favelier`_)
+
 Bugs
 ~~~~
 

mne/gui/_coreg.py

@@ -18,8 +18,8 @@
 from ..coreg import Coregistration, _is_mri_subject
 from ..viz._3d import (_plot_head_surface, _plot_head_fiducials,
                        _plot_head_shape_points, _plot_mri_fiducials,
-                       _plot_hpi_coils, _plot_sensors)
-from ..transforms import (read_trans, write_trans, _ensure_trans,
+                       _plot_hpi_coils, _plot_sensors, _plot_helmet)
+from ..transforms import (read_trans, write_trans, _ensure_trans, _get_trans,
                           rotation_angles, _get_transforms_to_coord_frame)
 from ..utils import (get_subjects_dir, check_fname, _check_fname, fill_doc,
                      warn, verbose, logger)
@@ -105,6 +105,7 @@ class CoregistrationUI(HasTraits):
     _eeg_channels = Bool()
     _head_resolution = Bool()
     _head_transparency = Bool()
+    _helmet = Bool()
     _grow_hair = Float()
     _scale_mode = Unicode()
     _icp_fid_match = Unicode()
@@ -154,6 +155,7 @@ def _get_default(var, val):
             eeg_channels=_get_default(eeg_channels, True),
             head_resolution=_get_default(head_resolution, True),
             head_transparency=_get_default(head_transparency, False),
+            helmet=False,
             head_opacity=0.5,
             sensor_opacity=_get_default(sensor_opacity, 1.0),
             fiducials=("LPA", "Nasion", "RPA"),
@@ -213,6 +215,7 @@ def _get_default(var, val):
         self._set_eeg_channels(self._defaults["eeg_channels"])
         self._set_head_resolution(self._defaults["head_resolution"])
         self._set_head_transparency(self._defaults["head_transparency"])
+        self._set_helmet(self._defaults["helmet"])
         self._set_grow_hair(self._defaults["grow_hair"])
         self._set_omit_hsp_distance(self._defaults["omit_hsp_distance"])
         self._set_icp_n_iterations(self._defaults["icp_n_iterations"])
@@ -324,6 +327,9 @@ def _set_head_resolution(self, state):
     def _set_head_transparency(self, state):
         self._head_transparency = bool(state)
 
+    def _set_helmet(self, state):
+        self._helmet = bool(state)
+
     def _set_grow_hair(self, value):
         self._grow_hair = value
 
@@ -497,6 +503,10 @@ def _head_transparency_changed(self, change=None):
         self._actors["head"].GetProperty().SetOpacity(self._head_opacity)
         self._renderer._update()
 
+    @observe("_helmet")
+    def _helmet_changed(self, change=None):
+        self._update_plot("helmet")
+
     @observe("_grow_hair")
     def _grow_hair_changed(self, change=None):
         self._coreg.set_grow_hair(self._grow_hair)
@@ -533,6 +543,7 @@ def _redraw(self, verbose=None):
             hpi=self._add_hpi_coils,
             eeg=self._add_eeg_channels,
             head_fids=self._add_head_fiducials,
+            helmet=self._add_helmet,
         )
         with self._redraw_mutex:
             logger.debug(f'Redrawing {self._redraws_pending}')
@@ -632,6 +643,7 @@ def _update_plot(self, changes="all"):
             'head', 'mri_fids',  # MRI first
             'hair',  # then hair
             'hsp', 'hpi', 'eeg', 'head_fids',  # then dig
+            'helmet',
         )
         if changes == 'all':
             changes = list(all_keys)
@@ -828,6 +840,15 @@ def _add_head_hair(self):
             self._surfaces["head"].points = \
                 self._coreg._get_processed_mri_points(res)
 
+    def _add_helmet(self):
+        if self._helmet:
+            head_mri_t = _get_trans(self._coreg.trans, 'head', 'mri')[0]
+            helmet_actor, _, _ = _plot_helmet(
+                self._renderer, self._info, head_mri_t)
+        else:
+            helmet_actor = None
+        self._update_actor("helmet", helmet_actor)
+
     def _fit_fiducials(self):
         if not self._lock_fids:
             self._display_message(
@@ -1044,6 +1065,13 @@ def _configure_dock(self):
             tooltip="Enable/Disable high resolution head surface",
             layout=layout
         )
+        self._widgets["helmet"] = self._renderer._dock_add_check_box(
+            name="Show helmet",
+            value=self._helmet,
+            callback=self._set_helmet,
+            tooltip="Enable/Disable helmet",
+            layout=layout
+        )
         self._renderer._dock_add_stretch()
 
         self._renderer._dock_initialize(name="Parameters", area="right")

mne/gui/tests/test_coreg_gui.py

@@ -151,6 +151,9 @@ def test_coreg_gui_pyvista(tmp_path, renderer_interactive_pyvistaqt):
     assert coreg._grow_hair == 0.1
 
     # visualization
+    assert not coreg._helmet
+    coreg._set_helmet(True)
+    assert coreg._helmet
     assert coreg._orient_glyphs
     assert coreg._scale_by_distance
     assert coreg._mark_inside

mne/viz/_3d.py

@@ -654,10 +654,6 @@ def plot_alignment(info=None, trans=None, subject=None, subjects_dir=None,
         skull['name'] = skull_name  # set name for alpha
         surfs[skull_name] = skull
 
-    if 'helmet' in meg and pick_types(info, meg=True).size > 0:
-        surfs['helmet'] = get_meg_helmet_surf(info, head_mri_t)
-        assert surfs['helmet']['coord_frame'] == FIFF.FIFFV_COORD_MRI
-
     # we've looked through all of them, raise if some remain
     if len(surfaces) > 0:
         raise ValueError(f'Unknown surface type{_pl(surfaces)}: {surfaces}')
@@ -675,9 +671,8 @@ def plot_alignment(info=None, trans=None, subject=None, subjects_dir=None,
         head_alpha = max_alpha
     else:
         head_alpha = alpha_range[0]
-    alphas = dict(helmet=0.25, lh=hemi_val, rh=hemi_val)
-    colors = dict(helmet=DEFAULTS['coreg']['helmet_color'],
-                  lh=(0.5,) * 3, rh=(0.5,) * 3)
+    alphas = dict(lh=hemi_val, rh=hemi_val)
+    colors = dict(lh=(0.5,) * 3, rh=(0.5,) * 3)
     for idx, name in enumerate(skulls):
         alphas[name] = alpha_range[idx + 1]
         colors[name] = (0.95 - idx * 0.2, 0.85, 0.95 - idx * 0.2)
@@ -702,6 +697,11 @@ def plot_alignment(info=None, trans=None, subject=None, subjects_dir=None,
         renderer, head, subject, subjects_dir, bem, coord_frame,
         to_cf_t, alpha=head_alpha)
 
+    # plot helmet
+    if 'helmet' in meg and pick_types(info, meg=True).size > 0:
+        _, _, src_surf = _plot_helmet(renderer, info, head_mri_t)
+        assert src_surf['coord_frame'] == FIFF.FIFFV_COORD_MRI
+
     # plot surfaces
     if brain and 'lh' not in surfs:  # one layer sphere
         assert bem['coord_frame'] == FIFF.FIFFV_COORD_HEAD
@@ -901,6 +901,15 @@ def _plot_head_surface(renderer, head, subject, subjects_dir, bem,
     return actor, dst_surf, src_surf
 
 
+def _plot_helmet(renderer, info, head_mri_t, alpha=0.25, color=None):
+    color = DEFAULTS['coreg']['helmet_color'] if color is None else color
+    src_surf = get_meg_helmet_surf(info, head_mri_t)
+    actor, dst_surf = renderer.surface(
+        surface=src_surf, color=color, opacity=alpha,
+        backface_culling=False)
+    return actor, dst_surf, src_surf
+
+
 def _plot_axes(renderer, info, to_cf_t, head_mri_t):
     """Render different axes a 3D scene."""
     axes = [(to_cf_t['head'], (0.9, 0.3, 0.3))]  # always show head