src/preprocess/image_processor.py

"""基础图像处理：按场景执行"""
from typing import List
import cv2
import numpy as np
from PIL import Image

from .scene_classifier import Scene, SceneResult

MAX_LONG_EDGE = 3840  # 4K
MIN_SHORT_EDGE = 100
SLICE_HEIGHT = 2000   # 长图切片高度阈值


def _pil_to_cv2(img: Image.Image) -> np.ndarray:
    return cv2.cvtColor(np.array(img.convert("RGB")), cv2.COLOR_RGB2BGR)


def _cv2_to_pil(arr: np.ndarray) -> Image.Image:
    return Image.fromarray(cv2.cvtColor(arr, cv2.COLOR_BGR2RGB))


def _deskew(img: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    coords = np.column_stack(np.where(gray < 200))
    if len(coords) < 10:
        return img
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = 90 + angle
    if abs(angle) < 0.5:
        return img
    h, w = img.shape[:2]
    M = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0)
    return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)


def _normalize_resolution(img: Image.Image) -> tuple[Image.Image, list[str]]:
    warnings = []
    w, h = img.size
    long_edge = max(w, h)
    short_edge = min(w, h)

    if short_edge < MIN_SHORT_EDGE:
        warnings.append(f"low_resolution: {w}x{h}")

    if long_edge > MAX_LONG_EDGE:
        scale = MAX_LONG_EDGE / long_edge
        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)

    return img, warnings


def _slice_long_image(img: Image.Image) -> List[Image.Image]:
    w, h = img.size
    if h <= SLICE_HEIGHT:
        return [img]
    slices = []
    for y in range(0, h, SLICE_HEIGHT):
        slices.append(img.crop((0, y, w, min(y + SLICE_HEIGHT, h))))
    return slices


def process(img: Image.Image, scene_result: SceneResult) -> dict:
    """
    返回:
      images: List[PIL.Image]  (长图切片后可能多张)
      warnings: List[str]
    """
    img, warnings = _normalize_resolution(img)
    scene = scene_result.scene
    arr = _pil_to_cv2(img)

    if scene == Scene.DOCUMENT:
        arr = _deskew(arr)
        arr = cv2.fastNlMeansDenoisingColored(arr, None, 10, 10, 7, 21)
        arr = cv2.detailEnhance(arr, sigma_s=10, sigma_r=0.15)
        img = _cv2_to_pil(arr)

    elif scene == Scene.POSTER:
        arr = cv2.convertScaleAbs(arr, alpha=1.3, beta=20)
        img = _cv2_to_pil(arr)

    images = _slice_long_image(img)
    return {"images": images, "warnings": warnings}