Developing an End-to-End Document Intelligence Pipeline with docTR for OCR, Layout Analysis, KIE, Benchmarking, and Searchable PDFs

import os, sys, io, json, time, math, re, subprocess, warnings
from collections import Counter, defaultdict
warnings.filterwarnings("ignore")
os.environ.setdefault("USE_TORCH", "1")
def _pip(*pkgs):
   subprocess.run([sys.executable, "-m", "pip", "install", "-q", *pkgs], check=False)
try:
   import doctr
except ImportError:
   print(">> Installing python-doctr (this takes ~1-2 min on Colab)...")
   _pip("python-doctr[viz]")
try:
   import reportlab
except ImportError:
   _pip("reportlab")
import numpy as np
import torch
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.patches import Rectangle, Polygon as MplPolygon
from PIL import Image, ImageDraw, ImageFont
import doctr
from doctr.io import DocumentFile
from doctr.models import (
   ocr_predictor,
   kie_predictor,
   detection_predictor,
   recognition_predictor,
)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("=" * 78)
print(f"docTR      : {doctr.__version__}")
print(f"torch      : {torch.__version__}")
print(f"device     : {DEVICE}"
     + (f"  ({torch.cuda.get_device_name(0)})" if DEVICE == "cuda" else ""))
print(f"python     : {sys.version.split()[0]}")
print("=" * 78)
print("NOTE: if the import above failed, restart the runtime "
     "(Runtime > Restart session) and re-run this cell.\n")
CFG = dict(
   RUN_BENCHMARK   = True,
   RUN_SECOND_PASS = True,
   RUN_ROTATION    = True,
   RUN_LAYOUT      = True,
   RUN_KIE         = True,
   RUN_SYNTHESIS   = True,
   RUN_PDF_EXPORT  = True,
)
WORK = "/content/doctr_demo" if os.path.isdir("/content") else "./doctr_demo"
os.makedirs(WORK, exist_ok=True)
print(f"working dir: {WORK}\n")
_FONT = font_manager.findfont(font_manager.FontProperties(family="DejaVu Sans"))
_FONT_B = font_manager.findfont(
   font_manager.FontProperties(family="DejaVu Sans", weight="bold"))
A4 = (1240, 1754)
INVOICE_LINES = [
   ( 80,  70, "NORTHWIND TRADING CO.",                    38, True ),
   ( 80, 122, "42 Harbour Road, Bristol BS1 5TY",         22, False),
   ( 80, 152, "VAT GB 884 5521 09",                       22, False),
   (820,  70, "INVOICE",                                  44, True ),
   (820, 132, "Invoice No: INV-2024-00817",               22, False),
   (820, 162, "Date: 14/03/2024",                         22, False),
   (820, 192, "Due Date: 13/04/2024",                     22, False),
   ( 80, 260, "BILL TO",                                  24, True ),
   ( 80, 296, "Aurora Robotics Ltd",                      24, False),
   ( 80, 328, "Unit 7 Fenway Business Park",              22, False),
   ( 80, 358, "Cambridge CB4 0WS",                        22, False),
   ( 80, 388, "Contact: procurement@aurorarobotics.co.uk",22, False),
   ( 80, 470, "DESCRIPTION",                              24, True ),
   (640, 470, "QTY",                                      24, True ),
   (780, 470, "UNIT PRICE",                               24, True ),
   (1010,470, "AMOUNT",                                   24, True ),
   ( 80, 520, "Servo controller board Rev C",             22, False),
   (640, 520, "12",                                       22, False),
   (780, 520, "84.50",                                    22, False),
   (1010,520, "1014.00",                                  22, False),
   ( 80, 560, "Harmonic drive gearbox 50:1",              22, False),
   (640, 560, "4",                                        22, False),
   (780, 560, "312.75",                                   22, False),
   (1010,560, "1251.00",                                  22, False),
   ( 80, 600, "Shielded encoder cable 2m",                22, False),
   (640, 600, "20",                                       22, False),
   (780, 600, "11.40",                                    22, False),
   (1010,600, "228.00",                                   22, False),
   ( 80, 640, "Calibration service on-site",              22, False),
   (640, 640, "1",                                        22, False),
   (780, 640, "450.00",                                   22, False),
   (1010,640, "450.00",                                   22, False),
   (780, 720, "Subtotal",                                 22, False),
   (1010,720, "2943.00",                                  22, False),
   (780, 756, "VAT 20%",                                  22, False),
   (1010,756, "588.60",                                   22, False),
   (780, 796, "TOTAL DUE",                                26, True ),
   (1010,796, "3531.60",                                  26, True ),
   ( 80, 900, "PAYMENT TERMS",                            24, True ),
   ( 80, 936, "Net 30 days. Late payments accrue interest at 2% per month.", 20, False),
   ( 80, 968, "Bank: Lloyds  Sort Code: 30-96-26  Account: 41775302",       20, False),
   ( 80,1010, "Reference: INV-2024-00817",                20, False),
]
PAGE2_LINES = [
   ( 80,  70, "APPENDIX A - DELIVERY SCHEDULE",           34, True ),
   ( 80, 140, "All shipments leave the Bristol warehouse before 16:00 GMT.", 22, False),
   ( 80, 176, "Tracking numbers are emailed on the day of dispatch.",       22, False),
   ( 80, 240, "MILESTONE",                                24, True ),
   (700, 240, "TARGET DATE",                              24, True ),
   ( 80, 288, "Purchase order acknowledged",              22, False),
   (700, 288, "18/03/2024",                               22, False),
   ( 80, 328, "Controller boards shipped",                22, False),
   (700, 328, "25/03/2024",                               22, False),
   ( 80, 368, "Gearboxes shipped",                        22, False),
   (700, 368, "02/04/2024",                               22, False),
   ( 80, 408, "On-site calibration window",               22, False),
   (700, 408, "08/04/2024",                               22, False),
   ( 80, 480, "Questions? Call +44 117 496 0022 or email ops@northwind.example", 20, False),
]
def render_page(lines, size=A4, bg=250):
   """Draw a clean document page from a list of (x, y, text, size, bold)."""
   img = Image.new("RGB", size, (bg, bg, bg))
   d = ImageDraw.Draw(img)
   for x, y, text, sz, bold in lines:
       font = ImageFont.truetype(_FONT_B if bold else _FONT, sz)
       d.text((x, y), text, fill=(18, 18, 22), font=font)
   d.line([(80, 455), (1160, 455)], fill=(60, 60, 60), width=2)
   d.line([(80, 505), (1160, 505)], fill=(160, 160, 160), width=1)
   d.line([(760, 700), (1160, 700)], fill=(60, 60, 60), width=2)
   return img
def scanify(img, angle=0.0, noise=6.0, jpeg_quality=72, blur_shadow=True):
   """Degrade a clean render so it behaves like a phone photo / flatbed scan."""
   if angle:
       img = img.rotate(angle, expand=True, resample=Image.BICUBIC,
                        fillcolor=(250, 250, 250))
   arr = np.asarray(img).astype(np.float32)
   if blur_shadow:
       h, w = arr.shape[:2]
       gx = np.linspace(-1, 1, w)[None, :]
       gy = np.linspace(-1, 1, h)[:, None]
       shade = 1.0 - 0.10 * (gx ** 2 + 0.6 * gy ** 2)
       arr *= shade[..., None]
   if noise:
       arr += np.random.normal(0, noise, arr.shape)
   arr = np.clip(arr, 0, 255).astype(np.uint8)
   out = Image.fromarray(arr)
   if jpeg_quality:
       buf = io.BytesIO()
       out.save(buf, format="JPEG", quality=jpeg_quality)
       buf.seek(0)
       out = Image.open(buf).convert("RGB")
   return out
clean1 = render_page(INVOICE_LINES)
clean2 = render_page(PAGE2_LINES)
page1_path   = os.path.join(WORK, "invoice_p1.png")
page2_path   = os.path.join(WORK, "invoice_p2.png")
rotated_path = os.path.join(WORK, "invoice_rotated.png")
pdf_path     = os.path.join(WORK, "invoice.pdf")
scanify(clean1, angle=0.4).save(page1_path)
scanify(clean2, angle=-0.3).save(page2_path)
scanify(clean1, angle=13.0, noise=8.0).save(rotated_path)
clean1.save(pdf_path, save_all=True, append_images=[clean2], resolution=150)
GT_WORDS_P1 = [w for _, _, t, _, _ in INVOICE_LINES for w in t.split()]
print(f"generated: {page1_path}, {page2_path}, {rotated_path}, {pdf_path}")
print(f"ground-truth words on page 1: {len(GT_WORDS_P1)}\n")
fig, ax = plt.subplots(1, 3, figsize=(15, 7))
for a, im, t in zip(ax, [Image.open(page1_path), Image.open(page2_path),
                        Image.open(rotated_path)],
                   ["page 1 (scanified)", "page 2", "rotated 13 deg"]):
   a.imshow(im); a.set_title(t, fontsize=10); a.axis("off")
plt.tight_layout(); plt.show()
imgs_doc  = DocumentFile.from_images([page1_path, page2_path])
pdf_doc   = DocumentFile.from_pdf(pdf_path)
pdf_hi    = DocumentFile.from_pdf(pdf_path, scale=3)
rot_doc   = DocumentFile.from_images(rotated_path)
print("from_images :", [p.shape for p in imgs_doc], imgs_doc[0].dtype)
print("from_pdf    :", [p.shape for p in pdf_doc])
print("from_pdf x3 :", [p.shape for p in pdf_hi])
print("""
Rules of thumb for `scale`:
 * body text should be >= ~10 px tall for the recognition model to be happy
 * scale=2 (default) suits 150-300 dpi scans; bump to 3-4 for dense 8pt text
 * you can also pass raw numpy arrays straight to any predictor:
       predictor([np.asarray(pil_image)])
 * DocumentFile.from_url(...) exists too, but needs the [html] extra
""")
def build_ocr(det="db_resnet50", reco="crnn_vgg16_bn", **kw):
   """Construct an OCR predictor and move it to the GPU when there is one."""
   model = ocr_predictor(det_arch=det, reco_arch=reco, pretrained=True, **kw)
   if DEVICE == "cuda":
       try:
           model = model.cuda()
       except Exception as e:
           print(f"  (cuda placement skipped: {e})")
   return model
def timeit(fn, *args, warmup=1, runs=3, **kw):
   """Warm up (weight load / cudnn autotune / lazy init), then time properly."""
   for _ in range(warmup):
       fn(*args, **kw)
   if DEVICE == "cuda":
       torch.cuda.synchronize()
   t0 = time.perf_counter()
   out = None
   for _ in range(runs):
       out = fn(*args, **kw)
   if DEVICE == "cuda":
       torch.cuda.synchronize()
   return out, (time.perf_counter() - t0) / runs
predictor = build_ocr()
result, dt = timeit(predictor, imgs_doc, runs=2)
print(f"\nbaseline end-to-end: {dt:.2f}s for {len(imgs_doc)} pages "
     f"({dt/len(imgs_doc):.2f}s/page on {DEVICE})")
print(f"first 90 chars of page 1: {result.pages[0].render()[:90]!r}")

def norm(w):
   return re.sub(r"[^\w@:./+-]", "", w.lower())
def bag_accuracy(gt_words, pred_words):
   """Order-insensitive word recall — good enough to rank models quickly."""
   g, p = Counter(map(norm, gt_words)), Counter(map(norm, pred_words))
   return sum((g & p).values()) / max(len(gt_words), 1)
def page_words(page):
   return [w.value for b in page.blocks for l in b.lines for w in l.words]
if CFG["RUN_BENCHMARK"]:
   combos = [
       ("db_mobilenet_v3_large", "crnn_mobilenet_v3_small"),
       ("fast_base",             "crnn_vgg16_bn"),
       ("db_resnet50",           "crnn_vgg16_bn"),
       ("db_resnet50",           "parseq"),
   ]
   rows = []
   for det, reco in combos:
       try:
           m = build_ocr(det, reco)
           res, dt = timeit(m, [imgs_doc[0]], warmup=1, runs=2)
           pw = page_words(res.pages[0])
           rows.append((f"{det} + {reco}", dt, len(pw), bag_accuracy(GT_WORDS_P1, pw)))
           del m
           if DEVICE == "cuda":
               torch.cuda.empty_cache()
       except Exception as e:
           rows.append((f"{det} + {reco}", float("nan"), 0, float("nan")))
           print(f"  !! {det}+{reco} failed: {e}")
   print("\n" + "-" * 78)
   print(f"{'architecture':<46}{'sec/page':>10}{'#words':>9}{'word acc':>11}")
   print("-" * 78)
   for name, dt, n, acc in rows:
       print(f"{name:<46}{dt:>10.2f}{n:>9}{acc:>10.1%}")
   print("-" * 78)
   print("""
Reading the table:
 * detection choice drives RECALL (#words found); recognition drives accuracy
 * mobilenet variants are 5-10x cheaper and lose only a couple of points on
   clean documents — they are usually the right default for bulk pipelines
 * parseq / master are worth it on noisy, handwritten or curved text only
 * these numbers are for ONE synthetic page; always benchmark on your own data
""")
page = result.pages[0]
print(f"page dimensions   : {page.dimensions}   (H, W in px)")
print(f"page orientation  : {page.orientation}")
print(f"page language     : {page.language}")
print(f"blocks/lines/words: {len(page.blocks)}, "
     f"{sum(len(b.lines) for b in page.blocks)}, {len(page_words(page))}\n")
for b_i, block in enumerate(page.blocks[:1]):
   print(f"Block {b_i}  geometry={np.round(np.array(block.geometry), 3).tolist()}")
   for l_i, line in enumerate(block.lines[:2]):
       print(f"  Line {l_i}: {' '.join(w.value for w in line.words)}")
       for w in line.words[:4]:
           geo = np.round(np.array(w.geometry), 4).tolist()
           print(f"    Word {w.value!r:<22} conf={w.confidence:.3f} "
                 f"objectness={getattr(w, 'objectness_score', None)} "
                 f"crop_orient={getattr(w, 'crop_orientation', None)}")
           print(f"      geometry={geo}")
print("""
Key facts about geometry:
 * coordinates are RELATIVE (0-1), so multiply by (W, H) to get pixels
 * assume_straight_pages=True  -> ((xmin, ymin), (xmax, ymax))
 * assume_straight_pages=False -> a 4-point polygon [(x,y) x 4], clockwise
 * confidence       = recognition softmax confidence for the whole word
 * objectness_score = how sure the DETECTOR was that this is text
   -> filter on objectness to kill hallucinated boxes, on confidence to
      flag words a human should review. They fail differently.
""")
def geom_to_pixels(geom, w, h):
   g = np.asarray(geom, dtype=np.float32)
   if g.ndim == 2 and g.shape == (2, 2):
       (x0, y0), (x1, y1) = g
       return np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) * [w, h]
   return g[:4] * [w, h]
def draw_result(page_obj, image, title="", min_conf=0.0, figsize=(13, 18),
               label=True):
   img = np.asarray(image)
   h, w = img.shape[:2]
   cmap = matplotlib.colormaps["RdYlGn"]
   fig, ax = plt.subplots(figsize=figsize)
   ax.imshow(img); ax.axis("off"); ax.set_title(title)
   for block in page_obj.blocks:
       for line in block.lines:
           for word in line.words:
               if word.confidence < min_conf:
                   continue
               pts = geom_to_pixels(word.geometry, w, h)
               c = cmap(float(word.confidence))
               ax.add_patch(MplPolygon(pts, closed=True, fill=False,
                                       edgecolor=c, linewidth=1.4))
               if label and word.confidence < 0.85:
                   ax.text(pts[:, 0].min(), pts[:, 1].min() - 4,
                           f"{word.value} {word.confidence:.2f}",
                           fontsize=6, color="crimson")
   sm = matplotlib.cm.ScalarMappable(cmap=cmap,
                                     norm=matplotlib.colors.Normalize(0, 1))
   fig.colorbar(sm, ax=ax, fraction=0.025, label="recognition confidence")
   plt.tight_layout(); plt.show()
draw_result(page, imgs_doc[0], "page 1 — words coloured by confidence")
confs = [w.confidence for w in
        (wd for b in page.blocks for l in b.lines for wd in l.words)]
print(f"confidence: mean={np.mean(confs):.3f}  p10={np.percentile(confs,10):.3f}  "
     f"min={np.min(confs):.3f}   below 0.8: {sum(c < .8 for c in confs)} words")
det = detection_predictor("db_resnet50", pretrained=True,
                         assume_straight_pages=True, preserve_aspect_ratio=True)
if DEVICE == "cuda":
   det = det.cuda()
det_out = det([imgs_doc[0]])[0]
key = list(det_out.keys())[0]
boxes = det_out[key]
print(f"detection output: key={key!r} shape={boxes.shape}  "
     f"(last column is the objectness score)")
print("first 3 boxes (relative):\n", np.round(boxes[:3], 4))
def crop_words(image, boxes, pad=0.004):
   """Cut relative boxes out of an image, with a little padding."""
   img = np.asarray(image)
   h, w = img.shape[:2]
   crops = []
   for b in boxes:
       x0, y0, x1, y1 = b[:4]
       x0 = int(max(0, (x0 - pad)) * w); x1 = int(min(1, (x1 + pad)) * w)
       y0 = int(max(0, (y0 - pad)) * h); y1 = int(min(1, (y1 + pad)) * h)
       if x1 > x0 + 2 and y1 > y0 + 2:
           crops.append(img[y0:y1, x0:x1])
   return crops
crops = crop_words(imgs_doc[0], boxes)
print(f"\nextracted {len(crops)} crops")
reco = recognition_predictor("crnn_vgg16_bn", pretrained=True)
if DEVICE == "cuda":
   reco = reco.cuda()
reco_out = reco(crops[:24])
print("crop-level predictions (text, confidence):")
print(reco_out[:8])
print(f"\nmodel vocab ({len(reco.model.cfg['vocab'])} chars): "
     f"{reco.model.cfg['vocab'][:70]}...")
print("""
The vocab matters: the default checkpoints ship with a French/Latin vocab.
If your text contains characters outside it, the model literally cannot emit
them and you must fine-tune with a wider `vocab` (see doctr.datasets.VOCABS).
""")
fig, axes = plt.subplots(4, 3, figsize=(11, 5))
for a, c, (txt, cf) in zip(axes.ravel(), crops, reco_out):
   a.imshow(c); a.axis("off"); a.set_title(f"{txt} ({cf:.2f})", fontsize=8)
plt.tight_layout(); plt.show()

if CFG["RUN_SECOND_PASS"]:
   CONF_GATE = 0.85
   fast_model = build_ocr("db_resnet50", "crnn_mobilenet_v3_small")
   res_fast = fast_model([imgs_doc[0]])
   pg = res_fast.pages[0]
   weak = [(w, w.geometry) for b in pg.blocks for l in b.lines for w in l.words
           if w.confidence < CONF_GATE]
   print(f"pass 1 (crnn_mobilenet_v3_small): {len(page_words(pg))} words, "
         f"{len(weak)} below {CONF_GATE}")
   if weak:
       h, w_ = imgs_doc[0].shape[:2]
       rects = []
       for _, g in weak:
           pts = geom_to_pixels(g, 1.0, 1.0)
           rects.append([pts[:, 0].min(), pts[:, 1].min(),
                         pts[:, 0].max(), pts[:, 1].max()])
       weak_crops = crop_words(imgs_doc[0], np.array(rects), pad=0.006)
       strong = recognition_predictor("parseq", pretrained=True)
       if DEVICE == "cuda":
           strong = strong.cuda()
       redo = strong(weak_crops)
       print(f"\n{'before':<26}{'conf':>7}   {'after (parseq)':<26}{'conf':>7}")
       print("-" * 72)
       changed = 0
       for (word, _), (new_txt, new_cf) in zip(weak, redo):
           flag = "  <-- changed" if new_txt != word.value else ""
           changed += new_txt != word.value
           print(f"{word.value:<26}{word.confidence:>7.3f}   "
                 f"{new_txt:<26}{new_cf:>7.3f}{flag}")
       print(f"\n{changed}/{len(weak)} words revised, "
             f"but parseq only ran on {len(weak)/max(len(page_words(pg)),1):.0%} "
             f"of the crops.")
   del fast_model
tuner = build_ocr("db_resnet50", "crnn_vgg16_bn")
pp = tuner.det_predictor.model.postprocessor
orig = (pp.bin_thresh, pp.box_thresh)
print(f"defaults: bin_thresh={orig[0]}, box_thresh={orig[1]}\n")
print(f"{'bin':>6}{'box':>7}{'#words':>9}{'mean conf':>12}{'sec':>8}")
print("-" * 42)
for bin_t, box_t in [(0.1, 0.05), (0.3, 0.1), (0.5, 0.2), (0.7, 0.4), (0.9, 0.6)]:
   pp.bin_thresh, pp.box_thresh = bin_t, box_t
   t0 = time.perf_counter()
   r = tuner([imgs_doc[0]])
   dt = time.perf_counter() - t0
   ws = [w for b in r.pages[0].blocks for l in b.lines for w in l.words]
   mc = np.mean([w.confidence for w in ws]) if ws else 0
   print(f"{bin_t:>6}{box_t:>7}{len(ws):>9}{mc:>12.3f}{dt:>8.2f}")
pp.bin_thresh, pp.box_thresh = orig
print("""
How to tune in practice:
 * LOW thresholds  -> more boxes: faint stamps, dot-matrix, carbon copies.
                      Cost: noise boxes, which you then filter by objectness.
 * HIGH thresholds -> fewer, cleaner boxes for crisp born-digital scans.
 * Sweep against a small labelled set and optimise F1, not eyeballs.
""")
class PadBoxesHook:
   """Recognition often improves when crops aren't cut flush to the glyphs."""
   def __init__(self, dx=0.004, dy=0.006):
       self.dx, self.dy = dx, dy
   def _pad(self, arr):
       a = np.array(arr, copy=True, dtype=np.float32)
       if a.ndim == 2 and a.shape[-1] >= 4:
           a[:, 0] = np.clip(a[:, 0] - self.dx, 0, 1)
           a[:, 1] = np.clip(a[:, 1] - self.dy, 0, 1)
           a[:, 2] = np.clip(a[:, 2] + self.dx, 0, 1)
           a[:, 3] = np.clip(a[:, 3] + self.dy, 0, 1)
       elif a.ndim == 3:
           pts = a[:, :4, :]
           ctr = pts.mean(axis=1, keepdims=True)
           a[:, :4, :] = np.clip(ctr + (pts - ctr) * 1.06, 0, 1)
       return a
   def __call__(self, loc_preds):
       out = []
       for p in loc_preds:
           out.append({k: self._pad(v) for k, v in p.items()}
                      if isinstance(p, dict) else self._pad(p))
       return out
class DropTinyBoxesHook:
   """Kill speckle boxes before they waste a recognition forward pass."""
   def __init__(self, min_h=0.006, min_w=0.004):
       self.min_h, self.min_w = min_h, min_w
   def _filt(self, arr):
       a = np.asarray(arr)
       if a.ndim == 2 and a.shape[-1] >= 4:
           keep = ((a[:, 2] - a[:, 0]) > self.min_w) & \
                  ((a[:, 3] - a[:, 1]) > self.min_h)
           return a[keep]
       if a.ndim == 3:
           pts = a[:, :4, :]
           wd = pts[..., 0].max(1) - pts[..., 0].min(1)
           ht = pts[..., 1].max(1) - pts[..., 1].min(1)
           return a[(wd > self.min_w) & (ht > self.min_h)]
       return a
   def __call__(self, loc_preds):
       return [{k: self._filt(v) for k, v in p.items()} if isinstance(p, dict)
               else self._filt(p) for p in loc_preds]
hooked = build_ocr("db_resnet50", "crnn_vgg16_bn")
before = hooked([imgs_doc[0]]).pages[0]
hooked.add_hook(DropTinyBoxesHook())
hooked.add_hook(PadBoxesHook())
after = hooked([imgs_doc[0]]).pages[0]
bw, aw = page_words(before), page_words(after)
print(f"no hooks : {len(bw):>4} words  mean conf "
     f"{np.mean([w.confidence for b in before.blocks for l in b.lines for w in l.words]):.4f}")
print(f"hooked   : {len(aw):>4} words  mean conf "
     f"{np.mean([w.confidence for b in after.blocks for l in b.lines for w in l.words]):.4f}")
print(f"word accuracy vs GT: {bag_accuracy(GT_WORDS_P1, bw):.1%} -> "
     f"{bag_accuracy(GT_WORDS_P1, aw):.1%}")
print("""
Other things hooks are good for:
 * snapping boxes to a known form template / table grid
 * merging boxes that the detector split across a hyphen or thin space
 * masking a redacted region so its crops never reach the recogniser
""")
if CFG["RUN_ROTATION"]:
   print("Three strategies for non-straight pages:\n"
         "  A) assume_straight_pages=True   fastest, breaks past ~5 deg skew\n"
         "  B) assume_straight_pages=False  returns 4-point polygons\n"
         "  C) straighten_pages=True        de-skews the page first, then A\n")
   variants = {
       "A straight (default)": dict(assume_straight_pages=True),
       "B polygons":           dict(assume_straight_pages=False,
                                    preserve_aspect_ratio=True),
       "C straighten first":   dict(assume_straight_pages=False,
                                    straighten_pages=True,
                                    detect_orientation=True),
       "B' polygons -> boxes": dict(assume_straight_pages=False,
                                    export_as_straight_boxes=True),
   }
   rot_results = {}
   for name, kw in variants.items():
       try:
           m = build_ocr("db_resnet50", "crnn_vgg16_bn", **kw)
           t0 = time.perf_counter()
           r = m(rot_doc)
           dt = time.perf_counter() - t0
           p = r.pages[0]
           ws = page_words(p)
           rot_results[name] = (r, p)
           print(f"{name:<24} words={len(ws):>4}  acc={bag_accuracy(GT_WORDS_P1, ws):>6.1%}  "
                 f"{dt:>5.2f}s  orientation={p.orientation}")
           del m
       except Exception as e:
           print(f"{name:<24} failed: {e}")
   if "B polygons" in rot_results:
       draw_result(rot_results["B polygons"][1], rot_doc[0],
                   "rotated page — polygon boxes", figsize=(11, 14), label=False)
   print("""
Extra speed switches once you know your data:
 disable_page_orientation=True  skip the 0/90/180/270 page classifier
 disable_crop_orientation=True  skip the per-word orientation classifier
Both only matter when assume_straight_pages=False / straighten_pages=True.
""")

if CFG["RUN_LAYOUT"]:
   try:
       lay = ocr_predictor(pretrained=True, detect_layout=True)
       if DEVICE == "cuda":
           lay = lay.cuda()
       lres = lay(imgs_doc)
       lpage = lres.pages[0]
       regions = getattr(lpage, "layout", []) or []
       print(f"detected {len(regions)} layout regions on page 1:")
       counts = Counter()
       for r in regions:
           counts[r.type] += 1
           print(f"  {r.type:<16} conf={r.confidence:.3f}  "
                 f"geom={np.round(np.array(r.geometry), 3).tolist()}")
       print("\nregion histogram:", dict(counts))
       h, w = imgs_doc[0].shape[:2]
       colors = {"Title": "tab:red", "Text": "tab:blue", "Table": "tab:green",
                 "Page-header": "tab:orange", "Page-footer": "tab:purple"}
       fig, ax = plt.subplots(figsize=(10, 14))
       ax.imshow(imgs_doc[0]); ax.axis("off")
       ax.set_title("layout regions")
       for r in regions:
           pts = geom_to_pixels(r.geometry, w, h)
           ax.add_patch(MplPolygon(pts, closed=True, fill=False, linewidth=2.2,
                                   edgecolor=colors.get(r.type, "black")))
           ax.text(pts[:, 0].min(), pts[:, 1].min() - 6, r.type, fontsize=9,
                   color=colors.get(r.type, "black"))
       plt.tight_layout(); plt.show()
       print("""
Why layout matters: it gives you *document structure*, not just text. Route
Table regions to a table parser, drop Page-header/Page-footer before feeding
an LLM, and use Title regions to chunk long documents sensibly.
""")
       del lay
   except TypeError:
       print("detect_layout not supported by this docTR version "
             "(needs >= 1.0) — upgrade with: pip install -U python-doctr")
   except Exception as e:
       print(f"layout detection unavailable: {e}")
if CFG["RUN_KIE"]:
   kie = kie_predictor(det_arch="db_resnet50", reco_arch="crnn_vgg16_bn",
                       pretrained=True)
   if DEVICE == "cuda":
       kie = kie.cuda()
   kres = kie([imgs_doc[0]])
   preds = kres.pages[0].predictions
   for cls, items in preds.items():
       print(f"class {cls!r}: {len(items)} predictions")
       for p in items[:5]:
           print(f"   {p.value!r:<24} conf={p.confidence:.3f} "
                 f"geom={np.round(np.array(p.geometry), 3).tolist()}")
   print("""
To make this genuinely useful, train a detection model with several classes
(references/detection/train_pytorch.py with a multi-class label file), e.g.
classes = ["invoice_number", "total", "date"]. Then KIE returns exactly those
fields already transcribed — no regex layer required.
""")
   del kie
res = predictor(imgs_doc)
txt = res.render()
print("--- render() -------------------------------------------------------")
print(txt[:320], "...\n")
open(os.path.join(WORK, "output.txt"), "w").write(txt)
js = res.export()
print("--- export() keys --------------------------------------------------")
print("document:", list(js.keys()))
print("page    :", list(js["pages"][0].keys()))
print("word    :", list(js["pages"][0]["blocks"][0]["lines"][0]["words"][0].keys()))
with open(os.path.join(WORK, "output.json"), "w") as f:
   json.dump(js, f, indent=2, default=str)
xml_out = res.export_as_xml()
xml_bytes, xml_tree = xml_out[0]
print("\n--- export_as_xml() (hOCR) ----------------------------------------")
print(xml_bytes.decode()[:520], "...")
for i, (b, _) in enumerate(xml_out):
   open(os.path.join(WORK, f"page_{i+1}.hocr"), "wb").write(b)
if CFG["RUN_SYNTHESIS"]:
   synth = res.synthesize()
   fig, ax = plt.subplots(1, 2, figsize=(14, 10))
   ax[0].imshow(imgs_doc[0]); ax[0].set_title("original"); ax[0].axis("off")
   ax[1].imshow(synth[0]);    ax[1].set_title("synthesize()"); ax[1].axis("off")
   plt.tight_layout(); plt.show()
   print("synthesize() re-renders text into the detected boxes. If the "
         "reconstruction looks right, geometry AND transcription are both OK.")
pg = res.pages[0]
H, W = pg.dimensions
def word_rect(word):
   """Relative geometry -> (x0, y0, x1, y1) axis-aligned, works for polygons."""
   p = np.asarray(word.geometry, dtype=np.float32)
   if p.shape == (2, 2):
       return float(p[0, 0]), float(p[0, 1]), float(p[1, 0]), float(p[1, 1])
   return (float(p[:, 0].min()), float(p[:, 1].min()),
           float(p[:, 0].max()), float(p[:, 1].max()))
flat = []
for b in pg.blocks:
   for l in b.lines:
       for w in l.words:
           x0, y0, x1, y1 = word_rect(w)
           flat.append(dict(text=w.value, conf=w.confidence,
                            x0=x0, y0=y0, x1=x1, y1=y1,
                            cx=(x0 + x1) / 2, cy=(y0 + y1) / 2, h=y1 - y0))
def group_rows(words, tol_factor=0.6):
   ws = sorted(words, key=lambda d: d["cy"])
   rows, cur, ref = [], [], None
   for w in ws:
       tol = max(w["h"] * tol_factor, 0.004)
       if ref is None or abs(w["cy"] - ref) <= tol:
           cur.append(w); ref = np.mean([c["cy"] for c in cur])
       else:
           rows.append(sorted(cur, key=lambda d: d["x0"])); cur, ref = [w], w["cy"]
   if cur:
       rows.append(sorted(cur, key=lambda d: d["x0"]))
   return rows
rows = group_rows(flat)
print(f"--- reading order: {len(rows)} rows ---")
for r in rows[:8]:
   print("   " + " ".join(w["text"] for w in r))
full_text = "\n".join(" ".join(w["text"] for w in r) for r in rows)
FIELDS = {
   "invoice_no":  r"Invoice\s*No[:\s]*([A-Z0-9\-]+)",
   "date":        r"\bDate[:\s]*(\d{2}/\d{2}/\d{4})",
   "due_date":    r"Due\s*Date[:\s]*(\d{2}/\d{2}/\d{4})",
   "vat_id":      r"VAT\s*(GB[\s\d]{8,})",
   "total_due":   r"TOTAL\s*DUE\s*([\d.,]+)",
   "subtotal":    r"Subtotal\s*([\d.,]+)",
   "email":       r"([\w.+-]+@[\w-]+\.[\w.]+)",
   "sort_code":   r"Sort\s*Code[:\s]*([\d-]{6,10})",
}
print("\n--- extracted fields ---")
extracted = {}
for name, pat in FIELDS.items():
   m = re.search(pat, full_text, flags=re.IGNORECASE)
   extracted[name] = m.group(1).strip() if m else None
   print(f"  {name:<12}: {extracted[name]}")
def detect_columns(rows, y_lo, y_hi, gap=0.03):
   """1-D clustering of word left-edges inside a band -> column boundaries."""
   xs = sorted(w["x0"] for r in rows for w in r if y_lo <= w["cy"] <= y_hi)
   if not xs:
       return []
   cols, cur = [], [xs[0]]
   for x in xs[1:]:
       if x - cur[-1] < gap:
           cur.append(x)
       else:
           cols.append(cur)
           cur = [x]
   cols.append(cur)
   return [float(np.min(c)) for c in cols if c]
band_lo, band_hi = 0.25, 0.40
col_x = detect_columns(rows, band_lo, band_hi)
print(f"\n--- table: {len(col_x)} columns at x={np.round(col_x, 3).tolist()} ---")
table = []
for r in rows:
   if not (band_lo <= np.mean([w["cy"] for w in r]) <= band_hi):
       continue
   cells = [""] * len(col_x)
   for w in r:
       idx = int(np.argmin([abs(w["x0"] - cx) for cx in col_x]))
       cells[idx] = (cells[idx] + " " + w["text"]).strip()
   table.append(cells)
for row in table:
   print("  | " + " | ".join(f"{c:<28}" if i == 0 else f"{c:<10}"
                             for i, c in enumerate(row)))
print("""
Escalation path when this gets hairy:
 * per-page dict -> pandas.DataFrame for downstream joins
 * detect_layout=True to isolate Table regions before column clustering
 * or hand result.render() / the hOCR to an LLM for schema-guided extraction —
   docTR's job is faithful text + geometry, not semantics
""")

if CFG["RUN_PDF_EXPORT"]:
   from reportlab.pdfgen import canvas as rl_canvas
   from reportlab.lib.utils import ImageReader
   def make_searchable_pdf(pages_np, doc_result, out_path, dpi=150):
       c = rl_canvas.Canvas(out_path)
       for img_np, page_obj in zip(pages_np, doc_result.pages):
           h_px, w_px = img_np.shape[:2]
           w_pt, h_pt = w_px * 72.0 / dpi, h_px * 72.0 / dpi
           c.setPageSize((w_pt, h_pt))
           c.drawImage(ImageReader(Image.fromarray(img_np)), 0, 0,
                       width=w_pt, height=h_pt)
           c.setFillColorRGB(0, 0, 0)
           for b in page_obj.blocks:
               for l in b.lines:
                   for wd in l.words:
                       if not wd.value.strip():
                           continue
                       x0, y0, x1, y1 = word_rect(wd)
                       bx, by = x0 * w_pt, (1 - y1) * h_pt
                       bw_, bh_ = (x1 - x0) * w_pt, (y1 - y0) * h_pt
                       size = max(bh_ * 0.82, 1.0)
                       t = c.beginText()
                       t.setTextRenderMode(3)
                       t.setFont("Helvetica", size)
                       adv = c.stringWidth(wd.value, "Helvetica", size) or 1.0
                       t.setHorizScale(100.0 * bw_ / adv)
                       t.setTextOrigin(bx, by + bh_ * 0.18)
                       t.textOut(wd.value)
                       c.drawText(t)
           c.showPage()
       c.save()
       return out_path
   out_pdf = make_searchable_pdf(imgs_doc, res,
                                 os.path.join(WORK, "invoice_searchable.pdf"))
   print(f"searchable PDF written: {out_pdf} "
         f"({os.path.getsize(out_pdf)/1024:.0f} KB)")
   print("Open it and Ctrl+F for 'INV-2024-00817' — the scan is unchanged, "
         "but the text is selectable.")
   try:
       from google.colab import files
       print("Run  files.download(out_pdf)  to pull it down from Colab.")
   except ImportError:
       pass
print("""
=============================== PERFORMANCE ==================================
Batch sizes (biggest single lever on GPU):
   ocr_predictor(pretrained=True, det_bs=4, reco_bs=1024)
 Detection is memory-bound (1024x1024 feature maps) so det_bs stays small;
 recognition crops are tiny (32x128) so reco_bs can be huge. On a T4 start at
 det_bs=2, reco_bs=512 and raise reco_bs until you OOM.
Cheap wins, in rough order of payoff:
 1. swap to db_mobilenet_v3_large + crnn_mobilenet_v3_small   (5-10x)
 2. pass ALL pages in one call — predictor(list_of_pages) batches internally
 3. assume_straight_pages=True + disable_*_orientation when data allows
 4. lower the PDF `scale` if your text is already large
 5. half precision:  predictor = predictor.half()  (test accuracy first;
    some post-processors expect float32, so keep a fallback)
Structure knobs (handled by DocumentBuilder):
   resolve_lines=True      group words into lines            (default True)
   resolve_blocks=False    group lines into blocks           (default False)
   paragraph_break=0.035   relative gap that splits paragraphs
============================== FINE-TUNING ===================================
Stock checkpoints are trained on a French/Latin vocab and generic documents.
Fine-tune when you have a custom alphabet, a specialist font, or a domain-
specific layout. In the repo:
   references/detection/train_pytorch.py
   references/recognition/train_pytorch.py
   references/classification/train_pytorch.py   (orientation classifiers)
Recognition wants word crops + labels.json; detection wants full pages with
polygon labels (multi-class supported -> feeds kie_predictor).
Then load your weights:
   from doctr.models import db_resnet50, ocr_predictor
   det = db_resnet50(pretrained=False)
   det.load_state_dict(torch.load("my_det.pt", map_location="cpu"))
   model = ocr_predictor(det_arch=det, reco_arch="crnn_vgg16_bn",
                         pretrained=True)
docTR also pushes/pulls checkpoints from the Hugging Face Hub
(doctr.models.factory: push_to_hf_hub / from_hub).
============================== DEPLOYMENT ====================================
 * FastAPI template in api/ with /detection /recognition /ocr /kie routes
 * GPU-ready Docker images: ghcr.io/mindee/doctr
 * Streamlit demo: streamlit run demo/app.py
 * Live demo: huggingface.co/spaces/mindee/doctr
 * Full docs: mindee.github.io/doctr
==============================================================================
""")
print(f"\nAll artefacts are in {WORK}:")
for f in sorted(os.listdir(WORK)):
   print(f"   {f:<28}{os.path.getsize(os.path.join(WORK, f))/1024:>8.0f} KB")
print("\nDone.")