From 5ef8107f9bbe8744e0e00f21ec8832ec666b9a5d Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 3 Nov 2015 10:09:38 -0500 Subject: [PATCH] Faster image generation in WebAgg/NbAgg backends - 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 --- .../backends/backend_webagg_core.py | 18 +---- src/_png.cpp | 73 +++++++++++++++++-- 2 files changed, 71 insertions(+), 20 deletions(-) diff --git a/lib/matplotlib/backends/backend_webagg_core.py b/lib/matplotlib/backends/backend_webagg_core.py index 4c3b41e664fc..3a9fc7938a33 100644 --- 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 index 3121edbb712d..50b6bd375d40 100644 --- 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,7 +90,9 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds) numpy::array_view 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 @@ -77,12 +100,14 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds) // 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,10 +208,12 @@ 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); + png_set_write_fn(png_ptr, (void *)&py_file, &write_png_data, &flush_png_data); } png_set_IHDR(png_ptr, info_ptr, @@ -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