Added support for alpha transparent webp images #204

PIL/WebPImagePlugin.py

@@ -3,8 +3,38 @@ from PIL import ImageFile
 from io import BytesIO
 import _webp
 
+
+_VALID_WEBP_ENCODERS_BY_MODE = {
+    "RGB": _webp.WebPEncodeRGB,
+    "RGBA": _webp.WebPEncodeRGBA,
+    }
+
+
+_VALID_WEBP_DECODERS_BY_MODE = {
+    "RGB": _webp.WebPDecodeRGB,
+    "RGBA": _webp.WebPDecodeRGBA,
+    }
+
+
+_STRIDE_MULTIPLIERS_BY_MODE = {
+    "RGB": 3,
+    "RGBA": 4,
+    }
+
+
+_VP8_MODES_BY_IDENTIFIER = {
+    b"VP8 ": "RGB",
+    b"VP8X": "RGBA",
+    } 
+
+
 def _accept(prefix):
-    return prefix[:4] == b"RIFF" and prefix[8:16] == b"WEBPVP8 "
+    is_riff_file_format = prefix[:4] == b"RIFF"
+    is_webp_file = prefix[8:12] == b"WEBP"
+    is_valid_vp8_mode = prefix[12:16] in _VP8_MODES_BY_IDENTIFIER
+    
+    return is_riff_file_format and is_webp_file and is_valid_vp8_mode
+
 
 class WebPImageFile(ImageFile.ImageFile):
 
@@ -12,20 +42,44 @@ class WebPImageFile(ImageFile.ImageFile):
     format_description = "WebP image"
 
     def _open(self):
-        self.mode = "RGB"
-        data, width, height = _webp.WebPDecodeRGB(self.fp.read())
+        file_header = self.fp.read(16)
+        vp8_header = file_header[12:16]
+        try:
+            webp_file_mode = _VP8_MODES_BY_IDENTIFIER[vp8_header]
+        except KeyError:
+            raise IOError("Unknown webp file mode")
+        finally:
+            self.fp.seek(0)
+        
+        self.mode = webp_file_mode
+        webp_decoder = _VALID_WEBP_DECODERS_BY_MODE[webp_file_mode]
+        
+        data, width, height = webp_decoder(self.fp.read())
         self.size = width, height
         self.fp = BytesIO(data)
-        self.tile = [("raw", (0, 0) + self.size, 0, 'RGB')]
+        self.tile = [("raw", (0, 0) + self.size, 0, webp_file_mode)]
+
 
 def _save(im, fp, filename):
-    if im.mode != "RGB":
-        raise IOError("cannot write mode %s as WEBP" % im.mode)
+    image_mode = im.mode
+    if image_mode not in _VALID_WEBP_ENCODERS_BY_MODE:
+        raise IOError("cannot write mode %s as WEBP" % image_mode)
+
+    webp_encoder = _VALID_WEBP_ENCODERS_BY_MODE[image_mode]
+    
+    stride = im.size[0] * _STRIDE_MULTIPLIERS_BY_MODE[image_mode]
     quality = im.encoderinfo.get("quality", 80)
     
-    data = _webp.WebPEncodeRGB(im.tobytes(), im.size[0], im.size[1], im.size[0] * 3, float(quality))
+    data = webp_encoder(
+        im.tobytes(),
+        im.size[0],
+        im.size[1],
+        stride,
+        float(quality),
+        )
     fp.write(data)
 
+
 Image.register_open("WEBP", WebPImageFile, _accept)
 Image.register_save("WEBP", _save)
 

Tests/test_file_webp.py

@@ -6,54 +6,95 @@ try:
     import _webp
 except:
     skip('webp support not installed')
-		
-def test_read():
-    """ Can we write a webp without error. Does it have the bits we expect?"""
-    
-    file = "Images/lena.webp"
-    im = Image.open(file)
 
-    assert_equal(im.mode, "RGB")
-    assert_equal(im.size, (128, 128))
-    assert_equal(im.format, "WEBP")
-    assert_no_exception(lambda: im.load())
-    assert_no_exception(lambda: im.getdata())
-
-    orig_bytes  = im.tobytes()
 
+def test_read_rgb():
+     
+    file_path = "Images/lena.webp"
+    image = Image.open(file_path)
+ 
+    assert_equal(image.mode, "RGB")
+    assert_equal(image.size, (128, 128))
+    assert_equal(image.format, "WEBP")
+    assert_no_exception(lambda: image.load())
+    assert_no_exception(lambda: image.getdata())
+ 
     # generated with: dwebp -ppm ../../Images/lena.webp -o lena_webp_bits.ppm
     target = Image.open('Tests/images/lena_webp_bits.ppm')
-    assert_image_equal(im, target)
+    assert_image_equal(image, target)
     
 
-def test_write():
-    """ Can we write a webp without error. Does it have the bits we expect?"""
+def test_read_rgba():
+    # Generated with `cwebp transparent.png -o transparent.webp`
+    file_path = "Images/transparent.webp"
+    image = Image.open(file_path)
 
-    file = tempfile("temp.webp")
+    assert_equal(image.mode, "RGBA")
+    assert_equal(image.size, (200, 150))
+    assert_equal(image.format, "WEBP")
+    assert_no_exception(lambda: image.load())
+    assert_no_exception(lambda: image.getdata())
 
-    lena("RGB").save(file)
+    orig_bytes  = image.tobytes()
 
-    im= Image.open(file)
-    im.load()
-
-    assert_equal(im.mode, "RGB")
-    assert_equal(im.size, (128, 128))
-    assert_equal(im.format, "WEBP")
-    assert_no_exception(lambda: im.load())
-    assert_no_exception(lambda: im.getdata())
+    target = Image.open('Images/transparent.png')
+    assert_image_similar(image, target, 20.0)
+    
 
+def test_write_rgb():
+    """
+    Can we write a RGB mode file to webp without error. Does it have the bits we
+    expect?
+      
+    """
+  
+    temp_file = tempfile("temp.webp")
+  
+    lena("RGB").save(temp_file)
+  
+    image = Image.open(temp_file)
+    image.load()
+  
+    assert_equal(image.mode, "RGB")
+    assert_equal(image.size, (128, 128))
+    assert_equal(image.format, "WEBP")
+    assert_no_exception(lambda: image.load())
+    assert_no_exception(lambda: image.getdata())
+  
     # If we're using the exact same version of webp, this test should pass.
     # but it doesn't if the webp is generated on Ubuntu and tested on Fedora.
-    
+      
     # generated with: dwebp -ppm temp.webp -o lena_webp_write.ppm
     #target = Image.open('Tests/images/lena_webp_write.ppm')
-    #assert_image_equal(im, target)
-
+    #assert_image_equal(image, target)
+  
     # 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 included lena images
     # for webp are showing ~16 on Ubuntu, the jpegs are showing ~18. 
-    target = lena('RGB')
-    assert_image_similar(im, target, 20.0)
-
-
+    target = lena("RGB")
+    assert_image_similar(image, target, 20.0)
+ 
+ 
+def test_write_rgba():
+    """
+    Can we write a RGBA mode file to webp without error. Does it have the bits we
+    expect?
+     
+    """
+ 
+    temp_file = tempfile("temp.webp")
+ 
+    pil_image = Image.new("RGBA", (10, 10), (255, 0, 0, 20))
+    pil_image.save(temp_file)
+ 
+    image = Image.open(temp_file)
+    image.load()
+ 
+    assert_equal(image.mode, "RGBA")
+    assert_equal(image.size, (10, 10))
+    assert_equal(image.format, "WEBP")
+    assert_no_exception(image.load)
+    assert_no_exception(image.getdata)
+ 
+    assert_image_similar(image, pil_image, 1.0)

_webp.c

@@ -38,6 +38,43 @@ PyObject* WebPEncodeRGB_wrapper(PyObject* self, PyObject* args)
 
 }
 
+
+PyObject* WebPEncodeRGBA_wrapper(PyObject* self, PyObject* args)
+{
+    PyBytesObject *rgba_string;
+    int width;
+    int height;
+    int stride;
+    float quality_factor;
+    uint8_t *rgba;
+    uint8_t *output;
+    Py_ssize_t size;
+    size_t ret_size;
+
+    if (!PyArg_ParseTuple(args, "Siiif", &rgba_string, &width, &height, &stride, &quality_factor)) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+
+    PyBytes_AsStringAndSize((PyObject *) rgba_string, (char**)&rgba, &size);
+
+    if (stride * height > size) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+
+    ret_size = WebPEncodeRGBA(rgba, width, height, stride, quality_factor, &output);
+    if (ret_size > 0) {
+        PyObject *ret = PyBytes_FromStringAndSize((char*)output, ret_size);
+        free(output);
+        return ret;
+    }
+    Py_INCREF(Py_None);
+    return Py_None;
+
+}
+
+
 PyObject* WebPDecodeRGB_wrapper(PyObject* self, PyObject* args)
 {
     PyBytesObject *webp_string;
@@ -62,13 +99,42 @@ PyObject* WebPDecodeRGB_wrapper(PyObject* self, PyObject* args)
     return Py_BuildValue("Sii", ret, width, height);
 }
 
+
+PyObject* WebPDecodeRGBA_wrapper(PyObject* self, PyObject* args)
+{
+    PyBytesObject *webp_string;
+    int width;
+    int height;
+    uint8_t *webp;
+    uint8_t *output;
+    Py_ssize_t size;
+    PyObject *ret;
+
+    if (!PyArg_ParseTuple(args, "S", &webp_string)) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+
+    PyBytes_AsStringAndSize((PyObject *) webp_string, (char**)&webp, &size);
+
+    output = WebPDecodeRGBA(webp, size, &width, &height);
+
+    ret = PyBytes_FromStringAndSize((char*)output, width * height * 4);
+    free(output);
+    return Py_BuildValue("Sii", ret, width, height);
+}
+
+
 static PyMethodDef webpMethods[] =
 {
     {"WebPEncodeRGB", WebPEncodeRGB_wrapper, METH_VARARGS, "WebPEncodeRGB"},
-    {"WebPDecodeRGB", WebPDecodeRGB_wrapper, METH_VARARGS, "WebPEncodeRGB"},
+    {"WebPEncodeRGBA", WebPEncodeRGBA_wrapper, METH_VARARGS, "WebPEncodeRGBA"},
+    {"WebPDecodeRGB", WebPDecodeRGB_wrapper, METH_VARARGS, "WebPDecodeRGB"},
+    {"WebPDecodeRGBA", WebPDecodeRGBA_wrapper, METH_VARARGS, "WebPDecodeRGBA"},
     {NULL, NULL}
 };
 
+
 #if PY_VERSION_HEX >= 0x03000000
 PyMODINIT_FUNC
 PyInit__webp(void) {