mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-19 14:55:40 +08:00
8d20cbd0b3
Supersedes / folds in #18305. The score-based layer-2 rotation selection from #18305 now lives here, on top of `WarpCrop` (layer 1), applied to **all three** Go OCR paths, together with the Python score plumbing the Go side depends on. #18305 is closed in favor of this PR.
360 lines
15 KiB
Python
Executable File
360 lines
15 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""WarpCrop (Go) vs an independent perspective-crop reference (cv2) pixel
|
|
alignment on REAL document text, with no model weights required.
|
|
|
|
This exercises the perspective de-skew step (layer 1 of the OCR warp work). It
|
|
answers the geometric part of the regression question: for real text regions on
|
|
a real PDF page, does Go's util.WarpCrop reproduce the same crop geometry as an
|
|
independent perspective transform?
|
|
|
|
Detection and recognition need the deepdoc ONNX weights (downloaded from
|
|
HuggingFace, unavailable offline). They are intentionally NOT used here:
|
|
- quads come from pdfplumber word boxes (rotated to simulate skew) -- no model
|
|
- the reference crop is pure geometry (getPerspectiveTransform +
|
|
warpPerspective with BORDER_REPLICATE + INTER_CUBIC) -- no model
|
|
|
|
So this runs fully offline and isolates the warp geometry, which is the part
|
|
WarpCrop actually implements.
|
|
|
|
Pipeline
|
|
--------
|
|
1. genquads : from the Go-rendered page, derive real word boxes (pdfplumber),
|
|
rotate each by a sweep of angles to build skewed quads, write
|
|
quads.json and the reference crops (py_warp_*.png).
|
|
2. (Go) : TestWarpAlignGo pass 2 reads quads.json and writes go_warp_*.png
|
|
(WarpCrop, de-skewed) and go_fastcrop_*.png (FastCrop, axis-aligned bbox).
|
|
3. compare : pixel-MSE Go-WarpCrop vs reference and Go-FastCrop vs reference,
|
|
aggregated (layer-1 geometry; no model weights needed).
|
|
4. rec : run recognize_batch (the actual Go->service path, NO layer-2
|
|
rotation) on all three crops and compare the recognized text
|
|
(layer-1 + recognizer; needs the deepdoc ONNX weights).
|
|
|
|
All outputs go under --out-dir (default /tmp/render_diff/align); the worktree
|
|
is never touched and nothing is committed.
|
|
|
|
Run
|
|
---
|
|
WARP_PDF=test/benchmark/test_docs/Doc1.pdf \
|
|
bash build.sh --test-manual ./internal/deepdoc/parser/pdf/ -run TestWarpAlignGo
|
|
.venv/bin/python tools/render_diff/warp_align.py genquads \
|
|
--pdf test/benchmark/test_docs/Doc1.pdf --go-page-png /tmp/render_diff/align/page0.png
|
|
WARP_PDF=test/benchmark/test_docs/Doc1.pdf \
|
|
bash build.sh --test-manual ./internal/deepdoc/parser/pdf/ -run TestWarpAlignGo
|
|
.venv/bin/python tools/render_diff/warp_align.py compare
|
|
# (needs HF weights in rag/res/deepdoc)
|
|
.venv/bin/python tools/render_diff/warp_align.py rec
|
|
"""
|
|
|
|
import argparse
|
|
import difflib
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import sys
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pdfplumber
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
|
|
|
|
# Go renders the page at exactly 3x (216 DPI); pdfplumber coords are in points.
|
|
SCALE = 3.0
|
|
|
|
|
|
def _sim(a, b):
|
|
"""Normalized similarity in [0,1] between two recognized strings."""
|
|
return difflib.SequenceMatcher(None, a or "", b or "").ratio()
|
|
|
|
|
|
def _quad_from_word(w, angle_deg, W, H, margin, perturb=4.0):
|
|
"""Return [x0,y0,x1,y1,x2,y2,x3,y3] (TL,TR,BR,BL) rotated by angle_deg, or
|
|
None if the rotated quad leaves the page.
|
|
|
|
A small deterministic perturbation is applied to the bottom-left corner so
|
|
the quad is a genuine trapezoid rather than a parallelogram. A plain
|
|
rotated rectangle maps to a parallelogram, whose getPerspectiveTransform has
|
|
no projective term -- the 8-DOF (true perspective) path in WarpCrop would
|
|
then go untested against cv2. The perturbation makes the cv2 reference
|
|
exercise the full homography the Go code implements.
|
|
"""
|
|
x0, top, x1, bottom = w["x0"], w["top"], w["x1"], w["bottom"]
|
|
pts = [
|
|
(x0 * SCALE, top * SCALE),
|
|
(x1 * SCALE, top * SCALE),
|
|
(x1 * SCALE, bottom * SCALE),
|
|
(x0 * SCALE, bottom * SCALE),
|
|
]
|
|
cx = (x0 + x1) / 2 * SCALE
|
|
cy = (top + bottom) / 2 * SCALE
|
|
a = math.radians(angle_deg)
|
|
ca, sa = math.cos(a), math.sin(a)
|
|
rot = []
|
|
for x, y in pts:
|
|
dx, dy = x - cx, y - cy
|
|
rot.append((cx + dx * ca - dy * sa, cy + dx * sa + dy * ca))
|
|
# Break the parallelogram so cv2 exercises the non-zero projective terms.
|
|
bx, by = rot[3]
|
|
rot[3] = (bx + perturb, by - perturb)
|
|
# in-bounds check
|
|
for x, y in rot:
|
|
if not (margin <= x <= W - margin and margin <= y <= H - margin):
|
|
return None
|
|
return [v for p in rot for v in p]
|
|
|
|
|
|
def _cv2_warp(page, quad):
|
|
p = np.array([[quad[0], quad[1]], [quad[2], quad[3]], [quad[4], quad[5]], [quad[6], quad[7]]], dtype=np.float32)
|
|
w = int(max(np.linalg.norm(p[0] - p[1]), np.linalg.norm(p[2] - p[3])))
|
|
h = int(max(np.linalg.norm(p[0] - p[3]), np.linalg.norm(p[1] - p[2])))
|
|
std = np.array([[0, 0], [w, 0], [w, h], [0, h]], dtype=np.float32)
|
|
M = cv2.getPerspectiveTransform(p, std)
|
|
return cv2.warpPerspective(page, M, (w, h), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC)
|
|
|
|
|
|
def cmd_genquads(args):
|
|
page = cv2.imread(args.go_page_png)
|
|
if page is None:
|
|
raise SystemExit(f"cannot read {args.go_page_png}")
|
|
H, W = page.shape[:2]
|
|
with pdfplumber.open(args.pdf) as pdf:
|
|
pg = pdf.pages[args.page]
|
|
words = pg.extract_words()
|
|
angles = [float(x) for x in args.angles.split(",")]
|
|
boxes = []
|
|
for w in words:
|
|
if w["bottom"] - w["top"] < 6: # skip tiny fragments
|
|
continue
|
|
for ang in angles:
|
|
if len(boxes) >= args.max_quads:
|
|
break
|
|
q = _quad_from_word(w, ang, W, H, args.margin)
|
|
if q is None:
|
|
continue
|
|
idx = len(boxes)
|
|
boxes.append({"quad": q, "text": w.get("text", ""), "angle": ang})
|
|
cv2.imwrite(os.path.join(args.out_dir, f"py_warp_{idx}.png"), _cv2_warp(page, q))
|
|
if len(boxes) >= args.max_quads:
|
|
break
|
|
with open(os.path.join(args.out_dir, "quads.json"), "w") as f:
|
|
json.dump({"boxes": [{"quad": b["quad"]} for b in boxes]}, f, indent=2)
|
|
with open(os.path.join(args.out_dir, "quads_meta.json"), "w") as f:
|
|
json.dump(boxes, f, indent=2)
|
|
logging.info("[genquads] %d quads (from %d words, angles=%s) -> quads.json + py_warp_*.png", len(boxes), len(words), angles)
|
|
|
|
|
|
def pixel_mse(a_path, b_path):
|
|
"""Pixel MSE between two crops that MUST share identical dimensions.
|
|
|
|
Width/height are part of the WarpCrop <-> cv2 parity contract: a 1px
|
|
dimension difference (e.g. float64 vs float32 truncation in the norm) is a
|
|
real divergence and must surface, not be hidden by resizing. Callers that
|
|
compare the de-skewed WarpCrop crop against the cv2 reference use this and
|
|
treat a shape mismatch as a contract failure.
|
|
"""
|
|
A = np.asarray(cv2.imread(a_path), dtype=np.float64)
|
|
B = np.asarray(cv2.imread(b_path), dtype=np.float64)
|
|
if A.shape != B.shape:
|
|
raise ValueError(f"shape mismatch for {os.path.basename(a_path)} vs {os.path.basename(b_path)}: {A.shape} != {B.shape} (dimensions are part of the parity contract)")
|
|
diff = A - B
|
|
return float((diff**2).mean())
|
|
|
|
|
|
def pixel_mse_baseline(a_path, b_path):
|
|
"""Pixel MSE for the FastCrop baseline vs cv2.
|
|
|
|
FastCrop is axis-aligned while cv2 de-skews, so the two crops legitimately
|
|
differ in size. Resize to a common size to obtain a scalar divergence
|
|
estimate (the baseline is known-wrong; only its magnitude matters).
|
|
"""
|
|
A = np.asarray(cv2.imread(a_path), dtype=np.float64)
|
|
B = np.asarray(cv2.imread(b_path), dtype=np.float64)
|
|
if A.shape != B.shape:
|
|
B = cv2.resize(B, (A.shape[1], A.shape[0]))
|
|
diff = A - B
|
|
return float((diff**2).mean())
|
|
|
|
|
|
def cmd_compare(args):
|
|
with open(os.path.join(args.out_dir, "quads.json")) as f:
|
|
boxes = json.load(f)["boxes"]
|
|
try:
|
|
with open(os.path.join(args.out_dir, "quads_meta.json")) as f:
|
|
meta = json.load(f)
|
|
except FileNotFoundError:
|
|
meta = [{}] * len(boxes)
|
|
|
|
limit = args.limit if args.limit and args.limit > 0 else len(boxes)
|
|
mse_warp, mse_fc = [], []
|
|
warp_results = [] # (box_index, mw) so the per-angle summary stays aligned
|
|
better = 0 # cases where WarpCrop is closer to cv2 than FastCrop
|
|
warp_dim_mismatch = 0 # WarpCrop vs cv2 dimension contract failures
|
|
for i in range(min(limit, len(boxes))):
|
|
gw = os.path.join(args.out_dir, f"go_warp_{i}.png")
|
|
gf = os.path.join(args.out_dir, f"go_fastcrop_{i}.png")
|
|
pw = os.path.join(args.out_dir, f"py_warp_{i}.png")
|
|
if not (os.path.exists(gw) and os.path.exists(gf) and os.path.exists(pw)):
|
|
continue
|
|
# WarpCrop must match cv2 exactly in size; a mismatch is a contract
|
|
# failure (previously hidden by resizing) and is not counted as "closer".
|
|
try:
|
|
mw = pixel_mse(gw, pw)
|
|
except ValueError as e:
|
|
warp_dim_mismatch += 1
|
|
logging.warning("[compare] box %d: %s", i, e)
|
|
continue
|
|
mf = pixel_mse_baseline(gf, pw)
|
|
mse_warp.append(mw)
|
|
mse_fc.append(mf)
|
|
warp_results.append((i, mw))
|
|
if mw < mf:
|
|
better += 1
|
|
|
|
n = len(mse_warp)
|
|
if n == 0:
|
|
logging.info("[compare] no go crops found; run TestWarpAlignGo pass 2 first")
|
|
return
|
|
mean_w = sum(mse_warp) / n
|
|
mean_f = sum(mse_fc) / n
|
|
logging.info("\n[compare] %d boxes (limit=%s)", n, limit)
|
|
logging.info(" pixel MSE WarpCrop vs cv2 : mean=%.2f max=%.2f", mean_w, max(mse_warp))
|
|
logging.info(" pixel MSE FastCrop vs cv2 : mean=%.2f max=%.2f", mean_f, max(mse_fc))
|
|
logging.info(" WarpCrop closer to cv2 than FastCrop : %d/%d", better, n)
|
|
if warp_dim_mismatch:
|
|
logging.info(" WarpCrop vs cv2 DIMENSION MISMATCH (contract fail) : %d", warp_dim_mismatch)
|
|
# show per-angle summary if meta present; pair each MSE with its box index
|
|
# so a skipped/missing crop cannot misalign the angle buckets.
|
|
by_angle = {}
|
|
for idx, mw in warp_results:
|
|
ang = meta[idx].get("angle", 0) if idx < len(meta) else 0
|
|
by_angle.setdefault(ang, []).append(mw)
|
|
if by_angle:
|
|
logging.info("\n -- mean MSE(WarpCrop vs cv2) by rotation angle --")
|
|
for ang in sorted(by_angle):
|
|
vals = by_angle[ang]
|
|
logging.info(" angle %+6.1fdeg : n=%3d meanMSE=%.2f", ang, len(vals), sum(vals) / len(vals))
|
|
|
|
|
|
def cmd_rec(args):
|
|
"""Recognition-level alignment (the decisive "no regression" evidence).
|
|
|
|
The Go->Python OCR service feeds cropped image BYTES (no box) to the
|
|
recognizer via recognize_batch(), which does NOT perform layer-2 rotation
|
|
selection. So the apples-to-apples comparison is recognize_batch() on each
|
|
of the three crops already on disk:
|
|
|
|
go_warp_*.png : Go util.WarpCrop (de-skewed, layer-1 only)
|
|
go_fastcrop_*.png : Go util.FastCrop (axis-aligned bounding box)
|
|
py_warp_*.png : cv2.getPerspectiveTransform + warpPerspective
|
|
(the de-skewed crop Python produces before recognition)
|
|
|
|
If go_warp matches py_warp at the text level and both beat go_fastcrop,
|
|
WarpCrop reproduces the de-skewed crop geometry that FastCrop misses.
|
|
"""
|
|
from deepdoc.vision.ocr import OCR
|
|
|
|
ocr = OCR()
|
|
limit = args.limit if args.limit and args.limit > 0 else 10**9
|
|
n = 0
|
|
warp_eq_py = 0 # WarpCrop text == cv2 text (exact)
|
|
fc_eq_py = 0 # FastCrop text == cv2 text (exact)
|
|
warp_sim_total = 0.0
|
|
fc_sim_total = 0.0
|
|
warp_closer = 0 # WarpCrop text closer to cv2 than FastCrop text
|
|
by_angle = {}
|
|
worst = [] # (sim_warp, i, text_warp, text_py, text_fc)
|
|
|
|
# Per-box rotation angle, if the genquads metadata is present.
|
|
meta = []
|
|
meta_path = os.path.join(args.out_dir, "quads_meta.json")
|
|
if os.path.exists(meta_path):
|
|
with open(meta_path) as f:
|
|
meta = json.load(f)
|
|
|
|
for i in range(limit):
|
|
gw = os.path.join(args.out_dir, f"go_warp_{i}.png")
|
|
gf = os.path.join(args.out_dir, f"go_fastcrop_{i}.png")
|
|
pw = os.path.join(args.out_dir, f"py_warp_{i}.png")
|
|
if not (os.path.exists(gw) and os.path.exists(gf) and os.path.exists(pw)):
|
|
break
|
|
img_w = cv2.imread(gw)
|
|
img_f = cv2.imread(gf)
|
|
img_p = cv2.imread(pw)
|
|
tw = ocr.recognize_batch([img_w])[0]
|
|
tf = ocr.recognize_batch([img_f])[0]
|
|
tp = ocr.recognize_batch([img_p])[0]
|
|
|
|
sw = _sim(tw, tp)
|
|
sf = _sim(tf, tp)
|
|
warp_sim_total += sw
|
|
fc_sim_total += sf
|
|
if tw == tp:
|
|
warp_eq_py += 1
|
|
if tf == tp:
|
|
fc_eq_py += 1
|
|
if sw > sf:
|
|
warp_closer += 1
|
|
# per-angle bucket
|
|
ang = meta[i].get("angle", 0.0) if i < len(meta) else 0.0
|
|
by_angle.setdefault(ang, []).append(sw)
|
|
worst.append((sw, i, tw, tp, tf))
|
|
n += 1
|
|
|
|
if n == 0:
|
|
logging.info("[rec] no crops found; run genquads + TestWarpAlignGo pass 2 first")
|
|
return
|
|
|
|
logging.info("\n[rec] %d boxes (recognize_batch, no layer-2 rotation)", n)
|
|
logging.info(" exact text match WarpCrop vs cv2 : %d/%d (%.1f%%)", warp_eq_py, n, 100 * warp_eq_py / n)
|
|
logging.info(" exact text match FastCrop vs cv2 : %d/%d (%.1f%%)", fc_eq_py, n, 100 * fc_eq_py / n)
|
|
logging.info(" mean text similarity WarpCrop vs cv2 : %.3f", warp_sim_total / n)
|
|
logging.info(" mean text similarity FastCrop vs cv2 : %.3f", fc_sim_total / n)
|
|
logging.info(" WarpCrop text strictly closer to cv2 than FastCrop : %d/%d (%.1f%%)", warp_closer, n, 100 * warp_closer / n)
|
|
|
|
if by_angle:
|
|
logging.info("\n -- mean text similarity(WarpCrop vs cv2) by rotation angle --")
|
|
for ang in sorted(by_angle):
|
|
vals = by_angle[ang]
|
|
logging.info(" angle %+6.1fdeg : n=%3d sim=%.3f", ang, len(vals), sum(vals) / len(vals))
|
|
|
|
logging.info("\n -- 8 worst WarpCrop-vs-cv2 cases --")
|
|
worst.sort(key=lambda x: x[0])
|
|
for sw, i, tw, tp, tf in worst[:8]:
|
|
logging.info(" #%3d sim=%.2f | warp=%r py=%r fastcrop=%r", i, sw, tw, tp, tf)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
|
|
g = sub.add_parser("genquads")
|
|
g.add_argument("--pdf", required=True)
|
|
g.add_argument("--page", type=int, default=0)
|
|
g.add_argument("--go-page-png", required=True)
|
|
g.add_argument("--out-dir", default="/tmp/render_diff/align")
|
|
g.add_argument("--angles", default="0,15,-15,30,-30")
|
|
g.add_argument("--max-quads", type=int, default=120)
|
|
g.add_argument("--margin", type=int, default=24)
|
|
|
|
c = sub.add_parser("compare")
|
|
c.add_argument("--out-dir", default="/tmp/render_diff/align")
|
|
c.add_argument("--limit", type=int, default=0, help="max boxes to compare (0 = all)")
|
|
|
|
r = sub.add_parser("rec", help="recognition-level alignment (needs deepdoc weights)")
|
|
r.add_argument("--out-dir", default="/tmp/render_diff/align")
|
|
r.add_argument("--limit", type=int, default=0, help="max boxes to recognize (0 = all)")
|
|
|
|
args = ap.parse_args()
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
if args.cmd == "genquads":
|
|
cmd_genquads(args)
|
|
elif args.cmd == "rec":
|
|
cmd_rec(args)
|
|
else:
|
|
cmd_compare(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|