Faster image generation in WebAgg/NbAgg backends

mdboom · mdboom · commit 62b556d470ec · 2015-11-03T10:58:33.000-05:00
- Write the PNG data to string buffer, rather than an io.BytesIO
  to eliminate the cost of memory reallocation

- Add compression and filter arguments to `_png.write_png`

- Use compression=3 and filter=NONE, which seems to be a sweet spot
  for processing time vs. file size
diff --git a/lib/matplotlib/backends/backend_webagg_core.py b/lib/matplotlib/backends/backend_webagg_core.py
@@ -143,11 +143,6 @@ class FigureCanvasWebAggCore(backend_agg.FigureCanvasAgg):
     def __init__(self, *args, **kwargs):
         backend_agg.FigureCanvasAgg.__init__(self, *args, **kwargs)
 
-        # A buffer to hold the PNG data for the last frame.  This is
-        # retained so it can be resent to each client without
-        # regenerating it.
-        self._png_buffer = io.BytesIO()
-
         # Set to True when the renderer contains data that is newer
         # than the PNG buffer.
         self._png_is_old = True
@@ -225,24 +220,19 @@ def get_diff_image(self):
                 diff = buff != last_buffer
                 output = np.where(diff, buff, 0)
 
-            # Clear out the PNG data buffer rather than recreating it
-            # each time.  This reduces the number of memory
-            # (de)allocations.
-            self._png_buffer.truncate()
-            self._png_buffer.seek(0)
-
             # TODO: We should write a new version of write_png that
             # handles the differencing inline
-            _png.write_png(
+            buff = _png.write_png(
                 output.view(dtype=np.uint8).reshape(output.shape + (4,)),
-                self._png_buffer)
+                None, compression=6, filter=_png.PNG_FILTER_NONE)
 
             # Swap the renderer frames
             self._renderer, self._last_renderer = (
                 self._last_renderer, renderer)
             self._force_full = False
             self._png_is_old = False
-        return self._png_buffer.getvalue()
+
+        return buff
 
     def get_renderer(self, cleared=None):
         # Mirrors super.get_renderer, but caches the old one
diff --git a/src/_png.cpp b/src/_png.cpp
@@ -34,6 +34,27 @@ extern "C" {
 #undef jmpbuf
 #endif
 
+struct buffer_t {
+    PyObject *str;
+    size_t cursor;
+    size_t size;
+};
+
+
+static void write_png_data_buffer(png_structp png_ptr, png_bytep data, png_size_t length)
+{
+    buffer_t *buff = (buffer_t *)png_get_io_ptr(png_ptr);
+    if (buff->cursor + length < buff->size) {
+        memcpy(PyBytes_AS_STRING(buff->str) + buff->cursor, data, length);
+        buff->cursor += length;
+    }
+}
+
+static void flush_png_data_buffer(png_structp png_ptr)
+{
+
+}
+
 static void write_png_data(png_structp png_ptr, png_bytep data, png_size_t length)
 {
     PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);
@@ -69,20 +90,24 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     numpy::array_view<unsigned char, 3> buffer;
     PyObject *filein;
     double dpi = 0;
-    const char *names[] = { "buffer", "file", "dpi", NULL };
+    int compression = 6;
+    int filter = -1;
+    const char *names[] = { "buffer", "file", "dpi", "compression", "filter", 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",
+                                     "O&O|dii:write_png",
                                      (char **)names,
                                      &buffer.converter_contiguous,
                                      &buffer,
                                      &filein,
-                                     &dpi)) {
+                                     &dpi,
+                                     &compression,
+                                     &filter)) {
         return NULL;
     }
 
@@ -104,6 +129,8 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     png_infop info_ptr = NULL;
     struct png_color_8_struct sig_bit;
     int png_color_type;
+    buffer_t buff;
+    buff.str = NULL;
 
     switch (channels) {
     case 1:
@@ -122,6 +149,12 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
         goto exit;
     }
 
+    if (compression < 0 || compression > 9) {
+        PyErr_Format(PyExc_ValueError,
+                     "compression must be in range 0-9, got %d", compression);
+        goto exit;
+    }
+
     if (PyBytes_Check(filein) || PyUnicode_Check(filein)) {
         if ((py_file = mpl_PyFile_OpenFile(filein, (char *)"wb")) == NULL) {
             goto exit;
@@ -131,7 +164,14 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
         py_file = filein;
     }
 
-    if ((fp = mpl_PyFile_Dup(py_file, (char *)"wb", &offset))) {
+    if (filein == Py_None) {
+        buff.size = width * height * 4 + 1024;
+        buff.str = PyBytes_FromStringAndSize(NULL, buff.size);
+        if (buff.str == NULL) {
+            goto exit;
+        }
+        buff.cursor = 0;
+    } else if ((fp = mpl_PyFile_Dup(py_file, (char *)"wb", &offset))) {
         close_dup_file = true;
     } else {
         PyErr_Clear();
@@ -152,6 +192,11 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
         goto exit;
     }
 
+    png_set_compression_level(png_ptr, compression);
+    if (filter >= 0) {
+        png_set_filter(png_ptr, 0, filter);
+    }
+
     info_ptr = png_create_info_struct(png_ptr);
     if (info_ptr == NULL) {
         PyErr_SetString(PyExc_RuntimeError, "Could not create info struct");
@@ -163,7 +208,9 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
         goto exit;
     }
 
-    if (fp) {
+    if (buff.str) {
+        png_set_write_fn(png_ptr, (void *)&buff, &write_png_data_buffer, &flush_png_data_buffer);
+    } else if (fp) {
         png_init_io(png_ptr, fp);
     } else {
         png_set_write_fn(png_ptr, (void *)py_file, &write_png_data, &flush_png_data);
@@ -227,8 +274,13 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     }
 
     if (PyErr_Occurred()) {
+        Py_XDECREF(buff.str);
         return NULL;
     } else {
+        if (buff.str) {
+            _PyBytes_Resize(&buff.str, buff.cursor);
+            return buff.str;
+        }
         Py_RETURN_NONE;
     }
 }
@@ -558,6 +610,15 @@ extern "C" {
 
         import_array();
 
+        if (PyModule_AddIntConstant(m, "PNG_FILTER_NONE", PNG_FILTER_NONE) ||
+            PyModule_AddIntConstant(m, "PNG_FILTER_SUB", PNG_FILTER_SUB) ||
+            PyModule_AddIntConstant(m, "PNG_FILTER_UP", PNG_FILTER_UP) ||
+            PyModule_AddIntConstant(m, "PNG_FILTER_AVG", PNG_FILTER_AVG) ||
+            PyModule_AddIntConstant(m, "PNG_FILTER_PAETH", PNG_FILTER_PAETH)) {
+            INITERROR;
+        }
+
+
 #if PY3K
         return m;
 #endif
