Python使用FastAPI提供图片缩略图生成接口

使用pillow的thumbnail生成缩略图时，会保持原图的宽高比；使用的opencv的resize则不会

具体代码如下：

#!/usr/bin/env python
import re
import sys
from enum import Enum
from io import BytesIO
from pathlib import Path
from typing import Annotated, Literal, Optional, Tuple, Union

# pip install python-multipart fastapi-cdn-host fastapi uvicorn pillow opencv-python
import cv2  # type:ignore[import-untyped]
import fastapi_cdn_host
import numpy as np
import uvicorn
from fastapi import FastAPI, File, HTTPException, Query
from fastapi.responses import RedirectResponse, Response
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)


class ValidationError(HTTPException):
    def __init__(self, detail: str, status_code=400) -> None:
        super().__init__(status_code=status_code, detail=detail)


app = FastAPI(title="Thumbnail Generator")
fastapi_cdn_host.monkey_patch_for_docs_ui(app)


@app.get("/", include_in_schema=False)
async def to_docs():
    return RedirectResponse("/docs")


class FmtEnum(str, Enum):
    jpg = "JPEG"
    png = "PNG"


class ImageResponse(Response):
    media_type: Literal["image/jpeg", "image/png"] = "image/jpeg"

    def __init__(self, content: bytes, status_code=200, **kw) -> None:
        super().__init__(content=content, status_code=status_code, **kw)

    @classmethod
    def docs_schema(cls) -> dict:
        example = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x06f..."
        return {
            "content": {
                "image/png": {"example": str(example)},
                "image/jpeg": {"example": str(example).replace("PNG", "JFIF")},
            },
            "description": "返回二进制JPEG/PNG图片.",
        }


class JpegResponse(ImageResponse):
    media_type = "image/jpeg"


class PngResponse(ImageResponse):
    media_type = "image/png"


@app.post(
    "/thumbnail",
    response_class=ImageResponse,
    responses={200: ImageResponse.docs_schema()},
)
async def generate_thumbnail(
    size: Annotated[
        Union[str, None],
        Query(
            max_length=50,
            description="缩略图尺寸，默认为: {}x{}".format(*Picture.default_size),
        ),
    ] = None,
    fmt: Annotated[FmtEnum, Query(description="缩略图格式")] = FmtEnum.jpg,
    uploaded_image: bytes = File(),
) -> ImageResponse:
    width, height = Picture.default_size
    if size:
        try:
            width, height = map(int, re.findall(r"\d+", size))
        except (ValueError, TypeError):
            raise ValidationError(
                f"Invalid size type! Expected: {width}x{height}, Example: 1080x720"
            )
    img_bytes = Picture.thumbnail(uploaded_image, size=(width, height), fmt=fmt)
    return PngResponse(img_bytes) if fmt == FmtEnum.png else JpegResponse(img_bytes)


def runserver() -> None:
    port = int(sys.argv[1]) if sys.argv[1:] and sys.argv[1].isdigit() else 8000
    uvicorn.run(f"{Path(__file__).stem}:app", port=port, reload=True)


if __name__ == "__main__":
    runserver()

效果如下：