311 lines
16 KiB
Python
311 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
overlay.py — human-viewable debug visualisation: draw the extractor's geometry over the rendered
|
|
exam page. Shows WHERE each question/part label was located and where Docling regions
|
|
(figures/tables/MCQ checkboxes) sit, so a reviewer can eyeball whether the structure landed in the
|
|
right place. This is the same geometry the exam-marker app uses to place regions on its canvas.
|
|
|
|
Coordinates: Docling/RapidOCR bboxes are PDF points with a BOTTOM-LEFT origin. We render the page
|
|
at DPI D (scale = D/72) and flip y against the rendered image height, so we never need the page's
|
|
point-height explicitly: y_top_px = H_px - t*scale.
|
|
|
|
With --docling, also draws every raw Docling text block (the body/question content the thin
|
|
extractor model discards) so a reviewer can see the FULL detection, not just what we persist.
|
|
Granite tables carry cells but no coordinates; we derive their box by locating the cell-texts in
|
|
the Docling text layer (content+geometry fusion).
|
|
|
|
Usage:
|
|
python scripts/overlay.py <structured.json> <source_pdf> [--pages 3,4,5] [--dpi 150] [--out DIR]
|
|
python scripts/overlay.py <structured.json> <pdf> --docling results/E_tess_full.json --pages 5
|
|
"""
|
|
import os, sys, json, re, argparse, subprocess, tempfile
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
PART_COLOR = (211, 47, 47) # red — question/part labels
|
|
BODY_COLOR = (150, 150, 150) # grey — raw Docling body-text blocks (--docling)
|
|
GRANITE_COLOR = (0, 150, 136) # teal — Granite table (geometry derived from cells)
|
|
REGION_COLORS = { # docling region taxonomy -> colour
|
|
"context_figure": (25, 118, 210), # blue
|
|
"context_data": (56, 142, 60), # green (tables)
|
|
"context_caption": (123, 31, 162), # purple
|
|
"mcq_option": (245, 124, 0), # orange (checkboxes)
|
|
}
|
|
|
|
|
|
def _norm(s):
|
|
return re.sub(r"[^a-z0-9]", "", (s or "").lower())
|
|
|
|
|
|
def docling_texts_by_page(doc):
|
|
"""All raw Docling text items -> {page: [(bbox, text, label)]}. The body content we discard."""
|
|
out = {}
|
|
for t in doc.get("texts", []):
|
|
prov = t.get("prov") or []
|
|
bb = prov[0].get("bbox") if prov else None
|
|
pg = prov[0].get("page_no") if prov else None
|
|
if bb and pg:
|
|
out.setdefault(pg, []).append((bb, t.get("text") or "", t.get("label") or "text"))
|
|
return out
|
|
|
|
|
|
def derive_table_bbox(grid, page_texts):
|
|
"""Granite tables have cells but no coordinates. Locate the cell-texts in the Docling text
|
|
layer and union their bboxes -> the table's on-page extent.
|
|
|
|
Two traps (seen on physics p5): (1) border/maths glyphs ('|','+') normalise to '' and an
|
|
empty string is a substring of everything; (2) cell WORDS recur in nearby content — the rock
|
|
names reappear in the MCQ options below the table ('Basalt or chalk'), far left and lower.
|
|
So we match only blocks whose normalised text is CONTAINED IN a cell (keeps fragments like
|
|
'2.90'/'Type', rejects the longer 'basaltorchalk'), require length >= 2, then keep the
|
|
dominant vertical cluster to drop any stray cell-word elsewhere on the page."""
|
|
import statistics
|
|
cells = {c for c in (_norm(x) for row in grid for x in row) if len(c) > 1}
|
|
hit = [bb for bb, txt, _ in page_texts
|
|
if len(_norm(txt)) > 1 and any(_norm(txt) in c for c in cells)]
|
|
if len(hit) < 3:
|
|
return None
|
|
med = statistics.median(sorted((b["t"] + b["b"]) / 2 for b in hit))
|
|
hit = [b for b in hit if abs((b["t"] + b["b"]) / 2 - med) <= 120] # table band only
|
|
return {"l": min(b["l"] for b in hit), "r": max(b["r"] for b in hit),
|
|
"t": max(b["t"] for b in hit), "b": min(b["b"] for b in hit)}
|
|
|
|
|
|
def _font(sz):
|
|
for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
|
|
if os.path.exists(p):
|
|
return ImageFont.truetype(p, sz)
|
|
return ImageFont.load_default()
|
|
|
|
|
|
MAIN_LINE = (25, 118, 210) # blue — main-question y-markers
|
|
PART_LINE = (211, 47, 47) # red — part y-markers
|
|
|
|
|
|
def _hline(draw, y_pdf, scale, H, W, color, label, width, font, dashed=False, inset=0):
|
|
"""Full-width horizontal marker line at a PDF-point y (BOTTOM-LEFT origin)."""
|
|
y = H - y_pdf * scale
|
|
if dashed:
|
|
x = inset
|
|
while x < W:
|
|
draw.line([x, y, min(x + 9, W), y], fill=color, width=width); x += 16
|
|
else:
|
|
draw.line([inset, y, W, y], fill=color, width=width)
|
|
if label:
|
|
tw = draw.textlength(label, font=font)
|
|
draw.rectangle([inset, y - 16, inset + tw + 6, y], fill=color)
|
|
draw.text((inset + 3, y - 15), label, fill=(255, 255, 255), font=font)
|
|
|
|
|
|
def _rect(draw, bb, scale, H, color, label, width=3, font=None):
|
|
"""Draw one bbox (BOTTOM-LEFT origin -> image space) + its label."""
|
|
x0, x1 = bb["l"] * scale, bb["r"] * scale
|
|
y0, y1 = H - bb["t"] * scale, H - bb["b"] * scale # t is the higher edge -> smaller y_px
|
|
draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
|
|
if label:
|
|
tw = draw.textlength(label, font=font)
|
|
draw.rectangle([x0, y0 - 17, x0 + tw + 6, y0], fill=color)
|
|
draw.text((x0 + 3, y0 - 16), label, fill=(255, 255, 255), font=font)
|
|
|
|
|
|
def draw_template(draw, tpl, pg, scale, H, W, font):
|
|
"""Render the editable template for one page: margins/bands as LINES, footprints as BOXES.
|
|
A confirmed element is drawn solid; an unconfirmed (auto) suggestion is drawn dashed."""
|
|
MARGIN, MAIN, PART = (0, 150, 136), (25, 118, 210), (211, 47, 47)
|
|
page = tpl["pages"].get(str(pg)) or tpl["pages"].get(pg) or {}
|
|
# role banner (top-left); margins suppressed entirely on no-margin pages (cover/blank)
|
|
role = page.get("role", "question")
|
|
draw.rectangle([0, 0, 8 + len(role) * 8, 16], fill=(70, 70, 70))
|
|
draw.text((4, 1), f"role:{role}", fill=(255, 255, 255), font=font)
|
|
margins_on = page.get("margins_enabled", True)
|
|
# margins: axis-locked lines (document scope on every page + this page's page-scope lines)
|
|
for m in (tpl.get("margins", []) if margins_on else []):
|
|
if m["scope"] == "page" and m.get("page") != pg:
|
|
continue
|
|
solid = m.get("confirmed")
|
|
if m["axis"] == "x":
|
|
x = m["value"] * scale
|
|
draw.line([x, 0, x, H], fill=MARGIN, width=2) if solid else _dash_v(draw, x, 0, H, MARGIN, 2)
|
|
else:
|
|
y = H - m["value"] * scale
|
|
draw.line([0, y, W, y], fill=MARGIN, width=2) if solid else _dash_h(draw, 0, W, y, MARGIN, 2)
|
|
for m in page.get("main_bands", []):
|
|
if not m.get("is_start", True): # continuation page: no spurious second "start" line
|
|
continue
|
|
_hline(draw, m["y_start"], scale, H, W, MAIN, f"Q{m['question']}", 3, font,
|
|
dashed=not m.get("confirmed"))
|
|
for p in page.get("part_bands", []):
|
|
_hline(draw, p["y_start"], scale, H, W, PART, p["label"], 2, font, inset=90,
|
|
dashed=not p.get("confirmed"))
|
|
for f in page.get("furniture", []):
|
|
if f.get("box"):
|
|
_rect(draw, f["box"], scale, H, (130, 130, 130), f"furniture:{f.get('kind','')}", 2, font)
|
|
for g in page.get("figures", []):
|
|
if g.get("box"):
|
|
_rect(draw, g["box"], scale, H, (56, 142, 60), "figure", 3, font)
|
|
for t in page.get("tables", []):
|
|
if t.get("box"):
|
|
_rect(draw, t["box"], scale, H, (0, 150, 136),
|
|
f"table {t.get('n_rows')}x{t.get('n_cols')}", 3, font)
|
|
|
|
|
|
def render_page(pdf, pg, dpi, td):
|
|
"""Render page `pg` and return an image in DOCLING's coordinate space. Docling reports bbox
|
|
relative to the CropBox, but pdftoppm renders the MediaBox — when CropBox != MediaBox (e.g. the
|
|
Edexcel 1MA1 papers: media 652x899, crop inset 28.35pt) that mismatch magnifies + shifts every
|
|
overlaid shape toward a corner. Fix: crop the rendered image to the CropBox so it matches Docling.
|
|
No-op when CropBox == MediaBox (h556) or when poppler already rendered the CropBox."""
|
|
base = os.path.join(td, f"p{pg}")
|
|
subprocess.run(["pdftoppm", "-png", "-r", str(dpi), "-f", str(pg), "-l", str(pg), pdf, base],
|
|
check=True)
|
|
png = next(p for p in (f"{base}-{pg:02d}.png", f"{base}-{pg}.png", f"{base}-{pg:03d}.png")
|
|
if os.path.exists(p))
|
|
img = Image.open(png).convert("RGB")
|
|
try:
|
|
import pypdf
|
|
page = pypdf.PdfReader(pdf).pages[pg - 1]
|
|
mb, cb = page.mediabox, page.cropbox
|
|
scale = dpi / 72.0
|
|
mbl, mbt = float(mb.left), float(mb.top)
|
|
dcrop = any(abs(a - b) > 0.5 for a, b in
|
|
((cb.left, mb.left), (cb.bottom, mb.bottom), (cb.right, mb.right), (cb.top, mb.top)))
|
|
rendered_mediabox = abs(img.width - (float(mb.right) - mbl) * scale) < 3
|
|
if dcrop and rendered_mediabox:
|
|
img = img.crop((round((float(cb.left) - mbl) * scale), round((mbt - float(cb.top)) * scale),
|
|
round((float(cb.right) - mbl) * scale), round((mbt - float(cb.bottom)) * scale)))
|
|
except Exception:
|
|
pass
|
|
return img
|
|
|
|
|
|
def _dash_v(draw, x, y0, y1, color, w):
|
|
y = y0
|
|
while y < y1:
|
|
draw.line([x, y, x, min(y + 9, y1)], fill=color, width=w); y += 16
|
|
|
|
|
|
def _dash_h(draw, x0, x1, y, color, w):
|
|
x = x0
|
|
while x < x1:
|
|
draw.line([x, y, min(x + 9, x1), y], fill=color, width=w); x += 16
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("structured"); ap.add_argument("pdf")
|
|
ap.add_argument("--docling", help="raw Docling doc JSON: also draw every body-text block "
|
|
"(the content the thin model discards) + derive Granite-table boxes")
|
|
ap.add_argument("--bands", help="bands.py JSON: draw main-question + part start/end y-marker lines")
|
|
ap.add_argument("--furniture", help="furniture.py JSON: mark recurring furniture vs real figures "
|
|
"+ draw the content x-margins")
|
|
ap.add_argument("--template", help="template.py JSON: render the editable first-pass template "
|
|
"(margins+bands as lines, furniture/figures as boxes). "
|
|
"When set, draws ONLY the template (the human-review view).")
|
|
ap.add_argument("--pages", help="comma list, e.g. 3,4,5 (default: all pages with geometry)")
|
|
ap.add_argument("--dpi", type=int, default=150)
|
|
ap.add_argument("--out", default="results/overlay")
|
|
a = ap.parse_args()
|
|
os.makedirs(a.out, exist_ok=True)
|
|
scale = a.dpi / 72.0
|
|
font = _font(14)
|
|
|
|
res = json.load(open(a.structured))
|
|
doc_texts = docling_texts_by_page(json.load(open(a.docling))) if a.docling else {}
|
|
bands = json.load(open(a.bands))["pages"] if a.bands else {}
|
|
furn = json.load(open(a.furniture)) if a.furniture else None
|
|
tpl = json.load(open(a.template)) if a.template else None
|
|
# gather geometry by page
|
|
parts_by_pg, regions_by_pg = {}, {}
|
|
for q in res.get("questions", []):
|
|
for p in q["parts"]:
|
|
if p.get("bbox") and p.get("page"):
|
|
parts_by_pg.setdefault(p["page"], []).append((p["label"], p["bbox"]))
|
|
for r in res.get("regions", []):
|
|
if r.get("bbox") and r.get("page"):
|
|
regions_by_pg.setdefault(r["page"], []).append((r["type"], r["bbox"]))
|
|
# tables: standard ones carry a bbox; Granite ones don't -> derive from the text layer
|
|
tables_by_pg = {}
|
|
for t in res.get("tables", []):
|
|
pg = t.get("page")
|
|
if not pg:
|
|
continue
|
|
bb = t.get("bbox") or (derive_table_bbox(t.get("grid", []), doc_texts.get(pg, []))
|
|
if a.docling else None)
|
|
if bb:
|
|
tables_by_pg.setdefault(pg, []).append(
|
|
(f"table {t.get('source','')} {t.get('n_rows')}x{t.get('n_cols')}", bb))
|
|
|
|
want = ([int(x) for x in a.pages.split(",")] if a.pages
|
|
else (sorted(int(p) for p in tpl["pages"]) if tpl
|
|
else sorted(set(parts_by_pg) | set(regions_by_pg) | set(doc_texts))))
|
|
if not want:
|
|
sys.exit("no bbox geometry in this result (born-digital text path carries no geometry; "
|
|
"use an OCR/rapid-path structured.json)")
|
|
|
|
written = []
|
|
with tempfile.TemporaryDirectory() as td:
|
|
for pg in want:
|
|
img = render_page(a.pdf, pg, a.dpi, td)
|
|
H = img.height
|
|
draw = ImageDraw.Draw(img)
|
|
if tpl: # template-only render = the human-review view
|
|
draw_template(draw, tpl, pg, scale, H, img.width, font)
|
|
out = os.path.join(a.out, f"p{pg:02d}.png")
|
|
img.save(out); written.append(out)
|
|
pgd = tpl["pages"].get(str(pg), {})
|
|
print(f"p{pg}: template — {len(pgd.get('main_bands',[]))} main, "
|
|
f"{len(pgd.get('part_bands',[]))} part, {len(pgd.get('furniture',[]))} furn, "
|
|
f"{len(pgd.get('figures',[]))} fig -> {out}")
|
|
continue
|
|
# layer 0: raw Docling body-text blocks (faint, no label) — the discarded content
|
|
for bb, txt, lab in doc_texts.get(pg, []):
|
|
_rect(draw, bb, scale, H, BODY_COLOR, None, 1, font)
|
|
# layer 1: taxonomy regions
|
|
for typ, bb in regions_by_pg.get(pg, []):
|
|
_rect(draw, bb, scale, H, REGION_COLORS.get(typ, (120, 120, 120)), typ, 2, font)
|
|
# layer 2: tables (Granite-derived boxes in teal)
|
|
for lab, bb in tables_by_pg.get(pg, []):
|
|
_rect(draw, bb, scale, H, GRANITE_COLOR, lab, 3, font)
|
|
# layer 3: part labels on top
|
|
for lab, bb in parts_by_pg.get(pg, []):
|
|
_rect(draw, bb, scale, H, PART_COLOR, lab, 3, font)
|
|
# layer 4: band y-marker lines (main-question = blue, part = red dashed; end = dashed)
|
|
pb = bands.get(str(pg)) or bands.get(pg)
|
|
nb = 0
|
|
if pb:
|
|
W = img.width
|
|
for m in pb["main"]:
|
|
if not m.get("is_start", True): # skip continuation-page duplicate
|
|
continue
|
|
_hline(draw, m["y_start"], scale, H, W, MAIN_LINE,
|
|
f"Q{m['question']} ▸ start", 3, font); nb += 1
|
|
_hline(draw, m["y_end"], scale, H, W, MAIN_LINE, None, 1, font, dashed=True)
|
|
for p in pb["part"]:
|
|
_hline(draw, p["y_start"], scale, H, W, PART_LINE,
|
|
f"{p['label']} start", 2, font, inset=90); nb += 1
|
|
# layer 5: furniture mask — green=real figure, grey=masked furniture; + content margins
|
|
if furn:
|
|
W = img.width
|
|
for it in furn["items"]:
|
|
if it["page"] != pg or it["kind"] != "picture":
|
|
continue
|
|
if it["furniture"]:
|
|
_rect(draw, it["bbox"], scale, H, (130, 130, 130), "furniture", 2, font)
|
|
else:
|
|
_rect(draw, it["bbox"], scale, H, (56, 142, 60), "figure ✓", 3, font)
|
|
band = (furn.get("content_margins") or {}).get("content_x_band")
|
|
if band:
|
|
for xk in ("x_left", "x_right"):
|
|
x = band[xk] * scale
|
|
draw.line([x, 0, x, H], fill=(0, 150, 136), width=2)
|
|
out = os.path.join(a.out, f"p{pg:02d}.png")
|
|
img.save(out); written.append(out)
|
|
print(f"p{pg}: {len(parts_by_pg.get(pg,[]))} part-labels, "
|
|
f"{len(regions_by_pg.get(pg,[]))} regions, {len(tables_by_pg.get(pg,[]))} tables, "
|
|
f"{len(doc_texts.get(pg,[]))} body-text blocks, {nb} band-lines -> {out}")
|
|
print(f"-> {len(written)} page(s) in {a.out}/")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|