TST: Calculate RMS and diff image in C++

QuLogic · QuLogic · commit 62a96ffdc5d1 · 2024-11-08T04:01:39.000-05:00
The current implementation is not slow, but uses a lot of memory per
image.

In `compare_images`, we have:

- one actual and one expected image as uint8 (2×image)
- both converted to int16 (though original is thrown away) (4×)

which adds up to 4× the image allocated in this function.

Then it calls `calculate_rms`, which has:

- a difference between them as int16 (2×)
- the difference cast to 64-bit float (8×)
- the square of the difference as 64-bit float (though possibly the
  original difference was thrown away) (8×)

which at its peak has 16× the image allocated in parallel.

If the RMS is over the desired tolerance, then `save_diff_image` is
called, which:

- loads the actual and expected images _again_ as uint8 (2× image)
- converts both to 64-bit float (throwing away the original) (16×)
- calculates the difference (8×)
- calculates the absolute value (8×)
- multiples that by 10 (in-place, so no allocation)
- clips to 0-255 (8×)
- casts to uint8 (1×)

which at peak uses 32× the image.

So at their peak, `compare_images`→`calculate_rms` will have 20× the
image allocated, and then `compare_images`→`save_diff_image` will have
36× the image allocated. This is generally not a problem, but on
resource-constrained places like WASM, it can sometimes run out of
memory just in `calculate_rms`.

This implementation in C++ always allocates the diff image, even when
not needed, but doesn't have all the temporaries, so it's a maximum of
3× the image size (plus a few scalar temporaries).
diff --git a/lib/matplotlib/testing/compare.py b/lib/matplotlib/testing/compare.py
@@ -19,7 +19,7 @@
 from PIL import Image
 
 import matplotlib as mpl
-from matplotlib import cbook
+from matplotlib import cbook, _image
 from matplotlib.testing.exceptions import ImageComparisonFailure
 
 _log = logging.getLogger(__name__)
@@ -398,7 +398,7 @@ def compare_images(expected, actual, tol, in_decorator=False):
 
     The two given filenames may point to files which are convertible to
     PNG via the `.converter` dictionary. The underlying RMS is calculated
-    with the `.calculate_rms` function.
+    in a similar way to the `.calculate_rms` function.
 
     Parameters
     ----------
@@ -469,17 +469,12 @@ def compare_images(expected, actual, tol, in_decorator=False):
         if np.array_equal(expected_image, actual_image):
             return None
 
-    # convert to signed integers, so that the images can be subtracted without
-    # overflow
-    expected_image = expected_image.astype(np.int16)
-    actual_image = actual_image.astype(np.int16)
-
-    rms = calculate_rms(expected_image, actual_image)
+    rms, abs_diff = _image.calculate_rms_and_diff(expected_image, actual_image)
 
     if rms <= tol:
         return None
 
-    save_diff_image(expected, actual, diff_image)
+    Image.fromarray(abs_diff).save(diff_image, format="png")
 
     results = dict(rms=rms, expected=str(expected),
                    actual=str(actual), diff=str(diff_image), tol=tol)
diff --git a/src/_image_wrapper.cpp b/src/_image_wrapper.cpp
@@ -1,6 +1,8 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/numpy.h>
 
+#include <algorithm>
+
 #include "_image_resample.h"
 #include "py_converters.h"
 
@@ -200,6 +202,76 @@ image_resample(py::array input_array,
 }
 
 
+// This is used by matplotlib.testing.compare to calculate RMS and a difference image.
+static py::tuple
+calculate_rms_and_diff(py::array_t<unsigned char> expected_image,
+                       py::array_t<unsigned char> actual_image)
+{
+    if (expected_image.ndim() != 3) {
+        auto exceptions = py::module_::import("matplotlib.testing.exceptions");
+        auto ImageComparisonFailure = exceptions.attr("ImageComparisonFailure");
+        py::set_error(
+            ImageComparisonFailure,
+            "Expected image must be 3-dimensional, but is {ndim}-dimensional"_s.format(
+                "ndim"_a=expected_image.ndim()));
+        throw py::error_already_set();
+    }
+
+    if (actual_image.ndim() != 3) {
+        auto exceptions = py::module_::import("matplotlib.testing.exceptions");
+        auto ImageComparisonFailure = exceptions.attr("ImageComparisonFailure");
+        py::set_error(
+            ImageComparisonFailure,
+            "Actual image must be 3-dimensional, but is {ndim}-dimensional"_s.format(
+                "ndim"_a=actual_image.ndim()));
+        throw py::error_already_set();
+    }
+
+    auto height = expected_image.shape(0);
+    auto width = expected_image.shape(1);
+    auto depth = expected_image.shape(2);
+
+    if (height != actual_image.shape(0) || width != actual_image.shape(1) ||
+            depth != actual_image.shape(2)) {
+        auto exceptions = py::module_::import("matplotlib.testing.exceptions");
+        auto ImageComparisonFailure = exceptions.attr("ImageComparisonFailure");
+        py::set_error(
+            ImageComparisonFailure,
+            "Image sizes do not match expected size: {expected_image.shape} "_s
+            "actual size {actual_image.shape}"_s.format(
+                "expected_image"_a=expected_image, "actual_image"_a=actual_image));
+        throw py::error_already_set();
+    }
+    auto expected = expected_image.unchecked<3>();
+    auto actual = actual_image.unchecked<3>();
+
+    py::ssize_t diff_dims[3] = {height, width, 3};
+    py::array_t<unsigned char> diff_image(diff_dims);
+    auto diff = diff_image.mutable_unchecked<3>();
+
+    double total = 0.0;
+    for (auto i = 0; i < height; i++) {
+        for (auto j = 0; j < width; j++) {
+            for (auto k = 0; k < depth; k++) {
+                auto pixel_diff = static_cast<double>(expected(i, j, k)) -
+                                  static_cast<double>(actual(i, j, k));
+
+                total += pixel_diff*pixel_diff;
+
+                if (k != 3) { // Hard-code a fully solid alpha channel by omitting it.
+                    diff(i, j, k) = static_cast<unsigned char>(std::clamp(
+                        abs(pixel_diff) * 10, // Expand differences in luminance domain.
+                        0.0, 255.0));
+                }
+            }
+        }
+    }
+    total = total / (width * height * depth);
+
+    return py::make_tuple(sqrt(total), diff_image);
+}
+
+
 PYBIND11_MODULE(_image, m, py::mod_gil_not_used())
 {
     py::enum_<interpolation_e>(m, "_InterpolationType")
@@ -232,4 +304,7 @@ PYBIND11_MODULE(_image, m, py::mod_gil_not_used())
         "norm"_a = false,
         "radius"_a = 1,
         image_resample__doc__);
+
+    m.def("calculate_rms_and_diff", &calculate_rms_and_diff,
+          "expected_image"_a, "actual_image"_a);
 }
