tacaswell · Jul 22, 2015 · Jul 8, 2015 · Jul 8, 2015 · Jul 8, 2015 · Jul 9, 2015
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:
@@ -1256,73 +1261,103 @@ def imageObject(self, image):
         self.images[image] = (name, ob)
         return name
 
-    def _rgb(self, im):
-        h, w, s = im.as_rgba_str()
+    def _unpack(self, im):
+        """
+        Unpack the image object im into height, width, data, alpha,
+        where data and alpha are HxWx3 (RGB) or HxWx1 (grayscale or alpha)
+        arrays, except alpha is None if the image is fully opaque.
+        """
 
+        h, w, s = im.as_rgba_str()
         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 = rgba[:, :, :3]
+        alpha = 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):
-        rgbat = im.as_rgba_str()
-        rgba = np.fromstring(rgbat[2], np.uint8)
-        rgba.shape = (rgbat[0], rgbat[1], 4)
-        rgba = rgba[::-1]
-        rgba_f = rgba.astype(np.float32)
-        r = rgba_f[:, :, 0]
-        g = rgba_f[:, :, 1]
-        b = rgba_f[:, :, 2]
-        a = rgba[:, :, 3]
-        if np.all(a == 255):
-            alpha = None
+            alpha = np.array(alpha, order='C')
+        if im.is_grayscale:
+            r, g, b = rgb.astype(np.float32).transpose(2, 0, 1)
+            gray = (0.3 * r + 0.59 * g + 0.11 * b).astype(np.uint8)[..., None]
+            return h, w, gray, alpha
         else:
-            alpha = a.tostring()
-        gray = (r*rc + g*gc + b*bc).astype(np.uint8).tostring()
-        return rgbat[0], rgbat[1], gray, alpha
+            rgb = np.array(rgb, order='C')
+            return h, w, rgb, alpha
 
-    def writeImages(self):
-        for img, pair in six.iteritems(self.images):
-            if img.is_grayscale:
-                height, width, data, adata = self._gray(img)
+    def _writePng(self, data):
+        """
+        Write the image *data* into the pdf file using png
+        predictors with Flate compression.
+        """
+
+        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(b'!L4s', bytes(header))
+            if type == b'IDAT':
+                data = bytearray(length)
+                n = buffer.readinto(data)
+                assert n == length
+                self.currentstream.write(bytes(data))
+                written += n
+            elif type == b'IEND':
+                break
             else:
-                height, width, data, adata = self._rgb(img)
+                buffer.seek(length, 1)
+            buffer.seek(4, 1)   # skip CRC
+
+    def _writeImg(self, data, height, width, grayscale, id, smask=None):
+        """
+        Write the image *data* of size *height* x *width*, as grayscale
+        if *grayscale* is true and RGB otherwise, as pdf object *id*
+        and with the soft mask (alpha channel) *smask*, which should be
+        either None or a *height* x *width* x 1 array.
+        """
 
-            colorspace = 'DeviceGray' if img.is_grayscale else 'DeviceRGB'
-            obj = {'Type': Name('XObject'),
-                   'Subtype': Name('Image'),
-                   'Width': width,
-                   'Height': height,
-                   'ColorSpace': Name(colorspace),
-                   'BitsPerComponent': 8}
+        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):
+            height, width, data, adata = self._unpack(img)
             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):

diff --git a/src/_png.cpp b/src/_png.cpp
@@ -71,11 +71,15 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     double dpi = 0;
     const char *names[] = { "buffer", "file", "dpi", NULL };
 
+    // We don't need strict contiguity, just for each row to be
+    // contiguous, and libpng has special handling for getting RGB out
+    // of RGBA, ARGB or BGR. But the simplest thing to do is to
+    // enforce contiguity using array_view::converter_contiguous.
     if (!PyArg_ParseTupleAndKeywords(args,
                                      kwds,
                                      "O&O|d:write_png",
                                      (char **)names,
-                                     &buffer.converter,
+                                     &buffer.converter_contiguous,
                                      &buffer,
                                      &filein,
                                      &dpi)) {
@@ -84,6 +88,7 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
 
     png_uint_32 width = (png_uint_32)buffer.dim(1);
     png_uint_32 height = (png_uint_32)buffer.dim(0);
+    int channels = buffer.dim(2);
     std::vector<png_bytep> row_pointers(height);
     for (png_uint_32 row = 0; row < (png_uint_32)height; ++row) {
         row_pointers[row] = (png_bytep)buffer[row].data();
@@ -98,9 +103,22 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     png_structp png_ptr = NULL;
     png_infop info_ptr = NULL;
     struct png_color_8_struct sig_bit;
-
-    if (buffer.dim(2) != 4) {
-        PyErr_SetString(PyExc_ValueError, "Buffer must be RGBA NxMx4 array");
+    int png_color_type;
+
+    switch (channels) {
+    case 1:
+	png_color_type = PNG_COLOR_TYPE_GRAY;
+	break;
+    case 3:
+	png_color_type = PNG_COLOR_TYPE_RGB;
+	break;
+    case 4:
+	png_color_type = PNG_COLOR_TYPE_RGB_ALPHA;
+	break;
+    default:
+        PyErr_SetString(PyExc_ValueError,
+			"Buffer must be an NxMxD array with D in 1, 3, 4 "
+			"(grayscale, RGB, RGBA)");
         goto exit;
     }
 
@@ -141,7 +159,7 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     }
 
     if (setjmp(png_jmpbuf(png_ptr))) {
-        PyErr_SetString(PyExc_RuntimeError, "Error setting jumps");
+        PyErr_SetString(PyExc_RuntimeError, "libpng signaled error");
         goto exit;
     }
 
@@ -155,7 +173,7 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
                  width,
                  height,
                  8,
-                 PNG_COLOR_TYPE_RGB_ALPHA,
+                 png_color_type,
                  PNG_INTERLACE_NONE,
                  PNG_COMPRESSION_TYPE_BASE,
                  PNG_FILTER_TYPE_BASE);
@@ -166,13 +184,27 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
         png_set_pHYs(png_ptr, info_ptr, dots_per_meter, dots_per_meter, PNG_RESOLUTION_METER);
     }
 
-    // this a a color image!
-    sig_bit.gray = 0;
-    sig_bit.red = 8;
-    sig_bit.green = 8;
-    sig_bit.blue = 8;
-    /* if the image has an alpha channel then */
-    sig_bit.alpha = 8;
+    sig_bit.alpha = 0;
+    switch (png_color_type) {
+    case PNG_COLOR_TYPE_GRAY:
+	sig_bit.gray = 8;
+	sig_bit.red = 0;
+	sig_bit.green = 0;
+	sig_bit.blue = 0;
+	break;
+    case PNG_COLOR_TYPE_RGB_ALPHA:
+	sig_bit.alpha = 8;
+	// fall through
+    case PNG_COLOR_TYPE_RGB:
+	sig_bit.gray = 0;
+	sig_bit.red = 8;
+	sig_bit.green = 8;
+	sig_bit.blue = 8;
+	break;
+    default:
+        PyErr_SetString(PyExc_RuntimeError, "internal error, bad png_color_type");
+        goto exit;
+    }
     png_set_sBIT(png_ptr, info_ptr, &sig_bit);
 
     png_write_info(png_ptr, info_ptr);