commit d8a8e887ef2cd528f61fc4ade6dd389b3dbb41b2 · cduck.me/drawsvg

+10

drawSvg/__init__.py

···

       48
        
       from .raster import Raster

     

       49
        
       from .drawing import Drawing

     

       50
        
       from .elements import *

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       51
        
       

     

       52
        
       # Make all elements available in the elements module

     

       53
        
       from . import defs

···

       48
        
       from .raster import Raster

     

       49
        
       from .drawing import Drawing

     

       50
        
       from .elements import *

     

       51
       +
       from .video import (

     

       52
       +
           render_svg_frames,

     

       53
       +
           save_video,

     

       54
       +
       )

     

       55
       +
       from .animation import (

     

       56
       +
           Animation,

     

       57
       +
           animate_video,

     

       58
       +
           animate_jupyter,

     

       59
       +
       )

     

       60
       +
       

     

       61
        
       

     

       62
        
       # Make all elements available in the elements module

     

       63
        
       from . import defs

+104

drawSvg/animation.py

···

       1
       +
       import time

     

       2
       +
       

     

       3
       +
       from . import video

     

       4
       +
       

     

       5
       +
       

     

       6
       +
       class Animation:

     

       7
       +
           def __init__(self, draw_func=None, callback=None):

     

       8
       +
               self.frames = []

     

       9
       +
               if draw_func is None:

     

       10
       +
                   draw_func = lambda d:d

     

       11
       +
               self.draw_func = draw_func

     

       12
       +
               if callback is None:

     

       13
       +
                   callback = lambda d:None

     

       14
       +
               self.callback = callback

     

       15
       +
       

     

       16
       +
           def append_frame(self, frame):

     

       17
       +
               self.frames.append(frame)

     

       18
       +
               self.callback(frame)

     

       19
       +
       

     

       20
       +
           def draw_frame(self, *args, **kwargs):

     

       21
       +
               frame = self.draw_func(*args, **kwargs)

     

       22
       +
               self.append_frame(frame)

     

       23
       +
               return frame

     

       24
       +
       

     

       25
       +
           def save_video(self, file, **kwargs):

     

       26
       +
               video.save_video(self.frames, file, **kwargs)

     

       27
       +
       

     

       28
       +
       

     

       29
       +
       class AnimationContext:

     

       30
       +
           def __init__(self, draw_func=None, out_file=None,

     

       31
       +
                        jupyter=False, pause=False, clear=True, delay=0, disable=False,

     

       32
       +
                        video_args=None, _patch_delay=0.05):

     

       33
       +
               self.jupyter = jupyter

     

       34
       +
               self.disable = disable

     

       35
       +
               if self.jupyter and not self.disable:

     

       36
       +
                   from IPython import display

     

       37
       +
                   self._jupyter_clear_output = display.clear_output

     

       38
       +
                   self._jupyter_display = display.display

     

       39
       +
                   callback = self.draw_jupyter_frame

     

       40
       +
               else:

     

       41
       +
                   callback = None

     

       42
       +
               self.anim = Animation(draw_func, callback=callback)

     

       43
       +
               self.out_file = out_file

     

       44
       +
               self.pause = pause

     

       45
       +
               self.clear = clear

     

       46
       +
               self.delay = delay

     

       47
       +
               if video_args is None:

     

       48
       +
                   video_args = {}

     

       49
       +
               self.video_args = video_args

     

       50
       +
               self._patch_delay = _patch_delay

     

       51
       +
       

     

       52
       +
           def draw_jupyter_frame(self, frame):

     

       53
       +
               if self.clear:

     

       54
       +
                   self._jupyter_clear_output(wait=True)

     

       55
       +
               self._jupyter_display(frame)

     

       56
       +
               if self.pause:

     

       57
       +
                   # Patch.  Jupyter sometimes clears the input field otherwise.

     

       58
       +
                   time.sleep(self._patch_delay)

     

       59
       +
                   input('Next?')

     

       60
       +
               elif self.delay != 0:

     

       61
       +
                   time.sleep(self.delay)

     

       62
       +
       

     

       63
       +
           def __enter__(self):

     

       64
       +
               return self.anim

     

       65
       +
       

     

       66
       +
           def __exit__(self, exc_type, exc_value, exc_traceback):

     

       67
       +
               if exc_value is None:

     

       68
       +
                   # No error

     

       69
       +
                   if self.out_file is not None and not self.disable:

     

       70
       +
                       self.anim.save_video(self.out_file, **self.video_args)

     

       71
       +
       

     

       72
       +
       

     

       73
       +
       def animate_video(out_file, draw_func=None, jupyter=False, **video_args):

     

       74
       +
           '''

     

       75
       +
           Returns a context manager that stores frames and saves a video when the

     

       76
       +
           context exits.

     

       77
       +
       

     

       78
       +
           Example:

     

       79
       +
           ```

     

       80
       +
           with animate_video('video.mp4') as draw_frame:

     

       81
       +
               while True:

     

       82
       +
                   ...

     

       83
       +
                   draw_frame(...)

     

       84
       +
           ```

     

       85
       +
           '''

     

       86
       +
           return AnimationContext(draw_func=draw_func, out_file=out_file,

     

       87
       +
                                   jupyter=jupyter, video_args=video_args)

     

       88
       +
       

     

       89
       +
       

     

       90
       +
       def animate_jupyter(draw_func=None, pause=False, clear=True, delay=0.1,

     

       91
       +
                           **kwargs):

     

       92
       +
           '''

     

       93
       +
           Returns a context manager that displays frames in a Jupyter notebook.

     

       94
       +
       

     

       95
       +
           Example:

     

       96
       +
           ```

     

       97
       +
           with animate_jupyter(delay=0.5) as draw_frame:

     

       98
       +
               while True:

     

       99
       +
                   ...

     

       100
       +
                   draw_frame(...)

     

       101
       +
           ```

     

       102
       +
           '''

     

       103
       +
           return AnimationContext(draw_func=draw_func, jupyter=True, pause=pause,

     

       104
       +
                                   clear=clear, delay=delay, **kwargs)

+54

drawSvg/video.py

···

       1
       +
       import numpy as np

     

       2
       +
       import imageio

     

       3
       +
       

     

       4
       +
       from .drawing import Drawing

     

       5
       +
       

     

       6
       +
       

     

       7
       +
       def render_svg_frames(frames, align_bottom=False, align_right=False,

     

       8
       +
                             bg=(255,)*4, **kwargs):

     

       9
       +
           arr_frames = [imageio.imread(d.rasterize().pngData)

     

       10
       +
                         for d in frames]

     

       11
       +
           max_width = max(map(lambda arr:arr.shape[1], arr_frames))

     

       12
       +
           max_height = max(map(lambda arr:arr.shape[0], arr_frames))

     

       13
       +
       

     

       14
       +
           def mod_frame(arr):

     

       15
       +
               new_arr = np.zeros((max_height, max_width) + arr.shape[2:],

     

       16
       +
                                  dtype=arr.dtype)

     

       17
       +
               new_arr[:,:] = bg[:new_arr.shape[-1]]

     

       18
       +
               if align_bottom:

     

       19
       +
                   slice0 = slice(-arr.shape[0], None)

     

       20
       +
               else:

     

       21
       +
                   slice0 = slice(None, arr.shape[0])

     

       22
       +
               if align_right:

     

       23
       +
                   slice1 = slice(-arr.shape[1], None)

     

       24
       +
               else:

     

       25
       +
                   slice1 = slice(None, arr.shape[1])

     

       26
       +
               new_arr[slice0, slice1] = arr

     

       27
       +
               return new_arr

     

       28
       +
           return list(map(mod_frame, arr_frames))

     

       29
       +
       

     

       30
       +
       def save_video(frames, file, **kwargs):

     

       31
       +
           '''

     

       32
       +
           Save a series of drawings as a GIF or video.

     

       33
       +
       

     

       34
       +
           Arguments:

     

       35
       +
               frames: A list of `Drawing`s or a list of `numpy.array`s.

     

       36
       +
               file: File name or file like object to write the video to.  The

     

       37
       +
                   extension determines the output format.

     

       38
       +
               align_bottom: If frames are different sizes, align the bottoms of each

     

       39
       +
                   frame in the video.

     

       40
       +
               align_right: If frames are different sizes, align the right edge of each

     

       41
       +
                   frame in the video.

     

       42
       +
               bg: If frames are different sizes, fill the background with this color.

     

       43
       +
                   (default is white: (255, 255, 255, 255))

     

       44
       +
               duration: If writing a GIF, sets the duration of each frame.

     

       45
       +
               fps: If writing a video, sets the frame rate in FPS.

     

       46
       +
               **kwargs: Other arguments to imageio.mimsave().

     

       47
       +
       

     

       48
       +
           '''

     

       49
       +
           if isinstance(frames[0], Drawing):

     

       50
       +
               frames = render_svg_frames(frames, **kwargs)

     

       51
       +
           kwargs.pop('align_bottom', None)

     

       52
       +
           kwargs.pop('align_right', None)

     

       53
       +
           kwargs.pop('bg', None)

     

       54
       +
           imageio.mimsave(file, frames, **kwargs)