541 lines
19 KiB
Python
541 lines
19 KiB
Python
"""
|
||
OCR 识别服务模块
|
||
支持多种 OCR 提供商：PaddleOCR（本地）/ 阿里云 / 腾讯云
|
||
auto 模式下本地优先 + 云端 fallback
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import abc
|
||
import asyncio
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import List, Optional
|
||
|
||
from app.config import OCRProvider, settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ──────────────────────────── 数据结构 ────────────────────────────
|
||
|
||
@dataclass
|
||
class OCRBlock:
|
||
"""OCR 识别的文本块"""
|
||
text: str
|
||
bbox: Optional[List[float]] = None # [x1, y1, x2, y2]
|
||
confidence: float = 0.0
|
||
|
||
|
||
@dataclass
|
||
class OCRResult:
|
||
"""OCR 识别结果"""
|
||
text: str
|
||
blocks: List[OCRBlock] = field(default_factory=list)
|
||
confidence: float = 0.0
|
||
provider: str = ""
|
||
keywords: List[str] = field(default_factory=list) # OCR 时提取的关键词标签
|
||
|
||
|
||
# ──────────────────────────── 抽象基类 ────────────────────────────
|
||
|
||
class OCRProviderBase(abc.ABC):
|
||
"""OCR 提供商抽象基类"""
|
||
|
||
@property
|
||
@abc.abstractmethod
|
||
def name(self) -> str:
|
||
"""提供商名称"""
|
||
...
|
||
|
||
@abc.abstractmethod
|
||
async def recognize(self, image_path: str) -> OCRResult:
|
||
"""
|
||
识别图片中的文字。
|
||
|
||
Args:
|
||
image_path: 图片文件路径
|
||
|
||
Returns:
|
||
OCRResult 包含识别文本、文本块和置信度
|
||
|
||
Raises:
|
||
Exception: 识别失败时抛出异常
|
||
"""
|
||
...
|
||
|
||
|
||
# ──────────────────────────── PaddleOCR 实现（本地） ────────────────────────────
|
||
|
||
class PaddleOCRProvider(OCRProviderBase):
|
||
"""
|
||
PaddleOCR 本地识别
|
||
优点：免费、离线可用、中文识别效果好
|
||
缺点：需要较多内存、首次加载模型较慢
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._ocr = None
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "paddleocr"
|
||
|
||
def _init_ocr(self):
|
||
"""延迟初始化 PaddleOCR 引擎"""
|
||
if self._ocr is not None:
|
||
return
|
||
|
||
logger.info("正在初始化 PaddleOCR 引擎...")
|
||
from paddleocr import PaddleOCR
|
||
|
||
self._ocr = PaddleOCR(
|
||
use_angle_cls=True,
|
||
lang="ch",
|
||
use_gpu=False, # 默认 CPU，生产环境可通过环境变量控制
|
||
show_log=False,
|
||
enable_mkldnn=True, # 启用 MKLDNN 加速
|
||
)
|
||
logger.info("PaddleOCR 引擎初始化完成")
|
||
|
||
async def recognize(self, image_path: str) -> OCRResult:
|
||
"""使用 PaddleOCR 识别图片"""
|
||
self._init_ocr()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def _run_ocr():
|
||
result = self._ocr.ocr(image_path, cls=True)
|
||
return result
|
||
|
||
raw_result = await loop.run_in_executor(None, _run_ocr)
|
||
|
||
# 解析 PaddleOCR 返回结果
|
||
blocks: List[OCRBlock] = []
|
||
all_text_parts: List[str] = []
|
||
total_confidence = 0.0
|
||
|
||
if raw_result and raw_result[0]:
|
||
for line in raw_result[0]:
|
||
bbox_points = line[0]
|
||
text = line[1][0]
|
||
confidence = float(line[1][1])
|
||
|
||
# 将四个角点转为 [x1, y1, x2, y2]
|
||
xs = [p[0] for p in bbox_points]
|
||
ys = [p[1] for p in bbox_points]
|
||
bbox = [min(xs), min(ys), max(xs), max(ys)]
|
||
|
||
blocks.append(OCRBlock(text=text, bbox=bbox, confidence=confidence))
|
||
all_text_parts.append(text)
|
||
total_confidence += confidence
|
||
|
||
full_text = "\n".join(all_text_parts)
|
||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||
|
||
return OCRResult(
|
||
text=full_text,
|
||
blocks=blocks,
|
||
confidence=avg_confidence,
|
||
provider=self.name,
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 阿里云 OCR 实现 ────────────────────────────
|
||
|
||
class AliyunOCRProvider(OCRProviderBase):
|
||
"""
|
||
阿里云 OCR 识别
|
||
使用阿里云文字识别 API（通用文字识别）
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
if not settings.ALIYUN_OCR_ACCESS_KEY or not settings.ALIYUN_OCR_SECRET:
|
||
raise ValueError("阿里云 OCR 需要配置 ALIYUN_OCR_ACCESS_KEY 和 ALIYUN_OCR_SECRET")
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "aliyun"
|
||
|
||
async def recognize(self, image_path: str) -> OCRResult:
|
||
"""调用阿里云 OCR API 识别图片"""
|
||
import base64
|
||
import json
|
||
|
||
import asyncio
|
||
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def _call_api():
|
||
from alibabacloud_tea_openapi.models import Config
|
||
from alibabacloud_ocr_api20210707.client import Client
|
||
from alibabacloud_ocr_api20210707.models import RecognizeGeneralRequest
|
||
from alibabacloud_tea_util.models import RuntimeOptions
|
||
|
||
config = Config(
|
||
access_key_id=settings.ALIYUN_OCR_ACCESS_KEY,
|
||
access_key_secret=settings.ALIYUN_OCR_SECRET,
|
||
endpoint="ocr-api.cn-hangzhou.aliyuncs.com",
|
||
)
|
||
client = Client(config)
|
||
|
||
# 读取图片并 Base64 编码
|
||
with open(image_path, "rb") as f:
|
||
image_bytes = f.read()
|
||
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||
|
||
request = RecognizeGeneralRequest(
|
||
body=json.dumps({"url": f"data:image/png;base64,{image_b64}"}),
|
||
)
|
||
runtime = RuntimeOptions()
|
||
|
||
response = client.recognize_general_with_options(request, runtime)
|
||
return response
|
||
|
||
response = await loop.run_in_executor(None, _call_api)
|
||
|
||
if response.body and response.body.data:
|
||
data = json.loads(response.body.data)
|
||
blocks: List[OCRBlock] = []
|
||
all_text_parts: List[str] = []
|
||
total_confidence = 0.0
|
||
|
||
# 阿里云返回格式可能包含多个文本块
|
||
if isinstance(data, dict) and "content" in data:
|
||
text = data["content"]
|
||
all_text_parts.append(text)
|
||
blocks.append(OCRBlock(text=text, confidence=1.0))
|
||
total_confidence = 1.0
|
||
elif isinstance(data, list):
|
||
for item in data:
|
||
text = item.get("text", item.get("word", ""))
|
||
score = float(item.get("score", item.get("confidence", 1.0)))
|
||
blocks.append(OCRBlock(text=text, confidence=score))
|
||
all_text_parts.append(text)
|
||
total_confidence += score
|
||
|
||
full_text = "\n".join(all_text_parts)
|
||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||
|
||
return OCRResult(
|
||
text=full_text,
|
||
blocks=blocks,
|
||
confidence=avg_confidence,
|
||
provider=self.name,
|
||
)
|
||
|
||
return OCRResult(text="", blocks=[], confidence=0.0, provider=self.name)
|
||
|
||
|
||
# ──────────────────────────── 腾讯云 OCR 实现 ────────────────────────────
|
||
|
||
class TencentOCRProvider(OCRProviderBase):
|
||
"""
|
||
腾讯云 OCR 识别
|
||
使用腾讯云文字识别 API（通用印刷体识别）
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
if not settings.TENCENT_OCR_SECRET_ID or not settings.TENCENT_OCR_SECRET_KEY:
|
||
raise ValueError("腾讯云 OCR 需要配置 TENCENT_OCR_SECRET_ID 和 TENCENT_OCR_SECRET_KEY")
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "tencent"
|
||
|
||
async def recognize(self, image_path: str) -> OCRResult:
|
||
"""调用腾讯云 OCR API 识别图片"""
|
||
import asyncio
|
||
import base64
|
||
import json
|
||
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def _call_api():
|
||
from tencentcloud.common import credential
|
||
from tencentcloud.common.profile.client_profile import ClientProfile
|
||
from tencentcloud.common.profile.http_profile import HttpProfile
|
||
from tencentcloud.ocr.v20181119 import ocr_client, models
|
||
|
||
cred = credential.Credential(
|
||
settings.TENCENT_OCR_SECRET_ID,
|
||
settings.TENCENT_OCR_SECRET_KEY,
|
||
)
|
||
http_profile = HttpProfile()
|
||
http_profile.endpoint = "ocr.tencentcloudapi.com"
|
||
|
||
client_profile = ClientProfile()
|
||
client_profile.httpProfile = http_profile
|
||
|
||
client = ocr_client.OcrClient(cred, "", client_profile)
|
||
|
||
# 读取图片并 Base64 编码
|
||
with open(image_path, "rb") as f:
|
||
image_bytes = f.read()
|
||
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||
|
||
req = models.GeneralBasicOCRRequest()
|
||
req.ImageBase64 = image_b64
|
||
|
||
resp = client.GeneralBasicOCR(req)
|
||
return resp
|
||
|
||
resp = await loop.run_in_executor(None, _call_api)
|
||
|
||
blocks: List[OCRBlock] = []
|
||
all_text_parts: List[str] = []
|
||
total_confidence = 0.0
|
||
|
||
if resp.TextDetections:
|
||
for item in resp.TextDetections:
|
||
text = item.DetectedText or ""
|
||
confidence = item.Confidence / 100.0 if item.Confidence else 0.0
|
||
|
||
# 解析多边形顶点
|
||
bbox = None
|
||
if item.Polygon:
|
||
xs = [p.X for p in item.Polygon]
|
||
ys = [p.Y for p in item.Polygon]
|
||
bbox = [min(xs), min(ys), max(xs), max(ys)]
|
||
|
||
blocks.append(OCRBlock(text=text, bbox=bbox, confidence=confidence))
|
||
all_text_parts.append(text)
|
||
total_confidence += confidence
|
||
|
||
full_text = "\n".join(all_text_parts)
|
||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||
|
||
return OCRResult(
|
||
text=full_text,
|
||
blocks=blocks,
|
||
confidence=avg_confidence,
|
||
provider=self.name,
|
||
)
|
||
|
||
|
||
# ──────────────────────────── DeepSeek OCR 实现 ────────────────────────────
|
||
|
||
class DeepSeekOCRProvider(OCRProviderBase):
|
||
"""
|
||
DeepSeek Vision OCR
|
||
通过 OpenAI 兼容接口发送图片 base64，让视觉模型识别文字。
|
||
支持 DeepSeek 官方 API 和 Ollama 自部署（如 deepseek-ocr 模型）。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
if not settings.DEEPSEEK_API_KEY:
|
||
raise ValueError("DeepSeek OCR 需要配置 DEEPSEEK_API_KEY")
|
||
import httpx
|
||
from openai import AsyncOpenAI
|
||
self._client = AsyncOpenAI(
|
||
api_key=settings.DEEPSEEK_API_KEY,
|
||
base_url=settings.DEEPSEEK_BASE_URL,
|
||
http_client=httpx.AsyncClient(timeout=120.0), # OCR 较慢，超时 120 秒
|
||
)
|
||
self._model = settings.DEEPSEEK_OCR_MODEL
|
||
# 判断是否为 Ollama 部署
|
||
self._is_ollama = "ollama" in settings.DEEPSEEK_API_KEY.lower() or \
|
||
"11434" in settings.DEEPSEEK_BASE_URL
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "deepseek"
|
||
|
||
async def recognize(self, image_path: str) -> OCRResult:
|
||
"""使用 DeepSeek Vision 模型识别图片中的文字"""
|
||
import base64
|
||
|
||
# 读取图片并编码为 base64
|
||
with open(image_path, "rb") as f:
|
||
image_bytes = f.read()
|
||
base64_image = base64.b64encode(image_bytes).decode("utf-8")
|
||
|
||
# 根据图片格式确定 MIME 类型
|
||
suffix = Path(image_path).suffix.lower()
|
||
mime_map = {
|
||
".png": "image/png",
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".bmp": "image/bmp",
|
||
".webp": "image/webp",
|
||
}
|
||
mime_type = mime_map.get(suffix, "image/png")
|
||
|
||
system_prompt = (
|
||
"你是一个 OCR 文字识别工具。"
|
||
"请识别图片中的所有文字，逐行输出，不要添加任何解释。"
|
||
)
|
||
|
||
response = await self._client.chat.completions.create(
|
||
model=self._model,
|
||
messages=[
|
||
{
|
||
"role": "system",
|
||
"content": system_prompt,
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:{mime_type};base64,{base64_image}"
|
||
},
|
||
},
|
||
{
|
||
"type": "text",
|
||
"text": "请识别这张图片中的所有文字内容。",
|
||
},
|
||
],
|
||
},
|
||
],
|
||
max_tokens=4096,
|
||
temperature=0.0,
|
||
)
|
||
|
||
text = response.choices[0].message.content or ""
|
||
|
||
# 清理特殊 token（deepseek-ocr 模型可能输出 <｜begin of sentence｜> 等）
|
||
import re
|
||
text = re.sub(r'<｜[^>]*｜>', '', text)
|
||
text = re.sub(r'<\|[^>]*\|>', '', text)
|
||
|
||
# 清理可能的 markdown 包裹
|
||
text = text.strip()
|
||
if text.startswith("```"):
|
||
text = text.split("\n", 1)[-1]
|
||
if text.endswith("```"):
|
||
text = text[:-3]
|
||
text = text.strip()
|
||
|
||
lines = [line for line in text.split("\n") if line.strip()]
|
||
blocks = [OCRBlock(text=line, confidence=1.0) for line in lines]
|
||
|
||
return OCRResult(
|
||
text=text,
|
||
blocks=blocks,
|
||
confidence=1.0,
|
||
provider=self.name,
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 工厂类 ────────────────────────────
|
||
|
||
class OCRService:
|
||
"""
|
||
OCR 服务工厂类
|
||
根据配置自动选择 OCR 提供商。
|
||
auto 模式下：本地 PaddleOCR 优先，失败后 fallback 到云端 OCR。
|
||
"""
|
||
|
||
_instance: Optional[OCRProviderBase] = None
|
||
|
||
@classmethod
|
||
def _create_provider(cls, provider: OCRProvider) -> OCRProviderBase:
|
||
"""根据提供商枚举创建对应的 OCR 实例"""
|
||
provider_map = {
|
||
OCRProvider.PADDLEOCR: PaddleOCRProvider,
|
||
OCRProvider.ALIYUN: AliyunOCRProvider,
|
||
OCRProvider.TENCENT: TencentOCRProvider,
|
||
OCRProvider.DEEPSEEK: DeepSeekOCRProvider,
|
||
}
|
||
provider_cls = provider_map.get(provider)
|
||
if provider_cls is None:
|
||
raise ValueError(f"不支持的 OCR 提供商: {provider}")
|
||
return provider_cls()
|
||
|
||
@classmethod
|
||
def get_instance(cls) -> OCRProviderBase:
|
||
"""获取 OCR 服务单例实例"""
|
||
if cls._instance is not None:
|
||
return cls._instance
|
||
|
||
if settings.OCR_PROVIDER == OCRProvider.AUTO:
|
||
# auto 模式返回 PaddleOCR，fallback 逻辑在 recognize_auto 中处理
|
||
cls._instance = PaddleOCRProvider()
|
||
else:
|
||
cls._instance = cls._create_provider(settings.OCR_PROVIDER)
|
||
|
||
logger.info("OCR 服务初始化完成: provider=%s", cls._instance.name)
|
||
return cls._instance
|
||
|
||
@classmethod
|
||
async def recognize(cls, image_path: str) -> OCRResult:
|
||
"""
|
||
识别图片文字。
|
||
auto 模式下：PaddleOCR 优先，失败后依次尝试阿里云、腾讯云。
|
||
"""
|
||
if settings.OCR_PROVIDER == OCRProvider.AUTO:
|
||
return await cls._recognize_auto(image_path)
|
||
|
||
instance = cls.get_instance()
|
||
return await instance.recognize(image_path)
|
||
|
||
@classmethod
|
||
async def _recognize_auto(cls, image_path: str) -> OCRResult:
|
||
"""
|
||
auto 模式识别逻辑：
|
||
1. 先尝试本地 PaddleOCR
|
||
2. 如果失败且配置了阿里云，尝试阿里云 OCR
|
||
3. 如果仍失败且配置了腾讯云，尝试腾讯云 OCR
|
||
"""
|
||
# 第一步：本地 PaddleOCR
|
||
try:
|
||
provider = PaddleOCRProvider()
|
||
result = await provider.recognize(image_path)
|
||
if result.text.strip():
|
||
logger.info("PaddleOCR 识别成功，文本长度: %d", len(result.text))
|
||
return result
|
||
logger.warning("PaddleOCR 返回空文本，尝试 fallback")
|
||
except Exception as exc:
|
||
logger.warning("PaddleOCR 识别失败: %s，尝试 fallback", exc)
|
||
|
||
# 第二步：DeepSeek OCR（如果配置了）
|
||
if settings.DEEPSEEK_API_KEY:
|
||
try:
|
||
provider = DeepSeekOCRProvider()
|
||
result = await provider.recognize(image_path)
|
||
if result.text.strip():
|
||
logger.info("DeepSeek OCR 识别成功，文本长度: %d", len(result.text))
|
||
return result
|
||
logger.warning("DeepSeek OCR 返回空文本，尝试 fallback")
|
||
except Exception as exc:
|
||
logger.warning("DeepSeek OCR 识别失败: %s，尝试 fallback", exc)
|
||
|
||
# 第三步：阿里云 OCR
|
||
if settings.ALIYUN_OCR_ACCESS_KEY and settings.ALIYUN_OCR_SECRET:
|
||
try:
|
||
provider = AliyunOCRProvider()
|
||
result = await provider.recognize(image_path)
|
||
if result.text.strip():
|
||
logger.info("阿里云 OCR 识别成功，文本长度: %d", len(result.text))
|
||
return result
|
||
logger.warning("阿里云 OCR 返回空文本，继续 fallback")
|
||
except Exception as exc:
|
||
logger.warning("阿里云 OCR 识别失败: %s，继续 fallback", exc)
|
||
|
||
# 第四步：腾讯云 OCR
|
||
if settings.TENCENT_OCR_SECRET_ID and settings.TENCENT_OCR_SECRET_KEY:
|
||
try:
|
||
provider = TencentOCRProvider()
|
||
result = await provider.recognize(image_path)
|
||
if result.text.strip():
|
||
logger.info("腾讯云 OCR 识别成功，文本长度: %d", len(result.text))
|
||
return result
|
||
logger.warning("腾讯云 OCR 返回空文本")
|
||
except Exception as exc:
|
||
logger.warning("腾讯云 OCR 识别失败: %s", exc)
|
||
|
||
# 所有 OCR 都失败
|
||
return OCRResult(
|
||
text="",
|
||
blocks=[],
|
||
confidence=0.0,
|
||
provider="none",
|
||
)
|
||
|
||
@classmethod
|
||
def reset(cls) -> None:
|
||
"""重置单例（用于测试或切换配置）"""
|
||
cls._instance = None
|