Compress raster images in PDF output using libpng

jkseppan · jkseppan · commit 60161fa19d4f · 2015-07-08T18:23:40.000+03:00
The PDF format allows for PNG predictors in Flate-compressed
streams, so we can use libpng to encode a png file and extract
the row data from that.
diff --git a/lib/matplotlib/backends/backend_pdf.py b/lib/matplotlib/backends/backend_pdf.py
@@ -12,6 +12,7 @@
 import codecs
 import os
 import re
+import struct
 import sys
 import time
 import warnings
@@ -43,6 +44,7 @@
 from matplotlib.transforms import Affine2D, BboxBase
 from matplotlib.path import Path
 from matplotlib import _path
+from matplotlib import _png
 from matplotlib import ttconv
 
 # Overview
@@ -87,7 +89,6 @@
 
 # TODOs:
 #
-# * image compression could be improved (PDF supports png-like compression)
 # * encoding of fonts, including mathtext fonts and unicode support
 # * TTF support has lots of small TODOs, e.g., how do you know if a font
 #   is serif/sans-serif, or symbolic/non-symbolic?
@@ -341,11 +342,12 @@ class Stream(object):
     """
     __slots__ = ('id', 'len', 'pdfFile', 'file', 'compressobj', 'extra', 'pos')
 
-    def __init__(self, id, len, file, extra=None):
+    def __init__(self, id, len, file, extra=None, png=None):
         """id: object id of stream; len: an unused Reference object for the
         length of the stream, or None (to use a memory buffer); file:
         a PdfFile; extra: a dictionary of extra key-value pairs to
-        include in the stream header """
+        include in the stream header; png: if the data is already
+        png compressed, the decode parameters"""
         self.id = id            # object id
         self.len = len          # id of length object
         self.pdfFile = file
@@ -354,10 +356,13 @@ def __init__(self, id, len, file, extra=None):
         if extra is None:
             self.extra = dict()
         else:
-            self.extra = extra
+            self.extra = extra.copy()
+        if png is not None:
+            self.extra.update({'Filter':      Name('FlateDecode'),
+                               'DecodeParms': png})
 
         self.pdfFile.recordXref(self.id)
-        if rcParams['pdf.compression']:
+        if rcParams['pdf.compression'] and not png:
             self.compressobj = zlib.compressobj(rcParams['pdf.compression'])
         if self.len is None:
             self.file = BytesIO()
@@ -590,9 +595,9 @@ def output(self, *data):
         self.write(fill([pdfRepr(x) for x in data]))
         self.write(b'\n')
 
-    def beginStream(self, id, len, extra=None):
+    def beginStream(self, id, len, extra=None, png=None):
         assert self.currentstream is None
-        self.currentstream = Stream(id, len, self, extra)
+        self.currentstream = Stream(id, len, self, extra, png)
 
     def endStream(self):
         if self.currentstream is not None:
@@ -1262,12 +1267,10 @@ def _rgb(self, im):
         rgba = np.fromstring(s, np.uint8)
         rgba.shape = (h, w, 4)
         rgba = rgba[::-1]
-        rgb = rgba[:, :, :3].tostring()
-        a = rgba[:, :, 3]
-        if np.all(a == 255):
+        rgb = np.ascontiguousarray(rgba[:, :, :3])
+        alpha = np.ascontiguousarray(rgba[:, :, 3][..., None])
+        if np.all(alpha == 255):
             alpha = None
-        else:
-            alpha = a.tostring()
         return h, w, rgb, alpha
 
     def _gray(self, im, rc=0.3, gc=0.59, bc=0.11):
@@ -1279,50 +1282,76 @@ def _gray(self, im, rc=0.3, gc=0.59, bc=0.11):
         r = rgba_f[:, :, 0]
         g = rgba_f[:, :, 1]
         b = rgba_f[:, :, 2]
-        a = rgba[:, :, 3]
-        if np.all(a == 255):
+        alpha = np.ascontiguousarray(rgba[:, :, 3][..., None])
+        if np.all(alpha == 255):
             alpha = None
-        else:
-            alpha = a.tostring()
-        gray = (r*rc + g*gc + b*bc).astype(np.uint8).tostring()
+        gray = (r*rc + g*gc + b*bc).astype(np.uint8)[..., None]
         return rgbat[0], rgbat[1], gray, alpha
 
+    def _writePng(self, data):
+        buffer = BytesIO()
+        _png.write_png(data, buffer)
+        buffer.seek(8)
+        written = 0
+        header = bytearray(8)
+        while True:
+            n = buffer.readinto(header)
+            assert n == 8
+            length, type = struct.unpack('!L4s', header)
+            if type == b'IDAT':
+                data = bytearray(length)
+                n = buffer.readinto(data)
+                assert n == length
+                self.currentstream.write(data)
+                written += n
+            elif type == b'IEND':
+                break
+            else:
+                buffer.seek(length, 1)
+            buffer.seek(4, 1)   # skip CRC
+
+    def _writeImg(self, data, height, width, grayscale, id, smask=None):
+        obj = {'Type':             Name('XObject'),
+               'Subtype':          Name('Image'),
+               'Width':            width,
+               'Height':           height,
+               'ColorSpace':       Name('DeviceGray' if grayscale
+                                        else 'DeviceRGB'),
+               'BitsPerComponent': 8}
+        if smask:
+            obj['Smask'] = smask
+        if rcParams['pdf.compression']:
+            png = {'Predictor': 10,
+                   'Colors':    1 if grayscale else 3,
+                   'Columns':   width}
+        else:
+            png = None
+        self.beginStream(
+            id,
+            self.reserveObject('length of image stream'),
+            obj,
+            png=png
+            )
+        if png:
+            self._writePng(data)
+        else:
+            self.currentstream.write(data.tostring())
+        self.endStream()
+
     def writeImages(self):
         for img, pair in six.iteritems(self.images):
             if img.is_grayscale:
                 height, width, data, adata = self._gray(img)
             else:
                 height, width, data, adata = self._rgb(img)
 
-            colorspace = 'DeviceGray' if img.is_grayscale else 'DeviceRGB'
-            obj = {'Type': Name('XObject'),
-                   'Subtype': Name('Image'),
-                   'Width': width,
-                   'Height': height,
-                   'ColorSpace': Name(colorspace),
-                   'BitsPerComponent': 8}
-
             if adata is not None:
                 smaskObject = self.reserveObject("smask")
-                self.beginStream(
-                    smaskObject.id,
-                    self.reserveObject('length of smask stream'),
-                    {'Type': Name('XObject'), 'Subtype': Name('Image'),
-                     'Width': width, 'Height': height,
-                     'ColorSpace': Name('DeviceGray'), 'BitsPerComponent': 8})
-                # TODO: predictors (i.e., output png)
-                self.currentstream.write(adata)
-                self.endStream()
-                obj['SMask'] = smaskObject
-
-            self.beginStream(
-                pair[1].id,
-                self.reserveObject('length of image stream'),
-                obj
-                )
-            # TODO: predictors (i.e., output png)
-            self.currentstream.write(data)
-            self.endStream()
+                self._writeImg(adata, height, width, True, smaskObject.id)
+            else:
+                smaskObject = None
+            self._writeImg(data, height, width, img.is_grayscale,
+                           pair[1].id, smaskObject)
 
     def markerObject(self, path, trans, fillp, strokep, lw, joinstyle,
                      capstyle):
