Move flatten_rgba to image.py and implment flatten_rgba for jpeg files.

cbreeden · cbreeden · commit 1759556f2d36 · 2015-11-06T10:41:26.000-06:00
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -39,6 +39,7 @@
 from matplotlib.path import Path
 from matplotlib.transforms import Bbox, BboxBase
 from matplotlib import colors as mcolors
+from matplotlib import image as mimage
 
 from matplotlib.backends._backend_agg import RendererAgg as _RendererAgg
 from matplotlib import _png
@@ -514,7 +515,7 @@ def print_raw(self, filename_or_obj, *args, **kwargs):
         if kwargs.get('flatten', False):
             w, h = int(renderer.width), int(renderer.height)
             img = np.array(memoryview(img)).reshape((w, h, 4))
-            img = self.flatten_rgba(img)
+            img = mimage.flatten_rgba(img)
 
         try:
             fileobj.write(img)
@@ -524,34 +525,6 @@ def print_raw(self, filename_or_obj, *args, **kwargs):
         renderer.dpi = original_dpi
     print_rgba = print_raw
 
-    def flatten_rgba(self, src, bg=None):
-        """
-        Flatten an RGBA image `src` with a background color `bg`.
-        The resulting image will have an alpha channel, but no
-        transparency.  This is useful when interfacing with
-        file formats that don't support transparency or only support
-        boolean transparency.
-        """
-        if bg is None:
-            bg = mcolors.colorConverter.to_rgb(
-                     rcParams.get('savefig.facecolor', 'white'))
-            bg = tuple([int(x * 255.0) for x in bg]) 
-       
-        # Numpy images have dtype=uint8 which will overflow 
-        src = src.astype(dtype = np.uint16)
-
-        alpha = src[:,:,3]
-        src_rgb = src[:,:,:3]
-        w, h, _ = src.shape
-
-        result = np.empty((w, h, 4))
-        result[:,:,0] = (255 - alpha)*bg[0] + alpha*src_rgb[:,:,0]
-        result[:,:,1] = (255 - alpha)*bg[1] + alpha*src_rgb[:,:,1]
-        result[:,:,2] = (255 - alpha)*bg[2] + alpha*src_rgb[:,:,2]
-        result = (result/255).astype(np.uint8)
-        result[:,:,3] = 255
-        return result
-        
     def print_png(self, filename_or_obj, *args, **kwargs):
         FigureCanvasAgg.draw(self)
         renderer = self.get_renderer()
@@ -562,15 +535,18 @@ def print_png(self, filename_or_obj, *args, **kwargs):
             close = True
         else:
             close = False
+
+        # Flatten RGBA if used as intermediate fileformat for something
+        # that doesn't support transparency (ie: Animations)
+        img = renderer._renderer
+        if kwargs.get('flatten', False):
+            img = img.buffer_rgba()
+            w, h = int(renderer.width), int(renderer.height)
+            img = renderer._renderer.buffer_rgba()
+            img = np.array(memoryview(img)).reshape((h, w, 4))
+            img = mimage.flatten_rgba(img)
         try:
-            if kwargs.get('flatten', False):
-                w, h = int(renderer.width), int(renderer.height)
-                img = renderer._renderer.buffer_rgba()
-                img = np.array(memoryview(img)).reshape((h,w,4))
-                img = self.flatten_rgba(img)
-                _png.write_png(img, filename_or_obj, self.figure.dpi)
-            else:
-                _png.write_png(renderer._renderer, filename_or_obj, self.figure.dpi)
+            _png.write_png(img, filename_or_obj, self.figure.dpi)
         finally:
             if close:
                 filename_or_obj.close()
@@ -605,24 +581,22 @@ def print_jpg(self, filename_or_obj, *args, **kwargs):
             *progressive*: If present, indicates that this image
                 should be stored as a progressive JPEG file.
             """
-            buf, size = self.print_to_buffer()
+            buf, (w, h) = self.print_to_buffer()
+            buf = np.array(memoryview(buf)).reshape((w, h, 4))
+
             if kwargs.pop("dryrun", False):
                 return
             # The image is "pasted" onto a white background image to safely
             # handle any transparency
-            image = Image.frombuffer('RGBA', size, buf, 'raw', 'RGBA', 0, 1)
-            color = mcolors.colorConverter.to_rgb(
-                rcParams.get('savefig.facecolor', 'white'))
-            color = tuple([int(x * 255.0) for x in color])
-            background = Image.new('RGB', size, color)
-            background.paste(image, image)
+            buf = mimage.flatten_rgba(buf)
+            img = Image.frombuffer('RGBA', (w, h), buf, 'raw', 'RGBA', 0, 1)
+
             options = restrict_dict(kwargs, ['quality', 'optimize',
                                              'progressive'])
-
             if 'quality' not in options:
                 options['quality'] = rcParams['savefig.jpeg_quality']
 
-            return background.save(filename_or_obj, format='jpeg', **options)
+            return img.save(filename_or_obj, format='jpeg', **options)
         print_jpeg = print_jpg
 
         # add TIFF support
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1502,3 +1502,45 @@ def thumbnail(infile, thumbfile, scale=0.1, interpolation='bilinear',
     ax.imshow(im, aspect='auto', resample=True, interpolation=interpolation)
     fig.savefig(thumbfile, dpi=dpi)
     return fig
+
+def flatten_rgba(src, bg=None):
+    """
+    Flatten an RGBA image *src* with a background color *bg*.
+    The resulting image will have an alpha channel, but no transparency.
+    This can be useful when interfacing with file formats that don't support
+    transparency or only support boolean transparency.
+
+    Parameters
+    ----------
+    src : MxNx4 Numpy array, dtype=uint8
+        Image source in RGBA to be flattened.
+
+    bg  : Tuple(int,int,int), optional
+        Background color to merge *src* with.  If no bg color is provided
+        the color from the rcParam 'savefig.facecolor' will be use.
+
+    Returns
+    -------
+    dest : MxNx4 Numpy array, dtype=uint8
+    """
+
+    if bg is None:
+        bg = mcolors.colorConverter.to_rgb(
+                 rcParams.get('savefig.facecolor', 'white'))
+        bg = tuple([int(x * 255.0) for x in bg])
+
+    # Numpy images have dtype=uint8 which will overflow
+    src = src.astype(np.uint16)
+
+    alpha = src[:,:,3]
+    src_rgb = src[:,:,:3]
+    w, h, _ = src.shape
+
+    dest = np.empty((w, h, 4))
+    dest[:,:,0] = (255 - alpha)*bg[0] + alpha*src_rgb[:,:,0]
+    dest[:,:,1] = (255 - alpha)*bg[1] + alpha*src_rgb[:,:,1]
+    dest[:,:,2] = (255 - alpha)*bg[2] + alpha*src_rgb[:,:,2]
+    dest = (dest/255).astype(np.uint8)
+    dest[:,:,3] = 255
+
+    return dest
