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: [email protected]",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 [email protected]", 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}")