Skip to content

Commit 125394d

Browse files
committed
Correctly handle high dpi in Pillow animation writer.
The idea is just to pass self.dpi to savefig() and then use self.frame_size as the frame size rather than looking up renderer internals, but this is made easier by properly moving some attributes and methods (some of `__init__`, `setup`, `frame_size`) from MovieWriter (which should really be called SubprocessMovieWriter, as that's what it does) to AbstractMovieWriter, so that PillowWriter can reuse them without having to carefully prevent any attempt of starting a subprocess.
1 parent e23b2aa commit 125394d

File tree

1 file changed

+41
-55
lines changed

1 file changed

+41
-55
lines changed

lib/matplotlib/animation.py

Lines changed: 41 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,14 @@ class AbstractMovieWriter(abc.ABC):
188188
``writer`` argument of `Animation.save()`.
189189
'''
190190

191+
def __init__(self, fps=5, metadata=None, codec=None, bitrate=None):
192+
self.fps = fps
193+
self.metadata = metadata if metadata is not None else {}
194+
self.codec = (
195+
mpl.rcParams['animation.codec'] if codec is None else codec)
196+
self.bitrate = (
197+
mpl.rcParams['animation.bitrate'] if bitrate is None else bitrate)
198+
191199
@abc.abstractmethod
192200
def setup(self, fig, outfile, dpi=None):
193201
'''
@@ -203,6 +211,17 @@ def setup(self, fig, outfile, dpi=None):
203211
The DPI (or resolution) for the file. This controls the size
204212
in pixels of the resulting movie file. Default is ``fig.dpi``.
205213
'''
214+
self.outfile = outfile
215+
self.fig = fig
216+
if dpi is None:
217+
dpi = self.fig.dpi
218+
self.dpi = dpi
219+
220+
@property
221+
def frame_size(self):
222+
'''A tuple ``(width, height)`` in pixels of a movie frame.'''
223+
w, h = self.fig.get_size_inches()
224+
return int(w * self.dpi), int(h * self.dpi)
206225

207226
@abc.abstractmethod
208227
def grab_frame(self, **savefig_kwargs):
@@ -275,7 +294,7 @@ def __init__(self, fps=5, codec=None, bitrate=None, extra_args=None,
275294
output file. Some keys that may be of use include:
276295
title, artist, genre, subject, copyright, srcform, comment.
277296
"""
278-
if self.__class__ is MovieWriter:
297+
if type(self) is MovieWriter:
279298
# TODO MovieWriter is still an abstract class and needs to be
280299
# extended with a mixin. This should be clearer in naming
281300
# and description. For now, just give a reasonable error
@@ -284,35 +303,15 @@ def __init__(self, fps=5, codec=None, bitrate=None, extra_args=None,
284303
'MovieWriter cannot be instantiated directly. Please use one '
285304
'of its subclasses.')
286305

287-
self.fps = fps
288-
self.frame_format = 'rgba'
306+
super().__init__(fps=fps, metadata=metadata)
289307

290-
if codec is None:
291-
self.codec = mpl.rcParams['animation.codec']
292-
else:
293-
self.codec = codec
294-
295-
if bitrate is None:
296-
self.bitrate = mpl.rcParams['animation.bitrate']
297-
else:
298-
self.bitrate = bitrate
308+
self.frame_format = 'rgba'
299309

300310
if extra_args is None:
301311
self.extra_args = list(mpl.rcParams[self.args_key])
302312
else:
303313
self.extra_args = extra_args
304314

305-
if metadata is None:
306-
self.metadata = dict()
307-
else:
308-
self.metadata = metadata
309-
310-
@property
311-
def frame_size(self):
312-
'''A tuple ``(width, height)`` in pixels of a movie frame.'''
313-
w, h = self.fig.get_size_inches()
314-
return int(w * self.dpi), int(h * self.dpi)
315-
316315
def _adjust_frame_size(self):
317316
if self.codec == 'h264':
318317
wo, ho = self.fig.get_size_inches()
@@ -340,13 +339,8 @@ def setup(self, fig, outfile, dpi=None):
340339
The DPI (or resolution) for the file. This controls the size
341340
in pixels of the resulting movie file. Default is fig.dpi.
342341
'''
343-
self.outfile = outfile
344-
self.fig = fig
345-
if dpi is None:
346-
dpi = self.fig.dpi
347-
self.dpi = dpi
342+
super().setup(fig, outfile, dpi=dpi)
348343
self._w, self._h = self._adjust_frame_size()
349-
350344
# Run here so that grab_frame() can write the data to a pipe. This
351345
# eliminates the need for temp files.
352346
self._run()
@@ -540,35 +534,27 @@ def cleanup(self):
540534

541535

542536
@writers.register('pillow')
543-
class PillowWriter(MovieWriter):
537+
class PillowWriter(AbstractMovieWriter):
544538
@classmethod
545539
def isAvailable(cls):
546540
return True
547541

548-
def __init__(self, *args, **kwargs):
549-
if kwargs.get("extra_args") is None:
550-
kwargs["extra_args"] = ()
551-
super().__init__(*args, **kwargs)
552-
553542
def setup(self, fig, outfile, dpi=None):
543+
super().setup(fig, outfile, dpi=dpi)
554544
self._frames = []
555-
self._outfile = outfile
556-
self._dpi = dpi
557-
self._fig = fig
558545

559546
def grab_frame(self, **savefig_kwargs):
560547
from PIL import Image
561548
buf = BytesIO()
562-
self._fig.savefig(buf, **dict(savefig_kwargs, format="rgba"))
563-
renderer = self._fig.canvas.get_renderer()
549+
self.fig.savefig(
550+
buf, **{**savefig_kwargs, "format": "rgba", "dpi": self.dpi})
551+
renderer = self.fig.canvas.get_renderer()
564552
self._frames.append(Image.frombuffer(
565-
"RGBA",
566-
(int(renderer.width), int(renderer.height)), buf.getbuffer(),
567-
"raw", "RGBA", 0, 1))
553+
"RGBA", self.frame_size, buf.getbuffer(), "raw", "RGBA", 0, 1))
568554

569555
def finish(self):
570556
self._frames[0].save(
571-
self._outfile, save_all=True, append_images=self._frames[1:],
557+
self.outfile, save_all=True, append_images=self._frames[1:],
572558
duration=int(1000 / self.fps), loop=0)
573559

574560

@@ -1075,11 +1061,15 @@ def func(current_frame: int, total_frames: int) -> Any
10751061
if dpi == 'figure':
10761062
dpi = self._fig.dpi
10771063

1078-
if codec is None:
1079-
codec = mpl.rcParams['animation.codec']
1080-
1081-
if bitrate is None:
1082-
bitrate = mpl.rcParams['animation.bitrate']
1064+
writer_kwargs = {}
1065+
if codec is not None:
1066+
writer_kwargs['codec'] = codec
1067+
if bitrate is not None:
1068+
writer_kwargs['bitrate'] = bitrate
1069+
if extra_args is not None:
1070+
writer_kwargs['extra_args'] = extra_args
1071+
if metadata is not None:
1072+
writer_kwargs['metadata'] = metadata
10831073

10841074
all_anim = [self]
10851075
if extra_anim is not None:
@@ -1091,9 +1081,7 @@ def func(current_frame: int, total_frames: int) -> Any
10911081
# registered class.
10921082
if isinstance(writer, str):
10931083
if writers.is_available(writer):
1094-
writer = writers[writer](fps, codec, bitrate,
1095-
extra_args=extra_args,
1096-
metadata=metadata)
1084+
writer = writers[writer](fps, **writer_kwargs)
10971085
else:
10981086
alt_writer = next(writers, None)
10991087
if alt_writer is None:
@@ -1102,9 +1090,7 @@ def func(current_frame: int, total_frames: int) -> Any
11021090
"save animations.")
11031091
_log.warning("MovieWriter %s unavailable; trying to use %s "
11041092
"instead.", writer, alt_writer)
1105-
writer = alt_writer(
1106-
fps, codec, bitrate,
1107-
extra_args=extra_args, metadata=metadata)
1093+
writer = alt_writer(fps, **writer_kwargs)
11081094
_log.info('Animation.save using %s', type(writer))
11091095

11101096
if 'bbox_inches' in savefig_kwargs:

0 commit comments

Comments
 (0)