"""端到端识别流水：文件 → QR → (可选)OCR → 字段抽取 → InvoiceResult 整体有 OCR_TOTAL_TIMEOUT_S 兜底；单页有 OCR_PAGE_TIMEOUT_S 兜底。 任一超时立即返回失败，不再死等。 QR 流程 (config: QR_FULL_OCR): - 先扫二维码取 开票时间/发票号/金额 (国家税务总局 8 字段逗号规范) - 扫到 + qr_full_ocr=false → 直接返回这 3 个字段 (快路径, 跳过 OCR) - 扫到 + qr_full_ocr=true → 继续跑 OCR, QR 的 3 字段覆盖 OCR 抽取结果 - 没扫到 / 格式不合法 → 直接判定为非发票 (不再跑 OCR) """ from __future__ import annotations import shutil import tempfile import time from pathlib import Path from typing import List from loguru import logger from app.config import settings from app.core import OCRTimeout, auto_rotate, enhance, pdf_to_images, recognize from app.models import InvoiceResult, InvoiceFields, OCRLine from app.services import decode_qr, extract_invoice, qr_has_any_field _IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tif", ".tiff"} def _save_upload(upload_bytes: bytes, suffix: str) -> Path: tmp = Path(tempfile.mkstemp(suffix=suffix)[1]) tmp.write_bytes(upload_bytes) return tmp def _ocr_image(image_path: Path) -> List[OCRLine]: img = auto_rotate(image_path) img = enhance(img) return recognize(img) def _overlay_qr_fields(fields: InvoiceFields, qr) -> InvoiceFields: """QR 解出的 3 字段优先, 没解到的保持 OCR 结果""" data = fields.model_dump() if qr.invoice_no: data["invoice_no"] = qr.invoice_no if qr.amount is not None: data["amount"] = qr.amount if qr.invoice_date: data["invoice_date"] = qr.invoice_date return InvoiceFields(**data) def _build_qr_only_result(qr, elapsed_ms: int, page_count: int) -> InvoiceResult: """快路径: 只用 QR 字段, 跳过 OCR""" fields = InvoiceFields( invoice_no=qr.invoice_no, amount=qr.amount, invoice_date=qr.invoice_date, ) return InvoiceResult( success=True, is_invoice=True, raw_text=f"[QR only] {qr.raw}", lines=[], fields=fields, page_count=page_count, engine="qr", elapsed_ms=elapsed_ms, from_qr=True, qr_raw=qr.raw, ) def recognize_file(filename: str, content: bytes) -> InvoiceResult: """识别入口（从字节流）：自动判断 PDF / 图片；任何环节超时立即返回失败""" t0 = time.time() suffix = Path(filename).suffix.lower() tmp_path = _save_upload(content, suffix or ".bin") try: return recognize_path(tmp_path, delete_after=True) except Exception: raise def recognize_path(file_path: Path, *, delete_after: bool = False) -> InvoiceResult: """识别入口（从磁盘路径）：不做上传字节校验，直接读本地文件 Args: file_path: 已在白名单校验过的本地文件路径 delete_after: True 表示临时文件用完删除（上传流场景），False 保留（原文件场景） """ t0 = time.time() suffix = file_path.suffix.lower() total_deadline = t0 + settings.ocr_total_timeout_s def _elapsed_ms() -> int: return int((time.time() - t0) * 1000) try: # ---------- 1. PDF / 图片 → 图片列表 ---------- try: if suffix == ".pdf": page_imgs = pdf_to_images(file_path) elif suffix in _IMG_EXTS: page_imgs = [file_path] else: return InvoiceResult( success=False, is_invoice=False, error=f"不支持的文件类型: {suffix}（仅支持 PDF / 图片）", error_code="unsupported", engine="paddleocr", elapsed_ms=_elapsed_ms(), ) except Exception as e: logger.exception("PDF/图片处理失败: {}", e) return InvoiceResult( success=False, is_invoice=False, error=f"PDF/图片处理失败: {e}", error_code="process_failed", engine="paddleocr", elapsed_ms=_elapsed_ms(), ) # ---------- 2. QR 优先识别 ---------- qr = decode_qr(page_imgs[0]) qr_ok = qr_has_any_field(qr) if qr_ok and not settings.qr_full_ocr: # 快路径: 扫到 QR 且配置为 fast, 直接返回 logger.info("QR 快路径: {} 耗时 {}ms", file_path.name, _elapsed_ms()) return _build_qr_only_result(qr, _elapsed_ms(), len(page_imgs)) if not qr_ok: # 没扫到 / 格式不合法 → 直接认为非发票, 不再跑 OCR reason = "未识别到发票二维码" if not qr.raw else "二维码格式不合法" logger.info("非发票 (无有效 QR): {} 耗时 {}ms", file_path.name, _elapsed_ms()) return InvoiceResult( success=False, is_invoice=False, error=f"{reason}（可能不是发票图片）", error_code="not_invoice", page_count=len(page_imgs), engine="paddleocr", elapsed_ms=_elapsed_ms(), from_qr=False, qr_raw=qr.raw or None, qr_error="no_qr" if not qr.raw else "bad_format", ) # ---------- 3. 每页 OCR（单页超时 + 总超时双重保护）---------- # 到这里 qr_ok=True, OCR 仅为补充信息; QR 已保证是发票 all_lines: List[OCRLine] = [] for idx, img_path in enumerate(page_imgs, start=1): remaining = total_deadline - time.time() if remaining <= 0: logger.warning("达到总超时 ({:.1f}s), 中断 OCR", settings.ocr_total_timeout_s) return InvoiceResult( success=False, is_invoice=False, error=f"达到总超时 ({settings.ocr_total_timeout_s}秒), 已识别 {idx-1}/{len(page_imgs)} 页", error_code="timeout", raw_text="\n".join(l.text for l in all_lines), lines=all_lines, page_count=len(page_imgs), engine="paddleocr", elapsed_ms=_elapsed_ms(), from_qr=True, qr_raw=qr.raw, ) page_timeout = min(settings.ocr_page_timeout_s, remaining) try: all_lines.extend(_ocr_image(img_path)) except OCRTimeout as e: logger.warning("第 {} 页 OCR 超时: {}", idx, e) return InvoiceResult( success=False, is_invoice=False, error=f"第 {idx} 页识别超时 ({page_timeout:.1f}秒)", error_code="timeout", raw_text="\n".join(l.text for l in all_lines), lines=all_lines, page_count=len(page_imgs), engine="paddleocr", elapsed_ms=_elapsed_ms(), from_qr=True, qr_raw=qr.raw, ) except Exception as e: logger.exception("第 {} 页 OCR 失败: {}", idx, e) return InvoiceResult( success=False, is_invoice=False, error=f"第 {idx} 页识别失败: {e}", error_code="ocr_failed", raw_text="\n".join(l.text for l in all_lines), lines=all_lines, page_count=len(page_imgs), engine="paddleocr", elapsed_ms=_elapsed_ms(), from_qr=True, qr_raw=qr.raw, ) # ---------- 4. 字段抽取 + QR 字段覆盖 ---------- raw_text = "\n".join(l.text for l in all_lines) fields = extract_invoice(raw_text, all_lines) fields = _overlay_qr_fields(fields, qr) elapsed = _elapsed_ms() logger.info("识别完成: {} 页={}, from_qr=true, 字段数={}, 耗时={}ms", file_path.name, len(page_imgs), sum(1 for f in fields.model_dump().values() if f), elapsed) return InvoiceResult( success=True, is_invoice=True, raw_text=raw_text, lines=all_lines, fields=fields, page_count=len(page_imgs), engine="paddleocr", elapsed_ms=elapsed, from_qr=True, qr_raw=qr.raw, ) finally: if delete_after: try: file_path.unlink(missing_ok=True) if suffix == ".pdf": shutil.rmtree(file_path.parent / f".{file_path.stem}_pages", ignore_errors=True) except Exception: pass