Add private local OCR comparison package
This commit is contained in:
OCR Team
@@ -0,0 +1,310 @@
|
||||
"""Offline worker for the full official PaddleOCR-VL layout pipeline.
|
||||
|
||||
The caller sends JPEG bytes on stdin. This process uses explicit local model
|
||||
directories; it writes a mode-0600 temporary file and deletes it afterward.
|
||||
No private images or OCR results are saved to the model cache or logs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
ROOT = Path(os.environ.get("OCR_MODEL_HOME", Path(__file__).resolve().parents[1] / ".cache")).expanduser().resolve()
|
||||
os.environ["PADDLE_PDX_CACHE_HOME"] = str(ROOT)
|
||||
os.environ["HF_HOME"] = str(ROOT / "hf")
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
||||
os.environ["OMP_NUM_THREADS"] = "1"
|
||||
|
||||
|
||||
class _TableRows(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.rows = []
|
||||
self.row = None
|
||||
self.cell = None
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
if tag == "tr":
|
||||
self.row = []
|
||||
elif tag in ("td", "th") and self.row is not None:
|
||||
self.cell = []
|
||||
elif tag == "br" and self.cell is not None:
|
||||
self.cell.append(" ")
|
||||
|
||||
def handle_data(self, data):
|
||||
if self.cell is not None:
|
||||
self.cell.append(data)
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
if tag in ("td", "th") and self.cell is not None and self.row is not None:
|
||||
self.row.append(" ".join("".join(self.cell).split())[:500])
|
||||
self.cell = None
|
||||
elif tag == "tr" and self.row is not None:
|
||||
if self.row:
|
||||
self.rows.append(self.row)
|
||||
self.row = None
|
||||
|
||||
|
||||
def _table_rows(content: str) -> list[list[str]]:
|
||||
parser = _TableRows()
|
||||
parser.feed(content)
|
||||
parser.close()
|
||||
return parser.rows
|
||||
|
||||
|
||||
def _box(value, width: int, height: int, offset=(0, 0)):
|
||||
if value is None or len(value) != 4:
|
||||
return None
|
||||
try:
|
||||
x1, y1, x2, y2 = (float(item) for item in value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if not (0 <= x1 < x2 <= width and 0 <= y1 < y2 <= height):
|
||||
return None
|
||||
ox, oy = offset
|
||||
return [round(x1 + ox, 1), round(y1 + oy, 1),
|
||||
round(x2 + ox, 1), round(y2 + oy, 1)]
|
||||
|
||||
|
||||
def extract(payload: dict, offset=(0, 0)) -> dict:
|
||||
"""Keep original block text and its actual layout region, never fake line boxes."""
|
||||
width, height = int(payload["width"]), int(payload["height"])
|
||||
lines, regions = [], []
|
||||
for block in payload.get("parsing_res_list", []):
|
||||
box = _box(block.get("block_bbox"), width, height, offset)
|
||||
label = str(block.get("block_label", "region"))[:80]
|
||||
content = block.get("block_content", "")
|
||||
if not isinstance(content, str):
|
||||
content = str(content)
|
||||
if box:
|
||||
regions.append({"label": label, "box": box})
|
||||
if label == "table" and "<tr" in content.lower():
|
||||
for cells in _table_rows(content):
|
||||
text = " | ".join(cells)
|
||||
if text.strip():
|
||||
lines.append({"text": text[:1000], "cells": cells,
|
||||
"box": box, "source": "PaddleOCR-VL-Tabellenzeile; nur Tabellenregion als Bildbeleg"})
|
||||
else:
|
||||
for part in content.splitlines():
|
||||
text = part.strip()
|
||||
if text:
|
||||
lines.append({"text": text[:1000], "box": box,
|
||||
"source": f"PaddleOCR-VL-Layoutblock {label}; keine eigene Zeilenbox"})
|
||||
return {
|
||||
"model": "PaddleOCR-VL 1.6 · vollständige lokale Layoutpipeline",
|
||||
"lines": lines[:500], "regions": regions[:200],
|
||||
"layout_note": "Vision-Bildbelege markieren echte Layoutregionen. Mehrere Textzeilen können dieselbe Region teilen; die Region ist keine präzise Feldbox.",
|
||||
}
|
||||
|
||||
|
||||
def _retry_table(block: dict, width: int, height: int) -> bool:
|
||||
if block.get("block_label") != "table":
|
||||
return False
|
||||
box = _box(block.get("block_bbox"), width, height)
|
||||
if not box:
|
||||
return False
|
||||
x1, y1, x2, y2 = box
|
||||
area = (x2 - x1) * (y2 - y1) / (width * height)
|
||||
content = block.get("block_content", "")
|
||||
if not isinstance(content, str):
|
||||
return False
|
||||
visible_text = re.sub(r"<[^>]*>", "", content).strip()
|
||||
return area > .45 and len(visible_text) < 80
|
||||
|
||||
|
||||
def _split_table(box: list[float], width: int, height: int) -> list[tuple[int, int, int, int]]:
|
||||
x1, y1, x2, y2 = box
|
||||
left = max(0, int(x1) - 20)
|
||||
top = max(0, int(y1) - 20)
|
||||
right = min(width, int(x2) + 20)
|
||||
bottom = min(height, int(y2) + 20)
|
||||
middle = (left + right) // 2
|
||||
overlap = min(120, max(40, round((right - left) * .045)))
|
||||
return [(left, top, min(right, middle + overlap), bottom),
|
||||
(max(left, middle - overlap), top, right, bottom)]
|
||||
|
||||
|
||||
def _missing_neighbor_box(block: dict, blocks: list[dict], width: int, height: int):
|
||||
"""Find an uncovered page column directly left of a detected right table."""
|
||||
if block.get("block_label") != "table":
|
||||
return None
|
||||
box = _box(block.get("block_bbox"), width, height)
|
||||
if not box:
|
||||
return None
|
||||
x1, y1, x2, y2 = box
|
||||
table_width, table_height = x2 - x1, y2 - y1
|
||||
if x1 < .52 * width or table_width > .4 * width or table_height < .2 * height:
|
||||
return None
|
||||
left = max(0, round(x1 - 1.2 * table_width))
|
||||
right = min(width, round(x1 + .07 * table_width))
|
||||
top = max(0, round(y1 - .1 * table_height))
|
||||
bottom = min(height, round(y2 + .3 * table_height))
|
||||
if right - left < 300 or bottom - top < 300:
|
||||
return None
|
||||
for other in blocks:
|
||||
if other is block or other.get("block_label") != "table":
|
||||
continue
|
||||
other_box = _box(other.get("block_bbox"), width, height)
|
||||
if not other_box:
|
||||
continue
|
||||
a, b, c, d = other_box
|
||||
overlap = max(0, min(x1, c) - max(left, a)) * max(0, min(bottom, d) - max(top, b))
|
||||
if overlap > .15 * (x1 - left) * (bottom - top):
|
||||
return None
|
||||
return (left, top, right, bottom)
|
||||
|
||||
|
||||
def _run_tile(pipeline, image: Image.Image, crop_box, private: Path):
|
||||
descriptor, name = tempfile.mkstemp(prefix="table-", suffix=".jpg", dir=private)
|
||||
tile_path = Path(name)
|
||||
try:
|
||||
with os.fdopen(descriptor, "wb") as stream:
|
||||
image.crop(crop_box).save(stream, format="JPEG", quality=92)
|
||||
results = list(pipeline.predict(str(tile_path)))
|
||||
if len(results) != 1:
|
||||
return None
|
||||
payload = results[0].json
|
||||
return extract(payload.get("res", payload), crop_box[:2])
|
||||
except Exception:
|
||||
# Keep the original page result if a diagnostic retry fails.
|
||||
return None
|
||||
finally:
|
||||
tile_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def predict_tiled(pipeline, path: Path, private: Path) -> dict:
|
||||
"""Generic full-page coverage for a wide page whose full parsing timed out."""
|
||||
with Image.open(path) as image:
|
||||
width, height = image.size
|
||||
boxes = [(0, 0, round(width * .34), height),
|
||||
(round(width * .25), 0, round(width * .66), height),
|
||||
(round(width * .61), 0, width, height)]
|
||||
tiles = [_run_tile(pipeline, image, box, private) for box in boxes]
|
||||
valid = [tile for tile in tiles if tile and tile["lines"]]
|
||||
if not valid:
|
||||
raise RuntimeError("tile_inference_failed")
|
||||
result = extract({"width": width, "height": height, "parsing_res_list": []})
|
||||
for tile in valid:
|
||||
result["lines"].extend(tile["lines"])
|
||||
result["regions"].extend(tile["regions"])
|
||||
result["lines"] = result["lines"][:500]
|
||||
result["regions"] = result["regions"][:200]
|
||||
result["retry_used"] = True
|
||||
result["layout_note"] = (
|
||||
f"Die Gesamtansicht brauchte zu lange. {len(valid)} von 3 überlappenden "
|
||||
"Ansichten derselben Seite wurden lokal gelesen; widersprüchliche "
|
||||
"Feldwerte bleiben ohne Zuordnung. Tabellenbelege sind grobe Regionen.")
|
||||
return result
|
||||
|
||||
|
||||
def predict_adaptive(pipeline, path: Path, private: Path) -> dict:
|
||||
results = list(pipeline.predict(str(path)))
|
||||
if len(results) != 1:
|
||||
raise RuntimeError("invalid_page_count")
|
||||
original = results[0].json
|
||||
payload = original.get("res", original)
|
||||
width, height = int(payload["width"]), int(payload["height"])
|
||||
blocks = payload.get("parsing_res_list", [])
|
||||
large_table = next((block for block in blocks if _retry_table(block, width, height)), None)
|
||||
neighbor_box = next((box for block in blocks
|
||||
if (box := _missing_neighbor_box(block, blocks, width, height))), None)
|
||||
if large_table is None and neighbor_box is None:
|
||||
return extract(payload)
|
||||
|
||||
# A large, almost empty table is a measurable layout/scale failure. Retry
|
||||
# only that detected region in two smaller views; all pixels still come
|
||||
# from the same upload and all coordinates map back to the common image.
|
||||
tiles = []
|
||||
large_recovered = False
|
||||
neighbor_recovered = False
|
||||
with Image.open(path) as image:
|
||||
if large_table is not None:
|
||||
box = _box(large_table["block_bbox"], width, height)
|
||||
for crop_box in _split_table(box, width, height):
|
||||
parsed = _run_tile(pipeline, image, crop_box, private)
|
||||
if parsed and parsed["lines"]:
|
||||
tiles.append(parsed)
|
||||
large_recovered = True
|
||||
if neighbor_box is not None:
|
||||
parsed = _run_tile(pipeline, image, neighbor_box, private)
|
||||
if parsed and parsed["lines"]:
|
||||
tiles.append(parsed)
|
||||
neighbor_recovered = True
|
||||
if not tiles:
|
||||
return extract(payload)
|
||||
kept = [block for block in blocks if block is not large_table] if large_recovered else blocks
|
||||
result = extract({**payload, "parsing_res_list": kept})
|
||||
for tile in tiles:
|
||||
result["lines"].extend(tile["lines"])
|
||||
result["regions"].extend(tile["regions"])
|
||||
result["lines"] = result["lines"][:500]
|
||||
result["regions"] = result["regions"][:200]
|
||||
result["retry_used"] = True
|
||||
notes = []
|
||||
if large_recovered:
|
||||
notes.append("Ein übergroßer, fast leerer Tabellenblock wurde in zwei kleineren Ansichten erneut gelesen.")
|
||||
if neighbor_recovered:
|
||||
notes.append("Eine von der Layoutstufe ausgelassene Nachbarspalte wurde zusätzlich gelesen.")
|
||||
notes.append("Tabellenzeilen haben weiterhin nur einen groben Regionsbeleg.")
|
||||
result["layout_note"] = " ".join(notes)
|
||||
return result
|
||||
|
||||
|
||||
def main() -> None:
|
||||
data = sys.stdin.buffer.read(16 * 1024 * 1024 + 1)
|
||||
if not data or len(data) > 16 * 1024 * 1024:
|
||||
print(json.dumps({"error": "invalid_input"}))
|
||||
return
|
||||
models = ROOT / "official_models"
|
||||
layout = models / "PP-DocLayoutV3"
|
||||
vlm = models / "PaddleOCR-VL-1.6"
|
||||
if not (layout / "inference.pdiparams").is_file() or not (vlm / "model.safetensors").is_file():
|
||||
print(json.dumps({"error": "models_missing"}))
|
||||
return
|
||||
private = Path(os.environ.get("OCR_PRIVATE_DIR", ROOT / "private"))
|
||||
private.mkdir(mode=0o700, exist_ok=True)
|
||||
descriptor, name = tempfile.mkstemp(prefix="page-", suffix=".jpg", dir=private)
|
||||
path = Path(name)
|
||||
try:
|
||||
with os.fdopen(descriptor, "wb") as stream:
|
||||
stream.write(data)
|
||||
# Paddle logs should never expose text or private file paths through
|
||||
# the SSH result channel. The app suppresses stderr as well.
|
||||
stdout_copy = os.dup(1)
|
||||
null_fd = os.open(os.devnull, os.O_WRONLY)
|
||||
os.dup2(null_fd, 1)
|
||||
os.close(null_fd)
|
||||
try:
|
||||
from paddleocr import PaddleOCRVL
|
||||
pipeline = PaddleOCRVL(pipeline_version="v1.6", use_layout_detection=True,
|
||||
layout_detection_model_dir=str(layout), vl_rec_model_dir=str(vlm),
|
||||
use_doc_orientation_classify=False,
|
||||
use_doc_unwarping=False,
|
||||
device=os.environ.get("OCR_VISION_DEVICE", "cpu"))
|
||||
result = (predict_tiled(pipeline, path, private)
|
||||
if "--tiles-only" in sys.argv[1:] else
|
||||
predict_adaptive(pipeline, path, private))
|
||||
finally:
|
||||
os.dup2(stdout_copy, 1)
|
||||
os.close(stdout_copy)
|
||||
print(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
|
||||
except Exception:
|
||||
# No exception text: model libraries may include recognized content.
|
||||
print(json.dumps({"error": "inference_failed"}))
|
||||
finally:
|
||||
path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user