guoju0808/ry-ocr/app/services/recognize_service.py

"""端到端识别流水：文件 → 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