Spaces:
Sleeping
Sleeping
File size: 2,753 Bytes
c1b1f3a | 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 | from typing import Dict, Any, List
import numpy as np
from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS
class ProximityHeuristicsExtractor:
"""
Baseline B: Spatial Bounding-Box Proximity Extractor.
Associates values located directly to the right or below target key words.
"""
def __init__(self):
self.triggers = {
"invoice_number": ["invoice", "inv", "number", "no", "#"],
"invoice_date": ["date", "dated"],
"total_amount": ["total", "due", "balance", "amount"],
"company_name": ["store", "mart", "inc", "ltd", "cafe"],
"date": ["date", "time"],
"applicant_name": ["name", "applicant", "employee", "signature"],
"address": ["address", "street", "city", "place"]
}
def _find_words_near(self, key_word: OCRWord, all_words: List[OCRWord], max_dx=300, max_dy=80) -> List[OCRWord]:
kx0, ky0, kx1, ky1 = key_word.bbox
candidates = []
for w in all_words:
if w == key_word:
continue
wx0, wy0, wx1, wy1 = w.bbox
# 1. Directly to the right (same line)
is_right = (wx0 >= kx0) and (wx0 - kx1 <= max_dx) and (abs(wy0 - ky0) <= 25)
# 2. Directly below (within column)
is_below = (wy0 >= ky1) and (wy0 - ky1 <= max_dy) and (abs(wx0 - kx0) <= 100)
if is_right or is_below:
dist = np.sqrt((wx0 - kx1)**2 + (wy0 - ky0)**2)
candidates.append((dist, w))
candidates.sort(key=lambda x: x[0])
return [c[1] for c in candidates]
def extract(self, record: DocumentRecord, doc_class: str = None) -> Dict[str, Any]:
cls_name = doc_class or record.doc_class
schema_fields = CLASS_FIELD_SCHEMAS.get(cls_name, [])
extracted = {f: None for f in schema_fields}
all_words = record.ocr_words
for field in schema_fields:
trig_list = self.triggers.get(field, [])
matched_key_word = None
for w in all_words:
w_lower = w.text.lower()
if any(t in w_lower for t in trig_list):
matched_key_word = w
break
if matched_key_word:
near_words = self._find_words_near(matched_key_word, all_words)
if near_words:
# Collect tokens on same line
val_tokens = [near_words[0].text]
for nw in near_words[1:4]:
if abs(nw.bbox[1] - near_words[0].bbox[1]) <= 20:
val_tokens.append(nw.text)
extracted[field] = " ".join(val_tokens)
return extracted
|