python使用pillow和opencv生成圖片縮圖

waketzheng發表於2024-03-13

程式碼如下:

from io import BytesIO
from pathlib import Path
from typing import Annotated, Literal, Optional, Tuple, Union

# pip install pillow opencv-python
import cv2  # type:ignore[import-untyped]
import numpy as np
from PIL import Image

ImageSizeType = Annotated[Tuple[int, int], "圖片尺寸(寬,高),如:(1080, 720)"]


class Picture:
    default_size = (351, 190)

    @staticmethod
    def generate_thumbnail_pil(
        img_bytes: bytes, size, fmt: Literal["JPEG", "PNG"], *, verbose=False
    ) -> bytes:
        with Image.open(BytesIO(img_bytes)) as image:
            origin_size = image.size
            image.thumbnail(size)
            if verbose:
                print(f"pil[target_{size=}]: {origin_size} -> {image.size}")
            bio = BytesIO()
            if fmt == "JPEG":
                image = image.convert("RGB")
            image.save(bio, format=fmt)
            return bio.getvalue()

    @staticmethod
    def generate_thumbnail_opencv(
        img_bytes: bytes, size, fmt: Literal[".jpeg", ".png"], *, verbose=False
    ) -> bytes:
        img = cv2.imdecode(np.frombuffer(img_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        resized = cv2.resize(img, size)
        if verbose:
            origin_size, converted_size = img.shape[:2][::-1], resized.shape[:2][::-1]
            print(f"opencv[target_{size=}]: {origin_size} -> {converted_size}")
        _, img_encode = cv2.imencode(fmt, resized)
        return img_encode.tobytes()

    @classmethod
    def thumbnail(
        cls,
        img: Union[bytes, BytesIO, str, Path],
        size: Optional[ImageSizeType] = None,
        keep_scale=False,
        fmt: Optional[str] = None,
        *,
        verbose=False,
    ) -> bytes:
        """生成縮圖

        :param img: 圖片二進位制或路徑
        :param size: 縮圖的寬、高, 如果為None,則使用類的default_size
        :param keep_scale: 是否保持寬高比
        :param fmt: 縮圖格式(jpg或png)
        :param verbose: 除錯用引數,是否列印生成的縮圖尺寸

        Usage::
            >>> p = Path('a.jpg')
            >>> thumb = Picture.thumbnail(p.read_bytes(), fmt=p.suffix)
            >>> isinstance(thumb, bytes)
            True
        """
        if size is None:
            size = cls.default_size
        if fmt is None:
            if isinstance(img, (str, Path)) and str(img).lower().endswith(".png"):
                fmt = "png"
        else:
            fmt = fmt.strip(".").lower()
            assert fmt in ("jpg", "jpeg", "png"), "Invalid `fmt`: only support png/jpg"
        if isinstance(img, BytesIO):
            img = img.getvalue()
        elif isinstance(img, (str, Path)):
            img = Path(img).read_bytes()
        if keep_scale:
            fmt1: Literal["PNG", "JPEG"] = "PNG" if fmt == "png" else "JPEG"
            return cls.generate_thumbnail_pil(img, size, fmt=fmt1, verbose=verbose)
        else:
            fmt2: Literal[".png", ".jpeg"] = ".png" if fmt == "png" else ".jpeg"
            return cls.generate_thumbnail_opencv(img, size, fmt=fmt2, verbose=verbose)


def main():
    import sys
    import doctest
    filepath = sys.argv[1]
    if not (p := Path(filepath)).exists():
        raise FileNotFoundError(filepath)
    b1 = Picture.thumbnail(p, keep_scale=True, verbose=True, fmt=p.suffix)
    b2 = Picture.thumbnail(p, verbose=True, fmt=p.suffix)
    p1 = p.with_name(p.stem + "-pil" + p.suffix)
    p2 = p.with_name(p.stem + "-opencv" + p.suffix)
    size = p1.write_bytes(b1)
    print(f"Write to {p1} with {size=}")
    size = p2.write_bytes(b2)
    print(f"Write to {p2} with {size=}")
    # verify param type
    content = p.read_bytes()
    assert isinstance(Picture.thumbnail(content, keep_scale=True, fmt=".png"), bytes)
    assert isinstance(Picture.thumbnail(str(p), keep_scale=True), bytes)
    assert isinstance(
        Picture.thumbnail(BytesIO(content), keep_scale=True, fmt="jpg"), bytes
    )
    assert isinstance(Picture.thumbnail(content, fmt=".jpg"), bytes)
    assert isinstance(Picture.thumbnail(str(p)), bytes)
    assert isinstance(Picture.thumbnail(BytesIO(content), fmt="png"), bytes)
    # doctest
    if Path("a.jpg").exists():
        doctest.testmod(verbose=True)
    # doctest
    if Path('a.jpg').exists():
        doctest.testmod(verbose=True)


if __name__ == "__main__":
    main()

相關文章