cduck.me
1import time
2
3from . import video
4
5
6class 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
29class 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
73def 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
90def 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)