File size: 5,664 Bytes
ba8abef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | #!/usr/bin/env python3
"""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
|