Merge pull request #4547 from radarhere/webp

Added method argument to single frame WebP saving

Author: hugovk

Tests/test_file_webp.py

@@ -1,3 +1,5 @@
+import io
+
 import pytest
 from PIL import Image, WebPImagePlugin
 
@@ -54,72 +56,65 @@ def test_read_rgb(self):
             # dwebp -ppm ../../Tests/images/hopper.webp -o hopper_webp_bits.ppm
             assert_image_similar_tofile(image, "Tests/images/hopper_webp_bits.ppm", 1.0)
 
-    def test_write_rgb(self, tmp_path):
-        """
-        Can we write a RGB mode file to webp without error.
-        Does it have the bits we expect?
-        """
-
+    def _roundtrip(self, tmp_path, mode, epsilon, args={}):
         temp_file = str(tmp_path / "temp.webp")
 
-        hopper(self.rgb_mode).save(temp_file)
+        hopper(mode).save(temp_file, **args)
         with Image.open(temp_file) as image:
             assert image.mode == self.rgb_mode
             assert image.size == (128, 128)
             assert image.format == "WEBP"
             image.load()
             image.getdata()
 
-            # generated with: dwebp -ppm temp.webp -o hopper_webp_write.ppm
-            assert_image_similar_tofile(
-                image, "Tests/images/hopper_webp_write.ppm", 12.0
-            )
+            if mode == self.rgb_mode:
+                # generated with: dwebp -ppm temp.webp -o hopper_webp_write.ppm
+                assert_image_similar_tofile(
+                    image, "Tests/images/hopper_webp_write.ppm", 12.0
+                )
 
             # This test asserts that the images are similar. If the average pixel
             # difference between the two images is less than the epsilon value,
             # then we're going to accept that it's a reasonable lossy version of
-            # the image. The old lena images for WebP are showing ~16 on
-            # Ubuntu, the jpegs are showing ~18.
-            target = hopper(self.rgb_mode)
-            assert_image_similar(image, target, 12.0)
+            # the image.
+            target = hopper(mode)
+            if mode != self.rgb_mode:
+                target = target.convert(self.rgb_mode)
+            assert_image_similar(image, target, epsilon)
+
+    def test_write_rgb(self, tmp_path):
+        """
+        Can we write a RGB mode file to webp without error?
+        Does it have the bits we expect?
+        """
+
+        self._roundtrip(tmp_path, self.rgb_mode, 12.5)
+
+    def test_write_method(self, tmp_path):
+        self._roundtrip(tmp_path, self.rgb_mode, 12.0, {"method": 6})
+
+        buffer_no_args = io.BytesIO()
+        hopper().save(buffer_no_args, format="WEBP")
+
+        buffer_method = io.BytesIO()
+        hopper().save(buffer_method, format="WEBP", method=6)
+        assert buffer_no_args.getbuffer() != buffer_method.getbuffer()
 
     def test_write_unsupported_mode_L(self, tmp_path):
         """
         Saving a black-and-white file to WebP format should work, and be
         similar to the original file.
         """
 
-        temp_file = str(tmp_path / "temp.webp")
-        hopper("L").save(temp_file)
-        with Image.open(temp_file) as image:
-            assert image.mode == self.rgb_mode
-            assert image.size == (128, 128)
-            assert image.format == "WEBP"
-
-            image.load()
-            image.getdata()
-            target = hopper("L").convert(self.rgb_mode)
-
-            assert_image_similar(image, target, 10.0)
+        self._roundtrip(tmp_path, "L", 10.0)
 
     def test_write_unsupported_mode_P(self, tmp_path):
         """
         Saving a palette-based file to WebP format should work, and be
         similar to the original file.
         """
 
-        temp_file = str(tmp_path / "temp.webp")
-        hopper("P").save(temp_file)
-        with Image.open(temp_file) as image:
-            assert image.mode == self.rgb_mode
-            assert image.size == (128, 128)
-            assert image.format == "WEBP"
-
-            image.load()
-            image.getdata()
-            target = hopper("P").convert(self.rgb_mode)
-
-            assert_image_similar(image, target, 50.0)
+        self._roundtrip(tmp_path, "P", 50.0)
 
     def test_WebPEncode_with_invalid_args(self):
         """

src/PIL/WebPImagePlugin.py

@@ -314,6 +314,7 @@ def _save(im, fp, filename):
     if isinstance(exif, Image.Exif):
         exif = exif.tobytes()
     xmp = im.encoderinfo.get("xmp", "")
+    method = im.encoderinfo.get("method", 0)
 
     if im.mode not in _VALID_WEBP_LEGACY_MODES:
         alpha = (
@@ -331,6 +332,7 @@ def _save(im, fp, filename):
         float(quality),
         im.mode,
         icc_profile,
+        method,
         exif,
         xmp,
     )

src/_webp.c

@@ -545,6 +545,7 @@ PyObject* WebPEncode_wrapper(PyObject* self, PyObject* args)
     int height;
     int lossless;
     float quality_factor;
+    int method;
     uint8_t* rgb;
     uint8_t* icc_bytes;
     uint8_t* exif_bytes;
@@ -556,49 +557,75 @@ PyObject* WebPEncode_wrapper(PyObject* self, PyObject* args)
     Py_ssize_t exif_size;
     Py_ssize_t xmp_size;
     size_t ret_size;
+    int rgba_mode;
+    int channels;
+    int ok;
     ImagingSectionCookie cookie;
+    WebPConfig config;
+    WebPMemoryWriter writer;
+    WebPPicture pic;
 
-    if (!PyArg_ParseTuple(args, "y#iiifss#s#s#",
+    if (!PyArg_ParseTuple(args, "y#iiifss#is#s#",
                 (char**)&rgb, &size, &width, &height, &lossless, &quality_factor, &mode,
-                &icc_bytes, &icc_size, &exif_bytes, &exif_size, &xmp_bytes, &xmp_size)) {
+                &icc_bytes, &icc_size, &method, &exif_bytes, &exif_size, &xmp_bytes, &xmp_size)) {
         return NULL;
     }
-    if (strcmp(mode, "RGBA")==0){
-        if (size < width * height * 4){
-            Py_RETURN_NONE;
-        }
-        #if WEBP_ENCODER_ABI_VERSION >= 0x0100
-        if (lossless) {
-            ImagingSectionEnter(&cookie);
-            ret_size = WebPEncodeLosslessRGBA(rgb, width, height, 4 * width, &output);
-            ImagingSectionLeave(&cookie);
-        } else
-        #endif
-        {
-            ImagingSectionEnter(&cookie);
-            ret_size = WebPEncodeRGBA(rgb, width, height, 4 * width, quality_factor, &output);
-            ImagingSectionLeave(&cookie);
-        }
-    } else if (strcmp(mode, "RGB")==0){
-        if (size < width * height * 3){
-            Py_RETURN_NONE;
-        }
-        #if WEBP_ENCODER_ABI_VERSION >= 0x0100
-        if (lossless) {
-            ImagingSectionEnter(&cookie);
-            ret_size = WebPEncodeLosslessRGB(rgb, width, height, 3 * width, &output);
-            ImagingSectionLeave(&cookie);
-        } else
-        #endif
-        {
-            ImagingSectionEnter(&cookie);
-            ret_size = WebPEncodeRGB(rgb, width, height, 3 * width, quality_factor, &output);
-            ImagingSectionLeave(&cookie);
-        }
-    } else {
+
+    rgba_mode = strcmp(mode, "RGBA") == 0;
+    if (!rgba_mode && strcmp(mode, "RGB") != 0) {
         Py_RETURN_NONE;
     }
 
+    channels = rgba_mode ? 4 : 3;
+    if (size < width * height * channels) {
+        Py_RETURN_NONE;
+    }
+
+    // Setup config for this frame
+    if (!WebPConfigInit(&config)) {
+        PyErr_SetString(PyExc_RuntimeError, "failed to initialize config!");
+        return NULL;
+    }
+    config.lossless = lossless;
+    config.quality = quality_factor;
+    config.method = method;
+
+    // Validate the config
+    if (!WebPValidateConfig(&config)) {
+        PyErr_SetString(PyExc_ValueError, "invalid configuration");
+        return NULL;
+    }
+
+    if (!WebPPictureInit(&pic)) {
+        PyErr_SetString(PyExc_ValueError, "could not initialise picture");
+        return NULL;
+    }
+    pic.width = width;
+    pic.height = height;
+    pic.use_argb = 1; // Don't convert RGB pixels to YUV
+
+    if (rgba_mode) {
+        WebPPictureImportRGBA(&pic, rgb, channels * width);
+    } else {
+        WebPPictureImportRGB(&pic, rgb, channels * width);
+    }
+
+    WebPMemoryWriterInit(&writer);
+    pic.writer = WebPMemoryWrite;
+    pic.custom_ptr = &writer;
+
+    ImagingSectionEnter(&cookie);
+    ok = WebPEncode(&config, &pic);
+    ImagingSectionLeave(&cookie);
+
+    WebPPictureFree(&pic);
+    if (!ok) {
+        PyErr_SetString(PyExc_ValueError, "encoding error");
+        return NULL;
+    }
+    output = writer.mem;
+    ret_size = writer.size;
+
 #ifndef HAVE_WEBPMUX
     if (ret_size > 0) {
         PyObject *ret = PyBytes_FromStringAndSize((char*)output, ret_size);