File: shrinkpix/piexif/_webp.py

File: shrinkpix/piexif/_webp.py
⇨ Raw text view and save
⇨ Page script and template
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