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}")