ENH: new coreg gui features (part 2)

mne/coreg.py

@@ -1632,6 +1632,11 @@ def _nearest_transformed_high_res_mri_idx_hsp(self):
         return self._nearest_calc.query(
             apply_trans(self._head_mri_t, self._filtered_extra_points))[1]
 
+    @property
+    def _has_hsp_data(self):
+        return (self._has_mri_data and
+                len(self._nearest_transformed_high_res_mri_idx_hsp) > 0)
+
     @property
     def _has_hpi_data(self):
         return (self._has_mri_data and
@@ -1963,3 +1968,67 @@ def reset(self):
         self._extra_points_filter = None
         self._update_nearest_calc()
         return self
+
+    def _get_fiducials_distance(self):
+        transformed_mri_lpa = apply_trans(self._mri_trans, self._lpa)
+        transformed_hsp_lpa = apply_trans(
+            self._head_mri_t, self._dig_dict['lpa'])
+        lpa_distance = np.linalg.norm(
+            np.ravel(transformed_mri_lpa - transformed_hsp_lpa))
+
+        transformed_mri_nasion = apply_trans(self._mri_trans, self._nasion)
+        transformed_hsp_nasion = apply_trans(
+            self._head_mri_t, self._dig_dict['nasion'])
+        nasion_distance = np.linalg.norm(
+            np.ravel(transformed_mri_nasion - transformed_hsp_nasion))
+
+        transformed_mri_rpa = apply_trans(self._mri_trans, self._rpa)
+        transformed_hsp_rpa = apply_trans(
+            self._head_mri_t, self._dig_dict['rpa'])
+        rpa_distance = np.linalg.norm(
+            np.ravel(transformed_mri_rpa - transformed_hsp_rpa))
+
+        return (lpa_distance * 1e3, nasion_distance * 1e3, rpa_distance * 1e3)
+
+    def _get_fiducials_distance_str(self):
+        dists = self._get_fiducials_distance()
+        return f"Fiducials: {dists[0]:.1f}, {dists[1]:.1f}, {dists[2]:.1f} mm"
+
+    def _get_point_distance(self):
+        mri_points = list()
+        hsp_points = list()
+        if self._hsp_weight > 0 and self._has_hsp_data:
+            mri_points.append(self._transformed_high_res_mri_points[
+                self._nearest_transformed_high_res_mri_idx_hsp])
+            hsp_points.append(self._transformed_dig_extra)
+            assert len(mri_points[-1]) == len(hsp_points[-1])
+        if self._eeg_weight > 0 and self._has_eeg_data:
+            mri_points.append(self._transformed_high_res_mri_points[
+                self._nearest_transformed_high_res_mri_idx_eeg])
+            hsp_points.append(self._transformed_dig_eeg)
+            assert len(mri_points[-1]) == len(hsp_points[-1])
+        if self._hpi_weight > 0 and self._has_hpi_data:
+            mri_points.append(self._transformed_high_res_mri_points[
+                self._nearest_transformed_high_res_mri_idx_hpi])
+            hsp_points.append(self._transformed_dig_hpi)
+            assert len(mri_points[-1]) == len(hsp_points[-1])
+        if all(len(h) == 0 for h in hsp_points):
+            return None
+        mri_points = np.concatenate(mri_points)
+        hsp_points = np.concatenate(hsp_points)
+        return np.linalg.norm(mri_points - hsp_points, axis=-1)
+
+    def _get_point_distance_str(self):
+        point_distance = self._get_point_distance()
+        if point_distance is None:
+            return ""
+        dists = 1e3 * point_distance
+        av_dist = np.mean(dists)
+        std_dist = np.std(dists)
+        kinds = [kind for kind, check in
+                 (('HSP', self._hsp_weight > 0 and self._has_hsp_data),
+                  ('EEG', self._eeg_weight > 0 and self._has_eeg_data),
+                  ('HPI', self._hpi_weight > 0 and self._has_hpi_data))
+                 if check]
+        kinds = '+'.join(kinds)
+        return f"{len(dists)} {kinds}: {av_dist:.1f} ± {std_dist:.1f} mm"

mne/gui/_coreg.py

@@ -219,6 +219,7 @@ def _get_default(var, val):
 
         # configure UI
         self._reset_fitting_parameters()
+        self._configure_status_bar()
         self._configure_dock()
         self._configure_picking()
 
@@ -242,9 +243,11 @@ def _get_default(var, val):
                  False: dict(azimuth=180, elevation=90)}  # left
         self._renderer.set_camera(distance=None, **views[self._lock_fids])
         self._redraw()
+        # XXX: internal plotter/renderer should not be exposed
         if not self._immediate_redraw:
             self._renderer.plotter.add_callback(
                 self._redraw, self._refresh_rate_ms)
+        self._renderer.plotter.show_axes()
         if standalone:
             _qt_app_exec(self._renderer.figure.store["app"])
 
@@ -373,6 +376,8 @@ def _set_point_weight(self, weight, point):
         if point in funcs.keys():
             getattr(self, funcs[point])(weight > 0)
         setattr(self, f"_{point}_weight", weight)
+        setattr(self._coreg, f"_{point}_weight", weight)
+        self._update_distance_estimation()
 
     @observe("_subjects_dir")
     def _subjects_dir_changed(self, change=None):
@@ -402,6 +407,7 @@ def _lock_fids_changed(self, change=None):
         if self._lock_fids:
             self._forward_widget_command(view_widgets, "set_enabled", True)
             self._display_message()
+            self._update_distance_estimation()
         else:
             self._forward_widget_command(view_widgets, "set_enabled", False)
             self._display_message("Picking fiducials - "
@@ -415,6 +421,7 @@ def _lock_fids_changed(self, change=None):
     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)
 
@@ -541,7 +548,7 @@ def _on_button_press(self, vtk_picker, event):
     def _on_button_release(self, vtk_picker, event):
         if self._mouse_no_mvt > 0:
             x, y = vtk_picker.GetEventPosition()
-            # XXX: plotter/renderer should not be exposed if possible
+            # XXX: internal plotter/renderer should not be exposed
             plotter = self._renderer.figure.plotter
             picked_renderer = self._renderer.figure.plotter.renderer
             # trigger the pick
@@ -589,11 +596,19 @@ def _reset_fiducials(self):
 
     def _omit_hsp(self):
         self._coreg.omit_head_shape_points(self._omit_hsp_distance / 1e3)
+        n_omitted = np.sum(~self._coreg._extra_points_filter)
+        n_remaining = len(self._coreg._dig_dict['hsp']) - n_omitted
         self._update_plot("hsp")
+        self._renderer._status_bar_show_message(
+            f"{n_omitted} head shape points omitted, "
+            f"{n_remaining} remaining.")
 
     def _reset_omit_hsp_filter(self):
         self._coreg._extra_points_filter = None
         self._update_plot("hsp")
+        n_total = len(self._coreg._dig_dict['hsp'])
+        self._renderer._status_bar_show_message(
+            f"No head shape point is omitted, the total is {n_total}.")
 
     def _update_plot(self, changes="all"):
         # Update list of things that need to be updated/plotted (and maybe
@@ -659,6 +674,11 @@ def _update_fiducials(self):
             self._forward_widget_command(
                 ["fid_X", "fid_Y", "fid_Z"], "set_value", val)
 
+    def _update_distance_estimation(self):
+        value = self._coreg._get_fiducials_distance_str() + '\n' + \
+            self._coreg._get_point_distance_str()
+        self._forward_widget_command("fit_label", "set_value", value)
+
     def _update_parameters(self):
         with self._lock_plot():
             # rotation
@@ -697,6 +717,7 @@ def _set_sensors_visibility(self, state):
         self._renderer._update()
 
     def _update_actor(self, actor_name, actor):
+        # XXX: internal plotter/renderer should not be exposed
         self._renderer.plotter.remove_actor(self._actors.get(actor_name))
         self._actors[actor_name] = actor
 
@@ -798,6 +819,7 @@ def _fit_fiducials(self):
             f"Fitting fiducials finished in {end - start:.2f} seconds.")
         self._update_plot("sensors")
         self._update_parameters()
+        self._update_distance_estimation()
 
     def _fit_icp(self):
         self._current_icp_iterations = 0
@@ -806,6 +828,16 @@ def callback(iteration, n_iterations):
             self._display_message(f"Fitting ICP - iteration {iteration + 1}")
             self._update_plot("sensors")
             self._current_icp_iterations = iteration
+            dists = self._coreg.compute_dig_mri_distances() * 1e3
+            self._status_msg.set_value(
+                "Distance between HSP and MRI (mean/min/max): "
+                f"{np.mean(dists):.2f} mm "
+                f"/ {np.min(dists):.2f} mm / {np.max(dists):.2f} mm"
+            )
+            self._status_msg.show()
+            self._status_msg.update()
+            self._renderer._status_bar_update()
+            self._update_distance_estimation()
             self._renderer._process_events()  # allow a draw or cancel
 
         start = time.time()
@@ -818,6 +850,8 @@ def callback(iteration, n_iterations):
             verbose=self._verbose,
         )
         end = time.time()
+        self._status_msg.set_value("")
+        self._status_msg.hide()
         self._display_message()
         self._renderer._status_bar_show_message(
             f"Fitting ICP finished in {end - start:.2f} seconds and "
@@ -827,6 +861,8 @@ def callback(iteration, n_iterations):
 
     def _save_trans(self, fname):
         write_trans(fname, self._coreg.trans)
+        self._renderer._status_bar_show_message(
+            f"{fname} transform file is saved.")
 
     def _load_trans(self, fname):
         mri_head_t = _ensure_trans(read_trans(fname, return_all=True),
@@ -838,6 +874,8 @@ def _load_trans(self, fname):
             tra=np.array([x, y, z]),
         )
         self._update_parameters()
+        self._renderer._status_bar_show_message(
+            f"{fname} transform file is loaded.")
 
     def _get_subjects(self, sdir=None):
         # XXX: would be nice to move this function to util
@@ -931,15 +969,15 @@ def _configure_dock(self):
             layout=layout,
         )
         self._widgets["grow_hair"] = self._renderer._dock_add_spin_box(
-            name="Grow Hair",
+            name="Grow Hair (mm)",
             value=self._grow_hair,
             rng=[0.0, 10.0],
             callback=self._set_grow_hair,
             layout=layout,
         )
         hlayout = self._renderer._dock_add_layout(vertical=False)
         self._widgets["omit_distance"] = self._renderer._dock_add_spin_box(
-            name="Omit Distance",
+            name="Omit Distance (mm)",
             value=self._omit_hsp_distance,
             rng=[0.0, 100.0],
             callback=self._set_omit_hsp_distance,
@@ -1032,6 +1070,10 @@ def _configure_dock(self):
             layout=hlayout,
         )
         self._renderer._layout_add_widget(layout, hlayout)
+        self._widgets["fit_label"] = self._renderer._dock_add_label(
+            value="",
+            layout=layout,
+        )
         self._widgets["icp_n_iterations"] = self._renderer._dock_add_spin_box(
             name="Number Of ICP Iterations",
             value=self._defaults["icp_n_iterations"],
@@ -1110,6 +1152,10 @@ def _configure_dock(self):
         self._renderer._layout_add_widget(layout, hlayout)
         self._renderer._dock_add_stretch()
 
+    def _configure_status_bar(self):
+        self._status_msg = self._renderer._status_bar_add_label("", stretch=1)
+        self._status_msg.hide()
+
     def _clean(self):
         self._renderer = None
         self._coreg = None
@@ -1118,6 +1164,8 @@ def _clean(self):
         self._surfaces.clear()
         self._defaults.clear()
         self._head_geo = None
+        self._redraw = None
+        self._status_msg = None
 
     def close(self):
         """Close interface and cleanup data structure."""