Skip to content

Commit 176680a

Browse files
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
1 parent 3bc8f4e commit 176680a

File tree

7 files changed

+224
-168
lines changed

7 files changed

+224
-168
lines changed

mne/viz/_brain/_brain.py

Lines changed: 63 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@
2727
from .callback import (ShowView, TimeCallBack, SmartCallBack,
2828
UpdateLUT, UpdateColorbarScale)
2929

30-
from ..utils import _show_help_fig, _get_color_list, concatenate_images
30+
from ..utils import (_show_help_fig, _get_color_list, concatenate_images,
31+
_generate_default_filename, _save_ndarray_img)
3132
from .._3d import _process_clim, _handle_time, _check_views
3233

3334
from ...externals.decorator import decorator
@@ -704,7 +705,7 @@ def _clean(self):
704705
self.plotter.picker = None
705706
# XXX end PyVista
706707
for key in ('plotter', 'window', 'dock', 'tool_bar', 'menu_bar',
707-
'status_bar', 'interactor', 'mpl_canvas', 'time_actor',
708+
'interactor', 'mpl_canvas', 'time_actor',
708709
'picked_renderer', 'act_data_smooth', '_scalar_bar',
709710
'actions', 'widgets', 'geo', '_data'):
710711
setattr(self, key, None)
@@ -1225,22 +1226,19 @@ def _configure_picking(self):
12251226
self._on_pick
12261227
)
12271228

1228-
def _save_movie_noname(self):
1229-
return self.save_movie(None)
1230-
12311229
def _configure_tool_bar(self):
12321230
self._renderer._tool_bar_load_icons()
12331231
self._renderer._tool_bar_set_theme(self.theme)
12341232
self._renderer._tool_bar_initialize()
1235-
self._renderer._tool_bar_add_screenshot_button(
1233+
self._renderer._tool_bar_add_file_button(
12361234
name="screenshot",
12371235
desc="Take a screenshot",
1238-
func=partial(self.screenshot, time_viewer=True),
1236+
func=self.save_image,
12391237
)
1240-
self._renderer._tool_bar_add_button(
1238+
self._renderer._tool_bar_add_file_button(
12411239
name="movie",
12421240
desc="Save movie...",
1243-
func=self._save_movie_noname,
1241+
func=self.save_movie,
12441242
shortcut="ctrl+shift+s",
12451243
)
12461244
self._renderer._tool_bar_add_button(
@@ -2574,7 +2572,7 @@ def reset_view(self):
25742572
self._renderer.set_camera(**views_dicts[h][v],
25752573
reset_camera=False)
25762574

2577-
def save_image(self, filename, mode='rgb'):
2575+
def save_image(self, filename=None, mode='rgb'):
25782576
"""Save view from all panels to disk.
25792577
25802578
Parameters
@@ -2584,7 +2582,10 @@ def save_image(self, filename, mode='rgb'):
25842582
mode : str
25852583
Either 'rgb' or 'rgba' for values to return.
25862584
"""
2587-
self._renderer.screenshot(mode=mode, filename=filename)
2585+
if filename is None:
2586+
filename = _generate_default_filename(".png")
2587+
_save_ndarray_img(
2588+
filename, self.screenshot(mode=mode, time_viewer=True))
25882589

25892590
@fill_doc
25902591
def screenshot(self, mode='rgb', time_viewer=False):
@@ -2991,8 +2992,52 @@ def _save_movie(self, filename, time_dilation=4., tmin=None, tmax=None,
29912992
kwargs['bitrate'] = bitrate
29922993
imageio.mimwrite(filename, images, **kwargs)
29932994

2995+
def _save_movie_tv(self, filename, time_dilation=4., tmin=None, tmax=None,
2996+
framerate=24, interpolation=None, codec=None,
2997+
bitrate=None, callback=None, time_viewer=False,
2998+
**kwargs):
2999+
def frame_callback(frame, n_frames):
3000+
if frame == n_frames:
3001+
# On the ImageIO step
3002+
self.status_msg.set_value(
3003+
"Saving with ImageIO: %s"
3004+
% filename
3005+
)
3006+
self.status_msg.show()
3007+
self.status_progress.hide()
3008+
self._renderer._status_bar_update()
3009+
else:
3010+
self.status_msg.set_value(
3011+
"Rendering images (frame %d / %d) ..."
3012+
% (frame + 1, n_frames)
3013+
)
3014+
self.status_msg.show()
3015+
self.status_progress.show()
3016+
self.status_progress.set_range([0, n_frames - 1])
3017+
self.status_progress.set_value(frame)
3018+
self.status_progress.update()
3019+
self.status_msg.update()
3020+
self._renderer._status_bar_update()
3021+
3022+
# set cursor to busy
3023+
default_cursor = self._renderer._window_get_cursor()
3024+
self._renderer._window_set_cursor(
3025+
self._renderer._window_new_cursor("WaitCursor"))
3026+
3027+
try:
3028+
self._save_movie(
3029+
filename=filename,
3030+
time_dilation=(1. / self.playback_speed),
3031+
callback=frame_callback,
3032+
**kwargs
3033+
)
3034+
except (Exception, KeyboardInterrupt):
3035+
warn('Movie saving aborted:\n' + traceback.format_exc())
3036+
finally:
3037+
self._renderer._window_set_cursor(default_cursor)
3038+
29943039
@fill_doc
2995-
def save_movie(self, filename, time_dilation=4., tmin=None, tmax=None,
3040+
def save_movie(self, filename=None, time_dilation=4., tmin=None, tmax=None,
29963041
framerate=24, interpolation=None, codec=None,
29973042
bitrate=None, callback=None, time_viewer=False, **kwargs):
29983043
"""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,
30443089
dialog : object
30453090
The opened dialog is returned for testing purpose only.
30463091
"""
3047-
if self.time_viewer:
3048-
try:
3049-
from pyvista.plotting.qt_plotting import FileDialog
3050-
except ImportError:
3051-
from pyvistaqt.plotting import FileDialog
3052-
3053-
if filename is None:
3054-
self.status_msg.setText("Choose movie path ...")
3055-
self.status_msg.show()
3056-
self.status_progress.setValue(0)
3057-
3058-
def _post_setup(unused):
3059-
del unused
3060-
self.status_msg.hide()
3061-
self.status_progress.hide()
3062-
3063-
dialog = FileDialog(
3064-
self.plotter.app_window,
3065-
callback=partial(self._save_movie, **kwargs)
3066-
)
3067-
dialog.setDirectory(os.getcwd())
3068-
dialog.finished.connect(_post_setup)
3069-
return dialog
3070-
else:
3071-
from PyQt5.QtCore import Qt
3072-
from PyQt5.QtGui import QCursor
3073-
3074-
def frame_callback(frame, n_frames):
3075-
if frame == n_frames:
3076-
# On the ImageIO step
3077-
self.status_msg.setText(
3078-
"Saving with ImageIO: %s"
3079-
% filename
3080-
)
3081-
self.status_msg.show()
3082-
self.status_progress.hide()
3083-
self.status_bar.layout().update()
3084-
else:
3085-
self.status_msg.setText(
3086-
"Rendering images (frame %d / %d) ..."
3087-
% (frame + 1, n_frames)
3088-
)
3089-
self.status_msg.show()
3090-
self.status_progress.show()
3091-
self.status_progress.setRange(0, n_frames - 1)
3092-
self.status_progress.setValue(frame)
3093-
self.status_progress.update()
3094-
self.status_progress.repaint()
3095-
self.status_msg.update()
3096-
self.status_msg.parent().update()
3097-
self.status_msg.repaint()
3098-
3099-
# set cursor to busy
3100-
default_cursor = self._renderer._window_get_cursor()
3101-
self._renderer._window_set_cursor(QCursor(Qt.WaitCursor))
3102-
3103-
try:
3104-
self._save_movie(
3105-
filename=filename,
3106-
time_dilation=(1. / self.playback_speed),
3107-
callback=frame_callback,
3108-
**kwargs
3109-
)
3110-
except (Exception, KeyboardInterrupt):
3111-
warn('Movie saving aborted:\n' + traceback.format_exc())
3112-
finally:
3113-
self._renderer._window_set_cursor(default_cursor)
3114-
else:
3115-
self._save_movie(filename, time_dilation, tmin, tmax,
3116-
framerate, interpolation, codec,
3117-
bitrate, callback, time_viewer, **kwargs)
3092+
if filename is None:
3093+
filename = _generate_default_filename(".mp4")
3094+
func = self._save_movie_tv if self.time_viewer else self._save_movie
3095+
func(filename, time_dilation, tmin, tmax,
3096+
framerate, interpolation, codec,
3097+
bitrate, callback, time_viewer, **kwargs)
31183098

31193099
def _make_movie_frames(self, time_dilation, tmin, tmax, framerate,
31203100
interpolation, callback, time_viewer):

mne/viz/_brain/tests/test_notebook.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ def test_notebook_alignment(renderer_notebook, brain_gc, nbexec):
3838
def test_notebook_interactive(renderer_notebook, brain_gc, nbexec):
3939
"""Test interactive modes."""
4040
import os
41+
import tempfile
4142
from contextlib import contextmanager
4243
from numpy.testing import assert_allclose
4344
from ipywidgets import Button
@@ -72,6 +73,11 @@ def interactive(on):
7273
assert brain._renderer.figure.notebook
7374
assert brain._renderer.figure.display is not None
7475
brain._renderer._update()
76+
tmp_path = tempfile.mkdtemp()
77+
movie_path = os.path.join(tmp_path, 'test.gif')
78+
screenshot_path = os.path.join(tmp_path, 'test.png')
79+
brain._renderer.actions['movie_field'].value = movie_path
80+
brain._renderer.actions['screenshot_field'].value = screenshot_path
7581
total_number_of_buttons = sum(
7682
'_field' not in k for k in brain._renderer.actions.keys())
7783
number_of_buttons = 0
@@ -80,6 +86,8 @@ def interactive(on):
8086
action.click()
8187
number_of_buttons += 1
8288
assert number_of_buttons == total_number_of_buttons
89+
assert os.path.isfile(movie_path)
90+
assert os.path.isfile(screenshot_path)
8391
img_nv = brain.screenshot()
8492
assert img_nv.shape == (300, 300, 3), img_nv.shape
8593
img_v = brain.screenshot(time_viewer=True)

mne/viz/backends/_abstract.py

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -473,7 +473,7 @@ def _tool_bar_add_spacer(self):
473473
pass
474474

475475
@abstractmethod
476-
def _tool_bar_add_screenshot_button(self, name, desc, func):
476+
def _tool_bar_add_file_button(self, name, desc, func, shortcut=None):
477477
pass
478478

479479
@abstractmethod
@@ -565,6 +565,10 @@ def _status_bar_add_label(self, value, stretch=0):
565565
def _status_bar_add_progress_bar(self, stretch=0):
566566
pass
567567

568+
@abstractmethod
569+
def _status_bar_update(self):
570+
pass
571+
568572

569573
class _AbstractPlayback(ABC):
570574
@abstractmethod
@@ -578,7 +582,7 @@ def _layout_initialize(self, max_width):
578582
pass
579583

580584
@abstractmethod
581-
def _layout_add_widget(self, layout, widget):
585+
def _layout_add_widget(self, layout, widget, stretch=0):
582586
pass
583587

584588

@@ -598,6 +602,22 @@ def set_value(self, value):
598602
def get_value(self):
599603
pass
600604

605+
@abstractmethod
606+
def set_range(self, rng):
607+
pass
608+
609+
@abstractmethod
610+
def show(self):
611+
pass
612+
613+
@abstractmethod
614+
def hide(self):
615+
pass
616+
617+
@abstractmethod
618+
def update(self, repaint=True):
619+
pass
620+
601621

602622
class _AbstractMplInterface(ABC):
603623
@abstractmethod
@@ -768,6 +788,10 @@ def _window_get_cursor(self):
768788
def _window_set_cursor(self, cursor):
769789
pass
770790

791+
@abstractmethod
792+
def _window_new_cursor(self, name):
793+
pass
794+
771795
@abstractmethod
772796
def _window_ensure_minimum_sizes(self):
773797
pass

0 commit comments

Comments
 (0)