| |
| """PaddleOCR-VL client: parallel layout crops + merge. |
| |
| Proven on RTX 5090 + vLLM next to Qwen3-4B. Prompt is ``OCR:``, temperature 0. |
| """ |
| from __future__ import annotations |
|
|
| import base64 |
| import io |
| import re |
| import time |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| from openai import OpenAI |
| from PIL import Image |
|
|
| DEFAULT_BASE_URL = "http://127.0.0.1:18000/v1" |
| DEFAULT_MODEL = "PaddlePaddle/PaddleOCR-VL" |
| PROMPT = "OCR:" |
| _MIN_SIDE = 420 |
|
|
| _REPEAT_TOKEN = re.compile(r"(\S+)(?:\s+\1){6,}") |
| _REPEAT_CHAR = re.compile(r"(.{6,}?)\1{5,}", re.DOTALL) |
| _REFUSAL = re.compile( |
| r"(not fully clear|cannot extract|cannot准确|无法|没有文字)", |
| re.I, |
| ) |
|
|
|
|
| @dataclass |
| class RegionResult: |
| name: str |
| text: str |
| latency_s: float |
|
|
|
|
| def _to_data_url(img: Image.Image) -> str: |
| buf = io.BytesIO() |
| img.convert("RGB").save(buf, format="PNG", optimize=True) |
| return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii") |
|
|
|
|
| def _pad_min(img: Image.Image, min_side: int = _MIN_SIDE) -> Image.Image: |
| w, h = img.size |
| if w >= min_side and h >= min_side: |
| return img |
| nw, nh = max(w, min_side), max(h, min_side) |
| canvas = Image.new("RGB", (nw, nh), (255, 255, 255)) |
| canvas.paste(img, ((nw - w) // 2, (nh - h) // 2)) |
| return canvas |
|
|
|
|
| def _band(img: Image.Image, top: float, bot: float, left: float = 0.0, right: float = 1.0) -> Image.Image: |
| w, h = img.size |
| box = ( |
| int(w * left), |
| int(h * top), |
| max(int(w * left) + 8, int(w * right)), |
| max(int(h * top) + 8, int(h * bot)), |
| ) |
| return img.crop(box) |
|
|
|
|
| def regions_for(img: Image.Image) -> list[tuple[str, Image.Image]]: |
| w, h = img.size |
| out: list[tuple[str, Image.Image]] = [ |
| ("full", img), |
| ("header", _pad_min(_band(img, 0.0, 0.16))), |
| ("top_right", _pad_min(_band(img, 0.0, 0.22, 0.48, 1.0))), |
| ("right", _pad_min(_band(img, 0.0, 0.62, 0.52, 1.0))), |
| ("totals", _pad_min(_band(img, 0.34, 0.72))), |
| ] |
| if w < 900: |
| out.append(("masthead", img.crop((0, 0, w, min(h, 130))))) |
| return out |
|
|
|
|
| def strip_collapse(text: str) -> str: |
| if not text: |
| return text |
| m = _REPEAT_TOKEN.search(text) |
| if m: |
| text = text[: m.start()] + m.group(1) |
| m = _REPEAT_CHAR.search(text) |
| if m: |
| text = text[: m.start()] + m.group(1) |
| lines = text.splitlines() |
| cleaned: list[str] = [] |
| run = 0 |
| prev = None |
| short_run = 0 |
| for ln in lines: |
| key = ln.strip() |
| if len(key) <= 2: |
| short_run += 1 |
| if short_run >= 8: |
| break |
| continue |
| short_run = 0 |
| if key and key == prev: |
| run += 1 |
| if run >= 6: |
| break |
| continue |
| prev = key |
| run = 0 |
| cleaned.append(ln) |
| return "\n".join(cleaned).strip() |
|
|
|
|
| def is_junk_region(text: str) -> bool: |
| if not text or not text.strip(): |
| return True |
| if _REFUSAL.search(text): |
| return True |
| cjk = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff") |
| return cjk > 12 |
|
|
|
|
| class OcrEngine: |
| def __init__( |
| self, |
| base_url: str = DEFAULT_BASE_URL, |
| model: str = DEFAULT_MODEL, |
| timeout: float = 180.0, |
| ) -> None: |
| self.model = model |
| self.client = OpenAI(api_key="EMPTY", base_url=base_url, timeout=timeout) |
|
|
| def _complete(self, img: Image.Image, *, max_tokens: int, penalty: bool) -> str: |
| extra = {"repetition_penalty": 1.05} |
| if penalty: |
| extra = {"repetition_penalty": 1.12, "frequency_penalty": 0.22} |
| resp = self.client.chat.completions.create( |
| model=self.model, |
| temperature=0.0, |
| max_tokens=max_tokens, |
| extra_body=extra, |
| messages=[ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image_url", "image_url": {"url": _to_data_url(img)}}, |
| {"type": "text", "text": PROMPT}, |
| ], |
| } |
| ], |
| ) |
| return strip_collapse(resp.choices[0].message.content or "") |
|
|
| def ocr_image(self, image: Image.Image | Path | str) -> tuple[str, float, list[RegionResult]]: |
| if not isinstance(image, Image.Image): |
| image = Image.open(image).convert("RGB") |
| else: |
| image = image.convert("RGB") |
| jobs = regions_for(image) |
| t0 = time.perf_counter() |
| results: list[RegionResult] = [] |
|
|
| def run(name: str, crop: Image.Image) -> RegionResult: |
| t1 = time.perf_counter() |
| light = name in {"full", "masthead"} |
| text = self._complete( |
| crop, |
| max_tokens=1024 if name == "full" else 512, |
| penalty=not light, |
| ) |
| return RegionResult(name, text, time.perf_counter() - t1) |
|
|
| with ThreadPoolExecutor(max_workers=len(jobs)) as pool: |
| futs = [pool.submit(run, name, crop) for name, crop in jobs] |
| for fut in as_completed(futs): |
| results.append(fut.result()) |
|
|
| order = {"full": 0, "header": 1, "masthead": 2, "top_right": 3, "right": 4, "totals": 5} |
| results.sort(key=lambda r: order.get(r.name, 9)) |
| parts = [] |
| for r in results: |
| chunk = strip_collapse(r.text) |
| if chunk.strip() and not is_junk_region(chunk): |
| parts.append(chunk) |
| return "\n\n".join(parts), time.perf_counter() - t0, results |
|
|
