import_frame for anim_encoder_add

uploadcare · Sep 3, 2024 · b144bfd · b144bfd
1 parent 721ff57
commit b144bfd
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 65 deletions.
diff --git a/src/PIL/WebPImagePlugin.py b/src/PIL/WebPImagePlugin.py
@@ -248,11 +248,8 @@ def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
 
                 # Append the frame to the animation encoder
                 enc.add(
-                    frame.tobytes(),
+                    frame.im.ptr,
                     round(timestamp),
-                    frame.size[0],
-                    frame.size[1],
-                    frame.mode,
                     lossless,
                     quality,
                     alpha_quality,
@@ -270,7 +267,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)

diff --git a/src/_webp.c b/src/_webp.c
@@ -83,6 +83,44 @@ HandleMuxError(WebPMuxError err, char *chunk) {
     return NULL;
 }
 
+/* -------------------------------------------------------------------- */
+/* Frame import                                                         */
+/* -------------------------------------------------------------------- */
+
+static int
+import_frame_libwebp(WebPPicture *frame, Imaging im) {
+    UINT32 mask = 0;
+
+    if (strcmp(im->mode, "RGBA") && strcmp(im->mode, "RGB") && strcmp(im->mode, "RGBX")) {
+        PyErr_SetString(PyExc_ValueError, "unsupported image mode");
+        return -1;
+    }
+
+    if (strcmp(im->mode, "RGBA")) {
+        mask = MASK_UINT32_CHANNEL_3;
+    }
+
+    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;
+    }
+
+    for (int y = 0; y < im->ysize; ++y) {
+        UINT8 *src = (UINT8 *)im->image32[y];
+        UINT32 *dst = frame->argb + frame->argb_stride * y;
+        for (int x = 0; x < im->xsize; ++x) {
+            UINT32 pix = MAKE_UINT32(src[x * 4 + 2], src[x * 4 + 1], src[x * 4], src[x * 4 + 3]);
+            dst[x] = pix | mask;
+        }
+    }
+
+    return 0;
+}
+
 /* -------------------------------------------------------------------- */
 /* WebP Animation Support                                               */
 /* -------------------------------------------------------------------- */
@@ -180,12 +218,9 @@ _anim_encoder_dealloc(PyObject *self) {
 
 PyObject *
 _anim_encoder_add(PyObject *self, PyObject *args) {
-    uint8_t *rgb;
-    Py_ssize_t size;
+    PyObject *i0;
+    Imaging im;
     int timestamp;
-    int width;
-    int height;
-    char *mode;
     int lossless;
     float quality_factor;
     float alpha_quality_factor;
@@ -198,13 +233,9 @@ _anim_encoder_add(PyObject *self, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "z#iiisiffi",
-            (char **)&rgb,
-            &size,
+            "Oiiffi",
+            &i0,
             &timestamp,
-            &width,
-            &height,
-            &mode,
             &lossless,
             &quality_factor,
             &alpha_quality_factor,
@@ -214,11 +245,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!");
@@ -235,16 +273,8 @@ _anim_encoder_add(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    // Populate the frame with raw bytes passed to us
-    frame->width = width;
-    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);
+    if (import_frame_libwebp(frame, im)) {
+        return NULL;
     }
 
     ImagingSectionEnter(&cookie);
@@ -570,44 +600,6 @@ static PyTypeObject WebPAnimDecoder_Type = {
     0,                                 /*tp_getset*/
 };
 
-/* -------------------------------------------------------------------- */
-/* Frame import                                                         */
-/* -------------------------------------------------------------------- */
-
-static int
-import_frame_libwebp(WebPPicture *frame, Imaging im) {
-    UINT32 mask = 0;
-
-    if (strcmp(im->mode, "RGBA") && strcmp(im->mode, "RGB") && strcmp(im->mode, "RGBX")) {
-        PyErr_SetString(PyExc_ValueError, "unsupported image mode");
-        return -1;
-    }
-
-    if (strcmp(im->mode, "RGBA")) {
-        mask = MASK_UINT32_CHANNEL_3;
-    }
-
-    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;
-    }
-
-    for (int y = 0; y < im->ysize; ++y) {
-        UINT8 *src = (UINT8 *)im->image32[y];
-        UINT32 *dst = frame->argb + frame->argb_stride * y;
-        for (int x = 0; x < im->xsize; ++x) {
-            UINT32 pix = MAKE_UINT32(src[x * 4 + 2], src[x * 4 + 1], src[x * 4], src[x * 4 + 3]);
-            dst[x] = pix | mask;
-        }
-    }
-
-    return 0;
-}
-
 /* -------------------------------------------------------------------- */
 /* Legacy WebP Support                                                  */
 /* -------------------------------------------------------------------- */