299 lines
9.4 KiB
Python
299 lines
9.4 KiB
Python
"""
|
||
感知服务编排。
|
||
|
||
职责：
|
||
1. 接收图片 → 预处理（灰度增强/对比度调整）
|
||
2. 调用天目 OCR（带缓存和重试）
|
||
3. 后处理 → 结构化
|
||
4. 写入事件系统（可选）
|
||
|
||
使用示例::
|
||
|
||
service = PerceptionService()
|
||
result = await service.recognize(image_bytes)
|
||
if result.success:
|
||
print(f"识别结果: {result.text}")
|
||
await service.close()
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import io
|
||
from collections import OrderedDict
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from app.core.logging import logger
|
||
from app.integrations.ocr import OCRClient, OCRResult
|
||
|
||
|
||
class LRUCache:
|
||
"""线程安全的 LRU 缓存（基于 OrderedDict）。"""
|
||
|
||
def __init__(self, maxsize: int = 128) -> None:
|
||
self._data: OrderedDict[str, Any] = OrderedDict()
|
||
self._maxsize = maxsize
|
||
|
||
def get(self, key: str) -> Any | None:
|
||
if key not in self._data:
|
||
return None
|
||
self._data.move_to_end(key)
|
||
return self._data[key]
|
||
|
||
def put(self, key: str, value: Any) -> None:
|
||
self._data[key] = value
|
||
self._data.move_to_end(key)
|
||
if len(self._data) > self._maxsize:
|
||
self._data.popitem(last=False)
|
||
|
||
def clear(self) -> None:
|
||
self._data.clear()
|
||
|
||
|
||
class PerceptionService:
|
||
"""感知服务：图片处理、OCR 识别、结构化输出。
|
||
|
||
支持缓存（相同图片 hash 命中时跳过重复识别）、
|
||
图片预处理（灰度增强有助于提高 OCR 精度）、
|
||
自动重试（主引擎失败时降级到备选引擎）。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
ocr_client: OCRClient | None = None,
|
||
cache_size: int = 128,
|
||
enable_preprocessing: bool = True,
|
||
enable_cache: bool = True,
|
||
) -> None:
|
||
self._ocr = ocr_client or OCRClient()
|
||
self._cache = LRUCache(maxsize=cache_size) if enable_cache else None
|
||
self._enable_preprocessing = enable_preprocessing
|
||
|
||
async def close(self) -> None:
|
||
"""释放 OCR 客户端和缓存。"""
|
||
await self._ocr.close()
|
||
if self._cache:
|
||
self._cache.clear()
|
||
|
||
# ── 通用 OCR ────────────────────────────────────────
|
||
|
||
async def recognize(
|
||
self,
|
||
image: bytes | str | Path,
|
||
*,
|
||
engine: str = "auto",
|
||
save_md: bool = False,
|
||
use_cache: bool = True,
|
||
) -> OCRResult:
|
||
"""通用 OCR 识别（带缓存和自动降级）。
|
||
|
||
Args:
|
||
image: 图片数据（bytes）、文件路径（str/Path）、HTTP URL。
|
||
engine: 引擎选择。
|
||
save_md: 是否保存 Markdown。
|
||
use_cache: 是否使用缓存（相同图片跳过重复识别）。
|
||
|
||
Returns:
|
||
OCRResult: 识别结果。
|
||
"""
|
||
image_data = self._resolve_image(image)
|
||
|
||
# 缓存命中
|
||
cache_key = self._cache_key(image_data) if use_cache and self._cache else None
|
||
if cache_key:
|
||
cached = self._cache.get(cache_key) if self._cache else None
|
||
if cached is not None:
|
||
logger.info("OCR 缓存命中 | key={}", cache_key[:12])
|
||
return cached
|
||
|
||
# 预处理
|
||
if self._enable_preprocessing and isinstance(image_data, bytes):
|
||
preprocessed = self._preprocess_image(image_data)
|
||
if preprocessed is not None:
|
||
image_data = preprocessed
|
||
|
||
# 主引擎识别（自动降级）
|
||
result = await self._recognize_with_fallback(image_data, engine, save_md)
|
||
|
||
# 写入缓存
|
||
if cache_key and self._cache and result.success:
|
||
self._cache.put(cache_key, result)
|
||
|
||
if result.success:
|
||
logger.info(
|
||
"OCR 识别完成 | engine={} blocks={} ms={}",
|
||
result.engine,
|
||
len(result.blocks),
|
||
result.duration_ms,
|
||
)
|
||
else:
|
||
logger.warning("OCR 识别失败 | engine={} error={}", engine, result.error)
|
||
|
||
return result
|
||
|
||
# ── 复杂文档解析 ────────────────────────────────────
|
||
|
||
async def parse_document(
|
||
self,
|
||
image: bytes | str | Path,
|
||
*,
|
||
output_format: str = "markdown",
|
||
save_md: bool = False,
|
||
) -> OCRResult:
|
||
"""复杂文档解析（PaddleOCR-VL）。
|
||
|
||
Args:
|
||
image: 图片数据（bytes）、文件路径（str/Path）、HTTP URL。
|
||
output_format: 输出格式。
|
||
save_md: 是否保存 Markdown。
|
||
|
||
Returns:
|
||
OCRResult: 结构化解析结果。
|
||
"""
|
||
image_data = self._resolve_image(image)
|
||
result = await self._ocr.ocr_vl(
|
||
image_data,
|
||
output_format=output_format,
|
||
save_md=save_md,
|
||
)
|
||
|
||
if result.success:
|
||
logger.info(
|
||
"文档解析完成 | engine={} format={} ms={}",
|
||
result.engine,
|
||
output_format,
|
||
result.duration_ms,
|
||
)
|
||
else:
|
||
logger.warning("文档解析失败 | error={}", result.error)
|
||
|
||
return result
|
||
|
||
# ── 批处理 ───────────────────────────────────────────
|
||
|
||
async def recognize_batch(
|
||
self,
|
||
images: list[bytes | str | Path],
|
||
*,
|
||
engine: str = "auto",
|
||
concurrency: int = 3,
|
||
) -> list[OCRResult]:
|
||
"""批量 OCR 识别（限制并发数）。
|
||
|
||
Args:
|
||
images: 图片列表。
|
||
engine: 引擎选择。
|
||
concurrency: 最大并发数。
|
||
|
||
Returns:
|
||
OCRResult 列表，顺序与输入一致。
|
||
"""
|
||
import asyncio
|
||
|
||
semaphore = asyncio.Semaphore(concurrency)
|
||
|
||
async def _limited(img: bytes | str | Path) -> OCRResult:
|
||
async with semaphore:
|
||
return await self.recognize(img, engine=engine)
|
||
|
||
tasks = [_limited(img) for img in images]
|
||
return await asyncio.gather(*tasks)
|
||
|
||
# ── 辅助 ─────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _resolve_image(image: bytes | str | Path) -> bytes | str:
|
||
"""将图片输入统一为 bytes 或 URL 字符串。
|
||
|
||
Args:
|
||
image: bytes 数据、文件路径或 URL。
|
||
|
||
Returns:
|
||
bytes（文件/二进制）或 str（URL）。
|
||
"""
|
||
if isinstance(image, Path):
|
||
return image.read_bytes()
|
||
if isinstance(image, str):
|
||
if image.startswith(("http://", "https://")):
|
||
return image
|
||
return Path(image).read_bytes()
|
||
return image
|
||
|
||
@staticmethod
|
||
def _cache_key(image_data: bytes | str) -> str | None:
|
||
"""生成缓存键（bytes 用 SHA256，URL 直接作为键）。"""
|
||
if isinstance(image_data, bytes):
|
||
return hashlib.sha256(image_data).hexdigest()
|
||
if isinstance(image_data, str) and image_data.startswith(("http://", "https://")):
|
||
return f"url:{image_data}"
|
||
return None
|
||
|
||
@staticmethod
|
||
def _preprocess_image(image_data: bytes) -> bytes | None:
|
||
"""图片预处理：灰度化 + 对比度增强。
|
||
|
||
对某些低对比度或偏色图片，预处理能显著提高 OCR 精度。
|
||
仅对 bytes 输入生效（URL 跳过）。
|
||
|
||
Returns:
|
||
预处理后的图片 bytes，失败时返回 None。
|
||
"""
|
||
try:
|
||
from PIL import Image, ImageEnhance, ImageFilter
|
||
|
||
img = Image.open(io.BytesIO(image_data))
|
||
|
||
# 1. 转灰度
|
||
if img.mode != "L":
|
||
img = img.convert("L")
|
||
|
||
# 2. 增强对比度
|
||
enhancer = ImageEnhance.Contrast(img)
|
||
img = enhancer.enhance(1.5)
|
||
|
||
# 3. 轻微锐化
|
||
img = img.filter(ImageFilter.SHARPEN)
|
||
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
return buf.getvalue()
|
||
except Exception as e:
|
||
logger.warning("图片预处理失败（不影响原图识别）: {}", e)
|
||
return None
|
||
|
||
async def _recognize_with_fallback(
|
||
self,
|
||
image_data: bytes | str,
|
||
engine: str,
|
||
save_md: bool,
|
||
) -> OCRResult:
|
||
"""主引擎识别 + 降级逻辑。
|
||
|
||
ppocrv6 失败时自动降级到 paddleocr-vl（如果配置了懒加载），
|
||
paddleocr-vl 失败时不再降级。
|
||
"""
|
||
result = await self._ocr.ocr(
|
||
image_data,
|
||
engine=engine,
|
||
save_md=save_md,
|
||
)
|
||
|
||
# 主引擎成功 → 直接返回
|
||
if result.success or engine != "auto":
|
||
return result
|
||
|
||
# auto 模式下 ppocrv6 失败 → 尝试 paddleocr-vl 降级
|
||
# 只对某些可恢复的错误降级（空结果而非服务器错误）
|
||
if result.error and ("引擎未加载" in result.error or "未识别到文字" in result.error):
|
||
logger.info("ppocrv6 结果为空，降级到 paddleocr-vl")
|
||
fallback = await self._ocr.ocr(
|
||
image_data,
|
||
engine="paddleocr-vl",
|
||
save_md=save_md,
|
||
)
|
||
if fallback.success:
|
||
return fallback
|
||
|
||
return result
|