parallel-matplotlib-animation 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- parallel_animate/__init__.py +1 -1
- parallel_animate/animator.py +150 -96
- parallel_animate/examples/nondeterministic_video_loader.py +46 -0
- parallel_animate/util.py +64 -0
- {parallel_matplotlib_animation-0.1.0.dist-info → parallel_matplotlib_animation-0.1.2.dist-info}/METADATA +30 -20
- parallel_matplotlib_animation-0.1.2.dist-info/RECORD +13 -0
- {parallel_matplotlib_animation-0.1.0.dist-info → parallel_matplotlib_animation-0.1.2.dist-info}/WHEEL +1 -1
- parallel_matplotlib_animation-0.1.0.dist-info/RECORD +0 -11
- {parallel_matplotlib_animation-0.1.0.dist-info → parallel_matplotlib_animation-0.1.2.dist-info}/licenses/LICENSE +0 -0
parallel_animate/__init__.py
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
from .animator import Animator
|
|
1
|
+
from .animator import Animator, IndexedFrameParams
|
parallel_animate/animator.py
CHANGED
|
@@ -3,17 +3,34 @@ import matplotlib
|
|
|
3
3
|
matplotlib.use("Agg")
|
|
4
4
|
|
|
5
5
|
import tempfile
|
|
6
|
+
import itertools
|
|
6
7
|
import logging
|
|
7
|
-
import time
|
|
8
8
|
import multiprocessing as mp
|
|
9
|
-
from pathlib import Path
|
|
10
|
-
from abc import ABC, abstractmethod
|
|
11
|
-
from typing import Any
|
|
12
|
-
|
|
13
9
|
import matplotlib.pyplot as plt
|
|
14
10
|
import av
|
|
15
11
|
from PIL import Image
|
|
16
|
-
from tqdm import tqdm
|
|
12
|
+
from tqdm import tqdm
|
|
13
|
+
from pathlib import Path
|
|
14
|
+
from abc import ABC, abstractmethod
|
|
15
|
+
from typing import Any, Iterable
|
|
16
|
+
from dataclasses import dataclass
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
_logger = logging.getLogger(__name__)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@dataclass(frozen=True)
|
|
23
|
+
class IndexedFrameParams:
|
|
24
|
+
"""Special dataclass to hold parameters for each frame **if frames arrive out of
|
|
25
|
+
order in the `param_by_frame` iterator**. The frame_id from this class will
|
|
26
|
+
override the index in the iterator."""
|
|
27
|
+
|
|
28
|
+
frame_id: int
|
|
29
|
+
params: Any
|
|
30
|
+
|
|
31
|
+
def __post_init__(self):
|
|
32
|
+
if self.frame_id < 0:
|
|
33
|
+
raise ValueError(f"frame_id must be non-negative, got {self.frame_id}")
|
|
17
34
|
|
|
18
35
|
|
|
19
36
|
class Animator(ABC):
|
|
@@ -21,10 +38,6 @@ class Animator(ABC):
|
|
|
21
38
|
Base class for creating matplotlib animations with efficient parallel rendering.
|
|
22
39
|
"""
|
|
23
40
|
|
|
24
|
-
def __init__(self):
|
|
25
|
-
"""Initialize the animator."""
|
|
26
|
-
self.logger = logging.getLogger(__name__)
|
|
27
|
-
|
|
28
41
|
@abstractmethod
|
|
29
42
|
def setup(self):
|
|
30
43
|
"""
|
|
@@ -60,24 +73,30 @@ class Animator(ABC):
|
|
|
60
73
|
def make_video(
|
|
61
74
|
self,
|
|
62
75
|
output_file: Path | str,
|
|
63
|
-
param_by_frame:
|
|
76
|
+
param_by_frame: Iterable[Any],
|
|
64
77
|
fps: int,
|
|
78
|
+
n_frames: int | None = None,
|
|
65
79
|
num_workers: int = -1,
|
|
66
80
|
disable_progress_bar: bool | None = None,
|
|
67
81
|
plotting_log_interval: int | None = None,
|
|
68
82
|
saving_log_interval: int | None = None,
|
|
69
|
-
savefig_params: dict[str, Any] =
|
|
83
|
+
savefig_params: dict[str, Any] | None = None,
|
|
70
84
|
video_codec: str = "libx264",
|
|
71
|
-
|
|
85
|
+
video_pixfmt: str = "yuv420p",
|
|
86
|
+
video_params: dict[str, Any] | None = None,
|
|
72
87
|
reuse_figure_object: bool = True,
|
|
88
|
+
preload_factor: int = 8,
|
|
73
89
|
) -> None:
|
|
74
90
|
"""
|
|
75
91
|
Render the animation to a video file.
|
|
76
92
|
|
|
77
93
|
Args:
|
|
78
94
|
output_file (Path | str): Path to output video file
|
|
79
|
-
param_by_frame (
|
|
95
|
+
param_by_frame (Iterable[Any]): Iterable of parameters, one per frame
|
|
80
96
|
fps (int): Frames per second for output video
|
|
97
|
+
n_frames (int | None): Number of frames to render. If None, use the length
|
|
98
|
+
of param_by_frame. If param_by_frame does not have __len__ implemented
|
|
99
|
+
and n_frames is None, the progress bar won't show completion percentage.
|
|
81
100
|
num_workers (int): Number of parallel workers. 1 for serial processing
|
|
82
101
|
(in the main thread), -1 for all CPU cores, -2 for all but one CPU core,
|
|
83
102
|
etc.
|
|
@@ -90,37 +109,51 @@ class Animator(ABC):
|
|
|
90
109
|
savefig_params (dict[str, Any]): Additional keyword arguments to
|
|
91
110
|
pass to plt.Figure.savefig() when saving frames (default: {}).
|
|
92
111
|
video_codec (str): Codec to use for video encoding (default: "libx264").
|
|
112
|
+
video_pixfmt (str): Pixel format for video encoding (default: "yuv420p").
|
|
93
113
|
video_params (dict[str, Any]): Additional parameters to set on the video
|
|
94
|
-
stream (default: {"
|
|
114
|
+
stream (default: {"crf": "23", "preset": "slow"}).
|
|
95
115
|
reuse_figure_object (bool): If False, the figure will be re-created for each
|
|
96
116
|
frame (i.e. setup() called every frame). There is basically no reason to
|
|
97
117
|
set this to False. Use only for testing and benchmarking.
|
|
118
|
+
preload_factor (int): Number of workloads to prefetch in the task queue per
|
|
119
|
+
worker (approximately). I.e. each worker will have up to this many
|
|
120
|
+
frames (on average) queued ahead of time to work on.
|
|
98
121
|
"""
|
|
99
|
-
#
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
num_frames = len(params_list)
|
|
122
|
+
# Convert param_by_frame to list if it's not already
|
|
123
|
+
if n_frames is None:
|
|
124
|
+
try:
|
|
125
|
+
n_frames = len(param_by_frame)
|
|
126
|
+
except TypeError:
|
|
127
|
+
_logger.warning(
|
|
128
|
+
"param_by_frame has no length and n_frames is not specified. "
|
|
129
|
+
"Progress bar won't show completion percentage."
|
|
130
|
+
)
|
|
109
131
|
|
|
110
132
|
# Determine number of workers
|
|
111
|
-
if num_workers ==
|
|
133
|
+
if num_workers == 0:
|
|
134
|
+
raise ValueError(
|
|
135
|
+
"num_workers cannot be 0. Use 1 for serial mode or -1 for all cores."
|
|
136
|
+
)
|
|
137
|
+
elif num_workers == -1:
|
|
112
138
|
num_workers = mp.cpu_count()
|
|
113
139
|
elif num_workers < -1:
|
|
114
140
|
num_workers = max(1, mp.cpu_count() + num_workers + 1)
|
|
115
141
|
|
|
116
|
-
|
|
142
|
+
_logger.info(f"Rendering {n_frames} frames at {fps} fps")
|
|
117
143
|
with tempfile.TemporaryDirectory(prefix="animator_frames_") as frames_dir:
|
|
118
|
-
|
|
144
|
+
_logger.info(f"Using temporary directory: {frames_dir}")
|
|
145
|
+
|
|
146
|
+
# Avoid mutable default arguments by creating defaults here
|
|
147
|
+
if savefig_params is None:
|
|
148
|
+
savefig_params = {}
|
|
149
|
+
if video_params is None:
|
|
150
|
+
video_params = {"crf": "23", "preset": "slow"}
|
|
119
151
|
|
|
120
152
|
if num_workers == 1:
|
|
121
|
-
|
|
153
|
+
_logger.info("Running in serial mode")
|
|
122
154
|
self._render_serial(
|
|
123
|
-
|
|
155
|
+
param_by_frame,
|
|
156
|
+
n_frames,
|
|
124
157
|
frames_dir,
|
|
125
158
|
disable_progress_bar,
|
|
126
159
|
plotting_log_interval,
|
|
@@ -128,30 +161,33 @@ class Animator(ABC):
|
|
|
128
161
|
reuse_figure_object,
|
|
129
162
|
)
|
|
130
163
|
else:
|
|
131
|
-
|
|
164
|
+
_logger.info(f"Running in parallel mode with {num_workers} workers")
|
|
132
165
|
self._render_parallel(
|
|
133
|
-
|
|
166
|
+
param_by_frame,
|
|
167
|
+
n_frames,
|
|
134
168
|
frames_dir,
|
|
135
169
|
num_workers,
|
|
136
170
|
disable_progress_bar,
|
|
137
171
|
plotting_log_interval,
|
|
138
172
|
savefig_params,
|
|
139
173
|
reuse_figure_object,
|
|
174
|
+
preload_factor,
|
|
140
175
|
)
|
|
141
176
|
|
|
142
|
-
|
|
177
|
+
_logger.info("Creating video with PyAV")
|
|
143
178
|
_merge_frames_into_video(
|
|
144
179
|
frames_dir,
|
|
145
180
|
output_file,
|
|
146
181
|
fps,
|
|
147
182
|
video_codec,
|
|
183
|
+
video_pixfmt,
|
|
148
184
|
video_params,
|
|
149
185
|
disable_progress_bar,
|
|
150
|
-
|
|
186
|
+
_logger,
|
|
151
187
|
log_interval=saving_log_interval,
|
|
152
188
|
)
|
|
153
189
|
|
|
154
|
-
|
|
190
|
+
_logger.info(f"Animation complete: {output_file}")
|
|
155
191
|
|
|
156
192
|
def _setup_and_check(self) -> plt.Figure:
|
|
157
193
|
"""Call setup() and validate its return type."""
|
|
@@ -164,7 +200,8 @@ class Animator(ABC):
|
|
|
164
200
|
|
|
165
201
|
def _render_serial(
|
|
166
202
|
self,
|
|
167
|
-
|
|
203
|
+
param_by_frame: Iterable[Any],
|
|
204
|
+
n_frames: int,
|
|
168
205
|
frames_dir: Path | str,
|
|
169
206
|
disable_progress_bar: bool | None,
|
|
170
207
|
log_interval: int | None,
|
|
@@ -172,53 +209,61 @@ class Animator(ABC):
|
|
|
172
209
|
reuse_figure_object: bool,
|
|
173
210
|
) -> None:
|
|
174
211
|
"""Render frames serially."""
|
|
175
|
-
|
|
212
|
+
_logger.info("Serial rendering")
|
|
176
213
|
|
|
214
|
+
fig = None
|
|
177
215
|
if reuse_figure_object:
|
|
178
|
-
# Setup once and get figure
|
|
179
216
|
fig = self._setup_and_check()
|
|
180
217
|
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
for frame_idx in
|
|
184
|
-
|
|
218
|
+
seen_frame_ids: set[int] = set()
|
|
219
|
+
frames_processed = 0
|
|
220
|
+
for frame_idx, params in tqdm(
|
|
221
|
+
enumerate(itertools.islice(param_by_frame, n_frames)),
|
|
222
|
+
total=n_frames,
|
|
223
|
+
disable=disable_progress_bar,
|
|
185
224
|
):
|
|
225
|
+
if isinstance(params, IndexedFrameParams):
|
|
226
|
+
if params.frame_id in seen_frame_ids:
|
|
227
|
+
raise ValueError(
|
|
228
|
+
f"Duplicate frame_id {params.frame_id} in param_by_frame"
|
|
229
|
+
)
|
|
230
|
+
seen_frame_ids.add(params.frame_id)
|
|
231
|
+
frame_idx = params.frame_id
|
|
232
|
+
params = params.params
|
|
233
|
+
|
|
186
234
|
if not reuse_figure_object:
|
|
187
235
|
fig = self._setup_and_check()
|
|
188
236
|
|
|
189
|
-
params = params_list[frame_idx]
|
|
190
237
|
self.update(frame_idx, params)
|
|
191
|
-
fig.canvas.draw()
|
|
192
|
-
|
|
193
238
|
frame_path = Path(frames_dir) / f"frame_{frame_idx:09d}.png"
|
|
194
239
|
fig.savefig(frame_path, **savefig_params)
|
|
195
240
|
if not reuse_figure_object:
|
|
196
241
|
plt.close(fig)
|
|
197
242
|
|
|
198
|
-
|
|
199
|
-
if log_interval and
|
|
200
|
-
|
|
243
|
+
frames_processed += 1
|
|
244
|
+
if log_interval and frames_processed % log_interval == 0:
|
|
245
|
+
_logger.info(f"Frame {frames_processed}/{n_frames} rendered")
|
|
201
246
|
|
|
202
|
-
|
|
247
|
+
if reuse_figure_object and fig is not None:
|
|
248
|
+
plt.close(fig)
|
|
203
249
|
|
|
204
250
|
def _render_parallel(
|
|
205
251
|
self,
|
|
206
|
-
|
|
252
|
+
param_by_frame: Iterable[Any],
|
|
253
|
+
n_frames: int,
|
|
207
254
|
frames_dir: Path | str,
|
|
208
255
|
num_workers: int,
|
|
209
256
|
disable_progress_bar: bool | None,
|
|
210
257
|
log_interval: int | None,
|
|
211
258
|
savefig_params: dict[str, Any],
|
|
212
259
|
reuse_figure_object: bool,
|
|
260
|
+
preload_factor: int,
|
|
213
261
|
) -> None:
|
|
214
262
|
"""Render frames in parallel using dynamic work distribution."""
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
self.logger.info(f"Using dynamic work distribution with {num_workers} workers")
|
|
263
|
+
_logger.info(f"Using dynamic work distribution with {num_workers} workers")
|
|
218
264
|
|
|
219
265
|
# Create queues for task distribution and atomic counter for progress
|
|
220
|
-
task_queue = mp.Queue()
|
|
221
|
-
num_frames_completed = mp.Value("i", 0) # atomic integer counter
|
|
266
|
+
task_queue = mp.Queue(maxsize=num_workers * preload_factor)
|
|
222
267
|
|
|
223
268
|
# Start worker processes
|
|
224
269
|
workers = []
|
|
@@ -229,7 +274,6 @@ class Animator(ABC):
|
|
|
229
274
|
self,
|
|
230
275
|
worker_id,
|
|
231
276
|
task_queue,
|
|
232
|
-
num_frames_completed,
|
|
233
277
|
frames_dir,
|
|
234
278
|
log_interval,
|
|
235
279
|
savefig_params,
|
|
@@ -239,36 +283,52 @@ class Animator(ABC):
|
|
|
239
283
|
p.start()
|
|
240
284
|
workers.append(p)
|
|
241
285
|
|
|
242
|
-
# Populate task queue with individual frames (batch_size = 1)
|
|
243
|
-
|
|
244
|
-
|
|
286
|
+
# Populate task queue with individual frames (batch_size = 1).
|
|
287
|
+
# The bounded queue naturally throttles this loop so it only runs ahead of
|
|
288
|
+
# workers by ~preload_factor frames. Note: tqdm here tracks enqueue speed,
|
|
289
|
+
# which approximates render progress but hits 100% while workers finish the
|
|
290
|
+
# last queued frames.
|
|
291
|
+
seen_frame_ids: set[int] = set()
|
|
292
|
+
for frame_idx, params in tqdm(
|
|
293
|
+
enumerate(itertools.islice(param_by_frame, n_frames)),
|
|
294
|
+
total=n_frames,
|
|
295
|
+
disable=disable_progress_bar,
|
|
296
|
+
):
|
|
297
|
+
if isinstance(params, IndexedFrameParams):
|
|
298
|
+
if params.frame_id in seen_frame_ids:
|
|
299
|
+
raise ValueError(
|
|
300
|
+
f"Duplicate frame_id {params.frame_id} in param_by_frame"
|
|
301
|
+
)
|
|
302
|
+
seen_frame_ids.add(params.frame_id)
|
|
303
|
+
frame_idx = params.frame_id
|
|
304
|
+
params = params.params
|
|
305
|
+
|
|
306
|
+
task_queue.put((frame_idx, params))
|
|
245
307
|
|
|
246
308
|
# Send sentinel values to signal workers to exit
|
|
247
309
|
for _ in range(num_workers):
|
|
248
310
|
task_queue.put(None)
|
|
249
311
|
|
|
250
|
-
#
|
|
251
|
-
|
|
252
|
-
while num_frames_completed.value < num_frames:
|
|
253
|
-
current_progress = num_frames_completed.value
|
|
254
|
-
if current_progress > pbar.n:
|
|
255
|
-
pbar.update(current_progress - pbar.n)
|
|
256
|
-
time.sleep(0.1)
|
|
257
|
-
pbar.update(num_frames_completed.value - pbar.n)
|
|
258
|
-
pbar.close()
|
|
259
|
-
|
|
260
|
-
# Wait for all workers to finish
|
|
312
|
+
# Wait for all workers to finish and check for errors
|
|
313
|
+
failed = []
|
|
261
314
|
for p in workers:
|
|
262
315
|
p.join()
|
|
316
|
+
if p.exitcode != 0:
|
|
317
|
+
failed.append((p.pid, p.exitcode))
|
|
318
|
+
|
|
319
|
+
if failed:
|
|
320
|
+
details = ", ".join(
|
|
321
|
+
f"pid {pid} exited with code {code}" for pid, code in failed
|
|
322
|
+
)
|
|
323
|
+
raise RuntimeError(f"One or more worker processes failed: {details}")
|
|
263
324
|
|
|
264
|
-
|
|
325
|
+
_logger.info("All workers completed")
|
|
265
326
|
|
|
266
327
|
|
|
267
328
|
def _worker_process(
|
|
268
329
|
animator: Animator,
|
|
269
330
|
worker_id: int,
|
|
270
331
|
task_queue: mp.Queue,
|
|
271
|
-
progress_counter,
|
|
272
332
|
frames_dir: Path | str,
|
|
273
333
|
log_interval: int | None,
|
|
274
334
|
savefig_params: dict[str, Any],
|
|
@@ -281,43 +341,34 @@ def _worker_process(
|
|
|
281
341
|
1. Calls setup() once to initialize the figure (unless reuse_figure_object is False)
|
|
282
342
|
2. Repeatedly pulls individual frames from the task queue
|
|
283
343
|
3. Renders each frame
|
|
284
|
-
4. Atomically increments the progress counter
|
|
285
344
|
"""
|
|
286
|
-
|
|
345
|
+
fig = None
|
|
287
346
|
if reuse_figure_object:
|
|
288
347
|
fig = animator._setup_and_check()
|
|
289
348
|
|
|
290
|
-
# Process frames until we get a sentinel value (None)
|
|
291
349
|
frames_processed = 0
|
|
292
350
|
while True:
|
|
293
351
|
task = task_queue.get()
|
|
294
352
|
if task is None:
|
|
295
|
-
break
|
|
353
|
+
break
|
|
296
354
|
frame_idx, params = task
|
|
297
355
|
|
|
298
|
-
# Render the frame
|
|
299
356
|
if not reuse_figure_object:
|
|
300
357
|
fig = animator._setup_and_check()
|
|
301
358
|
animator.update(frame_idx, params)
|
|
302
|
-
fig.canvas.draw()
|
|
303
359
|
frame_path = Path(frames_dir) / f"frame_{frame_idx:09d}.png"
|
|
304
360
|
fig.savefig(frame_path, **savefig_params)
|
|
305
361
|
frames_processed += 1
|
|
306
362
|
if not reuse_figure_object:
|
|
307
363
|
plt.close(fig)
|
|
308
364
|
|
|
309
|
-
# Logging
|
|
310
365
|
if log_interval and frames_processed % log_interval == 0:
|
|
311
|
-
|
|
312
|
-
f"Worker {worker_id}: processed {frames_processed} frames"
|
|
313
|
-
)
|
|
366
|
+
_logger.info(f"Worker {worker_id}: processed {frames_processed} frames")
|
|
314
367
|
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
progress_counter.value += 1
|
|
368
|
+
if fig is not None:
|
|
369
|
+
plt.close(fig)
|
|
318
370
|
|
|
319
|
-
|
|
320
|
-
animator.logger.info(f"Worker {worker_id}: completed {frames_processed} frames")
|
|
371
|
+
_logger.info(f"Worker {worker_id}: completed {frames_processed} frames")
|
|
321
372
|
|
|
322
373
|
|
|
323
374
|
def _merge_frames_into_video(
|
|
@@ -325,6 +376,7 @@ def _merge_frames_into_video(
|
|
|
325
376
|
output_file: Path | str,
|
|
326
377
|
fps: int,
|
|
327
378
|
video_codec: str,
|
|
379
|
+
video_pixfmt: str,
|
|
328
380
|
video_params: dict[str, Any],
|
|
329
381
|
disable_progress_bar: bool | None,
|
|
330
382
|
logger: logging.Logger,
|
|
@@ -350,8 +402,8 @@ def _merge_frames_into_video(
|
|
|
350
402
|
stream = container.add_stream(video_codec, rate=fps)
|
|
351
403
|
stream.width = width
|
|
352
404
|
stream.height = height
|
|
353
|
-
|
|
354
|
-
|
|
405
|
+
stream.pix_fmt = video_pixfmt
|
|
406
|
+
stream.options.update(video_params)
|
|
355
407
|
|
|
356
408
|
# Encode each frame
|
|
357
409
|
for i, frame_path in tqdm(
|
|
@@ -360,12 +412,14 @@ def _merge_frames_into_video(
|
|
|
360
412
|
desc="Merging frames",
|
|
361
413
|
disable=disable_progress_bar,
|
|
362
414
|
):
|
|
363
|
-
|
|
364
|
-
|
|
365
|
-
|
|
366
|
-
img
|
|
367
|
-
|
|
368
|
-
|
|
415
|
+
with Image.open(frame_path) as img:
|
|
416
|
+
img = img.convert("RGB")
|
|
417
|
+
# Ensure image has the right size
|
|
418
|
+
if img.size != (width, height):
|
|
419
|
+
img = img.resize((width, height))
|
|
420
|
+
|
|
421
|
+
video_frame = av.VideoFrame.from_image(img)
|
|
422
|
+
video_frame = video_frame.reformat(width, height, format=video_pixfmt)
|
|
369
423
|
for packet in stream.encode(video_frame):
|
|
370
424
|
container.mux(packet)
|
|
371
425
|
|
|
@@ -384,4 +438,4 @@ def _merge_frames_into_video(
|
|
|
384
438
|
|
|
385
439
|
except Exception as e:
|
|
386
440
|
logger.critical(f"PyAV failed to create video: {e}")
|
|
387
|
-
raise
|
|
441
|
+
raise
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import matplotlib.pyplot as plt
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
|
|
5
|
+
from parallel_animate import Animator, IndexedFrameParams
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class VideoFrameAnimation(Animator):
|
|
9
|
+
def setup(self):
|
|
10
|
+
fig, ax = plt.subplots(figsize=(6, 6))
|
|
11
|
+
self.imshow_artist = ax.imshow(np.zeros((128, 128, 3), dtype=np.uint8))
|
|
12
|
+
self.title_artist = ax.set_title("Video Frame X")
|
|
13
|
+
ax.axis("off")
|
|
14
|
+
return fig
|
|
15
|
+
|
|
16
|
+
def update(self, frame_idx, params):
|
|
17
|
+
self.imshow_artist.set_data(params["frame"])
|
|
18
|
+
self.title_artist.set_text(f"Video Frame {frame_idx}")
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def fake_video_loader(n_frames=64, frame_size=(128, 128)):
|
|
22
|
+
"""Emulates a video loader that yields in nondeterministic order (e.g. because
|
|
23
|
+
loading is parallelized)."""
|
|
24
|
+
frame_ids = np.arange(n_frames)
|
|
25
|
+
np.random.shuffle(frame_ids)
|
|
26
|
+
for frame_id in frame_ids:
|
|
27
|
+
frame = np.zeros((*frame_size, 3), dtype=np.uint8)
|
|
28
|
+
x = int((frame_id % frame_size[1]) * frame_size[1] / n_frames)
|
|
29
|
+
frame[:, x : x + 20] = (0, 255, 0) # moving green bar
|
|
30
|
+
yield IndexedFrameParams(frame_id=frame_id, params={"frame": frame})
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
if __name__ == "__main__":
|
|
34
|
+
# Create a "parallel video loader" that yields inputs out of order
|
|
35
|
+
frame_loader = fake_video_loader(n_frames=64)
|
|
36
|
+
|
|
37
|
+
# Create animation
|
|
38
|
+
anim = VideoFrameAnimation()
|
|
39
|
+
output_path = Path("example_output/nondeterministic_video_loader.mp4")
|
|
40
|
+
anim.make_video(
|
|
41
|
+
output_file=output_path,
|
|
42
|
+
param_by_frame=frame_loader,
|
|
43
|
+
n_frames=64,
|
|
44
|
+
fps=30,
|
|
45
|
+
num_workers=4,
|
|
46
|
+
)
|
parallel_animate/util.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
import matplotlib
|
|
2
|
+
import numpy as np
|
|
3
|
+
import logging
|
|
4
|
+
from matplotlib import pyplot as plt
|
|
5
|
+
from fractions import Fraction
|
|
6
|
+
|
|
7
|
+
_logger = logging.getLogger(__name__)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def configure_matplotlib_style():
|
|
11
|
+
"""Use sans serif font and export text as texts (not shapes) in PDFs."""
|
|
12
|
+
matplotlib.style.use("fast")
|
|
13
|
+
plt.rcParams["font.family"] = "Arial"
|
|
14
|
+
plt.rcParams["pdf.fonttype"] = 42
|
|
15
|
+
_logger.info("Configured matplotlib style.")
|
|
16
|
+
# suppress matplotlib font manager warnings
|
|
17
|
+
logging.getLogger("matplotlib.font_manager").setLevel(logging.ERROR)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def get_rendered_frame_ids(
|
|
21
|
+
data_fps: Fraction | int,
|
|
22
|
+
play_speed: float,
|
|
23
|
+
rendered_fps: Fraction | int,
|
|
24
|
+
n_data_frames: int,
|
|
25
|
+
) -> np.ndarray:
|
|
26
|
+
"""Get list of indices of input frames that should be rendered based on fps specs.
|
|
27
|
+
|
|
28
|
+
Example: if data is recorded at 330 FPS, and we want to play it back at 0.1x speed
|
|
29
|
+
at 30 FPS, then we need to render every `stride` frames in the original data, where
|
|
30
|
+
stride = data_fps / (rendered_fps / play_speed) = 1.1 frames.
|
|
31
|
+
|
|
32
|
+
Parameters
|
|
33
|
+
----------
|
|
34
|
+
data_fps : Fraction | int
|
|
35
|
+
The frame rate of the original data.
|
|
36
|
+
play_speed : float
|
|
37
|
+
The desired playback speed (e.g., 0.1 for 10% speed).
|
|
38
|
+
rendered_fps : Fraction | int
|
|
39
|
+
The frame rate at which the video will be rendered.
|
|
40
|
+
n_data_frames : int
|
|
41
|
+
The total number of data frames.
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
Returns
|
|
45
|
+
-------
|
|
46
|
+
np.ndarray of int
|
|
47
|
+
The indices of data frames that should be rendered.
|
|
48
|
+
"""
|
|
49
|
+
if n_data_frames <= 0:
|
|
50
|
+
return np.array([], dtype=int)
|
|
51
|
+
|
|
52
|
+
stride = Fraction(data_fps) / (Fraction(rendered_fps) / play_speed)
|
|
53
|
+
if stride < 1:
|
|
54
|
+
_logger.warning(
|
|
55
|
+
f"Calculated stride {stride} < 1. This will lead to repeated frames."
|
|
56
|
+
)
|
|
57
|
+
|
|
58
|
+
n_rendered_frames = max(1, int(n_data_frames / stride))
|
|
59
|
+
# Use floor so we never map to a future data frame index, then clip to valid range.
|
|
60
|
+
target_data_frame_ids = np.floor(
|
|
61
|
+
np.arange(n_rendered_frames) * float(stride)
|
|
62
|
+
).astype(int)
|
|
63
|
+
target_data_frame_ids = np.clip(target_data_frame_ids, 0, n_data_frames - 1)
|
|
64
|
+
return target_data_frame_ids
|
|
@@ -1,22 +1,16 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: parallel-matplotlib-animation
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Summary: Animate matplotlib figures into videos, in parallel but with efficient caching
|
|
5
|
+
Project-URL: Repository, https://github.com/sibocw/parallel-matplotlib-animation
|
|
6
|
+
Author-email: Sibo Wang-Chen <sibo.wang@epfl.ch>
|
|
5
7
|
License-File: LICENSE
|
|
6
|
-
Author: Sibo Wang-Chen
|
|
7
|
-
Author-email: sibo.wang@epfl.ch
|
|
8
8
|
Requires-Python: >=3.10
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
Classifier: Programming Language :: Python :: 3.14
|
|
15
|
-
Requires-Dist: Pillow (>=11.0)
|
|
16
|
-
Requires-Dist: av (>=14.0)
|
|
17
|
-
Requires-Dist: matplotlib (>=3.10.0)
|
|
18
|
-
Requires-Dist: numpy (>=2,<3)
|
|
19
|
-
Requires-Dist: tqdm (>=4.67)
|
|
9
|
+
Requires-Dist: av>=14.0
|
|
10
|
+
Requires-Dist: matplotlib>=3.10.0
|
|
11
|
+
Requires-Dist: numpy<3,>=2
|
|
12
|
+
Requires-Dist: pillow>=11.0
|
|
13
|
+
Requires-Dist: tqdm>=4.67
|
|
20
14
|
Description-Content-Type: text/markdown
|
|
21
15
|
|
|
22
16
|
# parallel-matplotlib-animation
|
|
@@ -84,7 +78,7 @@ anim = WaveAnimation()
|
|
|
84
78
|
anim.make_video("wave.mp4", param_by_frame=params, fps=30, num_workers=4)
|
|
85
79
|
```
|
|
86
80
|
|
|
87
|
-
|
|
81
|
+
<img src="assets/simple_wave_animation.gif" width="480"/>
|
|
88
82
|
|
|
89
83
|
## Usage
|
|
90
84
|
This library has a single class: `parallel_animate.Animator`. To make an animation, you must create your own class inheriting from it and define the following methods:
|
|
@@ -96,18 +90,34 @@ This library has a single class: `parallel_animate.Animator`. To make an animati
|
|
|
96
90
|
Once you have defined your animator class, there is a single method that you need to call that makes the video: **`.make_video(...)`**. It accepts the following arguments:
|
|
97
91
|
|
|
98
92
|
- `output_file` (Path or str): Output video path
|
|
99
|
-
- `param_by_frame` (
|
|
93
|
+
- `param_by_frame` (Iterable): Iterable of parameters. Each element is the `params` argument to be given to the `.update` call for the corresponding frame. Can be a list, tuple, generator, or any other iterable. Using generators is particularly useful for large data (e.g., bitmaps) to avoid loading everything into memory at once.
|
|
100
94
|
- `fps` (int): Frame rate of the output video
|
|
95
|
+
- `n_frames` (int or None): Number of frames to render. If None, use the length of `param_by_frame`. If param_by_frame does not have `__len__` implemented and `n_frames` is None, the progress bar won't show completion percentage.
|
|
101
96
|
- `num_workers` (int): Number of worker processes to be spawned. If -1, use all CPU cores. If -2, use all but one CPU cores, etc. If 1, no child process is created and the video is made in the main process itself. Default is -1.
|
|
102
97
|
- See the docstring for `parallel_animate.animator` directly for less commonly used, optional parameters. These control logging, rendering quality, etc.
|
|
103
98
|
|
|
99
|
+
### Special case: frame params arriving out-of-order in `param_by_frame`
|
|
100
|
+
In some cases, frames in `param_by_frame` might be out of order. We can handle these scenarios by populating `param_by_frame` with a special `parallel_animate.IndexedFrameParams` dataclass, which specifies the frame index that overrides the ordering in `param_by_frame`. This can be useful when, for example, the animator needs to draw frames that are decoded from a video, and the dataloader for that video might return frames in nondeterministic order because it's parallelized.
|
|
101
|
+
|
|
102
|
+
See `src/parallel_animate/examples/nondeterministic_video_loader.py` for details.
|
|
104
103
|
|
|
105
104
|
## Examples
|
|
106
105
|
|
|
107
106
|
See [`src/parallel_animate/examples/`](https://github.com/sibocw/parallel-matplotlib-animation/blob/main/src/parallel_animate/examples/):
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
107
|
+
|
|
108
|
+
`simple_wave_animation.py`: The example above
|
|
109
|
+
|
|
110
|
+
`multi_panel_animation.py`: 5 subplots with different plot types
|
|
111
|
+
|
|
112
|
+
<img src="assets/multi_panel_animation.gif" width="480"/>
|
|
113
|
+
|
|
114
|
+
`very_complex_animation.py`: 14 subplots with GridSpec layout
|
|
115
|
+
|
|
116
|
+
<img src="assets/very_complex_animation.gif" width="480"/>
|
|
117
|
+
|
|
118
|
+
`nondeterministic_video_loader.py`: handling frames that arrive out of order
|
|
119
|
+
|
|
120
|
+
<img src="assets/nondeterministic_video_loader.gif" width="240"/>
|
|
111
121
|
|
|
112
122
|
|
|
113
123
|
## Performance test
|
|
@@ -122,4 +132,4 @@ The left-most blue dot indicates serial processing with resources reuse. The bla
|
|
|
122
132
|
## Unit tests
|
|
123
133
|
```bash
|
|
124
134
|
python -m unittest discover -s tests
|
|
125
|
-
```
|
|
135
|
+
```
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
parallel_animate/__init__.py,sha256=cFUvpCvfACdcmJSSwxjlASCSeaoJQ-5g6BCojZ8xzHs,51
|
|
2
|
+
parallel_animate/animator.py,sha256=MUQJMRIYGwemO4_hlPNwejqOryMB7tdNaXg460matHc,16120
|
|
3
|
+
parallel_animate/util.py,sha256=Ym0OyySuU23AE9AWxjBe9pZHizkuBYhlo7wCn7AdzB0,2157
|
|
4
|
+
parallel_animate/examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
5
|
+
parallel_animate/examples/multi_panel_animation.py,sha256=stUoAQn0h4Hgg0BBJQtRmXuW9sj5j-FQBjUmqYA-5mI,5127
|
|
6
|
+
parallel_animate/examples/nondeterministic_video_loader.py,sha256=3lDo65327dOERI1pFBICM9b8V24LlY30rStWUSPKvcY,1566
|
|
7
|
+
parallel_animate/examples/scaling_test.py,sha256=aHvKjlzsUPGAqCWdmysllesdhJccHlcMa0cap9Z836Q,3396
|
|
8
|
+
parallel_animate/examples/simple_wave_animation.py,sha256=O_SyXltGJ5iovquwGnu4KicTucUgBZY78oDzTW-31Gc,1043
|
|
9
|
+
parallel_animate/examples/very_complex_animation.py,sha256=X37-vDlbxUGVApvRJuO9nhGpzAvf3WjCOAEKbDolXeA,24101
|
|
10
|
+
parallel_matplotlib_animation-0.1.2.dist-info/METADATA,sha256=VpLv0bLihzTZ4pL4mWFal7RDxD1tNdcFbJKfWLKYhs8,7287
|
|
11
|
+
parallel_matplotlib_animation-0.1.2.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
|
|
12
|
+
parallel_matplotlib_animation-0.1.2.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
|
|
13
|
+
parallel_matplotlib_animation-0.1.2.dist-info/RECORD,,
|
|
@@ -1,11 +0,0 @@
|
|
|
1
|
-
parallel_animate/__init__.py,sha256=LhvDAwx0dS6zhxit2hZ7C47l7S2t4hj-p9262RxTjtY,31
|
|
2
|
-
parallel_animate/animator.py,sha256=EqzV70jqk_d1wIsrKa90FnMDWG32PcD4GiMRxQnA8Mw,13585
|
|
3
|
-
parallel_animate/examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
4
|
-
parallel_animate/examples/multi_panel_animation.py,sha256=stUoAQn0h4Hgg0BBJQtRmXuW9sj5j-FQBjUmqYA-5mI,5127
|
|
5
|
-
parallel_animate/examples/scaling_test.py,sha256=aHvKjlzsUPGAqCWdmysllesdhJccHlcMa0cap9Z836Q,3396
|
|
6
|
-
parallel_animate/examples/simple_wave_animation.py,sha256=O_SyXltGJ5iovquwGnu4KicTucUgBZY78oDzTW-31Gc,1043
|
|
7
|
-
parallel_animate/examples/very_complex_animation.py,sha256=X37-vDlbxUGVApvRJuO9nhGpzAvf3WjCOAEKbDolXeA,24101
|
|
8
|
-
parallel_matplotlib_animation-0.1.0.dist-info/METADATA,sha256=yaUGS13TYDYFwDuWdnj4YtFZT9z5SqngfsQxhLafhPg,6211
|
|
9
|
-
parallel_matplotlib_animation-0.1.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
|
|
10
|
-
parallel_matplotlib_animation-0.1.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
|
|
11
|
-
parallel_matplotlib_animation-0.1.0.dist-info/RECORD,,
|
|
File without changes
|