311 lines
12 KiB
Python
311 lines
12 KiB
Python
|
"""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()
|