File size: 3,471 Bytes
4dce182 e24bf15 4dce182 e24bf15 4dce182 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Per-segment features for the extraction model. Used identically in training and inference."""
import math
import re
import numpy as np
from .charset import CHAR_TO_ID, MATH_ID
MAX_CHARS = 40
MAX_SEGMENTS = 256
_MONEY = re.compile(r"^[-(]?\s*[$€£¥₹]?\s*-?\d{1,3}([,. ]\d{3})*([.,]\d{2})\s*-?\)?$|^[-(]?[$€£¥₹]\s*\d+([.,]\d{1,2})?\)?$")
_NUMBER = re.compile(r"^[-+(]?[$€£¥₹]?\s*\d[\d,. ]*[)%]?$")
_DATE = re.compile(r"\b(\d{1,4}[-/.]\d{1,2}[-/.]\d{1,4}|\d{1,2}\s*[A-Za-z]{3,9}\.?,?\s*\d{2,4}|[A-Za-z]{3,9}\.?\s+\d{1,2},?\s+\d{4})\b")
_TIME = re.compile(r"\b\d{1,2}:\d{2}(:\d{2})?\s*([AaPp][Mm])?\b")
_PHONE = re.compile(r"(\+?\d[\d\s().-]{7,}\d)")
_CURRENCY = re.compile(r"[$€£¥₹]|\b(USD|EUR|GBP|INR|JPY|CAD|AUD|SGD|AED|CHF)\b")
def text_flags(text: str) -> list[float]:
t = text.strip()
n = max(len(t), 1)
digits = sum(c.isdigit() for c in t)
alphas = sum(c.isalpha() for c in t)
words = t.split()
return [
float(bool(_NUMBER.match(t))),
float(bool(_MONEY.match(t))),
float(bool(_DATE.search(t))),
float(bool(_TIME.search(t))),
float("%" in t),
float(bool(_CURRENCY.search(t))),
float(":" in t),
float(t.endswith(":")),
float(alphas > 1 and t.upper() == t),
float(bool(words) and all(w[:1].isupper() for w in words if w[:1].isalpha())),
digits / n,
alphas / n,
min(len(t), MAX_CHARS) / MAX_CHARS,
float("@" in t),
float(bool(_PHONE.search(t)) and digits >= 7),
float(len(set(t)) <= 2 and len(t) >= 3 and not t[0].isalnum()), # ----- / ===== separators
float(t[:1] in "•-*–" or bool(re.match(r"^\d+[.)]\s", t))),
min(len(words), 10) / 10.0,
float(t[:1] == "#" or "no." in t.lower()),
]
NUM_FLAGS = len(text_flags(""))
NUM_FEATS = 8 + 2 + 16 + NUM_FLAGS + 1
def segment_features(segments, page_w, page_h):
"""segments: list of dicts with 'bbox' [x0, y0, x1, y1] in pixels, 'text', optional 'conf'.
Returns dict of numpy arrays: char_ids (N, MAX_CHARS) int64, feats (N, F) float32,
rel_boxes (N, 4) float32 (box in units of median text height, for pairwise geometry).
"""
n = len(segments)
char_ids = np.zeros((n, MAX_CHARS), np.int64)
feats = np.zeros((n, NUM_FEATS), np.float32)
rel = np.zeros((n, 4), np.float32)
if n == 0:
return {"char_ids": char_ids, "feats": feats, "rel_boxes": rel}
boxes = np.array([s["bbox"] for s in segments], np.float32)
hs = np.maximum(boxes[:, 3] - boxes[:, 1], 1.0)
h_med = float(np.median(hs))
for i, s in enumerate(segments):
ids = [min(CHAR_TO_ID[c], MATH_ID) for c in s["text"][:MAX_CHARS] if c in CHAR_TO_ID]
char_ids[i, : len(ids)] = ids
x0, y0, x1, y1 = boxes[i]
nx0, nx1, ny0, ny1 = x0 / page_w, x1 / page_w, y0 / page_h, y1 / page_h
cx, cy = (nx0 + nx1) / 2, (ny0 + ny1) / 2
w, h = max(x1 - x0, 1.0), hs[i]
f = [nx0, ny0, nx1, ny1, cx, cy, nx1 - nx0, ny1 - ny0, math.log(h / h_med), math.log(w / h) / 4]
for freq in (1, 2, 4, 8):
f += [math.sin(math.pi * freq * cx), math.cos(math.pi * freq * cx),
math.sin(math.pi * freq * cy), math.cos(math.pi * freq * cy)]
f += text_flags(s["text"])
f.append(float(s.get("conf", 1.0)))
feats[i] = f
rel[:] = boxes / h_med
return {"char_ids": char_ids, "feats": feats, "rel_boxes": rel}
|