Merge pull request #46 from hMatoba/dev_webp

Dev webp

.gitignore

@@ -100,4 +100,5 @@ load_sample.py
 *.suo
 doc/_build/
 run_coverage.bat
-up2pypi.bat
+up2pypi.bat
+out/

.travis.yml

@@ -11,7 +11,7 @@ python:
  - 3.6
 
 install:
- - pip install pillow==4.0.0
+ - pip install pillow==5.0.0
  - pip install coveralls
 
 script:

doc/changes.rst

@@ -1,6 +1,11 @@
 Changelog
 =========
 
+1.1.0b
+------
+
+- "load", "insert", and "remove" support WebP format.
+
 1.0.13
 ------
 

piexif.pyproj

@@ -29,6 +29,7 @@
     <Folder Include="tests\" />
     <Folder Include="piexif\" />
     <Folder Include="tests\images\" />
+    <Folder Include="tests\images\out\" />
   </ItemGroup>
   <ItemGroup>
     <Compile Include="piexif\_common.py" />
@@ -46,6 +47,9 @@
     </Compile>
     <Compile Include="piexif\_remove.py" />
     <Compile Include="piexif\_transplant.py" />
+    <Compile Include="piexif\_webp.py">
+      <SubType>Code</SubType>
+    </Compile>
     <Compile Include="piexif\__init__.py" />
     <Compile Include="tests\s_test.py" />
   </ItemGroup>
@@ -57,6 +61,12 @@
     <Content Include="tests\images\L02.jpg" />
     <Content Include="tests\images\noapp01.jpg" />
     <Content Include="tests\images\noexif.jpg" />
+    <Content Include="tests\images\pil1.webp" />
+    <Content Include="tests\images\pil2.webp" />
+    <Content Include="tests\images\pil3.webp" />
+    <Content Include="tests\images\pil_rgb.webp" />
+    <Content Include="tests\images\pil_rgba.webp" />
+    <Content Include="tests\images\tool1.webp" />
   </ItemGroup>
   <!-- Uncomment the CoreCompile target to enable the Build command in
        Visual Studio and specify your pre- and post-build commands in

piexif/__init__.py

@@ -7,4 +7,5 @@
 from ._exceptions import *
 
 
-VERSION = '1.0.13'
+
+VERSION = '1.1.0b'

piexif/_insert.py

@@ -3,7 +3,7 @@
 
 from ._common import *
 from ._exceptions import InvalidImageDataError
-
+from piexif import _webp
 
 def insert(exif, image, new_file=None):
     """
@@ -21,14 +21,26 @@ def insert(exif, image, new_file=None):
     output_file = False
     if image[0:2] == b"\xff\xd8":
         image_data = image
+        file_type = "jpeg"
+    elif image[0:4] == b"RIFF" and image[8:12] == b"WEBP":
+        image_data = image
+        file_type = "webp"
     else:
         with open(image, 'rb') as f:
             image_data = f.read()
-        if image_data[0:2] != b"\xff\xd8":
+        if image_data[0:2] == b"\xff\xd8":
+            file_type = "jpeg"
+        elif image_data[0:4] == b"RIFF" and image_data[8:12] == b"WEBP":
+            file_type = "webp"
+        else:
             raise InvalidImageDataError
         output_file = True
-    segments = split_into_segments(image_data)
-    new_data = merge_segments(segments, exif)
+
+    if file_type == "jpeg":
+        segments = split_into_segments(image_data)
+        new_data = merge_segments(segments, exif)
+    elif file_type == "webp":
+        new_data = _webp.insert(image_data, exif)
 
     if isinstance(new_file, io.BytesIO):
         new_file.write(new_data)

piexif/_load.py

@@ -3,7 +3,7 @@
 from ._common import *
 from ._exceptions import InvalidImageDataError
 from ._exif import *
-
+from piexif import _webp
 
 LITTLE_ENDIAN = b"\x49\x49"
 
@@ -73,14 +73,13 @@ def __init__(self, data):
                 self.tiftag = None
         elif data[0:2] in (b"\x49\x49", b"\x4d\x4d"):  # TIFF
             self.tiftag = data
+        elif data[0:4] == b"RIFF" and data[8:12] == b"WEBP":
+            self.tiftag = _webp.get_exif(data)
         elif data[0:4] == b"Exif":  # Exif
             self.tiftag = data[6:]
         else:
-            try:
-                with open(data, 'rb') as f:
-                    magic_number = f.read(2)
-            except:
-                raise ValueError("Got invalid value.")
+            with open(data, 'rb') as f:
+                magic_number = f.read(2)
             if magic_number == b"\xff\xd8":  # JPEG
                 app1 = read_exif_from_file(data)
                 if app1:
@@ -91,7 +90,14 @@ def __init__(self, data):
                 with open(data, 'rb') as f:
                     self.tiftag = f.read()
             else:
-                raise InvalidImageDataError("Given file is neither JPEG nor TIFF.")
+                with open(data, 'rb') as f:
+                    header = f.read(12)
+                if header[0:4] == b"RIFF"and header[8:12] == b"WEBP":
+                    with open(data, 'rb') as f:
+                        file_data = f.read()
+                    self.tiftag = _webp.get_exif(file_data)
+                else:
+                    raise InvalidImageDataError("Given file is neither JPEG nor TIFF.")
 
     def get_ifd_dict(self, pointer, ifd_name, read_unknown=False):
         ifd_dict = {}

piexif/_remove.py

@@ -1,7 +1,7 @@
 import io
 
 from ._common import *
-
+from piexif import _webp
 
 def remove(src, new_file=None):
     """
@@ -14,17 +14,34 @@ def remove(src, new_file=None):
     output_is_file = False
     if src[0:2] == b"\xff\xd8":
         src_data = src
+        file_type = "jpeg"
+    elif src[0:4] == b"RIFF" and src[8:12] == b"WEBP":
+        src_data = src
+        file_type = "webp"
     else:
         with open(src, 'rb') as f:
             src_data = f.read()
         output_is_file = True
-    segments = split_into_segments(src_data)
-    exif = get_exif_seg(segments)
+        if src_data[0:2] == b"\xff\xd8":
+            file_type = "jpeg"
+        elif src_data[0:4] == b"RIFF" and src_data[8:12] == b"WEBP":
+            file_type = "webp"
 
-    if exif:
-        new_data = src_data.replace(exif, b"")
-    else:
-        new_data = src_data
+    if file_type == "jpeg":
+        segments = split_into_segments(src_data)
+        exif = get_exif_seg(segments)
+        if exif:
+            new_data = src_data.replace(exif, b"")
+        else:
+            new_data = src_data
+    elif file_type == "webp":
+        try:
+            new_data = _webp.remove(src_data)
+        except ValueError:
+            new_data = src_data
+        except e:
+            print(e.args)
+            raise ValueError("Error ocured.")
 
     if isinstance(new_file, io.BytesIO):
         new_file.write(new_data)
@@ -36,4 +53,4 @@ def remove(src, new_file=None):
         with open(src, "wb+") as f:
             f.write(new_data)
     else:
-        raise ValueError("Give a 2nd argment to 'remove' to output file")
+        raise ValueError("Give a second argment to 'remove' to output file")

piexif/_webp.py

@@ -0,0 +1,240 @@
+
+import struct
+
+
+def split(data):
+    if data[0:4] != b"RIFF" or data[8:12] != b"WEBP":
+        raise ValueError("Not WebP")
+
+    webp_length_bytes = data[4:8]
+    webp_length = struct.unpack("<L", webp_length_bytes)[0]
+    RIFF_HEADER_SIZE = 8
+    file_size = RIFF_HEADER_SIZE + webp_length
+
+    start = 12
+    pointer = start
+    CHUNK_FOURCC_LENGTH = 4
+    LENGTH_BYTES_LENGTH = 4
+
+    chunks = []
+    while pointer + CHUNK_FOURCC_LENGTH + LENGTH_BYTES_LENGTH < file_size:
+        fourcc = data[pointer:pointer + CHUNK_FOURCC_LENGTH]
+        pointer += CHUNK_FOURCC_LENGTH
+        chunk_length_bytes = data[pointer:pointer + LENGTH_BYTES_LENGTH]
+        chunk_length = struct.unpack("<L", chunk_length_bytes)[0]
+        pointer += LENGTH_BYTES_LENGTH
+
+        chunk_data = data[pointer:pointer + chunk_length]
+        chunks.append({"fourcc":fourcc, "length_bytes":chunk_length_bytes, "data":chunk_data})
+
+        padding = 1 if chunk_length % 2 else 0
+
+        pointer += chunk_length + padding
+    return chunks
+
+def merge_chunks(chunks):
+    merged = b"".join([chunk["fourcc"] 
+                       + chunk["length_bytes"]
+                       + chunk["data"]
+                       + (len(chunk["data"]) % 2) * b"\x00"
+                        for chunk in chunks])
+    return merged
+
+
+def _get_size_from_vp8x(chunk):
+    width_minus_one_bytes = chunk["data"][-6:-3] + b"\x00"
+    width_minus_one = struct.unpack("<L", width_minus_one_bytes)[0]
+    width = width_minus_one + 1
+    height_minus_one_bytes = chunk["data"][-3:] + b"\x00"
+    height_minus_one = struct.unpack("<L", height_minus_one_bytes)[0]
+    height = height_minus_one + 1
+    return (width, height)
+
+def _get_size_from_vp8(chunk):
+    BEGIN_CODE = b"\x9d\x01\x2a"
+    begin_index = chunk["data"].find(BEGIN_CODE)
+    if begin_index == -1:
+        ValueError("wrong VP8")
+    else:
+        BEGIN_CODE_LENGTH = len(BEGIN_CODE)
+        LENGTH_BYTES_LENGTH = 2
+        length_start = begin_index + BEGIN_CODE_LENGTH
+        width_bytes = chunk["data"][length_start:length_start + LENGTH_BYTES_LENGTH]
+        width = struct.unpack("<H", width_bytes)[0]
+        height_bytes = chunk["data"][length_start + LENGTH_BYTES_LENGTH:length_start + 2 * LENGTH_BYTES_LENGTH]
+        height = struct.unpack("<H", height_bytes)[0]
+        return (width, height)
+
+def _vp8L_contains_alpha(chunk_data):
+    flag = ord(chunk_data[4:5]) >> 5-1 & ord(b"\x01")
+    contains = 1 * flag
+    return contains
+
+def _get_size_from_vp8L(chunk):
+    b1 = chunk["data"][1:2]
+    b2 = chunk["data"][2:3]
+    b3 = chunk["data"][3:4]
+    b4 = chunk["data"][4:5]
+
+    width_minus_one = (ord(b2) & ord(b"\x3F")) << 8 | ord(b1)
+    width = width_minus_one + 1
+
+    height_minus_one = (ord(b4) & ord(b"\x0F")) << 10 | ord(b3) << 2 | (ord(b2) & ord(b"\xC0")) >> 6
+    height = height_minus_one + 1
+
+    return (width, height)
+
+def _get_size_from_anmf(chunk):
+    width_minus_one_bytes = chunk["data"][6:9] + b"\x00"
+    width_minus_one = struct.unpack("<L", width_minus_one_bytes)[0]
+    width = width_minus_one + 1
+    height_minus_one_bytes = chunk["data"][9:12] + b"\x00"
+    height_minus_one = struct.unpack("<L", height_minus_one_bytes)[0]
+    height = height_minus_one + 1
+    return (width, height)
+
+def set_vp8x(chunks):
+
+    width = None
+    height = None
+    flags = [b"0", b"0", b"0", b"0", b"0", b"0", b"0", b"0"]  # [0, 0, ICC, Alpha, EXIF, XMP, Anim, 0]
+
+    for chunk in chunks:
+        if chunk["fourcc"] == b"VP8X":
+            width, height = _get_size_from_vp8x(chunk)
+        elif chunk["fourcc"] == b"VP8 ":
+            width, height = _get_size_from_vp8(chunk)
+        elif chunk["fourcc"] == b"VP8L":
+            is_rgba = _vp8L_contains_alpha(chunk["data"])
+            if is_rgba:
+                flags[3] = b"1"
+            width, height = _get_size_from_vp8L(chunk)
+        elif chunk["fourcc"] == b"ANMF":
+            width, height = _get_size_from_anmf(chunk)
+        elif chunk["fourcc"] == b"ICCP":
+            flags[2] = b"1"
+        elif chunk["fourcc"] == b"ALPH":
+            flags[3] = b"1"
+        elif chunk["fourcc"] == b"EXIF":
+            flags[4] = b"1"
+        elif chunk["fourcc"] == b"XMP ":
+            flags[5] = b"1"
+        elif chunk["fourcc"] == b"ANIM":
+            flags[6] = b"1"
+    width_minus_one = width - 1
+    height_minus_one = height - 1
+
+    if chunks[0]["fourcc"] == b"VP8X":
+        chunks.pop(0)
+
+    header_bytes = b"VP8X"
+    length_bytes = b"\x0a\x00\x00\x00"
+    flags_bytes = struct.pack("B", int(b"".join(flags), 2))
+    padding_bytes = b"\x00\x00\x00"
+    width_bytes = struct.pack("<L", width_minus_one)[:3]
+    height_bytes = struct.pack("<L", height_minus_one)[:3]
+
+    data_bytes = flags_bytes + padding_bytes + width_bytes + height_bytes
+
+    vp8x_chunk = {"fourcc":header_bytes, "length_bytes":length_bytes, "data":data_bytes}
+    chunks.insert(0, vp8x_chunk)
+
+    return chunks
+
+def get_file_header(chunks):
+    WEBP_HEADER_LENGTH = 4
+    FOURCC_LENGTH = 4
+    LENGTH_BYTES_LENGTH = 4
+
+    length = WEBP_HEADER_LENGTH
+    for chunk in chunks:
+        data_length = struct.unpack("<L", chunk["length_bytes"])[0]
+        data_length += 1 if data_length % 2 else 0
+        length += FOURCC_LENGTH + LENGTH_BYTES_LENGTH + data_length
+    length_bytes = struct.pack("<L", length)
+    riff = b"RIFF"
+    webp_header = b"WEBP"
+    file_header = riff + length_bytes + webp_header
+    return file_header
+
+
+
+def get_exif(data):
+    if data[0:4] != b"RIFF" or data[8:12] != b"WEBP":
+        raise ValueError("Not WebP")
+
+    if data[12:16] != b"VP8X":
+        raise ValueError("doesnot have exif")
+
+    webp_length_bytes = data[4:8]
+    webp_length = struct.unpack("<L", webp_length_bytes)[0]
+    RIFF_HEADER_SIZE = 8
+    file_size = RIFF_HEADER_SIZE + webp_length
+
+    start = 12
+    pointer = start
+    CHUNK_FOURCC_LENGTH = 4
+    LENGTH_BYTES_LENGTH = 4
+
+    chunks = []
+    exif = b""
+    while pointer < file_size:
+        fourcc = data[pointer:pointer + CHUNK_FOURCC_LENGTH]
+        pointer += CHUNK_FOURCC_LENGTH
+        chunk_length_bytes = data[pointer:pointer + LENGTH_BYTES_LENGTH]
+        chunk_length = struct.unpack("<L", chunk_length_bytes)[0]
+        if chunk_length % 2:
+            chunk_length += 1
+        pointer += LENGTH_BYTES_LENGTH
+        if fourcc == b"EXIF":
+            return data[pointer:pointer + chunk_length]
+        pointer += chunk_length
+    return None  # if there isn't exif, return None.
+
+
+def insert_exif_into_chunks(chunks, exif_bytes):
+    EXIF_HEADER = b"EXIF"
+    exif_length_bytes = struct.pack("<L", len(exif_bytes))
+    exif_chunk = {"fourcc":EXIF_HEADER, "length_bytes":exif_length_bytes, "data":exif_bytes}
+
+    xmp_index = None
+    animation_index = None
+
+    for index, chunk in enumerate(chunks):
+        if chunk["fourcc"] == b"EXIF":
+            chunks.pop(index)
+
+    for index, chunk in enumerate(chunks):
+        if chunk["fourcc"] == b"XMP ":
+            xmp_index = index
+        elif chunk["fourcc"] == b"ANIM":
+            animation_index = index
+    if xmp_index is not None:
+        chunks.insert(xmp_index, exif_chunk)
+    elif animation_index is not None:
+        chunks.insert(animation_index, exif_chunk)
+    else:
+        chunks.append(exif_chunk)
+    return chunks
+
+
+def insert(webp_bytes, exif_bytes):
+    chunks = split(webp_bytes)
+    chunks = insert_exif_into_chunks(chunks, exif_bytes)
+    chunks = set_vp8x(chunks)
+    file_header = get_file_header(chunks)
+    merged = merge_chunks(chunks)
+    new_webp_bytes = file_header + merged
+    return new_webp_bytes
+
+
+def remove(webp_bytes):
+    chunks = split(webp_bytes)
+    for index, chunk in enumerate(chunks):
+        if chunk["fourcc"] == b"EXIF":
+            chunks.pop(index)
+    chunks = set_vp8x(chunks)
+    file_header = get_file_header(chunks)
+    merged = merge_chunks(chunks)
+    new_webp_bytes = file_header + merged
+    return new_webp_bytes

tests/s_test.py

@@ -12,6 +12,7 @@
 from PIL import Image
 import piexif
 from piexif import _common, ImageIFD, ExifIFD, GPSIFD, TAGS, InvalidImageDataError
+from piexif import _webp
 from piexif import helper
 
 
@@ -144,16 +145,6 @@ def test_load_tif_m(self):
         zeroth_ifd2 = exif2["0th"]
         self.assertDictEqual(zeroth_ifd, zeroth_ifd2)
 
-    def test_load_fail(self):
-        with self.assertRaises(ValueError):
-            exif = piexif.load(os.path.join("tests", "images", "note.txt"))
-
-        with self.assertRaises(ValueError):
-            exif = piexif.load(os.path.join("tests", "images", "notjpeg.jpg"))
-
-        with self.assertRaises(ValueError):
-            exif = piexif.load(os.path.join("Oh", "My", "God"))
-
     def test_load_from_pilImage_property(self):
         o = io.BytesIO()
         i = Image.open(INPUT_FILE1)
@@ -855,11 +846,203 @@ def test_decode_bad_encoding(self):
         self.assertRaises(ValueError, helper.UserComment.load, b'hello world')
 
 
+class WebpTests(unittest.TestCase):
+    def setUp(self):
+        try:
+            os.mkdir("tests/images/out")
+        except:
+            pass
+
+    def test_merge_chunks(self):
+        """Can PIL open our output WebP?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+            "pil1.webp",
+            "pil2.webp",
+            "pil3.webp",
+            "pil_rgb.webp",
+            "pil_rgba.webp",
+        ]
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+
+            with open(IMAGE_DIR + filename, "rb") as f:
+                data = f.read()
+
+            chunks = _webp.split(data)
+            file_header = _webp.get_file_header(chunks)
+            merged = _webp.merge_chunks(chunks)
+            new_webp_bytes = file_header + merged
+            with open(OUT_DIR + "raw_" + filename, "wb") as f:
+                f.write(new_webp_bytes)
+            Image.open(OUT_DIR + "raw_" + filename)
+
+    def test_insert_exif(self):
+        """Can PIL open WebP that is inserted exif?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+            "pil1.webp",
+            "pil2.webp",
+            "pil3.webp",
+            "pil_rgb.webp",
+            "pil_rgba.webp",
+        ]
+
+        exif_dict = {
+            "0th":{
+                piexif.ImageIFD.Software: b"PIL",
+                piexif.ImageIFD.Make: b"Make",
+            }
+        }
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+
+            with open(IMAGE_DIR + filename, "rb") as f:
+                data = f.read()
+            exif_bytes = piexif.dump(exif_dict)
+            exif_inserted = _webp.insert(data, exif_bytes)
+            with open(OUT_DIR + "i_" + filename, "wb") as f:
+                f.write(exif_inserted)
+            Image.open(OUT_DIR + "i_" + filename)
+
+    def test_remove_exif(self):
+        """Can PIL open WebP that is removed exif?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+            "pil1.webp",
+            "pil2.webp",
+            "pil3.webp",
+            "pil_rgb.webp",
+            "pil_rgba.webp",
+        ]
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+
+            with open(IMAGE_DIR + filename, "rb") as f:
+                data = f.read()
+            exif_removed = _webp.remove(data)
+            with open(OUT_DIR + "r_" + filename, "wb") as f:
+                f.write(exif_removed)
+            Image.open(OUT_DIR + "r_" + filename)
+
+    def test_get_exif(self):
+        """Can we get exif from WebP?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+        ]
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+
+            with open(IMAGE_DIR + filename, "rb") as f:
+                data = f.read()
+            exif_bytes = _webp.get_exif(data)
+            self.assertEqual(exif_bytes[0:2], b"MM")
+
+    def test_load(self):
+        """Can we get exif from WebP?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+        ]
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+            print(piexif.load(IMAGE_DIR + filename))
+
+    def test_remove(self):
+        """Can PIL open WebP that is removed exif?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+            "pil1.webp",
+            "pil2.webp",
+            "pil3.webp",
+            "pil_rgb.webp",
+            "pil_rgba.webp",
+        ]
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+            piexif.remove(IMAGE_DIR + filename, OUT_DIR + "rr_" + filename)
+            Image.open(OUT_DIR + "rr_" + filename)
+
+    def test_insert(self):
+        """Can PIL open WebP that is inserted exif?"""
+        IMAGE_DIR = "tests/images/"
+        OUT_DIR = "tests/images/out/"
+        files = [
+            "tool1.webp",
+            "pil1.webp",
+            "pil2.webp",
+            "pil3.webp",
+            "pil_rgb.webp",
+            "pil_rgba.webp",
+        ]
+
+        exif_dict = {
+            "0th":{
+                piexif.ImageIFD.Software: b"PIL",
+                piexif.ImageIFD.Make: b"Make",
+            }
+        }
+        exif_bytes = piexif.dump(exif_dict)
+
+        for filename in files:
+            try:
+                Image.open(IMAGE_DIR + filename)
+            except:
+                print("Pillow can't read {0}".format(filename))
+                continue
+            piexif.insert(exif_bytes, IMAGE_DIR + filename, OUT_DIR + "ii_" + filename)
+            Image.open(OUT_DIR + "ii_" + filename)
+
+
 def suite():
     suite = unittest.TestSuite()
-    suite.addTests([unittest.makeSuite(UTests),
-                    unittest.makeSuite(ExifTests),
-                    unittest.makeSuite(HelperTests)])
+    suite.addTests([
+        unittest.makeSuite(UTests),
+        unittest.makeSuite(ExifTests),
+        unittest.makeSuite(HelperTests),
+        unittest.makeSuite(WebpTests),
+    ])
     return suite