295 lines
11 KiB
Python
295 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
vl_compare.py — Run PaddleOCR-VL on one newspaper page and compare its
|
|
layout / reading-order output against the existing PaddleOCR LayoutDetection
|
|
regions (page_XXX.regions.json) produced by extractor.py.
|
|
|
|
Goal: see whether VL's reading-order blocks line up with article boundaries
|
|
better than the current "box soup" before deciding to swap engines.
|
|
|
|
USAGE
|
|
# auto-pick the most recent output/<job>/pages/page_001.png and its regions.json
|
|
python vl_compare.py
|
|
|
|
# or point at a specific page + regions file
|
|
python vl_compare.py output/20260529_004651/pages/page_001.png \
|
|
output/20260529_004651/pages/page_001.regions.json
|
|
|
|
OUTPUT (written next to the page image)
|
|
page_XXX.vl_layout.json normalized VL blocks (id, type, bbox, reading order, text)
|
|
page_XXX.vl_debug.png page with VL boxes + reading-order numbers drawn
|
|
page_XXX.vl_vs_paddle.json side-by-side summary (counts by type, totals)
|
|
page_XXX.vl_raw.json raw VL result (best-effort, for debugging)
|
|
|
|
NOTES
|
|
* PaddleOCR-VL is a ~0.9B vision-language model — expect a GPU and a model
|
|
download on first run. It is SEPARATE from the lightweight LayoutDetection
|
|
used in extractor.py.
|
|
* The VL entry point has moved across releases. This script tries several
|
|
known import paths and prints a clear hint if none are found so you can
|
|
adapt the one line to your installed version.
|
|
* Run inside your project venv: ./venv/bin/python vl_compare.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import glob
|
|
import json
|
|
import traceback
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 1. Resolve input page image + matching regions.json
|
|
# --------------------------------------------------------------------------
|
|
def resolve_inputs(argv):
|
|
if len(argv) >= 2:
|
|
page_png = argv[1]
|
|
regions_json = argv[2] if len(argv) >= 3 else os.path.splitext(page_png)[0] + ".regions.json"
|
|
return page_png, regions_json
|
|
|
|
# auto-discover: most recently modified page_001.png under output/
|
|
candidates = sorted(
|
|
glob.glob("output/*/pages/page_001.png"),
|
|
key=os.path.getmtime,
|
|
reverse=True,
|
|
)
|
|
if not candidates:
|
|
sys.exit("No page_001.png found under output/*/pages/. "
|
|
"Pass a page image path explicitly: python vl_compare.py <page.png> [regions.json]")
|
|
page_png = candidates[0]
|
|
regions_json = os.path.splitext(page_png)[0] + ".regions.json"
|
|
return page_png, regions_json
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 2. Load the VL model (defensive across releases)
|
|
# --------------------------------------------------------------------------
|
|
def load_vl_model():
|
|
import paddleocr
|
|
print(f"paddleocr version: {getattr(paddleocr, '__version__', 'unknown')}")
|
|
|
|
# Try the documented entry points, newest-first. Adjust if your release differs.
|
|
attempts = [
|
|
("paddleocr.PaddleOCRVL", lambda: paddleocr.PaddleOCRVL()),
|
|
("paddleocr.PPDocVL", lambda: paddleocr.PPDocVL()),
|
|
("paddleocr.DocVLM", lambda: paddleocr.DocVLM()),
|
|
]
|
|
for name, ctor in attempts:
|
|
try:
|
|
model = ctor()
|
|
print(f"Loaded VL model via: {name}")
|
|
return model
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" - {name} unavailable: {e}")
|
|
|
|
# Last resort: PaddleX pipeline name
|
|
try:
|
|
from paddlex import create_pipeline
|
|
model = create_pipeline(pipeline="PaddleOCR-VL")
|
|
print("Loaded VL model via: paddlex.create_pipeline('PaddleOCR-VL')")
|
|
return model
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" - paddlex pipeline unavailable: {e}")
|
|
|
|
sys.exit(
|
|
"\nCould not load PaddleOCR-VL with any known entry point.\n"
|
|
"Check the exact class name for your installed version, e.g.:\n"
|
|
" python -c \"import paddleocr; print([n for n in dir(paddleocr) if 'VL' in n or 'Doc' in n])\"\n"
|
|
"then edit the `attempts` list near the top of load_vl_model()."
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 3. Run VL + normalize its output to the same shape as regions.json
|
|
# --------------------------------------------------------------------------
|
|
def to_bbox(coord):
|
|
"""Accept [x1,y1,x2,y2] or polygon [x1,y1,...] -> [x1,y1,x2,y2] ints."""
|
|
if coord is None:
|
|
return [0, 0, 0, 0]
|
|
if hasattr(coord, "tolist"):
|
|
coord = coord.tolist()
|
|
coord = [float(v) for v in coord]
|
|
if len(coord) == 4:
|
|
return [int(v) for v in coord]
|
|
xs = coord[0::2]
|
|
ys = coord[1::2]
|
|
return [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))]
|
|
|
|
|
|
def run_vl(model, page_png):
|
|
if hasattr(model, "predict"):
|
|
raw = model.predict(page_png)
|
|
elif callable(model):
|
|
raw = model(page_png)
|
|
else:
|
|
raw = model.predict(input=page_png)
|
|
|
|
# raw is typically a list of result objects (dict-like). Iterate generically.
|
|
results = list(raw) if not isinstance(raw, list) else raw
|
|
|
|
blocks = []
|
|
raw_dump = []
|
|
|
|
for res in results:
|
|
# try to get a plain dict for dumping/parsing
|
|
d = None
|
|
for attr in ("json", "res", "_to_dict"):
|
|
try:
|
|
v = getattr(res, attr, None)
|
|
d = v() if callable(v) else v
|
|
if d:
|
|
break
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
if d is None and isinstance(res, dict):
|
|
d = res
|
|
raw_dump.append(_jsonable(d if d is not None else str(res)))
|
|
|
|
# VL layout blocks commonly live under one of these keys
|
|
layout = None
|
|
if isinstance(d, dict):
|
|
for key in ("layout_det_res", "parsing_res_list", "layout", "blocks", "boxes"):
|
|
if key in d and d[key]:
|
|
layout = d[key]
|
|
break
|
|
if layout is None:
|
|
continue
|
|
|
|
# layout may itself be a dict containing 'boxes'
|
|
if isinstance(layout, dict):
|
|
layout = layout.get("boxes", layout.get("parsing_res_list", []))
|
|
|
|
for i, b in enumerate(layout):
|
|
if not isinstance(b, dict):
|
|
continue
|
|
coord = b.get("coordinate", b.get("bbox", b.get("block_bbox", b.get("layout_bbox"))))
|
|
label = str(b.get("label", b.get("type", b.get("block_label", "unknown")))).lower()
|
|
score = float(b.get("score", b.get("confidence", 0)) or 0)
|
|
order = b.get("reading_order", b.get("order", b.get("index", i)))
|
|
text = b.get("text", b.get("block_content", b.get("rec_text", "")))
|
|
if isinstance(text, list):
|
|
text = " ".join(map(str, text))
|
|
blocks.append({
|
|
"id": len(blocks) + 1,
|
|
"type": label,
|
|
"bbox": to_bbox(coord),
|
|
"confidence": round(score, 4),
|
|
"reading_order": order,
|
|
"text_preview": (str(text)[:120] if text else ""),
|
|
})
|
|
|
|
# sort by reading order if available
|
|
try:
|
|
blocks.sort(key=lambda x: (x["reading_order"] if isinstance(x["reading_order"], (int, float)) else 0))
|
|
for i, b in enumerate(blocks, 1):
|
|
b["id"] = i
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
|
|
return blocks, raw_dump
|
|
|
|
|
|
def _jsonable(obj):
|
|
try:
|
|
json.dumps(obj)
|
|
return obj
|
|
except Exception: # noqa: BLE001
|
|
return str(obj)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# 4. Compare against existing PaddleOCR regions
|
|
# --------------------------------------------------------------------------
|
|
def type_counts(regions):
|
|
c = {}
|
|
for r in regions:
|
|
t = r.get("type", "unknown")
|
|
c[t] = c.get(t, 0) + 1
|
|
return dict(sorted(c.items(), key=lambda x: -x[1]))
|
|
|
|
|
|
def draw_debug(page_png, blocks, out_png):
|
|
try:
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
except Exception: # noqa: BLE001
|
|
print("Pillow not available — skipping debug image.")
|
|
return
|
|
img = Image.open(page_png).convert("RGB")
|
|
draw = ImageDraw.Draw(img)
|
|
try:
|
|
font = ImageFont.truetype("DejaVuSans-Bold.ttf", 60)
|
|
except Exception: # noqa: BLE001
|
|
font = ImageFont.load_default()
|
|
for b in blocks:
|
|
x1, y1, x2, y2 = b["bbox"]
|
|
draw.rectangle([x1, y1, x2, y2], outline=(220, 30, 30), width=6)
|
|
tag = f"{b['id']}:{b['type']}"
|
|
draw.text((x1 + 8, y1 + 8), tag, fill=(220, 30, 30), font=font)
|
|
img.save(out_png)
|
|
print(f" wrote {out_png}")
|
|
|
|
|
|
def main():
|
|
page_png, regions_json = resolve_inputs(sys.argv)
|
|
if not os.path.exists(page_png):
|
|
sys.exit(f"Page image not found: {page_png}")
|
|
print(f"Page image : {page_png}")
|
|
print(f"Regions JSON: {regions_json} (exists={os.path.exists(regions_json)})")
|
|
|
|
stem = os.path.splitext(page_png)[0]
|
|
|
|
# existing paddle regions
|
|
paddle_regions = []
|
|
if os.path.exists(regions_json):
|
|
with open(regions_json) as f:
|
|
paddle_regions = json.load(f).get("regions", [])
|
|
|
|
# run VL
|
|
print("\nLoading PaddleOCR-VL (first run downloads the model)...")
|
|
model = load_vl_model()
|
|
print("Running VL on the page...")
|
|
blocks, raw_dump = run_vl(model, page_png)
|
|
|
|
# write outputs
|
|
with open(stem + ".vl_layout.json", "w") as f:
|
|
json.dump({"page_image": os.path.basename(page_png), "blocks": blocks}, f, indent=2, ensure_ascii=False)
|
|
with open(stem + ".vl_raw.json", "w") as f:
|
|
json.dump(raw_dump, f, indent=2, ensure_ascii=False)
|
|
|
|
summary = {
|
|
"page_image": os.path.basename(page_png),
|
|
"paddle_layoutdetection": {
|
|
"total_regions": len(paddle_regions),
|
|
"by_type": type_counts(paddle_regions),
|
|
},
|
|
"paddleocr_vl": {
|
|
"total_blocks": len(blocks),
|
|
"by_type": type_counts(blocks),
|
|
"has_reading_order": any(isinstance(b.get("reading_order"), (int, float)) for b in blocks),
|
|
},
|
|
}
|
|
with open(stem + ".vl_vs_paddle.json", "w") as f:
|
|
json.dump(summary, f, indent=2, ensure_ascii=False)
|
|
|
|
draw_debug(page_png, blocks, stem + ".vl_debug.png")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("COMPARISON SUMMARY")
|
|
print("=" * 60)
|
|
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
|
print("\nWrote:")
|
|
for ext in (".vl_layout.json", ".vl_raw.json", ".vl_vs_paddle.json", ".vl_debug.png"):
|
|
print(" " + stem + ext)
|
|
print("\nNext: open page_XXX_regions_debug.png (existing) next to page_XXX.vl_debug.png")
|
|
print("and eyeball whether VL's blocks track article boundaries better.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except SystemExit:
|
|
raise
|
|
except Exception: # noqa: BLE001
|
|
traceback.print_exc()
|
|
sys.exit("\nFailed. If the error is about the VL class name, see load_vl_model().")
|