Merge pull request #8386 from uploadcare/webp-capsule

Use Capsule for WEBP saving

Tests/test_file_webp.py

@@ -72,7 +72,7 @@ class TestFileWebp:
     def _roundtrip(
         self, tmp_path: Path, mode: str, epsilon: float, args: dict[str, Any] = {}
     ) -> None:
-        temp_file = str(tmp_path / "temp.webp")
+        temp_file = tmp_path / "temp.webp"
 
         hopper(mode).save(temp_file, **args)
         with Image.open(temp_file) as image:
@@ -116,7 +116,7 @@ class TestFileWebp:
         assert buffer_no_args.getbuffer() != buffer_method.getbuffer()
 
     def test_save_all(self, tmp_path: Path) -> None:
-        temp_file = str(tmp_path / "temp.webp")
+        temp_file = tmp_path / "temp.webp"
         im = Image.new("RGB", (1, 1))
         im2 = Image.new("RGB", (1, 1), "#f00")
         im.save(temp_file, save_all=True, append_images=[im2])
@@ -127,6 +127,11 @@ class TestFileWebp:
             reloaded.seek(1)
             assert_image_similar(im2, reloaded, 1)
 
+    def test_unsupported_image_mode(self) -> None:
+        im = Image.new("1", (1, 1))
+        with pytest.raises(ValueError):
+            _webp.WebPEncode(im.getim(), False, 0, 0, "", 4, 0, b"", "")
+
     def test_icc_profile(self, tmp_path: Path) -> None:
         self._roundtrip(tmp_path, self.rgb_mode, 12.5, {"icc_profile": None})
         self._roundtrip(
@@ -151,18 +156,16 @@ class TestFileWebp:
 
     @pytest.mark.skipif(sys.maxsize <= 2**32, reason="Requires 64-bit system")
     def test_write_encoding_error_message(self, tmp_path: Path) -> None:
-        temp_file = str(tmp_path / "temp.webp")
         im = Image.new("RGB", (15000, 15000))
         with pytest.raises(ValueError) as e:
-            im.save(temp_file, method=0)
+            im.save(tmp_path / "temp.webp", method=0)
         assert str(e.value) == "encoding error 6"
 
     @pytest.mark.skipif(sys.maxsize <= 2**32, reason="Requires 64-bit system")
     def test_write_encoding_error_bad_dimension(self, tmp_path: Path) -> None:
-        temp_file = str(tmp_path / "temp.webp")
         im = Image.new("L", (16384, 16384))
         with pytest.raises(ValueError) as e:
-            im.save(temp_file)
+            im.save(tmp_path / "temp.webp")
         assert (
             str(e.value)
             == "encoding error 5: Image size exceeds WebP limit of 16383 pixels"
@@ -187,9 +190,8 @@ class TestFileWebp:
     def test_no_resource_warning(self, tmp_path: Path) -> None:
         file_path = "Tests/images/hopper.webp"
         with Image.open(file_path) as image:
-            temp_file = str(tmp_path / "temp.webp")
             with warnings.catch_warnings():
-                image.save(temp_file)
+                image.save(tmp_path / "temp.webp")
 
     def test_file_pointer_could_be_reused(self) -> None:
         file_path = "Tests/images/hopper.webp"
@@ -204,15 +206,16 @@ class TestFileWebp:
     def test_invalid_background(
         self, background: int | tuple[int, ...], tmp_path: Path
     ) -> None:
-        temp_file = str(tmp_path / "temp.webp")
+        temp_file = tmp_path / "temp.webp"
         im = hopper()
         with pytest.raises(OSError):
             im.save(temp_file, save_all=True, append_images=[im], background=background)
 
     def test_background_from_gif(self, tmp_path: Path) -> None:
+        out_webp = tmp_path / "temp.webp"
+
         # Save L mode GIF with background
         with Image.open("Tests/images/no_palette_with_background.gif") as im:
-            out_webp = str(tmp_path / "temp.webp")
             im.save(out_webp, save_all=True)
 
         # Save P mode GIF with background
@@ -220,11 +223,10 @@ class TestFileWebp:
             original_value = im.convert("RGB").getpixel((1, 1))
 
             # Save as WEBP
-            out_webp = str(tmp_path / "temp.webp")
             im.save(out_webp, save_all=True)
 
         # Save as GIF
-        out_gif = str(tmp_path / "temp.gif")
+        out_gif = tmp_path / "temp.gif"
         with Image.open(out_webp) as im:
             im.save(out_gif)
 
@@ -234,10 +236,10 @@ class TestFileWebp:
         assert difference < 5
 
     def test_duration(self, tmp_path: Path) -> None:
+        out_webp = tmp_path / "temp.webp"
+
         with Image.open("Tests/images/dispose_bgnd.gif") as im:
             assert im.info["duration"] == 1000
-
-            out_webp = str(tmp_path / "temp.webp")
             im.save(out_webp, save_all=True)
 
         with Image.open(out_webp) as reloaded:
@@ -245,7 +247,7 @@ class TestFileWebp:
             assert reloaded.info["duration"] == 1000
 
     def test_roundtrip_rgba_palette(self, tmp_path: Path) -> None:
-        temp_file = str(tmp_path / "temp.webp")
+        temp_file = tmp_path / "temp.webp"
         im = Image.new("RGBA", (1, 1)).convert("P")
         assert im.mode == "P"
         assert im.palette is not None

src/PIL/WebPImagePlugin.py

@@ -13,10 +13,6 @@ except ImportError:
     SUPPORTED = False
 
 
-_VALID_WEBP_MODES = {"RGBX": True, "RGBA": True, "RGB": True}
-
-_VALID_WEBP_LEGACY_MODES = {"RGB": True, "RGBA": True}
-
 _VP8_MODES_BY_IDENTIFIER = {
     b"VP8 ": "RGB",
     b"VP8X": "RGBA",
@@ -153,6 +149,13 @@ class WebPImageFile(ImageFile.ImageFile):
         return self.__logical_frame
 
 
+def _convert_frame(im: Image.Image) -> Image.Image:
+    # Make sure image mode is supported
+    if im.mode not in ("RGBX", "RGBA", "RGB"):
+        im = im.convert("RGBA" if im.has_transparency_data else "RGB")
+    return im
+
+
 def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     encoderinfo = im.encoderinfo.copy()
     append_images = list(encoderinfo.get("append_images", []))
@@ -243,31 +246,13 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
 
             for idx in range(nfr):
                 ims.seek(idx)
-                ims.load()
 
-                # Make sure image mode is supported
+                frame = _convert_frame(ims)
-                frame = ims
-                rawmode = ims.mode
-                if ims.mode not in _VALID_WEBP_MODES:
-                    alpha = (
-                        "A" in ims.mode
-                        or "a" in ims.mode
-                        or (ims.mode == "P" and "A" in ims.im.getpalettemode())
-                    )
-                    rawmode = "RGBA" if alpha else "RGB"
-                    frame = ims.convert(rawmode)
-
-                if rawmode == "RGB":
-                    # For faster conversion, use RGBX
-                    rawmode = "RGBX"
 
                 # Append the frame to the animation encoder
                 enc.add(
-                    frame.tobytes("raw", rawmode),
+                    frame.getim(),
                     round(timestamp),
-                    frame.size[0],
-                    frame.size[1],
-                    rawmode,
                     lossless,
                     quality,
                     alpha_quality,
@@ -285,7 +270,7 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
         im.seek(cur_idx)
 
     # Force encoder to flush frames
-    enc.add(None, round(timestamp), 0, 0, "", lossless, quality, alpha_quality, 0)
+    enc.add(None, round(timestamp), lossless, quality, alpha_quality, 0)
 
     # Get the final output from the encoder
     data = enc.assemble(icc_profile, exif, xmp)
@@ -310,17 +295,13 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
     method = im.encoderinfo.get("method", 4)
     exact = 1 if im.encoderinfo.get("exact") else 0
 
-    if im.mode not in _VALID_WEBP_LEGACY_MODES:
+    im = _convert_frame(im)
-        im = im.convert("RGBA" if im.has_transparency_data else "RGB")
 
     data = _webp.WebPEncode(
-        im.tobytes(),
+        im.getim(),
-        im.size[0],
-        im.size[1],
         lossless,
         float(quality),
         float(alpha_quality),
-        im.mode,
         icc_profile,
         method,
         exact,

src/_webp.c

@@ -83,6 +83,49 @@ HandleMuxError(WebPMuxError err, char *chunk) {
     return NULL;
 }
 
+/* -------------------------------------------------------------------- */
+/* Frame import                                                         */
+/* -------------------------------------------------------------------- */
+
+static int
+import_frame_libwebp(WebPPicture *frame, Imaging im) {
+    if (strcmp(im->mode, "RGBA") && strcmp(im->mode, "RGB") &&
+        strcmp(im->mode, "RGBX")) {
+        PyErr_SetString(PyExc_ValueError, "unsupported image mode");
+        return -1;
+    }
+
+    frame->width = im->xsize;
+    frame->height = im->ysize;
+    frame->use_argb = 1;  // Don't convert RGB pixels to YUV
+
+    if (!WebPPictureAlloc(frame)) {
+        PyErr_SetString(PyExc_MemoryError, "can't allocate picture frame");
+        return -2;
+    }
+
+    int ignore_fourth_channel = strcmp(im->mode, "RGBA");
+    for (int y = 0; y < im->ysize; ++y) {
+        UINT8 *src = (UINT8 *)im->image32[y];
+        UINT32 *dst = frame->argb + frame->argb_stride * y;
+        if (ignore_fourth_channel) {
+            for (int x = 0; x < im->xsize; ++x) {
+                dst[x] =
+                    ((UINT32)(src[x * 4 + 2]) | ((UINT32)(src[x * 4 + 1]) << 8) |
+                     ((UINT32)(src[x * 4]) << 16) | (0xff << 24));
+            }
+        } else {
+            for (int x = 0; x < im->xsize; ++x) {
+                dst[x] =
+                    ((UINT32)(src[x * 4 + 2]) | ((UINT32)(src[x * 4 + 1]) << 8) |
+                     ((UINT32)(src[x * 4]) << 16) | ((UINT32)(src[x * 4 + 3]) << 24));
+            }
+        }
+    }
+
+    return 0;
+}
+
 /* -------------------------------------------------------------------- */
 /* WebP Animation Support                                               */
 /* -------------------------------------------------------------------- */
@@ -180,16 +223,14 @@ _anim_encoder_dealloc(PyObject *self) {
 
 PyObject *
 _anim_encoder_add(PyObject *self, PyObject *args) {
-    uint8_t *rgb;
+    PyObject *i0;
-    Py_ssize_t size;
+    Imaging im;
     int timestamp;
-    int width;
-    int height;
-    char *mode;
     int lossless;
     float quality_factor;
     float alpha_quality_factor;
     int method;
+    ImagingSectionCookie cookie;
     WebPConfig config;
     WebPAnimEncoderObject *encp = (WebPAnimEncoderObject *)self;
     WebPAnimEncoder *enc = encp->enc;
@@ -197,13 +238,9 @@ _anim_encoder_add(PyObject *self, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "z#iiisiffi",
+            "Oiiffi",
-            (char **)&rgb,
+            &i0,
-            &size,
             &timestamp,
-            &width,
-            &height,
-            &mode,
             &lossless,
             &quality_factor,
             &alpha_quality_factor,
@@ -213,11 +250,18 @@ _anim_encoder_add(PyObject *self, PyObject *args) {
     }
 
     // Check for NULL frame, which sets duration of final frame
-    if (!rgb) {
+    if (i0 == Py_None) {
         WebPAnimEncoderAdd(enc, NULL, timestamp, NULL);
         Py_RETURN_NONE;
     }
 
+    if (!PyCapsule_IsValid(i0, IMAGING_MAGIC)) {
+        PyErr_Format(PyExc_TypeError, "Expected '%s' Capsule", IMAGING_MAGIC);
+        return NULL;
+    }
+
+    im = (Imaging)PyCapsule_GetPointer(i0, IMAGING_MAGIC);
+
     // Setup config for this frame
     if (!WebPConfigInit(&config)) {
         PyErr_SetString(PyExc_RuntimeError, "failed to initialize config!");
@@ -234,20 +278,15 @@ _anim_encoder_add(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    // Populate the frame with raw bytes passed to us
+    if (import_frame_libwebp(frame, im)) {
-    frame->width = width;
+        return NULL;
-    frame->height = height;
-    frame->use_argb = 1;  // Don't convert RGB pixels to YUV
-    if (strcmp(mode, "RGBA") == 0) {
-        WebPPictureImportRGBA(frame, rgb, 4 * width);
-    } else if (strcmp(mode, "RGBX") == 0) {
-        WebPPictureImportRGBX(frame, rgb, 4 * width);
-    } else {
-        WebPPictureImportRGB(frame, rgb, 3 * width);
     }
 
-    // Add the frame to the encoder
+    ImagingSectionEnter(&cookie);
-    if (!WebPAnimEncoderAdd(enc, frame, timestamp, &config)) {
+    int ok = WebPAnimEncoderAdd(enc, frame, timestamp, &config);
+    ImagingSectionLeave(&cookie);
+
+    if (!ok) {
         PyErr_SetString(PyExc_RuntimeError, WebPAnimEncoderGetError(enc));
         return NULL;
     }
@@ -572,26 +611,21 @@ static PyTypeObject WebPAnimDecoder_Type = {
 
 PyObject *
 WebPEncode_wrapper(PyObject *self, PyObject *args) {
-    int width;
-    int height;
     int lossless;
     float quality_factor;
     float alpha_quality_factor;
     int method;
     int exact;
-    uint8_t *rgb;
+    Imaging im;
+    PyObject *i0;
     uint8_t *icc_bytes;
     uint8_t *exif_bytes;
     uint8_t *xmp_bytes;
     uint8_t *output;
-    char *mode;
-    Py_ssize_t size;
     Py_ssize_t icc_size;
     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;
@@ -600,15 +634,11 @@ WebPEncode_wrapper(PyObject *self, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "y#iiiffss#iis#s#",
+            "Oiffs#iis#s#",
-            (char **)&rgb,
+            &i0,
-            &size,
-            &width,
-            &height,
             &lossless,
             &quality_factor,
             &alpha_quality_factor,
-            &mode,
             &icc_bytes,
             &icc_size,
             &method,
@@ -621,15 +651,12 @@ WebPEncode_wrapper(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    rgba_mode = strcmp(mode, "RGBA") == 0;
+    if (!PyCapsule_IsValid(i0, IMAGING_MAGIC)) {
-    if (!rgba_mode && strcmp(mode, "RGB") != 0) {
+        PyErr_Format(PyExc_TypeError, "Expected '%s' Capsule", IMAGING_MAGIC);
-        Py_RETURN_NONE;
+        return NULL;
     }
 
-    channels = rgba_mode ? 4 : 3;
+    im = (Imaging)PyCapsule_GetPointer(i0, IMAGING_MAGIC);
-    if (size < width * height * channels) {
-        Py_RETURN_NONE;
-    }
 
     // Setup config for this frame
     if (!WebPConfigInit(&config)) {
@@ -652,14 +679,9 @@ WebPEncode_wrapper(PyObject *self, PyObject *args) {
         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) {
+    if (import_frame_libwebp(&pic, im)) {
-        WebPPictureImportRGBA(&pic, rgb, channels * width);
+        return NULL;
-    } else {
-        WebPPictureImportRGB(&pic, rgb, channels * width);
     }
 
     WebPMemoryWriterInit(&writer);