baodan/api/insurance/document/ocr.py

"""按页预算执行 Tesseract OCR，并保留真实词块坐标。"""
from __future__ import annotations

import csv
import io
import os
import shutil
import subprocess
import tempfile
from collections import defaultdict


def default_ocr_provider():
    return tesseract_ocr_blocks if shutil.which("tesseract") else None


def tesseract_ocr_blocks(pdf_path: str, page_numbers: list[int]) -> dict[int, list[dict]]:
    executable = shutil.which("tesseract")
    if not executable or not page_numbers:
        return {}
    language = _available_language(executable)
    results: dict[int, list[dict]] = {}
    scale = 3.0
    with tempfile.TemporaryDirectory(prefix="insurance-document-ocr-") as temp_dir:
        rendered = _render_pages(pdf_path, page_numbers, temp_dir, scale)
        for page_number, image_path in rendered.items():
            completed = subprocess.run(
                [executable, image_path, "stdout", "-l", language, "--psm", "4", "--oem", "3", "tsv"],
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=90,
                check=False,
            )
            if completed.returncode != 0 or not completed.stdout.strip():
                continue
            blocks = _tsv_to_blocks(completed.stdout, scale)
            if blocks:
                results[page_number] = blocks
    return results


def _render_pages(pdf_path: str, page_numbers: list[int], temp_dir: str, scale: float) -> dict[int, str]:
    try:
        try:
            import fitz
        except ImportError:
            import pymupdf as fitz
        rendered = {}
        with fitz.open(pdf_path) as document:
            for page_number in page_numbers:
                if page_number < 1 or page_number > len(document):
                    continue
                page = document[page_number - 1]
                pixmap = page.get_pixmap(matrix=fitz.Matrix(scale, scale), colorspace=fitz.csGRAY)
                image_path = os.path.join(temp_dir, f"page-{page_number}.png")
                pixmap.save(image_path)
                rendered[page_number] = image_path
        return rendered
    except ImportError:
        pass
    try:
        import pypdfium2 as pdfium
    except ImportError:
        return {}
    rendered = {}
    document = pdfium.PdfDocument(pdf_path)
    try:
        for page_number in page_numbers:
            if page_number < 1 or page_number > len(document):
                continue
            page = document[page_number - 1]
            image = page.render(scale=scale, grayscale=True).to_pil()
            image_path = os.path.join(temp_dir, f"page-{page_number}.png")
            image.save(image_path)
            rendered[page_number] = image_path
            page.close()
    finally:
        document.close()
    return rendered


def _available_language(executable: str) -> str:
    try:
        completed = subprocess.run(
            [executable, "--list-langs"],
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=10,
            check=False,
        )
        available = set(completed.stdout.lower().split())
    except Exception:
        available = set()
    languages = [name for name in ("chi_tra", "chi_sim", "eng") if name in available]
    return "+".join(languages) or "eng"


def _tsv_to_blocks(payload: str, scale: float) -> list[dict]:
    grouped = defaultdict(list)
    reader = csv.DictReader(io.StringIO(payload), delimiter="\t")
    for row in reader:
        text = str(row.get("text") or "").strip()
        if not text or str(row.get("level") or "") != "5":
            continue
        try:
            confidence = float(row.get("conf") or -1)
            left = float(row.get("left") or 0) / scale
            top = float(row.get("top") or 0) / scale
            width = float(row.get("width") or 0) / scale
            height = float(row.get("height") or 0) / scale
        except (TypeError, ValueError):
            continue
        if confidence < 0 or width <= 0 or height <= 0:
            continue
        key = (row.get("block_num"), row.get("par_num"), row.get("line_num"))
        grouped[key].append((left, top, left + width, top + height, text, confidence))

    blocks = []
    for words in grouped.values():
        words.sort(key=lambda item: item[0])
        blocks.append({
            "bbox": [
                min(word[0] for word in words),
                min(word[1] for word in words),
                max(word[2] for word in words),
                max(word[3] for word in words),
            ],
            "text": " ".join(word[4] for word in words),
            "confidence": round(sum(word[5] for word in words) / len(words) / 100.0, 4),
        })
    return blocks