Skip to content

Commit

Permalink
Factor common parts of saving to different formats using pillow.
Browse files Browse the repository at this point in the history
We don't need to drop the alpha channel for jpeg anymore as imsave
handles that.

In imsave, also clarify that `sm` is not going to be used in the
already-usable-rgba case.

Also, spliting out pil_kwargs under "other parameters" seems overkill
for an effectively internal method (`savefig` is the user-facing API).
  • Loading branch information
anntzer committed Oct 18, 2021
1 parent 481701b commit 479084f
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 45 deletions.
59 changes: 18 additions & 41 deletions lib/matplotlib/backends/backend_agg.py
Original file line number Diff line number Diff line change
Expand Up @@ -487,6 +487,16 @@ def print_raw(self, filename_or_obj, *args):

print_rgba = print_raw

def _print_pil(self, filename_or_obj, fmt, pil_kwargs, metadata=None):
"""
Draw the canvas, then save it using `.image.imsave` (to which
*pil_kwargs* and *metadata* are forwarded).
"""
FigureCanvasAgg.draw(self)
mpl.image.imsave(
filename_or_obj, self.buffer_rgba(), format=fmt, origin="upper",
dpi=self.figure.dpi, metadata=metadata, pil_kwargs=pil_kwargs)

@_check_savefig_extra_args
@_api.delete_parameter("3.5", "args")
def print_png(self, filename_or_obj, *args,
Expand Down Expand Up @@ -537,10 +547,7 @@ def print_png(self, filename_or_obj, *args,
If the 'pnginfo' key is present, it completely overrides
*metadata*, including the default 'Software' key.
"""
FigureCanvasAgg.draw(self)
mpl.image.imsave(
filename_or_obj, self.buffer_rgba(), format="png", origin="upper",
dpi=self.figure.dpi, metadata=metadata, pil_kwargs=pil_kwargs)
self._print_pil(filename_or_obj, "png", pil_kwargs, metadata)

def print_to_buffer(self):
FigureCanvasAgg.draw(self)
Expand All @@ -555,68 +562,38 @@ def print_to_buffer(self):
@_check_savefig_extra_args()
@_api.delete_parameter("3.5", "args")
def print_jpg(self, filename_or_obj, *args, pil_kwargs=None, **kwargs):
"""
Write the figure to a JPEG file.
Parameters
----------
filename_or_obj : str or path-like or file-like
The file to write to.
Other Parameters
----------------
pil_kwargs : dict, optional
Additional keyword arguments that are passed to
`PIL.Image.Image.save` when saving the figure.
"""
# Remove transparency by alpha-blending on an assumed white background.
r, g, b, a = mcolors.to_rgba(self.figure.get_facecolor())
try:
self.figure.set_facecolor(a * np.array([r, g, b]) + 1 - a)
FigureCanvasAgg.draw(self)
self._print_pil(filename_or_obj, "jpeg", pil_kwargs)
finally:
self.figure.set_facecolor((r, g, b, a))
if pil_kwargs is None:
pil_kwargs = {}
pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
# Drop alpha channel now.
return (Image.fromarray(np.asarray(self.buffer_rgba())[..., :3])
.save(filename_or_obj, format='jpeg', **pil_kwargs))

print_jpeg = print_jpg

@_check_savefig_extra_args
def print_tif(self, filename_or_obj, *, pil_kwargs=None):
FigureCanvasAgg.draw(self)
if pil_kwargs is None:
pil_kwargs = {}
pil_kwargs.setdefault("dpi", (self.figure.dpi, self.figure.dpi))
return (Image.fromarray(np.asarray(self.buffer_rgba()))
.save(filename_or_obj, format='tiff', **pil_kwargs))
self._print_pil(filename_or_obj, "tiff", pil_kwargs)

print_tiff = print_tif

@_check_savefig_extra_args
def print_webp(self, filename_or_obj, *, pil_kwargs=None):
self._print_pil(filename_or_obj, "webp", pil_kwargs)

print_jpg.__doc__, print_tif.__doc__, print_webp.__doc__ = map(
"""
Write the figure to a WebP file.
Write the figure to a {} file.
Parameters
----------
filename_or_obj : str or path-like or file-like
The file to write to.
Other Parameters
----------------
pil_kwargs : dict, optional
Additional keyword arguments that are passed to
`PIL.Image.Image.save` when saving the figure.
"""
FigureCanvasAgg.draw(self)
if pil_kwargs is None:
pil_kwargs = {}
return (Image.fromarray(np.asarray(self.buffer_rgba()))
.save(filename_or_obj, format='webp', **pil_kwargs))
""".format, ["JPEG", "TIFF", "WebP"])


@_Backend.export
Expand Down
8 changes: 4 additions & 4 deletions lib/matplotlib/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -1601,20 +1601,20 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
else:
# Don't bother creating an image; this avoids rounding errors on the
# size when dividing and then multiplying by dpi.
sm = cm.ScalarMappable(cmap=cmap)
sm.set_clim(vmin, vmax)
if origin is None:
origin = mpl.rcParams["image.origin"]
if origin == "lower":
arr = arr[::-1]
if (isinstance(arr, memoryview) and arr.format == "B"
and arr.ndim == 3 and arr.shape[-1] == 4):
# Such an ``arr`` would also be handled fine by sm.to_rgba (after
# casting with asarray), but it is useful to special-case it
# Such an ``arr`` would also be handled fine by sm.to_rgba below
# (after casting with asarray), but it is useful to special-case it
# because that's what backend_agg passes, and can be in fact used
# as is, saving a few operations.
rgba = arr
else:
sm = cm.ScalarMappable(cmap=cmap)
sm.set_clim(vmin, vmax)
rgba = sm.to_rgba(arr, bytes=True)
if pil_kwargs is None:
pil_kwargs = {}
Expand Down

0 comments on commit 479084f

Please sign in to comment.