api/api/services/docling/bands.py
kcar 5938613893
Some checks failed
api-ci-deploy / test-build-deploy (push) Has been cancelled
[verified] add docling auto-map package wrapper
2026-06-07 20:03:06 +01:00

137 lines
6.6 KiB
Python

#!/usr/bin/env python3
"""
bands.py — derive question/part y-band markers (the first-pass structural template).
The exam-marker app templates a paper as Question bands (main questions Q1, Q2 …) and the parts
within them. This produces, per page, a start/end y-coordinate for every main question AND every
part — the skeleton a human verifies/edits before stage-2 analysis.
Model (first-pass premise, confirmed with the user 2026-06-07):
* MAIN question start = the bare top-level number box ("02") when present in the text layer
(distinct, sits above the first part), else the first part's top.
* PART start = the part label's top (we already carry this geometry).
* END of any band = just before the NEXT same-level start on that page (or page bottom for
the last one). Parts are nested: a part's end never exceeds its question's.
Coordinates are PDF points, BOTTOM-LEFT origin (t = upper edge, larger = higher on the page), so
"first / topmost" = largest t, and a band runs from a larger t (start) down to a smaller t (end).
Usage:
python bands.py <structured.json> [--docling results/E_tess_full.json] [--out results/bands/x.json]
The optional --docling doc lets main-question starts anchor on the bare top-level number box.
"""
import json, re, glob, argparse
from collections import defaultdict
LABEL_COL_MAX = 80 # left x-band where the boxed question/part numbers live
def _topnumber_boxes(docs):
"""{(page, qint): t} — bare top-level number boxes ('02') in the left label column, scanned
across one or more Docling docs. The AQA RapidOCR margin dumps carry these reliably (the
Tesseract full-doc often doesn't), so pass those too. Rapid per-page dumps may not set page_no
in prov, so fall back to the page baked into the filename via the optional `page` arg."""
out = {}
for doc, page_hint in docs:
for it in doc.get("texts", []):
prov = it.get("prov") or []
bb = prov[0].get("bbox") if prov else None
pg = (prov[0].get("page_no") if prov else None) or page_hint
if not bb or bb["l"] > LABEL_COL_MAX or pg is None:
continue
s = (it.get("text") or "").strip().replace(" ", "")
m = re.match(r"^(\d{1,2})$", s)
if m:
key = (pg, int(m.group(1)))
out[key] = max(bb["t"], out.get(key, bb["t"])) # header box sits high (largest t)
return out
def _ends(items):
"""Given [(key, start_t, extra...)] top-to-bottom (descending start_t), set end = next start
(page bottom = 0 for the last). Returns list of dicts with start/end."""
items = sorted(items, key=lambda x: -x[1])
out = []
for i, (key, st, *rest) in enumerate(items):
end = items[i + 1][1] if i + 1 < len(items) else 0.0
out.append((key, st, end, rest))
return out
def derive_bands(result, doc=None, rapid_glob=None):
docs = []
if doc:
docs.append((doc, None))
for fn in sorted(glob.glob(rapid_glob) if rapid_glob else []):
m = re.search(r"p(\d+)\.json", fn)
docs.append((json.load(open(fn)), int(m.group(1)) if m else None))
topnum = _topnumber_boxes(docs)
# gather parts with geometry, grouped by page
by_page = defaultdict(list) # page -> [(q, label, t, b)]
for q in result.get("questions", []):
for p in q["parts"]:
bb, pg = p.get("bbox"), p.get("page")
if bb and pg:
by_page[pg].append((q["question"], p["label"], bb["t"], bb["b"]))
# global first page each question appears on (to mark the true start vs continuation pages)
q_first_page = {}
for pg, parts in by_page.items():
for q, *_ in parts:
q_first_page[q] = min(pg, q_first_page.get(q, pg))
pages = {}
for pg, parts in by_page.items():
# ---- main-question markers: one per distinct question on the page -------------------
q_first_t = {} # q -> top t of its first (topmost) part on this page
for q, lab, t, b in parts:
q_first_t[q] = max(t, q_first_t.get(q, t))
main_starts = []
for q, ft in q_first_t.items():
tn = topnum.get((pg, int(re.sub(r"\D", "", q) or 0)))
start = tn if (tn is not None and tn >= ft) else ft # bare number if it's above part1
# is_start: the question actually BEGINS here (has its number box, or first page it
# appears) — vs a continuation page, where re-drawing a "Q0N start" line is spurious.
is_start = (tn is not None) or (pg == q_first_page.get(q))
main_starts.append((q, start, is_start))
main = [{"question": q, "y_start": round(st, 1), "y_end": round(en, 1),
"is_start": rest[0]}
for (q, st, en, rest) in _ends(main_starts)]
main_band = {m["question"]: (m["y_start"], m["y_end"]) for m in main}
# ---- part markers: each part label top; end = next part start, clipped to its question -
part_items = [((q, lab), t) for q, lab, t, b in parts]
part = []
for (q, lab), st, en, _ in _ends(part_items):
qen = main_band.get(q, (st, 0))[1] # don't run past the question end
part.append({"label": lab, "question": q,
"y_start": round(st, 1), "y_end": round(max(en, qen), 1)})
pages[pg] = {"main": main, "part": part}
return {"board": result.get("board"), "paper_code": result.get("paper_code"),
"coord_origin": "BOTTOMLEFT", "pages": pages}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("structured")
ap.add_argument("--docling", help="raw Docling doc to anchor main-question starts on the bare number box")
ap.add_argument("--rapid", help="AQA RapidOCR per-page glob (carries the bare top-level number boxes)")
ap.add_argument("--out", default="results/bands.json")
a = ap.parse_args()
res = json.load(open(a.structured))
doc = json.load(open(a.docling)) if a.docling else None
bands = derive_bands(res, doc, a.rapid)
json.dump(bands, open(a.out, "w"), indent=2)
nq = sum(len(p["main"]) for p in bands["pages"].values())
npt = sum(len(p["part"]) for p in bands["pages"].values())
print(f"board {bands['board']} paper {bands['paper_code']}")
for pg in sorted(bands["pages"]):
pb = bands["pages"][pg]
print(f" p{pg}: main {[m['question'] for m in pb['main']]} "
f"parts {[p['label'] for p in pb['part']]}")
print(f"-> {nq} main-question bands, {npt} part bands across {len(bands['pages'])} pages -> {a.out}")
if __name__ == "__main__":
main()