KIKIMR-19287: add task_stats_drawing script

1 files changed, 297 insertions, 0 deletions

diff --git a/contrib/python/matplotlib/py3/src/_image_wrapper.cpp b/contrib/python/matplotlib/py3/src/_image_wrapper.cpp
new file mode 100644
index 0000000000..ca6ae8b222
--- /dev/null
+++ b/contrib/python/matplotlib/py3/src/_image_wrapper.cpp
@@ -0,0 +1,297 @@
+#include "mplutils.h"
+#include "_image_resample.h"
+#include "numpy_cpp.h"
+#include "py_converters.h"
+
+
+/**********************************************************************
+ * Free functions
+ * */
+
+const char* image_resample__doc__ =
+"resample(input_array, output_array, transform, interpolation=NEAREST, resample=False, alpha=1.0, norm=False, radius=1.0)\n"
+"--\n\n"
+
+"Resample input_array, blending it in-place into output_array, using an\n"
+"affine transformation.\n\n"
+
+"Parameters\n"
+"----------\n"
+"input_array : 2-d or 3-d NumPy array of float, double or `numpy.uint8`\n"
+"    If 2-d, the image is grayscale.  If 3-d, the image must be of size\n"
+"    4 in the last dimension and represents RGBA data.\n\n"
+
+"output_array : 2-d or 3-d NumPy array of float, double or `numpy.uint8`\n"
+"    The dtype and number of dimensions must match `input_array`.\n\n"
+
+"transform : matplotlib.transforms.Transform instance\n"
+"    The transformation from the input array to the output array.\n\n"
+
+"interpolation : int, default: NEAREST\n"
+"    The interpolation method.  Must be one of the following constants\n"
+"    defined in this module:\n\n"
+
+"      NEAREST, BILINEAR, BICUBIC, SPLINE16, SPLINE36,\n"
+"      HANNING, HAMMING, HERMITE, KAISER, QUADRIC, CATROM, GAUSSIAN,\n"
+"      BESSEL, MITCHELL, SINC, LANCZOS, BLACKMAN\n\n"
+
+"resample : bool, optional\n"
+"    When `True`, use a full resampling method.  When `False`, only\n"
+"    resample when the output image is larger than the input image.\n\n"
+
+"alpha : float, default: 1\n"
+"    The transparency level, from 0 (transparent) to 1 (opaque).\n\n"
+
+"norm : bool, default: False\n"
+"    Whether to norm the interpolation function.\n\n"
+
+"radius: float, default: 1\n"
+"    The radius of the kernel, if method is SINC, LANCZOS or BLACKMAN.\n";
+
+
+static PyArrayObject *
+_get_transform_mesh(PyObject *py_affine, npy_intp *dims)
+{
+    /* TODO: Could we get away with float, rather than double, arrays here? */
+
+    /* Given a non-affine transform object, create a mesh that maps
+    every pixel in the output image to the input image.  This is used
+    as a lookup table during the actual resampling. */
+
+    PyObject *py_inverse = NULL;
+    npy_intp out_dims[3];
+
+    out_dims[0] = dims[0] * dims[1];
+    out_dims[1] = 2;
+
+    py_inverse = PyObject_CallMethod(py_affine, "inverted", NULL);
+    if (py_inverse == NULL) {
+        return NULL;
+    }
+
+    numpy::array_view<double, 2> input_mesh(out_dims);
+    double *p = (double *)input_mesh.data();
+
+    for (npy_intp y = 0; y < dims[0]; ++y) {
+        for (npy_intp x = 0; x < dims[1]; ++x) {
+            *p++ = (double)x;
+            *p++ = (double)y;
+        }
+    }
+
+    PyObject *output_mesh = PyObject_CallMethod(
+        py_inverse, "transform", "O", input_mesh.pyobj_steal());
+
+    Py_DECREF(py_inverse);
+
+    if (output_mesh == NULL) {
+        return NULL;
+    }
+
+    PyArrayObject *output_mesh_array =
+        (PyArrayObject *)PyArray_ContiguousFromAny(
+            output_mesh, NPY_DOUBLE, 2, 2);
+
+    Py_DECREF(output_mesh);
+
+    if (output_mesh_array == NULL) {
+        return NULL;
+    }
+
+    return output_mesh_array;
+}
+
+
+static PyObject *
+image_resample(PyObject *self, PyObject* args, PyObject *kwargs)
+{
+    PyObject *py_input = NULL;
+    PyObject *py_output = NULL;
+    PyObject *py_transform = NULL;
+    resample_params_t params;
+
+    PyArrayObject *input = NULL;
+    PyArrayObject *output = NULL;
+    PyArrayObject *transform_mesh = NULL;
+    int ndim;
+    int type;
+
+    params.interpolation = NEAREST;
+    params.transform_mesh = NULL;
+    params.resample = false;
+    params.norm = false;
+    params.radius = 1.0;
+    params.alpha = 1.0;
+
+    const char *kwlist[] = {
+        "input_array", "output_array", "transform", "interpolation",
+        "resample", "alpha", "norm", "radius", NULL };
+
+    if (!PyArg_ParseTupleAndKeywords(
+            args, kwargs, "OOO|iO&dO&d:resample", (char **)kwlist,
+            &py_input, &py_output, &py_transform,
+            &params.interpolation, &convert_bool, &params.resample,
+            &params.alpha, &convert_bool, &params.norm, &params.radius)) {
+        return NULL;
+    }
+
+    if (params.interpolation < 0 || params.interpolation >= _n_interpolation) {
+        PyErr_Format(PyExc_ValueError, "Invalid interpolation value %d",
+                     params.interpolation);
+        goto error;
+    }
+
+    input = (PyArrayObject *)PyArray_FromAny(
+        py_input, NULL, 2, 3, NPY_ARRAY_C_CONTIGUOUS, NULL);
+    if (!input) {
+        goto error;
+    }
+    ndim = PyArray_NDIM(input);
+    type = PyArray_TYPE(input);
+
+    if (!PyArray_Check(py_output)) {
+        PyErr_SetString(PyExc_ValueError, "Output array must be a NumPy array");
+        goto error;
+    }
+    output = (PyArrayObject *)py_output;
+    if (PyArray_NDIM(output) != ndim) {
+        PyErr_Format(
+            PyExc_ValueError,
+            "Input (%dD) and output (%dD) have different dimensionalities.",
+            ndim, PyArray_NDIM(output));
+        goto error;
+    }
+    // PyArray_FromAny above checks that input is 2D or 3D.
+    if (ndim == 3 && (PyArray_DIM(input, 2) != 4 || PyArray_DIM(output, 2) != 4)) {
+        PyErr_Format(
+            PyExc_ValueError,
+            "If 3D, input and output arrays must be RGBA with shape (M, N, 4); "
+            "got trailing dimensions of %" NPY_INTP_FMT " and %" NPY_INTP_FMT
+            " respectively", PyArray_DIM(input, 2), PyArray_DIM(output, 2));
+        goto error;
+    }
+    if (PyArray_TYPE(output) != type) {
+        PyErr_SetString(PyExc_ValueError, "Mismatched types");
+        goto error;
+    }
+    if (!PyArray_IS_C_CONTIGUOUS(output)) {
+        PyErr_SetString(PyExc_ValueError, "Output array must be C-contiguous");
+        goto error;
+    }
+
+    if (py_transform == NULL || py_transform == Py_None) {
+        params.is_affine = true;
+    } else {
+        PyObject *py_is_affine;
+        int py_is_affine2;
+        py_is_affine = PyObject_GetAttrString(py_transform, "is_affine");
+        if (!py_is_affine) {
+            goto error;
+        }
+
+        py_is_affine2 = PyObject_IsTrue(py_is_affine);
+        Py_DECREF(py_is_affine);
+
+        if (py_is_affine2 == -1) {
+            goto error;
+        } else if (py_is_affine2) {
+            if (!convert_trans_affine(py_transform, &params.affine)) {
+                goto error;
+            }
+            params.is_affine = true;
+        } else {
+            transform_mesh = _get_transform_mesh(
+                py_transform, PyArray_DIMS(output));
+            if (!transform_mesh) {
+                goto error;
+            }
+            params.transform_mesh = (double *)PyArray_DATA(transform_mesh);
+            params.is_affine = false;
+        }
+    }
+
+    if (auto resampler =
+            (ndim == 2) ? (
+                (type == NPY_UINT8) ? resample<agg::gray8> :
+                (type == NPY_INT8) ? resample<agg::gray8> :
+                (type == NPY_UINT16) ? resample<agg::gray16> :
+                (type == NPY_INT16) ? resample<agg::gray16> :
+                (type == NPY_FLOAT32) ? resample<agg::gray32> :
+                (type == NPY_FLOAT64) ? resample<agg::gray64> :
+                nullptr) : (
+            // ndim == 3
+                (type == NPY_UINT8) ? resample<agg::rgba8> :
+                (type == NPY_INT8) ? resample<agg::rgba8> :
+                (type == NPY_UINT16) ? resample<agg::rgba16> :
+                (type == NPY_INT16) ? resample<agg::rgba16> :
+                (type == NPY_FLOAT32) ? resample<agg::rgba32> :
+                (type == NPY_FLOAT64) ? resample<agg::rgba64> :
+                nullptr)) {
+        Py_BEGIN_ALLOW_THREADS
+        resampler(
+            PyArray_DATA(input), PyArray_DIM(input, 1), PyArray_DIM(input, 0),
+            PyArray_DATA(output), PyArray_DIM(output, 1), PyArray_DIM(output, 0),
+            params);
+        Py_END_ALLOW_THREADS
+    } else {
+        PyErr_SetString(
+            PyExc_ValueError,
+            "arrays must be of dtype byte, short, float32 or float64");
+        goto error;
+    }
+
+    Py_DECREF(input);
+    Py_XDECREF(transform_mesh);
+    Py_RETURN_NONE;
+
+ error:
+    Py_XDECREF(input);
+    Py_XDECREF(transform_mesh);
+    return NULL;
+}
+
+static PyMethodDef module_functions[] = {
+    {"resample", (PyCFunction)image_resample, METH_VARARGS|METH_KEYWORDS, image_resample__doc__},
+    {NULL}
+};
+
+static struct PyModuleDef moduledef = {
+    PyModuleDef_HEAD_INIT, "_image", NULL, 0, module_functions,
+};
+
+PyMODINIT_FUNC PyInit__image(void)
+{
+    PyObject *m;
+
+    import_array();
+
+    m = PyModule_Create(&moduledef);
+
+    if (m == NULL) {
+        return NULL;
+    }
+
+    if (PyModule_AddIntConstant(m, "NEAREST", NEAREST) ||
+        PyModule_AddIntConstant(m, "BILINEAR", BILINEAR) ||
+        PyModule_AddIntConstant(m, "BICUBIC", BICUBIC) ||
+        PyModule_AddIntConstant(m, "SPLINE16", SPLINE16) ||
+        PyModule_AddIntConstant(m, "SPLINE36", SPLINE36) ||
+        PyModule_AddIntConstant(m, "HANNING", HANNING) ||
+        PyModule_AddIntConstant(m, "HAMMING", HAMMING) ||
+        PyModule_AddIntConstant(m, "HERMITE", HERMITE) ||
+        PyModule_AddIntConstant(m, "KAISER", KAISER) ||
+        PyModule_AddIntConstant(m, "QUADRIC", QUADRIC) ||
+        PyModule_AddIntConstant(m, "CATROM", CATROM) ||
+        PyModule_AddIntConstant(m, "GAUSSIAN", GAUSSIAN) ||
+        PyModule_AddIntConstant(m, "BESSEL", BESSEL) ||
+        PyModule_AddIntConstant(m, "MITCHELL", MITCHELL) ||
+        PyModule_AddIntConstant(m, "SINC", SINC) ||
+        PyModule_AddIntConstant(m, "LANCZOS", LANCZOS) ||
+        PyModule_AddIntConstant(m, "BLACKMAN", BLACKMAN) ||
+        PyModule_AddIntConstant(m, "_n_interpolation", _n_interpolation)) {
+        Py_DECREF(m);
+        return NULL;
+    }
+
+    return m;
+}