27
27
from .callback import (ShowView , TimeCallBack , SmartCallBack ,
28
28
UpdateLUT , UpdateColorbarScale )
29
29
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 )
31
32
from .._3d import _process_clim , _handle_time , _check_views
32
33
33
34
from ...externals .decorator import decorator
@@ -704,7 +705,7 @@ def _clean(self):
704
705
self .plotter .picker = None
705
706
# XXX end PyVista
706
707
for key in ('plotter' , 'window' , 'dock' , 'tool_bar' , 'menu_bar' ,
707
- 'status_bar' , ' interactor' , 'mpl_canvas' , 'time_actor' ,
708
+ 'interactor' , 'mpl_canvas' , 'time_actor' ,
708
709
'picked_renderer' , 'act_data_smooth' , '_scalar_bar' ,
709
710
'actions' , 'widgets' , 'geo' , '_data' ):
710
711
setattr (self , key , None )
@@ -1225,22 +1226,19 @@ def _configure_picking(self):
1225
1226
self ._on_pick
1226
1227
)
1227
1228
1228
- def _save_movie_noname (self ):
1229
- return self .save_movie (None )
1230
-
1231
1229
def _configure_tool_bar (self ):
1232
1230
self ._renderer ._tool_bar_load_icons ()
1233
1231
self ._renderer ._tool_bar_set_theme (self .theme )
1234
1232
self ._renderer ._tool_bar_initialize ()
1235
- self ._renderer ._tool_bar_add_screenshot_button (
1233
+ self ._renderer ._tool_bar_add_file_button (
1236
1234
name = "screenshot" ,
1237
1235
desc = "Take a screenshot" ,
1238
- func = partial ( self .screenshot , time_viewer = True ) ,
1236
+ func = self .save_image ,
1239
1237
)
1240
- self ._renderer ._tool_bar_add_button (
1238
+ self ._renderer ._tool_bar_add_file_button (
1241
1239
name = "movie" ,
1242
1240
desc = "Save movie..." ,
1243
- func = self ._save_movie_noname ,
1241
+ func = self .save_movie ,
1244
1242
shortcut = "ctrl+shift+s" ,
1245
1243
)
1246
1244
self ._renderer ._tool_bar_add_button (
@@ -2574,7 +2572,7 @@ def reset_view(self):
2574
2572
self ._renderer .set_camera (** views_dicts [h ][v ],
2575
2573
reset_camera = False )
2576
2574
2577
- def save_image (self , filename , mode = 'rgb' ):
2575
+ def save_image (self , filename = None , mode = 'rgb' ):
2578
2576
"""Save view from all panels to disk.
2579
2577
2580
2578
Parameters
@@ -2584,7 +2582,10 @@ def save_image(self, filename, mode='rgb'):
2584
2582
mode : str
2585
2583
Either 'rgb' or 'rgba' for values to return.
2586
2584
"""
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 ))
2588
2589
2589
2590
@fill_doc
2590
2591
def screenshot (self , mode = 'rgb' , time_viewer = False ):
@@ -2991,8 +2992,52 @@ def _save_movie(self, filename, time_dilation=4., tmin=None, tmax=None,
2991
2992
kwargs ['bitrate' ] = bitrate
2992
2993
imageio .mimwrite (filename , images , ** kwargs )
2993
2994
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
+
2994
3039
@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 ,
2996
3041
framerate = 24 , interpolation = None , codec = None ,
2997
3042
bitrate = None , callback = None , time_viewer = False , ** kwargs ):
2998
3043
"""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,
3044
3089
dialog : object
3045
3090
The opened dialog is returned for testing purpose only.
3046
3091
"""
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 )
3118
3098
3119
3099
def _make_movie_frames (self , time_dilation , tmin , tmax , framerate ,
3120
3100
interpolation , callback , time_viewer ):
0 commit comments