Skip to content
Back to top

# Source code for PIL.WebPImagePlugin

    from __future__ import annotations

    from io import BytesIO

    from . import Image, ImageFile

    try:
        from . import _webp

        SUPPORTED = True
    except ImportError:
        SUPPORTED = False

    TYPE_CHECKING = False
    if TYPE_CHECKING:
        from typing import IO, Any

    _VP8_MODES_BY_IDENTIFIER = {
        b"VP8 ": "RGB",
        b"VP8X": "RGBA",
        b"VP8L": "RGBA",  # lossless
    }

    def _accept(prefix: bytes) -> bool | str:
        is_riff_file_format = prefix.startswith(b"RIFF")
        is_webp_file = prefix[8:12] == b"WEBP"
        is_valid_vp8_mode = prefix[12:16] in _VP8_MODES_BY_IDENTIFIER

        if is_riff_file_format and is_webp_file and is_valid_vp8_mode:
            if not SUPPORTED:
                return (
                    "image file could not be identified because WEBP support not installed"
                )
            return True
        return False

    class WebPImageFile(ImageFile.ImageFile):
        format = "WEBP"
        format_description = "WebP image"
        __loaded = 0
        __logical_frame = 0

        def _open(self) -> None:
            assert self.fp is not None
            s = self.fp.read()
            if not _accept(s):
                msg = "not a WEBP file"
                raise SyntaxError(msg)

            # Use the newer AnimDecoder API to parse the (possibly) animated file,
            # and access muxed chunks like ICC/EXIF/XMP.
            self._decoder = _webp.WebPAnimDecoder(s)

            # Get info from decoder
            self._size, self.info["loop"], bgcolor, self.n_frames, self.rawmode = (
                self._decoder.get_info()
            )
            self.info["background"] = (
                (bgcolor >> 16) & 0xFF,  # R
                (bgcolor >> 8) & 0xFF,  # G
                bgcolor & 0xFF,  # B
                (bgcolor >> 24) & 0xFF,  # A
            )
            self.is_animated = self.n_frames > 1
            self._mode = "RGB" if self.rawmode == "RGBX" else self.rawmode

            # Attempt to read ICC / EXIF / XMP chunks from file
            for key, chunk_name in {
                "icc_profile": "ICCP",
                "exif": "EXIF",
                "xmp": "XMP ",
            }.items():
                if value := self._decoder.get_chunk(chunk_name):
                    self.info[key] = value

            # Initialize seek state
            self._reset(reset=False)

        def _getexif(self) -> dict[int, Any] | None:
            if "exif" not in self.info:
                return None
            return self.getexif()._get_merged_dict()

        def seek(self, frame: int) -> None:
            if not self._seek_check(frame):
                return

            # Set logical frame to requested position
            self.__logical_frame = frame

        def _reset(self, reset: bool = True) -> None:
            if reset:
                self._decoder.reset()
            self.__physical_frame = 0
            self.__loaded = -1
            self.__timestamp = 0

        def _get_next(self) -> tuple[bytes, int, int]:
            # Get next frame
            ret = self._decoder.get_next()
            self.__physical_frame += 1

            # Check if an error occurred
            if ret is None:
                self._reset()  # Reset just to be safe
                self.seek(0)
                msg = "failed to decode next frame in WebP file"
                raise EOFError(msg)

            # Compute duration
            data, timestamp = ret
            duration = timestamp - self.__timestamp
            self.__timestamp = timestamp

            # libwebp gives frame end, adjust to start of frame
            timestamp -= duration
            return data, timestamp, duration

        def _seek(self, frame: int) -> None:
            if self.__physical_frame == frame:
                return  # Nothing to do
            if frame < self.__physical_frame:
                self._reset()  # Rewind to beginning
            while self.__physical_frame < frame:
                self._get_next()  # Advance to the requested frame

        def load(self) -> Image.core.PixelAccess | None:
            if self.__loaded != self.__logical_frame:
                self._seek(self.__logical_frame)

                # We need to load the image data for this frame
                data, self.info["timestamp"], self.info["duration"] = self._get_next()
                self.__loaded = self.__logical_frame

                # Set tile
                if self.fp and self._exclusive_fp:
                    self.fp.close()
                self.fp = BytesIO(data)
                self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.rawmode)]

            return super().load()

        def load_seek(self, pos: int) -> None:
            pass

        def tell(self) -> int:
            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", []))

        # If total frame count is 1, then save using the legacy API, which
        # will preserve non-alpha modes
        total = 0
        for ims in [im] + append_images:
            total += getattr(ims, "n_frames", 1)
        if total == 1:
            _save(im, fp, filename)
            return

        background: int | tuple[int, ...] = (0, 0, 0, 0)
        if "background" in encoderinfo:
            background = encoderinfo["background"]
        elif "background" in im.info:
            background = im.info["background"]
            if isinstance(background, int):
                # GifImagePlugin stores a global color table index in
                # info["background"]. So it must be converted to an RGBA value
                palette = im.getpalette()
                if palette:
                    r, g, b = palette[background * 3 : (background + 1) * 3]
                    background = (r, g, b, 255)
                else:
                    background = (background, background, background, 255)

        duration = im.encoderinfo.get("duration", im.info.get("duration", 0))
        loop = im.encoderinfo.get("loop", 0)
        minimize_size = im.encoderinfo.get("minimize_size", False)
        kmin = im.encoderinfo.get("kmin", None)
        kmax = im.encoderinfo.get("kmax", None)
        allow_mixed = im.encoderinfo.get("allow_mixed", False)
        verbose = False
        lossless = im.encoderinfo.get("lossless", False)
        quality = im.encoderinfo.get("quality", 80)
        alpha_quality = im.encoderinfo.get("alpha_quality", 100)
        method = im.encoderinfo.get("method", 0)
        icc_profile = im.encoderinfo.get("icc_profile") or ""
        exif = im.encoderinfo.get("exif", "")
        if isinstance(exif, Image.Exif):
            exif = exif.tobytes()
        xmp = im.encoderinfo.get("xmp", "")
        if allow_mixed:
            lossless = False

        # Sensible keyframe defaults are from gif2webp.c script
        if kmin is None:
            kmin = 9 if lossless else 3
        if kmax is None:
            kmax = 17 if lossless else 5

        # Validate background color
        if (
            not isinstance(background, (list, tuple))
            or len(background) != 4
            or not all(0 <= v < 256 for v in background)
        ):
            msg = f"Background color is not an RGBA tuple clamped to (0-255): {background}"
            raise OSError(msg)

        # Convert to packed uint
        bg_r, bg_g, bg_b, bg_a = background
        background = (bg_a << 24) | (bg_r << 16) | (bg_g << 8) | (bg_b << 0)

        # Setup the WebP animation encoder
        enc = _webp.WebPAnimEncoder(
            im.size,
            background,
            loop,
            minimize_size,
            kmin,
            kmax,
            allow_mixed,
            verbose,
        )

        # Add each frame
        frame_idx = 0
        timestamp = 0
        cur_idx = im.tell()
        try:
            for ims in [im] + append_images:
                # Get number of frames in this image
                nfr = getattr(ims, "n_frames", 1)

                for idx in range(nfr):
                    ims.seek(idx)

                    frame = _convert_frame(ims)

                    # Append the frame to the animation encoder
                    enc.add(
                        frame.getim(),
                        round(timestamp),
                        lossless,
                        quality,
                        alpha_quality,
                        method,
                    )

                    # Update timestamp and frame index
                    if isinstance(duration, (list, tuple)):
                        timestamp += duration[frame_idx]
                    else:
                        timestamp += duration
                    frame_idx += 1

        finally:
            im.seek(cur_idx)

        # Force encoder to flush frames
        enc.add(None, round(timestamp), lossless, quality, alpha_quality, 0)

        # Get the final output from the encoder
        data = enc.assemble(icc_profile, exif, xmp)
        if data is None:
            msg = "cannot write file as WebP (encoder returned None)"
            raise OSError(msg)

        fp.write(data)

    def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
        lossless = im.encoderinfo.get("lossless", False)
        quality = im.encoderinfo.get("quality", 80)
        alpha_quality = im.encoderinfo.get("alpha_quality", 100)
        icc_profile = im.encoderinfo.get("icc_profile") or ""
        exif = im.encoderinfo.get("exif", b"")
        if isinstance(exif, Image.Exif):
            exif = exif.tobytes()
        if exif.startswith(b"Exif\x00\x00"):
            exif = exif[6:]
        xmp = im.encoderinfo.get("xmp", "")
        method = im.encoderinfo.get("method", 4)
        exact = 1 if im.encoderinfo.get("exact") else 0

        im = _convert_frame(im)

        data = _webp.WebPEncode(
            im.getim(),
            lossless,
            float(quality),
            float(alpha_quality),
            icc_profile,
            method,
            exact,
            exif,
            xmp,
        )
        if data is None:
            msg = "cannot write file as WebP (encoder returned None)"
            raise OSError(msg)

        fp.write(data)

    Image.register_open(WebPImageFile.format, WebPImageFile, _accept)
    if SUPPORTED:
        Image.register_save(WebPImageFile.format, _save)
        Image.register_save_all(WebPImageFile.format, _save_all)
        Image.register_extension(WebPImageFile.format, ".webp")
        Image.register_mime(WebPImageFile.format, "image/webp")