ENH: button to save movie (#9190)

* First commit

* Fix icon

* Remove cruft

* Refactor

* Refactor

* Refactor

* Refactor

* Update conftest

* Refactor

* Fix layout

* Nitpick

* Update based on reviews

* Use tempfile

* Fix

* Update testing

Author: GuillaumeFavelier

mne/viz/_brain/_brain.py

@@ -27,7 +27,8 @@
 from .callback import (ShowView, TimeCallBack, SmartCallBack,
                        UpdateLUT, UpdateColorbarScale)
 
-from ..utils import _show_help_fig, _get_color_list, concatenate_images
+from ..utils import (_show_help_fig, _get_color_list, concatenate_images,
+                     _generate_default_filename, _save_ndarray_img)
 from .._3d import _process_clim, _handle_time, _check_views
 
 from ...externals.decorator import decorator
@@ -704,7 +705,7 @@ def _clean(self):
             self.plotter.picker = None
         # XXX end PyVista
         for key in ('plotter', 'window', 'dock', 'tool_bar', 'menu_bar',
-                    'status_bar', 'interactor', 'mpl_canvas', 'time_actor',
+                    'interactor', 'mpl_canvas', 'time_actor',
                     'picked_renderer', 'act_data_smooth', '_scalar_bar',
                     'actions', 'widgets', 'geo', '_data'):
             setattr(self, key, None)
@@ -1225,22 +1226,19 @@ def _configure_picking(self):
             self._on_pick
         )
 
-    def _save_movie_noname(self):
-        return self.save_movie(None)
-
     def _configure_tool_bar(self):
         self._renderer._tool_bar_load_icons()
         self._renderer._tool_bar_set_theme(self.theme)
         self._renderer._tool_bar_initialize()
-        self._renderer._tool_bar_add_screenshot_button(
+        self._renderer._tool_bar_add_file_button(
             name="screenshot",
             desc="Take a screenshot",
-            func=partial(self.screenshot, time_viewer=True),
+            func=self.save_image,
         )
-        self._renderer._tool_bar_add_button(
+        self._renderer._tool_bar_add_file_button(
             name="movie",
             desc="Save movie...",
-            func=self._save_movie_noname,
+            func=self.save_movie,
             shortcut="ctrl+shift+s",
         )
         self._renderer._tool_bar_add_button(
@@ -2574,7 +2572,7 @@ def reset_view(self):
                 self._renderer.set_camera(**views_dicts[h][v],
                                           reset_camera=False)
 
-    def save_image(self, filename, mode='rgb'):
+    def save_image(self, filename=None, mode='rgb'):
         """Save view from all panels to disk.
 
         Parameters
@@ -2584,7 +2582,10 @@ def save_image(self, filename, mode='rgb'):
         mode : str
             Either 'rgb' or 'rgba' for values to return.
         """
-        self._renderer.screenshot(mode=mode, filename=filename)
+        if filename is None:
+            filename = _generate_default_filename(".png")
+        _save_ndarray_img(
+            filename, self.screenshot(mode=mode, time_viewer=True))
 
     @fill_doc
     def screenshot(self, mode='rgb', time_viewer=False):
@@ -2991,8 +2992,52 @@ def _save_movie(self, filename, time_dilation=4., tmin=None, tmax=None,
             kwargs['bitrate'] = bitrate
         imageio.mimwrite(filename, images, **kwargs)
 
+    def _save_movie_tv(self, filename, time_dilation=4., tmin=None, tmax=None,
+                       framerate=24, interpolation=None, codec=None,
+                       bitrate=None, callback=None, time_viewer=False,
+                       **kwargs):
+        def frame_callback(frame, n_frames):
+            if frame == n_frames:
+                # On the ImageIO step
+                self.status_msg.set_value(
+                    "Saving with ImageIO: %s"
+                    % filename
+                )
+                self.status_msg.show()
+                self.status_progress.hide()
+                self._renderer._status_bar_update()
+            else:
+                self.status_msg.set_value(
+                    "Rendering images (frame %d / %d) ..."
+                    % (frame + 1, n_frames)
+                )
+                self.status_msg.show()
+                self.status_progress.show()
+                self.status_progress.set_range([0, n_frames - 1])
+                self.status_progress.set_value(frame)
+                self.status_progress.update()
+            self.status_msg.update()
+            self._renderer._status_bar_update()
+
+        # set cursor to busy
+        default_cursor = self._renderer._window_get_cursor()
+        self._renderer._window_set_cursor(
+            self._renderer._window_new_cursor("WaitCursor"))
+
+        try:
+            self._save_movie(
+                filename=filename,
+                time_dilation=(1. / self.playback_speed),
+                callback=frame_callback,
+                **kwargs
+            )
+        except (Exception, KeyboardInterrupt):
+            warn('Movie saving aborted:\n' + traceback.format_exc())
+        finally:
+            self._renderer._window_set_cursor(default_cursor)
+
     @fill_doc
-    def save_movie(self, filename, time_dilation=4., tmin=None, tmax=None,
+    def save_movie(self, filename=None, time_dilation=4., tmin=None, tmax=None,
                    framerate=24, interpolation=None, codec=None,
                    bitrate=None, callback=None, time_viewer=False, **kwargs):
         """Save a movie (for data with a time axis).
@@ -3044,77 +3089,12 @@ def save_movie(self, filename, time_dilation=4., tmin=None, tmax=None,
         dialog : object
             The opened dialog is returned for testing purpose only.
         """
-        if self.time_viewer:
-            try:
-                from pyvista.plotting.qt_plotting import FileDialog
-            except ImportError:
-                from pyvistaqt.plotting import FileDialog
-
-            if filename is None:
-                self.status_msg.setText("Choose movie path ...")
-                self.status_msg.show()
-                self.status_progress.setValue(0)
-
-                def _post_setup(unused):
-                    del unused
-                    self.status_msg.hide()
-                    self.status_progress.hide()
-
-                dialog = FileDialog(
-                    self.plotter.app_window,
-                    callback=partial(self._save_movie, **kwargs)
-                )
-                dialog.setDirectory(os.getcwd())
-                dialog.finished.connect(_post_setup)
-                return dialog
-            else:
-                from PyQt5.QtCore import Qt
-                from PyQt5.QtGui import QCursor
-
-                def frame_callback(frame, n_frames):
-                    if frame == n_frames:
-                        # On the ImageIO step
-                        self.status_msg.setText(
-                            "Saving with ImageIO: %s"
-                            % filename
-                        )
-                        self.status_msg.show()
-                        self.status_progress.hide()
-                        self.status_bar.layout().update()
-                    else:
-                        self.status_msg.setText(
-                            "Rendering images (frame %d / %d) ..."
-                            % (frame + 1, n_frames)
-                        )
-                        self.status_msg.show()
-                        self.status_progress.show()
-                        self.status_progress.setRange(0, n_frames - 1)
-                        self.status_progress.setValue(frame)
-                        self.status_progress.update()
-                        self.status_progress.repaint()
-                    self.status_msg.update()
-                    self.status_msg.parent().update()
-                    self.status_msg.repaint()
-
-                # set cursor to busy
-                default_cursor = self._renderer._window_get_cursor()
-                self._renderer._window_set_cursor(QCursor(Qt.WaitCursor))
-
-                try:
-                    self._save_movie(
-                        filename=filename,
-                        time_dilation=(1. / self.playback_speed),
-                        callback=frame_callback,
-                        **kwargs
-                    )
-                except (Exception, KeyboardInterrupt):
-                    warn('Movie saving aborted:\n' + traceback.format_exc())
-                finally:
-                    self._renderer._window_set_cursor(default_cursor)
-        else:
-            self._save_movie(filename, time_dilation, tmin, tmax,
-                             framerate, interpolation, codec,
-                             bitrate, callback, time_viewer, **kwargs)
+        if filename is None:
+            filename = _generate_default_filename(".mp4")
+        func = self._save_movie_tv if self.time_viewer else self._save_movie
+        func(filename, time_dilation, tmin, tmax,
+             framerate, interpolation, codec,
+             bitrate, callback, time_viewer, **kwargs)
 
     def _make_movie_frames(self, time_dilation, tmin, tmax, framerate,
                            interpolation, callback, time_viewer):

mne/viz/_brain/tests/test_notebook.py

@@ -38,6 +38,7 @@ def test_notebook_alignment(renderer_notebook, brain_gc, nbexec):
 def test_notebook_interactive(renderer_notebook, brain_gc, nbexec):
     """Test interactive modes."""
     import os
+    import tempfile
     from contextlib import contextmanager
     from numpy.testing import assert_allclose
     from ipywidgets import Button
@@ -72,6 +73,11 @@ def interactive(on):
         assert brain._renderer.figure.notebook
         assert brain._renderer.figure.display is not None
         brain._renderer._update()
+        tmp_path = tempfile.mkdtemp()
+        movie_path = os.path.join(tmp_path, 'test.gif')
+        screenshot_path = os.path.join(tmp_path, 'test.png')
+        brain._renderer.actions['movie_field'].value = movie_path
+        brain._renderer.actions['screenshot_field'].value = screenshot_path
         total_number_of_buttons = sum(
             '_field' not in k for k in brain._renderer.actions.keys())
         number_of_buttons = 0
@@ -80,6 +86,8 @@ def interactive(on):
                 action.click()
                 number_of_buttons += 1
         assert number_of_buttons == total_number_of_buttons
+        assert os.path.isfile(movie_path)
+        assert os.path.isfile(screenshot_path)
         img_nv = brain.screenshot()
         assert img_nv.shape == (300, 300, 3), img_nv.shape
         img_v = brain.screenshot(time_viewer=True)

mne/viz/backends/_abstract.py

@@ -473,7 +473,7 @@ def _tool_bar_add_spacer(self):
         pass
 
     @abstractmethod
-    def _tool_bar_add_screenshot_button(self, name, desc, func):
+    def _tool_bar_add_file_button(self, name, desc, func, shortcut=None):
         pass
 
     @abstractmethod
@@ -565,6 +565,10 @@ def _status_bar_add_label(self, value, stretch=0):
     def _status_bar_add_progress_bar(self, stretch=0):
         pass
 
+    @abstractmethod
+    def _status_bar_update(self):
+        pass
+
 
 class _AbstractPlayback(ABC):
     @abstractmethod
@@ -578,7 +582,7 @@ def _layout_initialize(self, max_width):
         pass
 
     @abstractmethod
-    def _layout_add_widget(self, layout, widget):
+    def _layout_add_widget(self, layout, widget, stretch=0):
         pass
 
 
@@ -598,6 +602,22 @@ def set_value(self, value):
     def get_value(self):
         pass
 
+    @abstractmethod
+    def set_range(self, rng):
+        pass
+
+    @abstractmethod
+    def show(self):
+        pass
+
+    @abstractmethod
+    def hide(self):
+        pass
+
+    @abstractmethod
+    def update(self, repaint=True):
+        pass
+
 
 class _AbstractMplInterface(ABC):
     @abstractmethod
@@ -768,6 +788,10 @@ def _window_get_cursor(self):
     def _window_set_cursor(self, cursor):
         pass
 
+    @abstractmethod
+    def _window_new_cursor(self, name):
+        pass
+
     @abstractmethod
     def _window_ensure_minimum_sizes(self):
         pass