夜雨聆风作者：MarkTechPost
编辑：浮世Talk
编辑：浮世Talk
引言
在本教程中，我们将用 docTR 搭建一条端到端的 OCR 工作流，并探讨现代文档理解流水线如何把文字检测、识别、几何定位、版面分析、结构化抽取与结果导出结合在一起。我们会生成贴近真实的合成发票文档，通过 DocumentFile 加载图片与 PDF，构建感知 GPU 的 OCR 预测器，并针对速度与精度对不同「检测-识别」架构组合做基准测试。接着我们会深入查看 docTR 内部的 Document 层级结构，可视化带置信度的检测框，单独使用检测与识别模型，对低置信词实现二次识别，调优检测阈值，并引入自定义流水线钩子来做检测框过滤与填充。我们还会处理旋转与倾斜文档，尝试版面检测与 KIE（关键信息抽取），重建阅读顺序与表格信息，抽取结构化发票字段，并把结果导出为文本、JSON、hOCR、合成文档图像以及可搜索 PDF。最后，我们会审视实际的性能表现、微调、批处理与部署要点，理解如何从一个基础的 OCR 示例，演进为面向生产的文档智能流水线。
1. 环境设置与数据生成
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: 1424-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}")
我们首先搭建 docTR 环境，安装所需依赖，检测 GPU 是否可用，并配置本教程的运行参数。我们生成合成发票页面，施加贴近真实的扫描退化（噪声、JPEG 压缩、阴影、轻微旋转），通过 DocumentFile 加载图片与 PDF，并准备用于评估的 ground-truth 文本。随后我们构建基线 OCR 预测器，并在真实生成的文档页面上测量端到端推理性能。
2. 架构基准测试与文档结构解析
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()
我们对多组检测与识别架构组合做基准测试，比较其处理速度、检出词数与识别准确率。我们深入查看 docTR 层级化的 Document 结构，并用几何信息与识别置信度可视化检测到的词。我们还将文字检测与识别拆开，单独取出每个词的裁剪图（crop），观察独立识别模型如何处理这些检测区域。
3. 二次识别、阈值调优、自定义钩子与旋转处理
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.
""")
我们实现了一套二次识别策略：先找出低置信词，再仅用更强的 PARSeq 识别器对这些裁剪图重新识别。我们调优检测后处理阈值，并引入自定义钩子在识别前过滤掉微小的检测框、对包围盒做填充。我们还评估了处理旋转与倾斜文档的多种策略，包括基于多边形的检测、页面拉直与方向检测。
4. 版面检测、KIE、导出与阅读顺序重建
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
""")
我们用版面检测与 KIE 能力扩展 OCR 流水线，以识别文档区域并支持结构化信息抽取。我们把 OCR 结果导出为纯文本、JSON、hOCR 与合成文档表示，同时保留文字与几何信息。接着我们重建阅读顺序，用正则表达式抽取发票字段，并依据空间坐标把检测到的词组织成表格结构。
5. 可搜索 PDF、性能优化、微调与部署
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.")
我们在原始扫描文档之上叠加一层不可见的 OCR 文本层，生成可搜索 PDF，同时保留其视觉外观。我们审视了批处理、轻量检测/识别模型、方向控制与 PDF 缩放等可落地的性能优化手段。我们还梳理了微调与部署路径，以便把 docTR 模型适配到专用数据集，并将最终 OCR 流水线集成进生产应用。
结论
总之，我们全面理解了 docTR 远不止「简单文字识别」：它能在同一条工作流中融合 OCR、文档几何、版面感知、结构化后处理与面向生产的优化。我们比较了模型架构，检查了检测与识别的置信度，通过有选择的二次识别改善了困难样本的预测，调优了后处理阈值，并用自定义钩子修改了中间检测框。我们还处理了旋转文档，探索了版面与 KIE 能力，把原始 OCR 输出转换为有序文本、抽取了字段、重建了表格，并生成了多种可复用的输出格式，包括带不可见文本层的可搜索 PDF。
引用链接