Merge pull request #5389 from mdboom/webagg-speed

jenshnielsen · jenshnielsen · commit cb3c6f361376 · 2015-11-04T15:21:19.000Z
Faster image generation in WebAgg/NbAgg backends
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);
@@ -62,27 +83,66 @@ static void flush_png_data(png_structp png_ptr)
     Py_XDECREF(result);
 }
 
-const char *Py_write_png__doc__ = "write_png(buffer, file, dpi=0)";
+const char *Py_write_png__doc__ =
+    "write_png(buffer, file, dpi=0, compression=6, filter=auto)\n"
+    "\n"
+    "Parameters\n"
+    "----------\n"
+    "buffer : numpy array of image data\n"
+    "    Must be an MxNxD array of dtype uint8.\n"
+    "    - if D is 1, the image is greyscale\n"
+    "    - if D is 3, the image is RGB\n"
+    "    - if D is 4, the image is RGBA\n"
+    "\n"
+    "file : str path, file-like object or None\n"
+    "    - If a str, must be a file path\n"
+    "    - If a file-like object, must write bytes\n"
+    "    - If None, a byte string containing the PNG data will be returned\n"
+    "\n"
+    "dpi : float\n"
+    "    The dpi to store in the file metadata.\n"
+    "\n"
+    "compression : int\n"
+    "    The level of lossless zlib compression to apply.  0 indicates no\n"
+    "    compression.  Values 1-9 indicate low/fast through high/slow\n"
+    "    compression.  Default is 6.\n"
+    "\n"
+    "filter : int\n"
+    "    Filter to apply.  Must be one of the constants: PNG_FILTER_NONE,\n"
+    "    PNG_FILTER_SUB, PNG_FILTER_UP, PNG_FILTER_AVG, PNG_FILTER_PAETH.\n"
+    "    See the PNG standard for more information.\n"
+    "    If not provided, libpng will try to automatically determine the\n"
+    "    best filter on a line-by-line basis.\n"
+    "\n"
+    "Returns\n"
+    "-------\n"
+    "buffer : bytes or None\n"
+    "    Byte string containing the PNG content if None was passed in for\n"
+    "    file, otherwise None is returned.\n";
 
 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 +164,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 +184,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 +199,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 +227,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 +243,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 +309,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;
     }
 }
@@ -494,21 +581,46 @@ static PyObject *_read_png(PyObject *filein, bool float_result)
     }
 }
 
-const char *Py_read_png_float__doc__ = "read_png_float(file)";
+const char *Py_read_png_float__doc__ =
+    "read_png_float(file)\n"
+    "\n"
+    "Read in a PNG file, converting values to floating-point doubles\n"
+    "in the range (0, 1)\n"
+    "\n"
+    "Parameters\n"
+    "----------\n"
+    "file : str path or file-like object\n";
 
 static PyObject *Py_read_png_float(PyObject *self, PyObject *args, PyObject *kwds)
 {
     return _read_png(args, true);
 }
 
-const char *Py_read_png_int__doc__ = "read_png_int(file)";
+const char *Py_read_png_int__doc__ =
+    "read_png_int(file)\n"
+    "\n"
+    "Read in a PNG file with original integer values.\n"
+    "\n"
+    "Parameters\n"
+    "----------\n"
+    "file : str path or file-like object\n";
 
 static PyObject *Py_read_png_int(PyObject *self, PyObject *args, PyObject *kwds)
 {
     return _read_png(args, false);
 }
 
-const char *Py_read_png__doc__ = "read_png(file)";
+const char *Py_read_png__doc__ =
+    "read_png(file)\n"
+    "\n"
+    "Read in a PNG file, converting values to floating-point doubles\n"
+    "in the range (0, 1)\n"
+    "\n"
+    "Alias for read_png_float()\n"
+    "\n"
+    "Parameters\n"
+    "----------\n"
+    "file : str path or file-like object\n";
 
 static PyMethodDef module_methods[] = {
     {"write_png", (PyCFunction)Py_write_png, METH_VARARGS|METH_KEYWORDS, Py_write_png__doc__},
@@ -558,6 +670,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
