Spaces:
Sleeping
Sleeping
| import re | |
| from typing import Dict, Any, List, Tuple, Optional | |
| import numpy as np | |
| from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS | |
| INV_NUM_PATTERNS = [ | |
| re.compile(r'(?i)\b(?:invoice\s*(?:number|num|no|#)?|inv(?:\s*#|\s*no|\.)?)\s*[:.\s#]*\s*([a-z0-9\-_/]+)'), | |
| re.compile(r'(?i)\b(?:order|po|bill)\s*(?:no|number|#)?\s*[:.\s#]*\s*([a-z0-9\-_/]+)') | |
| ] | |
| DATE_PATTERN = re.compile(r'(\d{1,4}[/\-.]\d{1,2}[/\-.]\d{1,4})') | |
| AMOUNT_PATTERNS = [ | |
| re.compile(r'(?i)(?:total|amount\s*due|net\s*amount|balance|grand\s*total)[\s:$]*([0-9]+[.,][0-9]{2,3})'), | |
| re.compile(r'(?i)(?:total|amount|due)[\s:$]*([0-9]+[.,]?[0-9]*)'), | |
| re.compile(r'[$€£¥]\s*([0-9]+[.,][0-9]{2})') | |
| ] | |
| class HybridFieldExtractor: | |
| """ | |
| Intelligent Hybrid Field Extractor: | |
| Combines spatial bounding-box proximity heuristics, regex pattern matching, | |
| and token-level confidence scoring to extract structured fields and bounding boxes. | |
| """ | |
| def __init__(self, doc_class: str): | |
| self.doc_class = doc_class | |
| self.schema_fields = CLASS_FIELD_SCHEMAS.get(doc_class, []) | |
| def _find_words_near(self, key_word: OCRWord, all_words: List[OCRWord], max_dx=350, max_dy=70) -> 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) <= 30) | |
| # 2. Directly below (within column) | |
| is_below = (wy0 >= ky1) and (wy0 - ky1 <= max_dy) and (abs(wx0 - kx0) <= 120) | |
| 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 _union_bboxes(self, words: List[OCRWord]) -> Optional[List[int]]: | |
| if not words: | |
| return None | |
| x0 = min(w.bbox[0] for w in words) | |
| y0 = min(w.bbox[1] for w in words) | |
| x1 = max(w.bbox[2] for w in words) | |
| y1 = max(w.bbox[3] for w in words) | |
| return [int(x0), int(y0), int(x1), int(y1)] | |
| def extract_with_details(self, record: DocumentRecord) -> Tuple[Dict[str, Any], Dict[str, List[int]], Dict[str, float]]: | |
| words = record.ocr_words | |
| full_text = " ".join([w.text for w in words]) | |
| extracted = {f: None for f in self.schema_fields} | |
| bboxes = {f: None for f in self.schema_fields} | |
| confidences = {f: 0.0 for f in self.schema_fields} | |
| if self.doc_class == "invoice": | |
| # 1. Invoice Number | |
| for pat in INV_NUM_PATTERNS: | |
| m = pat.search(full_text) | |
| if m: | |
| cand = m.group(1).strip() | |
| if cand.lower() not in ["number", "num", "no", "date", "due", "total", "ber", "page", "order"]: | |
| extracted["invoice_number"] = cand | |
| confidences["invoice_number"] = 0.92 | |
| break | |
| # Spatial fallback for invoice number | |
| if not extracted["invoice_number"]: | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["invoice", "inv", "number", "no.", "no:"]): | |
| near = self._find_words_near(w, words, max_dx=250, max_dy=50) | |
| val_words = [nw for nw in near if any(c.isalnum() for c in nw.text) and nw.text.lower() not in ["date", "tax", "total", "to", "number"]] | |
| if val_words: | |
| extracted["invoice_number"] = val_words[0].text | |
| bboxes["invoice_number"] = val_words[0].bbox | |
| confidences["invoice_number"] = 0.85 | |
| break | |
| # 2. Invoice Date | |
| for i, w in enumerate(words): | |
| m_date = DATE_PATTERN.search(w.text) | |
| if m_date: | |
| extracted["invoice_date"] = m_date.group(1) | |
| bboxes["invoice_date"] = w.bbox | |
| confidences["invoice_date"] = 0.95 | |
| break | |
| if not extracted["invoice_date"]: | |
| # Check for "date" trigger | |
| for w in words: | |
| if "date" in w.text.lower(): | |
| near = self._find_words_near(w, words) | |
| for nw in near: | |
| m_d = DATE_PATTERN.search(nw.text) | |
| if m_d: | |
| extracted["invoice_date"] = m_d.group(1) | |
| bboxes["invoice_date"] = nw.bbox | |
| confidences["invoice_date"] = 0.90 | |
| break | |
| if extracted["invoice_date"]: | |
| break | |
| # 3. Total Amount | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["total", "due", "balance", "amount"]): | |
| near = self._find_words_near(w, words, max_dx=350, max_dy=60) | |
| for nw in near: | |
| cleaned = re.sub(r'[^\d.-]', '', nw.text.replace(',', '')) | |
| if cleaned and any(c.isdigit() for c in cleaned): | |
| try: | |
| val = float(cleaned) | |
| if val > 0: | |
| extracted["total_amount"] = nw.text | |
| bboxes["total_amount"] = nw.bbox | |
| confidences["total_amount"] = 0.93 | |
| break | |
| except ValueError: | |
| pass | |
| if extracted["total_amount"]: | |
| break | |
| if not extracted["total_amount"]: | |
| m_amt = AMOUNT_PATTERNS[0].search(full_text) or AMOUNT_PATTERNS[1].search(full_text) | |
| if m_amt: | |
| extracted["total_amount"] = m_amt.group(1).strip() | |
| confidences["total_amount"] = 0.88 | |
| # 4. Seller / Vendor Name (typically top of document) | |
| top_words = [w for w in words if w.bbox[1] < 250 and len(w.text) > 2] | |
| if top_words: | |
| seller_tokens = [w for w in top_words[:4] if w.text.lower() not in ["invoice", "date", "number", "page", "to", "from", "bill"]] | |
| if seller_tokens: | |
| extracted["seller_name"] = " ".join([w.text for w in seller_tokens]) | |
| bboxes["seller_name"] = self._union_bboxes(seller_tokens) | |
| confidences["seller_name"] = 0.80 | |
| # 5. Buyer Name | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["bill to", "sold to", "customer", "attn", "client"]): | |
| near = self._find_words_near(w, words, max_dx=300, max_dy=50) | |
| if near: | |
| buyer_tokens = [nw for nw in near[:3] if nw.text.lower() not in [":", "to", "name"]] | |
| if buyer_tokens: | |
| extracted["buyer_name"] = " ".join([nw.text for nw in buyer_tokens]) | |
| bboxes["buyer_name"] = self._union_bboxes(buyer_tokens) | |
| confidences["buyer_name"] = 0.82 | |
| break | |
| elif self.doc_class == "receipt": | |
| # 1. Company Name (top header) | |
| top_words = sorted([w for w in words if w.bbox[1] < 200], key=lambda w: (w.bbox[1], w.bbox[0])) | |
| if top_words: | |
| company_tokens = [w for w in top_words[:3] if w.text.lower() not in ["receipt", "welcome", "tel", "phone", "date", "time"]] | |
| if company_tokens: | |
| extracted["company_name"] = " ".join([w.text for w in company_tokens]) | |
| bboxes["company_name"] = self._union_bboxes(company_tokens) | |
| confidences["company_name"] = 0.88 | |
| # 2. Date | |
| for w in words: | |
| m_d = DATE_PATTERN.search(w.text) | |
| if m_d: | |
| extracted["date"] = m_d.group(1) | |
| bboxes["date"] = w.bbox | |
| confidences["date"] = 0.94 | |
| break | |
| # 3. Address | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["street", "st.", "ave", "road", "rd", "blvd", "suite", "city", "mall", "plaza"]): | |
| near = self._find_words_near(w, words, max_dx=200, max_dy=30) | |
| addr_tokens = [w] + near[:4] | |
| extracted["address"] = " ".join([nw.text for nw in addr_tokens]) | |
| bboxes["address"] = self._union_bboxes(addr_tokens) | |
| confidences["address"] = 0.80 | |
| break | |
| # 4. Total Amount | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["total", "due", "cash", "amount", "paid", "subtotal", "charge"]): | |
| near = self._find_words_near(w, words, max_dx=350, max_dy=50) | |
| for nw in near: | |
| cleaned = re.sub(r'[^\d.-]', '', nw.text.replace(',', '')) | |
| if cleaned and any(c.isdigit() for c in cleaned): | |
| try: | |
| val = float(cleaned) | |
| if val > 0: | |
| extracted["total_amount"] = nw.text | |
| bboxes["total_amount"] = nw.bbox | |
| confidences["total_amount"] = 0.95 | |
| break | |
| except ValueError: | |
| pass | |
| if extracted["total_amount"]: | |
| break | |
| if not extracted["total_amount"]: | |
| # Try finding biggest currency/numeric number near bottom | |
| bottom_words = [w for w in words if w.bbox[1] > 300] | |
| for bw in reversed(bottom_words): | |
| cleaned = re.sub(r'[^\d.-]', '', bw.text.replace(',', '')) | |
| if cleaned and any(c.isdigit() for c in cleaned): | |
| try: | |
| val = float(cleaned) | |
| if val > 0: | |
| extracted["total_amount"] = bw.text | |
| bboxes["total_amount"] = bw.bbox | |
| confidences["total_amount"] = 0.85 | |
| break | |
| except ValueError: | |
| pass | |
| elif self.doc_class == "application_form": | |
| # 1. Applicant Name | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["applicant", "employee", "name", "signature", "client"]): | |
| near = self._find_words_near(w, words, max_dx=300, max_dy=40) | |
| cand_tokens = [nw for nw in near[:3] if nw.text.lower() not in [":", "name", "date", "is", "of", "title"]] | |
| if cand_tokens: | |
| extracted["applicant_name"] = " ".join([nw.text for nw in cand_tokens]) | |
| bboxes["applicant_name"] = self._union_bboxes(cand_tokens) | |
| confidences["applicant_name"] = 0.86 | |
| break | |
| # 2. Date | |
| for w in words: | |
| m_d = DATE_PATTERN.search(w.text) | |
| if m_d: | |
| extracted["date"] = m_d.group(1) | |
| bboxes["date"] = w.bbox | |
| confidences["date"] = 0.95 | |
| break | |
| # 3. Address | |
| for w in words: | |
| if any(t in w.text.lower() for t in ["address", "location", "residence", "city", "state", "street"]): | |
| near = self._find_words_near(w, words, max_dx=350, max_dy=50) | |
| cand_tokens = [nw for nw in near[:4] if nw.text.lower() not in [":", "address", "city", "phone"]] | |
| if cand_tokens: | |
| extracted["address"] = " ".join([nw.text for nw in cand_tokens]) | |
| bboxes["address"] = self._union_bboxes(cand_tokens) | |
| confidences["address"] = 0.83 | |
| break | |
| # 4. Other key value pairs | |
| kv_pairs = [] | |
| for w in words: | |
| if ":" in w.text and len(w.text) > 2: | |
| key_part = w.text.replace(":", "").strip() | |
| near = self._find_words_near(w, words, max_dx=200, max_dy=30) | |
| if near: | |
| val_part = " ".join([nw.text for nw in near[:3]]) | |
| kv_pairs.append(f"{key_part}: {val_part}") | |
| if kv_pairs: | |
| extracted["other_key_value_pairs"] = "; ".join(kv_pairs[:3]) | |
| confidences["other_key_value_pairs"] = 0.78 | |
| # Fallback bounding box search if text was extracted via regex without explicit bbox | |
| for f, val in extracted.items(): | |
| if val and not bboxes[f]: | |
| val_first = str(val).split()[0] | |
| for w in words: | |
| if val_first.lower() in w.text.lower(): | |
| bboxes[f] = w.bbox | |
| break | |
| return extracted, bboxes, confidences | |
| def extract(self, record: DocumentRecord) -> Dict[str, Any]: | |
| fields, _, _ = self.extract_with_details(record) | |
| return fields | |