"""基础图像处理:按场景执行""" 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}