Spaces:
Sleeping
Sleeping
| 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 | |