Spaces:
Sleeping
Sleeping
File size: 2,851 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 69 70 71 72 73 74 | import re
from typing import Dict, Any, List
from ..data.schema import DocumentRecord, CLASS_FIELD_SCHEMAS
INV_NUM_REGEX = re.compile(r'(?i)(?:invoice\s*(?:no|num|number|#)?[:.\s]*|inv\s*#?\s*|no[:.\s]+)([a-z0-9\-_/]+)')
DATE_REGEX = re.compile(r'(\d{1,4}[/\-.]\d{1,2}[/\-.]\d{1,4})')
AMOUNT_REGEX = re.compile(r'(?i)(?:total|amount|due|balance|subtotal)[\s:$]*([0-9]+[.,][0-9]{2,3})')
class RuleBasedExtractor:
"""Baseline A: Rule-based regular expression and keyword heuristic extractor."""
def __init__(self):
pass
def extract(self, record: DocumentRecord, doc_class: str = None) -> Dict[str, Any]:
cls_name = doc_class or record.doc_class
tokens = [w.text for w in record.ocr_words]
full_text = " ".join(tokens)
schema_fields = CLASS_FIELD_SCHEMAS.get(cls_name, [])
extracted = {f: None for f in schema_fields}
if cls_name == "invoice":
# 1. Invoice Number
m_num = INV_NUM_REGEX.search(full_text)
if m_num:
extracted["invoice_number"] = m_num.group(1).strip()
# 2. Invoice Date
m_date = DATE_REGEX.search(full_text)
if m_date:
extracted["invoice_date"] = m_date.group(1).strip()
# 3. Total Amount
m_amt = AMOUNT_REGEX.search(full_text)
if m_amt:
extracted["total_amount"] = m_amt.group(1).strip()
# 4. Buyer / Seller heuristics
m_to = re.search(r'(?i)(?:bill\s*to|sold\s*to|to)[:\s]+([A-Za-z0-9\s]{3,30})', full_text)
if m_to:
extracted["buyer_name"] = m_to.group(1).strip()
elif cls_name == "receipt":
# 1. Company Name: top token block
if tokens:
extracted["company_name"] = " ".join(tokens[:min(3, len(tokens))])
# 2. Date
m_date = DATE_REGEX.search(full_text)
if m_date:
extracted["date"] = m_date.group(1).strip()
# 3. Total Amount
m_amt = AMOUNT_REGEX.search(full_text)
if m_amt:
extracted["total_amount"] = m_amt.group(1).strip()
elif cls_name == "application_form":
# 1. Applicant Name
m_name = re.search(r'(?i)(?:name|applicant|signature)[:\s]+([A-Za-z\s]{3,25})', full_text)
if m_name:
extracted["applicant_name"] = m_name.group(1).strip()
# 2. Date
m_date = DATE_REGEX.search(full_text)
if m_date:
extracted["date"] = m_date.group(1).strip()
# 3. Address
m_addr = re.search(r'(?i)(?:address|location|city)[:\s]+([A-Za-z0-9\s,]{5,40})', full_text)
if m_addr:
extracted["address"] = m_addr.group(1).strip()
return extracted
|