bid-doctor/spikes/s5-ocr/ocr_pages.py

"""S5：对 PDF 里的扫描页整页 OCR、对文字页里的图片区域 OCR，结果换算回 PDF 坐标。

用法：.venv/bin/python -I ocr_pages.py <pdf> <out.json> [页码范围，如 16-33,104-131]
不给页码时按文字层规则自动挑：扫描页（字少图多）整页识别，图片占比 ≥ 0.15 的文字页只识别图片区域。
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np
import pdfplumber
import pypdfium2 as pdfium
from PIL import Image, ImageOps
from rapidocr import RapidOCR

DPI = 200
RCH = False  # main() 按 --rch 打开
SCALE = DPI / 72


def parse_ranges(s: str) -> set[int]:
    out: set[int] = set()
    for part in s.split(","):
        a, _, b = part.partition("-")
        out.update(range(int(a), int(b or a) + 1))
    return out


def image_cover(page) -> float:
    area = float(page.width) * float(page.height)
    total = 0.0
    for im in page.images:
        w = min(float(page.width), float(im["x1"])) - max(0.0, float(im["x0"]))
        h = min(float(page.height), float(im["bottom"])) - max(0.0, float(im["top"]))
        if w > 0 and h > 0:
            total += w * h
    return min(total / area, 1.0) if area else 0.0


def is_scanned(chars: int, cover: float) -> bool:
    return chars < 20 or (cover >= 0.2 and chars < 60) or (cover >= 0.45 and chars < 120)


def run_ocr(eng: RapidOCR, img: Image.Image) -> list[tuple[list[list[float]], str, float]]:
    r = eng(np.asarray(img.convert("RGB")))
    if r.boxes is None:
        return []
    return [(b.tolist(), t, float(s)) for b, t, s in zip(r.boxes, r.txts, r.scores)]


def red_ratio(img: Image.Image) -> float:
    a = np.asarray(img.convert("RGB").resize((img.width // 4, img.height // 4))).astype(np.int16)
    r, g, b = a[..., 0], a[..., 1], a[..., 2]
    return float(((r > 120) & (r - g > 40) & (r - b > 40)).mean())


def red_channel(img: Image.Image) -> Image.Image:
    """只取红色通道当灰度图：红章变浅、黑字不变（章压字的页用）。"""
    return ImageOps.autocontrast(img.convert("RGB").split()[0]).convert("RGB")


def best_orientation(eng: RapidOCR, img: Image.Image):
    """整页转正：0° 识别结果里竖长框多、置信度低时，再试 90° / 270°，取「字数 × 平均置信度」最高的。"""
    lines = run_ocr(eng, img)

    def quality(ls):
        n = sum(len(t) for _, t, _ in ls)
        return n * (sum(s for *_, s in ls) / len(ls) if ls else 0)

    tall = sum(1 for b, t, _ in lines if len(t) >= 3 and (b[2][1] - b[0][1]) > 1.5 * (b[1][0] - b[0][0]))
    if lines and tall < len(lines) * 0.3:
        return 0, img, lines
    best = (quality(lines), 0, img, lines)
    for rot in (90, 270):
        r_img = img.rotate(rot, expand=True)
        r_lines = run_ocr(eng, r_img)
        q = quality(r_lines)
        if q > best[0]:
            best = (q, rot, r_img, r_lines)
    return best[1], best[2], best[3]


def to_pdf_box(box, rot: int, img_w: int, img_h: int, ox: float, oy: float) -> list[float]:
    """OCR 框（转正后图片的像素）→ PDF 坐标（pt，左上原点），返回外接矩形 [x0, top, x1, bottom]。"""
    pts = []
    for x, y in box:
        # 逆转回原图像素：PIL rotate(rot) 是逆时针；原图宽高为 (img_w, img_h)
        if rot == 90:
            x, y = img_w - y, x
        elif rot == 270:
            x, y = y, img_h - x
        pts.append((ox + x / SCALE, oy + y / SCALE))
    xs, ys = [p[0] for p in pts], [p[1] for p in pts]
    return [round(min(xs), 1), round(min(ys), 1), round(max(xs), 1), round(max(ys), 1)]


def overlaps_text(bbox, chars) -> bool:
    """OCR 行和文字层已有的字重叠（扫描页上方的文字说明、图片上叠的文字），去掉，免得重复。"""
    x0, t, x1, b = bbox
    hit = sum(1 for c in chars if c["x0"] >= x0 - 2 and c["x1"] <= x1 + 2 and c["top"] >= t - 2 and c["bottom"] <= b + 2)
    return hit >= 2


def main() -> None:
    global RCH
    args = [a for a in sys.argv[1:] if a != "--rch"]
    RCH = "--rch" in sys.argv
    pdf_path, out_path = args[0], Path(args[1])
    only = parse_ranges(args[2]) if len(args) > 2 else None
    eng = RapidOCR()
    doc = pdfium.PdfDocument(pdf_path)
    results = []
    t_all = time.time()
    with pdfplumber.open(pdf_path) as pl:
        for i, page in enumerate(pl.pages):
            no = i + 1
            if only and no not in only:
                continue
            chars = [c for c in page.chars if str(c["text"]).strip()]
            cover = image_cover(page)
            scanned = is_scanned(len(chars), cover)
            if not scanned and cover < 0.15:
                continue
            t0 = time.time()
            full = doc[i].render(scale=SCALE).to_pil()
            regions = []
            if scanned:
                regions.append((0.0, 0.0, full))
            else:
                for im in page.images:
                    x0, top = max(0.0, float(im["x0"])), max(0.0, float(im["top"]))
                    x1, bot = min(float(page.width), float(im["x1"])), min(float(page.height), float(im["bottom"]))
                    if (x1 - x0) * (bot - top) < 0.05 * float(page.width) * float(page.height):
                        continue  # 小图：logo、印章、图标
                    crop = full.crop((int(x0 * SCALE), int(top * SCALE), int(x1 * SCALE), int(bot * SCALE)))
                    regions.append((x0, top, crop))
            lines, rots = [], []
            for ox, oy, img in regions:
                if RCH and scanned and red_ratio(img) > 0.002:
                    img = red_channel(img)
                rot, r_img, ls = best_orientation(eng, img) if scanned else (0, img, run_ocr(eng, img))
                rots.append(rot)
                for box, text, score in ls:
                    bbox = to_pdf_box(box, rot, img.width, img.height, ox, oy)
                    if overlaps_text(bbox, chars):
                        continue
                    lines.append({"bbox": bbox, "text": text, "score": round(score, 3)})
            results.append(
                {
                    "page": no,
                    "kind": "scan" if scanned else "figure",
                    "text_chars": len(chars),
                    "image_cover": round(cover, 3),
                    "rotation": rots,
                    "seconds": round(time.time() - t0, 2),
                    "lines": lines,
                }
            )
            print(f"p{no} {results[-1]['kind']} rot={rots} lines={len(lines)} {results[-1]['seconds']}s", flush=True)
    out_path.write_text(json.dumps(results, ensure_ascii=False, indent=1))
    print(f"total {len(results)} pages {round(time.time() - t_all, 1)}s")


if __name__ == "__main__":
    main()