315 lines
13 KiB
Python
315 lines
13 KiB
Python
"""
|
|
OCR服务模块
|
|
---------
|
|
提供OCR识别服务，协调OCR流程。
|
|
"""
|
|
|
|
from typing import Dict, List, Optional, Tuple, Union, Any, Callable
|
|
import os
|
|
|
|
from ..config.settings import ConfigManager
|
|
from ..core.utils.log_utils import get_logger
|
|
from ..core.ocr.table_ocr import OCRProcessor
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
class OCRService:
|
|
"""
|
|
OCR识别服务：协调OCR流程
|
|
"""
|
|
|
|
def __init__(self, config: Optional[ConfigManager] = None):
|
|
"""
|
|
初始化OCR服务
|
|
|
|
Args:
|
|
config: 配置管理器，如果为None则创建新的
|
|
"""
|
|
logger.info("初始化OCRService")
|
|
self.config = config or ConfigManager()
|
|
|
|
# 创建OCR处理器
|
|
self.ocr_processor = OCRProcessor(self.config)
|
|
|
|
logger.info("OCRService初始化完成")
|
|
|
|
def get_unprocessed_images(self) -> List[str]:
|
|
"""
|
|
获取待处理的图片列表
|
|
|
|
Returns:
|
|
待处理图片路径列表
|
|
"""
|
|
return self.ocr_processor.get_unprocessed_images()
|
|
|
|
def process_image(self, image_path: str) -> Optional[str]:
|
|
"""
|
|
处理单个图片文件
|
|
|
|
Args:
|
|
image_path: 图片文件路径
|
|
|
|
Returns:
|
|
生成的Excel文件路径，如果处理失败则返回None
|
|
"""
|
|
try:
|
|
# 检查文件是否存在
|
|
if not os.path.exists(image_path):
|
|
logger.error(f"文件不存在: {image_path}")
|
|
return None
|
|
|
|
# 检查文件类型
|
|
if not self._is_valid_image(image_path):
|
|
logger.error(f"不支持的文件类型: {image_path}")
|
|
return None
|
|
|
|
# 检查是否已处理
|
|
excel_file = self._get_excel_path(image_path)
|
|
if os.path.exists(excel_file):
|
|
logger.info(f"文件已处理过，跳过OCR识别: {image_path}")
|
|
# 即使 xlsx 已存在，仍补写 meta.json（首次 OCR 后可能未生成）
|
|
meta_path = str(excel_file).replace('.xlsx', '.meta.json')
|
|
if not os.path.exists(meta_path):
|
|
self._write_meta_json(str(excel_file), {}, image_path)
|
|
return excel_file
|
|
|
|
# 执行OCR识别
|
|
result = self.ocr_processor.process_image(image_path)
|
|
if not result:
|
|
logger.error(f"OCR识别失败: {image_path}")
|
|
return None
|
|
|
|
# 生成Excel文件
|
|
excel_file = self._generate_excel(result, image_path)
|
|
if not excel_file:
|
|
logger.error(f"生成Excel文件失败: {image_path}")
|
|
return None
|
|
|
|
# 写 meta.json（从 ocr_client 再做一次通用文字识别获取抬头/日期）
|
|
self._write_meta_json(excel_file, {}, image_path)
|
|
|
|
logger.info(f"处理完成: {image_path} -> {excel_file}")
|
|
return excel_file
|
|
|
|
except Exception as e:
|
|
logger.error(f"处理图片时发生错误: {e}", exc_info=True)
|
|
return None
|
|
|
|
def process_images_batch(self, batch_size: int = None, max_workers: int = None, progress_cb: Optional[Callable[[int], None]] = None) -> Tuple[int, int]:
|
|
"""
|
|
批量处理图片
|
|
|
|
Args:
|
|
batch_size: 批处理大小
|
|
max_workers: 最大线程数
|
|
|
|
Returns:
|
|
(总处理数, 成功处理数)元组
|
|
"""
|
|
logger.info(f"OCRService开始批量处理图片, batch_size={batch_size}, max_workers={max_workers}")
|
|
return self.ocr_processor.process_images_batch(batch_size, max_workers, progress_cb)
|
|
|
|
# 添加batch_process作为process_images_batch的别名，确保兼容性
|
|
def batch_process(self, batch_size: int = None, max_workers: int = None, progress_cb: Optional[Callable[[int], None]] = None) -> Tuple[int, int]:
|
|
"""
|
|
批量处理图片（别名方法，与process_images_batch功能相同）
|
|
|
|
Args:
|
|
batch_size: 批处理大小
|
|
max_workers: 最大线程数
|
|
|
|
Returns:
|
|
(总处理数, 成功处理数)元组
|
|
"""
|
|
logger.info(f"OCRService.batch_process被调用，转发到process_images_batch")
|
|
return self.process_images_batch(batch_size, max_workers, progress_cb)
|
|
|
|
def validate_image(self, image_path: str) -> bool:
|
|
"""
|
|
验证图片是否有效
|
|
|
|
Args:
|
|
image_path: 图片路径
|
|
|
|
Returns:
|
|
图片是否有效
|
|
"""
|
|
return self.ocr_processor.validate_image(image_path)
|
|
|
|
def _is_valid_image(self, image_path: str) -> bool:
|
|
"""
|
|
检查文件是否为有效的图片格式
|
|
|
|
Args:
|
|
image_path: 图片文件路径
|
|
|
|
Returns:
|
|
是否为有效图片格式
|
|
"""
|
|
return self.validate_image(image_path)
|
|
|
|
def _get_excel_path(self, image_path: str) -> str:
|
|
"""
|
|
根据图片路径生成对应的Excel文件路径
|
|
|
|
Args:
|
|
image_path: 图片文件路径
|
|
|
|
Returns:
|
|
Excel文件路径
|
|
"""
|
|
# 获取文件名（不含扩展名）
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
# 生成Excel文件路径
|
|
output_dir = self.config.get_path('Paths', 'output_folder', fallback='data/output', create=True) if hasattr(self.config, 'get_path') else os.path.abspath('data/output')
|
|
excel_path = os.path.join(output_dir, f"{base_name}.xlsx")
|
|
return excel_path
|
|
|
|
def _generate_excel(self, ocr_result: dict, image_path: str) -> Optional[str]:
|
|
"""
|
|
根据OCR结果生成Excel文件
|
|
|
|
Args:
|
|
ocr_result: OCR识别结果
|
|
image_path: 原始图片路径
|
|
|
|
Returns:
|
|
生成的Excel文件路径，失败返回None
|
|
"""
|
|
try:
|
|
excel_path = self._get_excel_path(image_path)
|
|
|
|
# 确保输出目录存在
|
|
os.makedirs(os.path.dirname(excel_path), exist_ok=True)
|
|
|
|
# 调用OCR处理器的Excel生成功能
|
|
if hasattr(self.ocr_processor, 'generate_excel'):
|
|
success = self.ocr_processor.generate_excel(ocr_result, excel_path)
|
|
if success:
|
|
return excel_path
|
|
else:
|
|
# OCR 处理器（table_ocr.OCRProcessor）已直接生成 xlsx
|
|
if os.path.exists(excel_path):
|
|
return excel_path
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"生成Excel文件时发生错误: {e}", exc_info=True)
|
|
return None
|
|
|
|
def _write_meta_json(self, excel_path: str, ocr_result: dict, image_path: str) -> None:
|
|
"""从 OCR 结果中抽取原始文本，写入与 Excel 同名的 .meta.json。
|
|
|
|
用途：供元信息提取器（OrderMetadataExtractor）使用，无需重新 OCR。
|
|
失败不影响主流程。
|
|
"""
|
|
try:
|
|
import json
|
|
import base64
|
|
from pathlib import Path
|
|
|
|
base = Path(excel_path)
|
|
meta_path = base.with_suffix('.meta.json')
|
|
|
|
# 0) 优先：调百度通用文字识别（/accurate），覆盖全图文字（含手写抬头/日期）
|
|
general_text = ''
|
|
general_lines = []
|
|
try:
|
|
# OCRService.ocr_processor = core.ocr.table_ocr.OCRProcessor
|
|
# 其 .ocr_client = BaiduOCRClient (process_image 时初始化)
|
|
client = getattr(self.ocr_processor, 'ocr_client', None)
|
|
if client is None:
|
|
# 显式触发一次 process_image 准备流程 (不会重复 OCR)
|
|
try:
|
|
self.ocr_processor._ensure_ocr_client()
|
|
except Exception:
|
|
pass
|
|
client = getattr(self.ocr_processor, 'ocr_client', None)
|
|
if client and hasattr(client, 'recognize_general') and image_path and os.path.exists(str(image_path)):
|
|
words = client.recognize_general(str(image_path))
|
|
if words:
|
|
# 按 location.top 排序（顶部先），方便后续提取供应商/日期
|
|
def _top(w):
|
|
loc = w.get('location') or {}
|
|
try:
|
|
return float(loc.get('top', 0))
|
|
except (TypeError, ValueError):
|
|
return 0.0
|
|
words_sorted = sorted(words, key=_top)
|
|
general_lines = [(w.get('words') or '').strip() for w in words_sorted if (w.get('words') or '').strip()]
|
|
general_text = '\n'.join(general_lines)
|
|
logger.info(f"通用识别获取 {len(general_lines)} 行文字")
|
|
except Exception as e:
|
|
logger.warning(f"通用识别失败（不影响主流程）: {e}")
|
|
|
|
# 1) 表格识别 header/body 拼表内文字（已在前面逻辑处理）
|
|
ocr_text = ''
|
|
ocr_rows = []
|
|
for key in ('text', 'raw_text', 'ocr_text', 'content'):
|
|
v = ocr_result.get(key)
|
|
if isinstance(v, str) and v.strip():
|
|
ocr_text = v
|
|
break
|
|
|
|
# 2) 如果没有，尝试从 tables_result 拼
|
|
if not ocr_text and isinstance(ocr_result, dict):
|
|
tables = ocr_result.get('tables_result') or []
|
|
if isinstance(tables, list):
|
|
lines = []
|
|
for t in tables:
|
|
if not isinstance(t, dict):
|
|
continue
|
|
# 关键：从 header / body / footer 都提取，header 含供应商/日期抬头
|
|
for region in ('header', 'body', 'footer'):
|
|
region_data = t.get(region) or []
|
|
if not isinstance(region_data, list):
|
|
continue
|
|
for cell in region_data:
|
|
if not isinstance(cell, dict):
|
|
continue
|
|
words = cell.get('words') or cell.get('word') or ''
|
|
if not words:
|
|
continue
|
|
row = cell.get('row') or []
|
|
col = cell.get('column') or []
|
|
# 按行归一化拼成文本
|
|
if isinstance(row, list) and row:
|
|
row_key = f"r{row[0]}"
|
|
else:
|
|
row_key = f"r{len(lines)}"
|
|
# 简化：直接每行一个 cell
|
|
lines.append(str(words).strip())
|
|
ocr_rows.append([str(words)])
|
|
if not ocr_text and lines:
|
|
ocr_text = '\n'.join(lines)
|
|
|
|
# 3) 兜底：从 Excel 文件读 cell 拼
|
|
if not ocr_text:
|
|
try:
|
|
import xlrd
|
|
rb = xlrd.open_workbook(str(excel_path))
|
|
ws = rb.sheet_by_index(0)
|
|
lines = []
|
|
for r in range(ws.nrows):
|
|
row_vals = [str(ws.cell_value(r, c)) for c in range(ws.ncols)]
|
|
ocr_rows.append(row_vals)
|
|
lines.append(' '.join(row_vals))
|
|
ocr_text = '\n'.join(lines)
|
|
except Exception as e:
|
|
logger.debug(f"从 xlsx 读 OCR 文本失败: {e}")
|
|
|
|
payload = {
|
|
'excel_path': str(excel_path),
|
|
'image_path': str(image_path),
|
|
'ocr_text': ocr_text,
|
|
'ocr_rows': ocr_rows,
|
|
'general_text': general_text,
|
|
'general_lines': general_lines,
|
|
'created_at': __import__('datetime').datetime.now().isoformat(timespec='seconds'),
|
|
}
|
|
meta_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding='utf-8')
|
|
logger.debug(f"meta.json 已写入: {meta_path}")
|
|
except Exception as e:
|
|
logger.warning(f"写 meta.json 失败（不影响主流程）: {e}")
|