MRG: Rework CoregistrationUI fiducial I/O widgets & add save button (#…

…10242)

* Rework CoregistrationUI fiducial I/O widgets & add save button

* Add docstring to & rename params of _forwared_widget_command()

* _core is now public

* Improve logging, make use of _forward_widget_command where possible

* Update mne/gui/_coreg.py

* Rework with Guillaume

* _update_mri_fiducials_label -> _update_fiducials_label

* Simplify

* Update mne/gui/_coreg.py

* Simplify

* Apply review suggestion

* Forgot one?

mne/coreg.py

@@ -1369,6 +1369,8 @@ def __init__(self, info, subject, subjects_dir=None, fiducials='auto', *,
 
         self._setup_digs()
         self._setup_bem()
+
+        self._fid_filename = None
         self._setup_fiducials(fiducials)
         self.reset()
 
@@ -1449,6 +1451,7 @@ def _setup_fiducials(self, fids):
                 logger.info(f'Using fiducials from: {fid_filename}.')
                 fids, _ = read_fiducials(fid_filename)
                 fid_accurate = True
+                self._fid_filename = fid_filename
             else:
                 fids = 'estimated'
 

mne/gui/_coreg.py

@@ -100,7 +100,6 @@ class CoregistrationUI(HasTraits):
     _subject = Unicode()
     _subjects_dir = Unicode()
     _lock_fids = Bool()
-    _fiducials_file = Unicode()
     _current_fiducial = Unicode()
     _info_file = Unicode()
     _orient_glyphs = Bool()
@@ -175,6 +174,7 @@ def _get_default(var, val):
         self._mouse_no_mvt = -1
         self._to_cf_t = None
         self._omit_hsp_distance = 0.0
+        self._fiducials_file = None
         self._fid_colors = tuple(
             DEFAULTS['coreg'][f'{key}_color'] for key in
             ('lpa', 'nasion', 'rpa'))
@@ -246,6 +246,8 @@ def _get_default(var, val):
             setattr(self, f"_{fid}_weight", self._defaults["weights"][fid])
 
         # set main traits
+        self._set_head_opacity(self._defaults["head_opacity"])
+        self._old_head_opacity = self._head_opacity
         self._set_subjects_dir(subjects_dir)
         self._set_subject(subject)
         self._set_info_file(info_file)
@@ -257,8 +259,6 @@ def _get_default(var, val):
         self._set_head_shape_points(self._defaults["head_shape_points"])
         self._set_eeg_channels(self._defaults["eeg_channels"])
         self._set_head_resolution(self._defaults["head_resolution"])
-        self._set_head_opacity(self._defaults["head_opacity"])
-        self._old_head_opacity = self._head_opacity
         self._set_helmet(self._defaults["helmet"])
         self._set_grow_hair(self._defaults["grow_hair"])
         self._set_skip_fiducials(self._defaults["skip_fiducials"])
@@ -282,10 +282,19 @@ def _get_default(var, val):
         if trans is not None:
             self._load_trans(trans)
         self._redraw()  # we need the elements to be present now
-        if not fid_accurate:
+
+        if fid_accurate:
+            assert self.coreg._fid_filename is not None
+            # _set_fiducials_file() calls _update_fiducials_label()
+            # internally
+            self._set_fiducials_file(self.coreg._fid_filename)
+        else:
             self._set_head_resolution('high')
             self._forward_widget_command('high_res_head', "set_value", True)
             self._set_lock_fids(True)  # hack to make the dig disappear
+            self._update_fiducials_label()
+            self._update_fiducials()
+
         self._set_lock_fids(fid_accurate)
 
         # configure worker
@@ -318,10 +327,34 @@ def _set_lock_fids(self, state):
         self._lock_fids = bool(state)
 
     def _set_fiducials_file(self, fname):
-        if not self._check_fif('fiducials', fname):
-            return
-        self._fiducials_file = _check_fname(
-            fname, overwrite='read', must_exist=True, need_dir=False)
+        if fname is None:
+            fids = 'auto'
+        else:
+            fname = _check_fname(
+                fname, overwrite='read', must_exist=True, need_dir=False
+            )
+            fids, _ = read_fiducials(fname)
+
+        self._fiducials_file = fname
+        self.coreg._setup_fiducials(fids)
+        self._update_distance_estimation()
+        self._update_fiducials_label()
+        self._update_fiducials()
+        self._reset()
+
+        if fname is None:
+            self._set_lock_fids(False)
+            self._forward_widget_command(
+                'reload_mri_fids', 'set_enabled', False
+            )
+        else:
+            self._set_lock_fids(True)
+            self._forward_widget_command(
+                'reload_mri_fids', 'set_enabled', True
+            )
+            self._display_message(
+                f"Loading MRI fiducials from {fname}... Done!"
+            )
 
     def _set_current_fiducial(self, fid):
         self._current_fiducial = fid.lower()
@@ -466,7 +499,10 @@ def _subjects_dir_changed(self, change=None):
         # XXX: add coreg.set_subjects_dir
         self.coreg._subjects_dir = self._subjects_dir
         subjects = self._get_subjects()
-        self._subject = subjects[0]
+
+        if self._subject not in subjects:  # Just pick the first available one
+            self._subject = subjects[0]
+
         self._reset()
 
     @observe("_subject")
@@ -478,11 +514,23 @@ def _subject_changed(self, changed=None):
         self._reset()
         self._update_projection_surface()
 
+        default_fid_fname = fid_fname.format(
+            subjects_dir=self._subjects_dir, subject=self._subject
+        )
+        if Path(default_fid_fname).exists():
+            fname = default_fid_fname
+        else:
+            fname = None
+
+        self._set_fiducials_file(fname)
+        self._reset_fiducials()
+
     @observe("_lock_fids")
     def _lock_fids_changed(self, change=None):
         locked_widgets = ["sX", "sY", "sZ", "tX", "tY", "tZ",
                           "rX", "rY", "rZ", "project_eeg",
-                          "fit_fiducials", "fit_icp"]
+                          "fit_fiducials", "fit_icp",
+                          "save_mri_fids"]
         fits_widgets = ["fits_fiducials", "fits_icp"]
         fid_widgets = ["fid_X", "fid_Y", "fid_Z", "fids_file", "fids"]
         if self._lock_fids:
@@ -498,19 +546,12 @@ def _lock_fids_changed(self, change=None):
             self._forward_widget_command(fits_widgets, "set_enabled", False)
             self._display_message("Placing MRI fiducials - "
                                   f"{self._current_fiducial.upper()}")
+
         self._set_sensors_visibility(self._lock_fids)
         self._forward_widget_command("lock_fids", "set_value", self._lock_fids)
         self._forward_widget_command(fid_widgets, "set_enabled",
                                      not self._lock_fids)
 
-    @observe("_fiducials_file")
-    def _fiducials_file_changed(self, change=None):
-        fids, _ = read_fiducials(self._fiducials_file)
-        self.coreg._setup_fiducials(fids)
-        self._update_distance_estimation()
-        self._reset()
-        self._set_lock_fids(True)
-
     @observe("_current_fiducial")
     def _current_fiducial_changed(self, change=None):
         self._update_fiducials()
@@ -816,6 +857,9 @@ def _update_projection_surface(self):
 
     def _update_fiducials(self):
         fid = self._current_fiducial
+        if not fid:
+            return
+
         idx = _map_fid_name_to_idx(name=fid)
         val = self.coreg.fiducials.dig[idx]['r'] * 1e3
 
@@ -1149,8 +1193,8 @@ def _save_mri_fiducials(self, fname):
         write_fiducials(
             fname=fname, pts=dig_montage.dig, coord_frame='mri', overwrite=True
         )
-        self._display_message(f"Saving {fname}... Done!")
         self._set_fiducials_file(fname)
+        self._display_message(f"Saving {fname}... Done!")
 
     def _save_trans(self, fname):
         write_trans(fname, self.coreg.trans, overwrite=True)
@@ -1184,14 +1228,29 @@ def _get_subjects(self, sdir=None):
             subjects = ['']
         return sorted(subjects)
 
-    def _check_fif(self, filetype, fname):
-        try:
-            check_fname(fname, filetype, ('.fif'), ('.fif'))
-        except IOError:
-            warn(f"The filename {fname} for {filetype} must end with '.fif'.")
-            self._widgets[f"{filetype}_file"].set_value(0, '')
-            return False
-        return True
+    def _update_fiducials_label(self):
+        if self._fiducials_file is None:
+            text = (
+                '<p><strong>No custom MRI fiducials loaded!</strong></p>'
+                '<p>MRI fiducials could not be found in the standard '
+                'location. The displayed initial MRI fiducial locations '
+                '(diamonds) were derived from fsaverage. Place, lock, and '
+                'save fiducials to discard this message.</p>'
+            )
+        else:
+            assert self._fiducials_file == fid_fname.format(
+                subjects_dir=self._subjects_dir, subject=self._subject
+            )
+            assert self.coreg._fid_accurate is True
+            text = (
+                f'<p><strong>MRI fiducials (diamonds) loaded from '
+                f'standard location:</strong></p>'
+                f'<p>{self._fiducials_file}</p>'
+            )
+
+        self._forward_widget_command(
+            'mri_fiducials_label', 'set_value', text
+        )
 
     def _configure_dock(self):
         self._renderer._dock_initialize(
@@ -1219,24 +1278,51 @@ def _configure_dock(self):
             layout=mri_subject_layout,
         )
 
-        mri_fiducials_layout = \
-            self._renderer._dock_add_group_box("MRI Fiducials")
+        mri_fiducials_layout =  self._renderer._dock_add_group_box(
+            "MRI Fiducials"
+        )
+        # Add MRI fiducials I/O widgets
+        self._widgets['mri_fiducials_label'] = self._renderer._dock_add_label(
+            value='',  # Will be filled via _update_fiducials_label()
+            layout=mri_fiducials_layout,
+            selectable=True
+        )
+        # Reload & Save buttons go into their own layout widget
+        mri_fiducials_button_layout =  self._renderer._dock_add_layout(
+            vertical=False
+        )
+        self._renderer._layout_add_widget(
+            layout=mri_fiducials_layout,
+            widget=mri_fiducials_button_layout
+        )
+        self._widgets["reload_mri_fids"] = self._renderer._dock_add_button(
+            name='Reload MRI Fid.',
+            callback=lambda: self._set_fiducials_file(self._fiducials_file),
+            tooltip="Reload MRI fiducials from the standard location",
+            layout=mri_fiducials_button_layout,
+        )
+        # Disable reload button until we've actually loaded a fiducial file
+        # (happens in _set_fiducials_file method)
+        self._forward_widget_command('reload_mri_fids', 'set_enabled', False)
+
+        self._widgets["save_mri_fids"] = self._renderer._dock_add_button(
+            name="Save MRI Fid.",
+            callback=lambda: self._save_mri_fiducials(
+                fid_fname.format(
+                    subjects_dir=self._subjects_dir, subject=self._subject
+                )
+            ),
+            tooltip="Save MRI fiducials to the standard location. Fiducials "
+                    "must be locked first!",
+            layout=mri_fiducials_button_layout,
+        )
         self._widgets["lock_fids"] = self._renderer._dock_add_check_box(
             name="Lock fiducials",
             value=self._lock_fids,
             callback=self._set_lock_fids,
             tooltip="Lock/Unlock interactive fiducial editing",
             layout=mri_fiducials_layout,
         )
-        self._widgets["fiducials_file"] = self._renderer._dock_add_file_button(
-            name="fiducials_file",
-            desc="Load",
-            func=self._set_fiducials_file,
-            value=self._fiducials_file,
-            placeholder="Path to fiducials",
-            tooltip="Load the fiducials from a FIFF file",
-            layout=mri_fiducials_layout,
-        )
         self._widgets["fids"] = self._renderer._dock_add_radio_buttons(
             value=self._defaults["fiducial"],
             rng=self._defaults["fiducials"],