Resample RGB images in C++

QuLogic · QuLogic · commit 2b73e7e33fd3 · 2025-01-14T21:02:43.000-05:00
Agg already has RGB resampling with output to RGBA builtin, so we just
need to correctly wire up the corresponding templates. With this RGB
resampling mode, we save the extra copy from RGB to RGBA in NumPy land
that was required for the previous always-RGBA resampling.
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -160,8 +160,7 @@ def flush_images():
         flush_images()
 
 
-def _resample(
-        image_obj, data, out_shape, transform, *, resample=None, alpha=1):
+def _resample(image_obj, data, out_shape, transform, *, resample=None, alpha=1):
     """
     Convenience wrapper around `._image.resample` to resample *data* to
     *out_shape* (with a third dimension if *data* is RGBA) that takes care of
@@ -204,7 +203,10 @@ def _resample(
             interpolation = 'nearest'
         else:
             interpolation = 'hanning'
-    out = np.zeros(out_shape + data.shape[2:], data.dtype)  # 2D->2D, 3D->3D.
+    if len(data.shape) == 3:
+        # Always output RGBA.
+        out_shape += (4, )
+    out = np.zeros(out_shape, data.dtype)
     if resample is None:
         resample = image_obj.get_resample()
     _image.resample(data, out, transform,
@@ -216,20 +218,6 @@ def _resample(
     return out
 
 
-def _rgb_to_rgba(A):
-    """
-    Convert an RGB image to RGBA, as required by the image resample C++
-    extension.
-    """
-    rgba = np.zeros((A.shape[0], A.shape[1], 4), dtype=A.dtype)
-    rgba[:, :, :3] = A
-    if rgba.dtype == np.uint8:
-        rgba[:, :, 3] = 255
-    else:
-        rgba[:, :, 3] = 1.0
-    return rgba
-
-
 class _ImageBase(mcolorizer.ColorizingArtist):
     """
     Base class for images.
@@ -508,10 +496,10 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
                     # alpha channel below.
                     output_alpha = (255 * alpha) if A.dtype == np.uint8 else alpha
                 else:
-                    output_alpha = _resample(  # resample alpha channel
-                        self, A[..., 3], out_shape, t, alpha=alpha)
-                output = _resample(  # resample rgb channels
-                    self, _rgb_to_rgba(A[..., :3]), out_shape, t, alpha=alpha)
+                    # resample alpha channel
+                    output_alpha = _resample(self, A[..., 3], out_shape, t, alpha=alpha)
+                # resample rgb channels
+                output = _resample(self, A[..., :3], out_shape, t, alpha=alpha)
                 output[..., 3] = output_alpha  # recombine rgb and alpha
 
             # output is now either a 2D array of normed (int or float) data
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -1577,8 +1577,8 @@ def test__resample_valid_output():
         resample(np.zeros((9, 9)), np.zeros((9, 9, 4)))
     with pytest.raises(ValueError, match="different dimensionalities"):
         resample(np.zeros((9, 9, 4)), np.zeros((9, 9)))
-    with pytest.raises(ValueError, match="3D input array must be RGBA"):
-        resample(np.zeros((9, 9, 3)), np.zeros((9, 9, 4)))
+    with pytest.raises(ValueError, match="3D input array must be RGB"):
+        resample(np.zeros((9, 9, 2)), np.zeros((9, 9, 4)))
     with pytest.raises(ValueError, match="3D output array must be RGBA"):
         resample(np.zeros((9, 9, 4)), np.zeros((9, 9, 3)))
     with pytest.raises(ValueError, match="mismatched types"):
diff --git a/src/_image_resample.h b/src/_image_resample.h
@@ -8,6 +8,7 @@
 #include "agg_image_accessors.h"
 #include "agg_path_storage.h"
 #include "agg_pixfmt_gray.h"
+#include "agg_pixfmt_rgb.h"
 #include "agg_pixfmt_rgba.h"
 #include "agg_renderer_base.h"
 #include "agg_renderer_scanline.h"
@@ -16,6 +17,7 @@
 #include "agg_span_allocator.h"
 #include "agg_span_converter.h"
 #include "agg_span_image_filter_gray.h"
+#include "agg_span_image_filter_rgb.h"
 #include "agg_span_image_filter_rgba.h"
 #include "agg_span_interpolator_adaptor.h"
 #include "agg_span_interpolator_linear.h"
@@ -496,16 +498,38 @@ typedef enum {
 } interpolation_e;
 
 
-// T is rgba if and only if it has an T::r field.
+// T is rgb(a) if and only if it has an T::r field.
 template<typename T, typename = void> struct is_grayscale : std::true_type {};
 template<typename T> struct is_grayscale<T, std::void_t<decltype(T::r)>> : std::false_type {};
 template<typename T> constexpr bool is_grayscale_v = is_grayscale<T>::value;
 
 
-template<typename color_type>
+template<typename color_type, bool input_has_alpha>
 struct type_mapping
 {
-    using blender_type = std::conditional_t<
+    using input_blender_type = std::conditional_t<
+        is_grayscale_v<color_type>,
+        agg::blender_gray<color_type>,
+        std::conditional_t<
+            input_has_alpha,
+            std::conditional_t<
+                std::is_same_v<color_type, agg::rgba8>,
+                fixed_blender_rgba_plain<color_type, agg::order_rgba>,
+                agg::blender_rgba_plain<color_type, agg::order_rgba>
+            >,
+            agg::blender_rgb<color_type, agg::order_rgb>
+        >
+    >;
+    using input_pixfmt_type = std::conditional_t<
+        is_grayscale_v<color_type>,
+        agg::pixfmt_alpha_blend_gray<input_blender_type, agg::rendering_buffer>,
+        std::conditional_t<
+            input_has_alpha,
+            agg::pixfmt_alpha_blend_rgba<input_blender_type, agg::rendering_buffer>,
+            agg::pixfmt_alpha_blend_rgb<input_blender_type, agg::rendering_buffer, 3>
+        >
+    >;
+    using output_blender_type = std::conditional_t<
         is_grayscale_v<color_type>,
         agg::blender_gray<color_type>,
         std::conditional_t<
@@ -514,25 +538,37 @@ struct type_mapping
             agg::blender_rgba_plain<color_type, agg::order_rgba>
         >
     >;
-    using pixfmt_type = std::conditional_t<
+    using output_pixfmt_type = std::conditional_t<
         is_grayscale_v<color_type>,
-        agg::pixfmt_alpha_blend_gray<blender_type, agg::rendering_buffer>,
-        agg::pixfmt_alpha_blend_rgba<blender_type, agg::rendering_buffer>
+        agg::pixfmt_alpha_blend_gray<output_blender_type, agg::rendering_buffer>,
+        agg::pixfmt_alpha_blend_rgba<output_blender_type, agg::rendering_buffer>
     >;
     template<typename A> using span_gen_affine_type = std::conditional_t<
         is_grayscale_v<color_type>,
         agg::span_image_resample_gray_affine<A>,
-        agg::span_image_resample_rgba_affine<A>
+        std::conditional_t<
+            input_has_alpha,
+            agg::span_image_resample_rgba_affine<A>,
+            agg::span_image_resample_rgb_affine<A>
+        >
     >;
     template<typename A, typename B> using span_gen_filter_type = std::conditional_t<
         is_grayscale_v<color_type>,
         agg::span_image_filter_gray<A, B>,
-        agg::span_image_filter_rgba<A, B>
+        std::conditional_t<
+            input_has_alpha,
+            agg::span_image_filter_rgba<A, B>,
+            agg::span_image_filter_rgb<A, B>
+        >
     >;
     template<typename A, typename B> using span_gen_nn_type = std::conditional_t<
         is_grayscale_v<color_type>,
         agg::span_image_filter_gray_nn<A, B>,
-        agg::span_image_filter_rgba_nn<A, B>
+        std::conditional_t<
+            input_has_alpha,
+            agg::span_image_filter_rgba_nn<A, B>,
+            agg::span_image_filter_rgb_nn<A, B>
+        >
     >;
 };
 
@@ -686,16 +722,16 @@ static void get_filter(const resample_params_t &params,
 }
 
 
-template<typename color_type>
+template<typename color_type, bool input_has_alpha = true>
 void resample(
-    const void *input, int in_width, int in_height,
-    void *output, int out_width, int out_height,
+    const void *input, int in_width, int in_height, int in_stride,
+    void *output, int out_width, int out_height, int out_stride,
     resample_params_t &params)
 {
-    using type_mapping_t = type_mapping<color_type>;
+    using type_mapping_t = type_mapping<color_type, input_has_alpha>;
 
-    using input_pixfmt_t = typename type_mapping_t::pixfmt_type;
-    using output_pixfmt_t = typename type_mapping_t::pixfmt_type;
+    using input_pixfmt_t = typename type_mapping_t::input_pixfmt_type;
+    using output_pixfmt_t = typename type_mapping_t::output_pixfmt_type;
 
     using renderer_t = agg::renderer_base<output_pixfmt_t>;
     using rasterizer_t = agg::rasterizer_scanline_aa<agg::rasterizer_sl_clip_dbl>;
@@ -711,11 +747,6 @@ void resample(
     using arbitrary_interpolator_t =
         agg::span_interpolator_adaptor<agg::span_interpolator_linear<>, lookup_distortion>;
 
-    size_t itemsize = sizeof(color_type);
-    if (is_grayscale<color_type>::value) {
-        itemsize /= 2;  // agg::grayXX includes an alpha channel which we don't have.
-    }
-
     if (params.interpolation != NEAREST &&
         params.is_affine &&
         fabs(params.affine.sx) == 1.0 &&
@@ -732,14 +763,12 @@ void resample(
     span_conv_alpha_t conv_alpha(params.alpha);
 
     agg::rendering_buffer input_buffer;
-    input_buffer.attach(
-        (unsigned char *)input, in_width, in_height, in_width * itemsize);
+    input_buffer.attach((unsigned char *)input, in_width, in_height, in_stride);
     input_pixfmt_t input_pixfmt(input_buffer);
     image_accessor_t input_accessor(input_pixfmt);
 
     agg::rendering_buffer output_buffer;
-    output_buffer.attach(
-        (unsigned char *)output, out_width, out_height, out_width * itemsize);
+    output_buffer.attach((unsigned char *)output, out_width, out_height, out_stride);
     output_pixfmt_t output_pixfmt(output_buffer);
     renderer_t renderer(output_pixfmt);
 
diff --git a/src/_image_wrapper.cpp b/src/_image_wrapper.cpp
@@ -106,10 +106,14 @@ image_resample(py::array input_array,
         throw std::invalid_argument("Input array must be a 2D or 3D array");
     }
 
-    if (ndim == 3 && input_array.shape(2) != 4) {
-        throw std::invalid_argument(
-            "3D input array must be RGBA with shape (M, N, 4), has trailing dimension of {}"_s.format(
-                input_array.shape(2)));
+    py::ssize_t ncomponents = 0;
+    if (ndim == 3) {
+        ncomponents = input_array.shape(2);
+        if (ncomponents != 3 && ncomponents != 4) {
+            throw std::invalid_argument(
+                "3D input array must be RGB with shape (M, N, 3) or RGBA with shape (M, N, 4), "
+                "has trailing dimension of {}"_s.format(ncomponents));
+        }
     }
 
     // Ensure input array is contiguous, regardless of dtype
@@ -173,25 +177,35 @@ image_resample(py::array input_array,
 
     if (auto resampler =
             (ndim == 2) ? (
-                (dtype.equal(py::dtype::of<std::uint8_t>())) ? resample<agg::gray8> :
-                (dtype.equal(py::dtype::of<std::int8_t>())) ? resample<agg::gray8> :
-                (dtype.equal(py::dtype::of<std::uint16_t>())) ? resample<agg::gray16> :
-                (dtype.equal(py::dtype::of<std::int16_t>())) ? resample<agg::gray16> :
-                (dtype.equal(py::dtype::of<float>())) ? resample<agg::gray32> :
-                (dtype.equal(py::dtype::of<double>())) ? resample<agg::gray64> :
+                dtype.equal(py::dtype::of<std::uint8_t>()) ? resample<agg::gray8> :
+                dtype.equal(py::dtype::of<std::int8_t>()) ? resample<agg::gray8> :
+                dtype.equal(py::dtype::of<std::uint16_t>()) ? resample<agg::gray16> :
+                dtype.equal(py::dtype::of<std::int16_t>()) ? resample<agg::gray16> :
+                dtype.equal(py::dtype::of<float>()) ? resample<agg::gray32> :
+                dtype.equal(py::dtype::of<double>()) ? resample<agg::gray64> :
                 nullptr) : (
-            // ndim == 3
-                (dtype.equal(py::dtype::of<std::uint8_t>())) ? resample<agg::rgba8> :
-                (dtype.equal(py::dtype::of<std::int8_t>())) ? resample<agg::rgba8> :
-                (dtype.equal(py::dtype::of<std::uint16_t>())) ? resample<agg::rgba16> :
-                (dtype.equal(py::dtype::of<std::int16_t>())) ? resample<agg::rgba16> :
-                (dtype.equal(py::dtype::of<float>())) ? resample<agg::rgba32> :
-                (dtype.equal(py::dtype::of<double>())) ? resample<agg::rgba64> :
-                nullptr)) {
+                // ndim == 3
+                (ncomponents == 4) ? (
+                    dtype.equal(py::dtype::of<std::uint8_t>()) ? resample<agg::rgba8, true> :
+                    dtype.equal(py::dtype::of<std::int8_t>()) ? resample<agg::rgba8, true> :
+                    dtype.equal(py::dtype::of<std::uint16_t>()) ? resample<agg::rgba16, true> :
+                    dtype.equal(py::dtype::of<std::int16_t>()) ? resample<agg::rgba16, true> :
+                    dtype.equal(py::dtype::of<float>()) ? resample<agg::rgba32, true> :
+                    dtype.equal(py::dtype::of<double>()) ? resample<agg::rgba64, true> :
+                    nullptr
+                ) : (
+                    dtype.equal(py::dtype::of<std::uint8_t>()) ? resample<agg::rgba8, false> :
+                    dtype.equal(py::dtype::of<std::int8_t>()) ? resample<agg::rgba8, false> :
+                    dtype.equal(py::dtype::of<std::uint16_t>()) ? resample<agg::rgba16, false> :
+                    dtype.equal(py::dtype::of<std::int16_t>()) ? resample<agg::rgba16, false> :
+                    dtype.equal(py::dtype::of<float>()) ? resample<agg::rgba32, false> :
+                    dtype.equal(py::dtype::of<double>()) ? resample<agg::rgba64, false> :
+                    nullptr)))
+    {
         Py_BEGIN_ALLOW_THREADS
         resampler(
-            input_array.data(), input_array.shape(1), input_array.shape(0),
-            output_array.mutable_data(), output_array.shape(1), output_array.shape(0),
+            input_array.data(), input_array.shape(1), input_array.shape(0), input_array.strides(0),
+            output_array.mutable_data(), output_array.shape(1), output_array.shape(0), output_array.strides(0),
             params);
         Py_END_ALLOW_THREADS
     } else {
