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