MRG, ENH: Add a screenshot button to the notebook 3d backend

mne/datasets/utils.py

@@ -384,7 +384,7 @@ def _data_path(path=None, force_update=False, update_path=True, download=True,
     want_version = _FAKE_VERSION if name == 'fake' else want_version
     if not need_download and want_version is not None:
         data_version = _dataset_version(folder_path[0], name)
-        need_download = data_version != want_version
+        need_download = LooseVersion(data_version) < LooseVersion(want_version)
         if need_download:
             logger.info(f'Dataset {name} version {data_version} out of date, '
                         f'latest version is {want_version}')

mne/label.py

@@ -1874,9 +1874,11 @@ def _cortex_parcellation(subject, n_parcel, hemis, vertices_, graphs,
     return labels
 
 
-def _read_annot_cands(dir_name):
+def _read_annot_cands(dir_name, raise_error=True):
     """List the candidate parcellations."""
     if not op.isdir(dir_name):
+        if not raise_error:
+            return list()
         raise IOError('Directory for annotation does not exist: %s',
                       dir_name)
     cands = os.listdir(dir_name)

mne/viz/__init__.py

@@ -6,7 +6,7 @@
 from .topo import plot_topo_image_epochs, iter_topography
 from .utils import (tight_layout, mne_analyze_colormap, compare_fiff,
                     ClickableImage, add_background_image, plot_sensors,
-                    centers_to_edges)
+                    centers_to_edges, concatenate_images)
 from ._3d import (plot_sparse_source_estimates, plot_source_estimates,
                   plot_vector_source_estimates, plot_evoked_field,
                   plot_dipole_locations, snapshot_brain_montage,

mne/viz/_brain/_brain.py

@@ -593,6 +593,7 @@ def setup_time_viewer(self, time_viewer=True, show_traces=True):
         self._configure_picking()
         self._configure_tool_bar()
         if self.notebook:
+            self._renderer._set_tool_bar(state=False)
             self.show()
         self._configure_trace_mode()
         self.toggle_interface()
@@ -1145,15 +1146,16 @@ def _set_annot(annot):
 
         from PyQt5.QtWidgets import QComboBox, QLabel
         dir_name = op.join(self._subjects_dir, self._subject_id, 'label')
-        cands = _read_annot_cands(dir_name)
+        cands = _read_annot_cands(dir_name, raise_error=False)
         self.tool_bar.addSeparator()
         self.tool_bar.addWidget(QLabel("Annotation"))
         self._annot_cands_widget = QComboBox()
         self.tool_bar.addWidget(self._annot_cands_widget)
-        self._annot_cands_widget.addItem('None')
+        cands = ['None'] + cands
         for cand in cands:
             self._annot_cands_widget.addItem(cand)
         self.annot = cands[0]
+        del cands
 
         # setup label extraction parameters
         def _set_label_mode(mode):
@@ -1215,22 +1217,28 @@ def _save_movie_noname(self):
         return self.save_movie(None)
 
     def _screenshot(self):
-        if not self.notebook:
+        if self.notebook:
+            from PIL import Image
+            fname = self.actions.get("screenshot_field").value
+            fname = self._renderer._get_screenshot_filename() \
+                if len(fname) == 0 else fname
+            img = self.screenshot(fname, time_viewer=True)
+            Image.fromarray(img).save(fname)
+        else:
             self.plotter._qt_screenshot()
 
     def _initialize_actions(self):
         if not self.notebook:
             self._load_icons()
             self.tool_bar = self.window.addToolBar("toolbar")
 
-    def _add_action(self, name, desc, func, icon_name, qt_icon_name=None,
+    def _add_button(self, name, desc, func, icon_name, qt_icon_name=None,
                     notebook=True):
         if self.notebook:
             if not notebook:
                 return
-            from ipywidgets import Button
-            self.actions[name] = Button(description=desc, icon=icon_name)
-            self.actions[name].on_click(lambda x: func())
+            self.actions[name] = self._renderer._add_button(
+                desc, func, icon_name)
         else:
             qt_icon_name = name if qt_icon_name is None else qt_icon_name
             self.actions[name] = self.tool_bar.addAction(
@@ -1239,61 +1247,71 @@ def _add_action(self, name, desc, func, icon_name, qt_icon_name=None,
                 func,
             )
 
+    def _add_text_field(self, name, value, placeholder):
+        if not self.notebook:
+            return
+        self.actions[name] = self._renderer._add_text_field(
+            value, placeholder)
+
     def _configure_tool_bar(self):
         self._initialize_actions()
-        self._add_action(
+        self._add_button(
             name="screenshot",
             desc="Take a screenshot",
             func=self._screenshot,
-            icon_name=None,
-            notebook=False,
+            icon_name="camera",
+        )
+        self._add_text_field(
+            name="screenshot_field",
+            value=None,
+            placeholder="Type a file name",
         )
-        self._add_action(
+        self._add_button(
             name="movie",
             desc="Save movie...",
             func=self._save_movie_noname,
             icon_name=None,
             notebook=False,
         )
-        self._add_action(
+        self._add_button(
             name="visibility",
             desc="Toggle Visibility",
             func=self.toggle_interface,
             icon_name="eye",
             qt_icon_name="visibility_on",
         )
-        self._add_action(
+        self._add_button(
             name="play",
             desc="Play/Pause",
             func=self.toggle_playback,
             icon_name=None,
             notebook=False,
         )
-        self._add_action(
+        self._add_button(
             name="reset",
             desc="Reset",
             func=self.reset,
             icon_name="history",
         )
-        self._add_action(
+        self._add_button(
             name="scale",
             desc="Auto-Scale",
             func=self.apply_auto_scaling,
             icon_name="magic",
         )
-        self._add_action(
+        self._add_button(
             name="restore",
             desc="Restore scaling",
             func=self.restore_user_scaling,
             icon_name="reply",
         )
-        self._add_action(
+        self._add_button(
             name="clear",
             desc="Clear traces",
             func=self.clear_glyphs,
             icon_name="trash",
         )
-        self._add_action(
+        self._add_button(
             name="help",
             desc="Help",
             func=self.help,
@@ -1302,10 +1320,7 @@ def _configure_tool_bar(self):
         )
 
         if self.notebook:
-            from IPython import display
-            from ipywidgets import HBox
-            self.tool_bar = HBox(tuple(self.actions.values()))
-            display.display(self.tool_bar)
+            self.tool_bar = self._renderer._show_tool_bar(self.actions)
         else:
             # Qt shortcuts
             self.actions["movie"].setShortcut("ctrl+shift+s")
@@ -1593,6 +1608,7 @@ def plot_time_course(self, hemi, vertex_id, color):
         if self.mpl_canvas is None:
             return
         time = self._data['time'].copy()  # avoid circular ref
+        mni = None
         if hemi == 'vol':
             hemi_str = 'V'
             xfm = read_talxfm(
@@ -1605,15 +1621,20 @@ def plot_time_course(self, hemi, vertex_id, color):
             mni = apply_trans(np.dot(xfm['trans'], src_mri_t), ijk)
         else:
             hemi_str = 'L' if hemi == 'lh' else 'R'
-            mni = vertex_to_mni(
-                vertices=vertex_id,
-                hemis=0 if hemi == 'lh' else 1,
-                subject=self._subject_id,
-                subjects_dir=self._subjects_dir
-            )
-        label = "{}:{} MNI: {}".format(
-            hemi_str, str(vertex_id).ljust(6),
-            ', '.join('%5.1f' % m for m in mni))
+            try:
+                mni = vertex_to_mni(
+                    vertices=vertex_id,
+                    hemis=0 if hemi == 'lh' else 1,
+                    subject=self._subject_id,
+                    subjects_dir=self._subjects_dir
+                )
+            except Exception:
+                mni = None
+        if mni is not None:
+            mni = ' MNI: ' + ', '.join('%5.1f' % m for m in mni)
+        else:
+            mni = ''
+        label = "{}:{}{}".format(hemi_str, str(vertex_id).ljust(6), mni)
         act_data, smooth = self.act_data_smooth[hemi]
         if smooth is not None:
             act_data = smooth[vertex_id].dot(act_data)[0]
@@ -2594,25 +2615,43 @@ def screenshot(self, mode='rgb', time_viewer=False):
                 not self.separate_canvas:
             canvas = self.mpl_canvas.fig.canvas
             canvas.draw_idle()
-            # In theory, one of these should work:
-            #
-            # trace_img = np.frombuffer(
-            #     canvas.tostring_rgb(), dtype=np.uint8)
-            # trace_img.shape = canvas.get_width_height()[::-1] + (3,)
-            #
-            # or
-            #
-            # trace_img = np.frombuffer(
-            #     canvas.tostring_rgb(), dtype=np.uint8)
-            # size = time_viewer.mpl_canvas.getSize()
-            # trace_img.shape = (size.height(), size.width(), 3)
-            #
-            # But in practice, sometimes the sizes does not match the
-            # renderer tostring_rgb() size. So let's directly use what
-            # matplotlib does in lib/matplotlib/backends/backend_agg.py
-            # before calling tobytes():
-            trace_img = np.asarray(
-                canvas.renderer._renderer).take([0, 1, 2], axis=2)
+            if self.notebook:
+                from io import BytesIO
+                # Here we need to save to a buffer without taking into
+                # account the dpi to keep the size used on creation
+                # of the canvas which is used for the 3D renderer.
+                # For some reason doing
+                # canvas.get_width_height() overestimates the size
+                # of the image while doing fig.savefig('fname.png')
+                # returns the right number of pixels.
+                fig = self.mpl_canvas.fig
+                output = BytesIO()
+                fig.savefig(output, format='raw')
+                output.seek(0)
+                trace_img = np.reshape(
+                    np.frombuffer(output.getvalue(), dtype=np.uint8),
+                    newshape=(-1, img.shape[1], 4))[:, :, :3]
+                output.close()
+            else:
+                # In theory, one of these should work:
+                #
+                # trace_img = np.frombuffer(
+                #     canvas.tostring_rgb(), dtype=np.uint8)
+                # trace_img.shape = canvas.get_width_height()[::-1] + (3,)
+                #
+                # or
+                #
+                # trace_img = np.frombuffer(
+                #     canvas.tostring_rgb(), dtype=np.uint8)
+                # size = time_viewer.mpl_canvas.getSize()
+                # trace_img.shape = (size.height(), size.width(), 3)
+                #
+                # But in practice, sometimes the sizes does not match the
+                # renderer tostring_rgb() size. So let's directly use what
+                # matplotlib does in lib/matplotlib/backends/backend_agg.py
+                # before calling tobytes():
+                trace_img = np.asarray(
+                    canvas.renderer._renderer).take([0, 1, 2], axis=2)
             # need to slice into trace_img because generally it's a bit
             # smaller
             delta = trace_img.shape[1] - img.shape[1]

mne/viz/_brain/tests/test.ipynb

@@ -29,6 +29,7 @@
    "source": [
     "import os\n",
     "import mne\n",
+    "from ipywidgets import Button\n",
     "import matplotlib.pyplot as plt\n",
     "from mne.datasets import testing\n",
     "data_path = testing.data_path()\n",
@@ -44,11 +45,22 @@
     "    brain = stc.plot(subjects_dir=subjects_dir, initial_time=initial_time,\n",
     "                     clim=dict(kind='value', pos_lims=[3, 6, 9]),\n",
     "                     time_viewer=True,\n",
+    "                     show_traces=True,\n",
     "                     hemi='split')\n",
     "    assert isinstance(brain, brain_class)\n",
     "    assert brain.notebook\n",
     "    assert brain._renderer.figure.display is not None\n",
     "    brain._update()\n",
+    "    total_number_of_buttons = len([k for k in brain.actions.keys() if '_field' not in k])\n",
+    "    number_of_buttons = 0\n",
+    "    for action in brain.actions.values():\n",
+    "        if isinstance(action, Button):\n",
+    "            action.click()\n",
+    "            number_of_buttons += 1\n",
+    "    assert number_of_buttons == total_number_of_buttons\n",
+    "    img1 = brain.screenshot(time_viewer=False)\n",
+    "    img2 = brain.screenshot(time_viewer=True)\n",
+    "    assert img1.shape[0] < img2.shape[0]\n",
     "    brain.close()"
    ]
   },
@@ -66,6 +78,13 @@
     "mne.viz.set_3d_view(fig, 200, 70, focalpoint=[0, 0, 0])\n",
     "assert fig.display is None\n",
     "rend.show()\n",
+    "total_number_of_buttons = len([k for k in rend.actions.keys() if '_field' not in k])\n",
+    "number_of_buttons = 0\n",
+    "for action in rend.actions.values():\n",
+    "    if isinstance(action, Button):\n",
+    "        action.click()\n",
+    "        number_of_buttons += 1\n",
+    "assert number_of_buttons == total_number_of_buttons\n",
     "assert fig.display is not None"
    ]
   }

mne/viz/_brain/tests/test_brain.py

@@ -17,7 +17,7 @@
 
 from mne import (read_source_estimate, read_evokeds, read_cov,
                  read_forward_solution, pick_types_forward,
-                 SourceEstimate, MixedSourceEstimate,
+                 SourceEstimate, MixedSourceEstimate, write_surface,
                  VolSourceEstimate)
 from mne.minimum_norm import apply_inverse, make_inverse_operator
 from mne.source_space import (read_source_spaces, vertex_to_mni,
@@ -369,6 +369,46 @@ def test_brain_save_movie(tmpdir, renderer, brain_gc):
     brain.close()
 
 
+@pytest.fixture()
+def tiny(tmpdir, request):
+    """Create a tiny fake brain."""
+    renderer_interactive = request.getfixturevalue('renderer_interactive')
+    if renderer_interactive._get_3d_backend() != 'pyvista':
+        pytest.skip('TimeViewer tests only supported on PyVista')
+    # This is a minimal version of what we need for our viz-with-timeviewer
+    # support currently
+    subject = 'test'
+    subject_dir = tmpdir.mkdir(subject)
+    surf_dir = subject_dir.mkdir('surf')
+    rng = np.random.RandomState(0)
+    rr = rng.randn(4, 3)
+    tris = np.array([[0, 1, 2], [2, 1, 3]])
+    curv = rng.randn(len(rr))
+    with open(surf_dir.join('lh.curv'), 'wb') as fid:
+        fid.write(np.array([255, 255, 255], dtype=np.uint8))
+        fid.write(np.array([len(rr), 0, 1], dtype='>i4'))
+        fid.write(curv.astype('>f4'))
+    write_surface(surf_dir.join('lh.white'), rr, tris)
+    write_surface(surf_dir.join('rh.white'), rr, tris)  # needed for vertex tc
+    vertices = [np.arange(len(rr)), []]
+    data = rng.randn(len(rr), 10)
+    stc = SourceEstimate(data, vertices, 0, 1, subject)
+    brain = stc.plot(subjects_dir=tmpdir, hemi='lh', surface='white',
+                     size=300)
+    ratio = brain.mpl_canvas.canvas.window().devicePixelRatio()
+    return brain, ratio
+
+
+def test_brain_screenshot(renderer_interactive, brain_gc, tiny):
+    """Test time viewer screenshot."""
+    tiny_brain, ratio = tiny
+    img1 = tiny_brain.screenshot(time_viewer=False)
+    img2 = tiny_brain.screenshot(time_viewer=True)
+    assert img1.shape[0] < img2.shape[0]
+    assert img1.shape[1] == 300 * ratio
+    tiny_brain.close()
+
+
 @testing.requires_testing_data
 @pytest.mark.slowtest
 def test_brain_time_viewer(renderer_interactive, pixel_ratio, brain_gc):