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}")
Credit: Source link
