services/ocr/image_io.py

from __future__ import annotations

import base64
import binascii
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

from .contracts import ImageSource
from .errors import OCRServiceError

MAX_IMAGE_BYTES = 25 * 1024 * 1024


@dataclass(frozen=True, slots=True)
class LoadedImage:
    content: bytes
    array: np.ndarray
    width: int
    height: int
    mime_type: str


def _decode_base64(value: str) -> tuple[bytes, str | None]:
    mime_type: str | None = None
    encoded = value.strip()
    if encoded.startswith("data:"):
        try:
            header, encoded = encoded.split(",", 1)
        except ValueError as exc:
            raise OCRServiceError("INVALID_BASE64", "Base64 data URL 缺少逗号") from exc
        if ";base64" not in header:
            raise OCRServiceError("INVALID_BASE64", "data URL 必须使用 base64 编码")
        mime_type = header[5:].split(";", 1)[0]
    try:
        return base64.b64decode(encoded, validate=True), mime_type
    except (binascii.Error, ValueError) as exc:
        raise OCRServiceError("INVALID_BASE64", "图片 Base64 编码无效") from exc


def _detect_mime(content: bytes, declared: str | None) -> str:
    if declared in {"image/png", "image/jpeg", "image/webp", "image/bmp"}:
        return declared
    if content.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if content.startswith(b"\xff\xd8\xff"):
        return "image/jpeg"
    if content.startswith(b"RIFF") and content[8:12] == b"WEBP":
        return "image/webp"
    if content.startswith(b"BM"):
        return "image/bmp"
    return "application/octet-stream"


def load_image(source: ImageSource) -> LoadedImage:
    declared_mime: str | None = None
    if source.path is not None:
        path = Path(source.path).expanduser()
        if not path.is_file():
            raise OCRServiceError(
                "IMAGE_NOT_FOUND", f"图片不存在: {path}", status_code=404
            )
        try:
            content = path.read_bytes()
        except OSError as exc:
            raise OCRServiceError("IMAGE_READ_ERROR", f"无法读取图片: {path}") from exc
    else:
        assert source.base64 is not None
        content, declared_mime = _decode_base64(source.base64)

    if not content:
        raise OCRServiceError("EMPTY_IMAGE", "图片内容为空")
    if len(content) > MAX_IMAGE_BYTES:
        raise OCRServiceError(
            "IMAGE_TOO_LARGE",
            f"图片超过 {MAX_IMAGE_BYTES // (1024 * 1024)} MiB 限制",
            status_code=413,
        )

    array = cv2.imdecode(np.frombuffer(content, dtype=np.uint8), cv2.IMREAD_COLOR)
    if array is None:
        raise OCRServiceError("IMAGE_DECODE_ERROR", "无法解码图片格式")
    height, width = array.shape[:2]
    return LoadedImage(
        content=content,
        array=array,
        width=int(width),
        height=int(height),
        mime_type=_detect_mime(content, declared_mime),
    )