""" Post-Processing and Key-Value Extraction Combines NER predictions with regex patterns and heuristics Implements confidence scoring and field validation """ import re from typing import Dict, List, Optional, Tuple import logging from datetime import datetime import numpy as np logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class PostProcessor: """ Post-processing layer for IDP Combines NER outputs with regex patterns and validation rules """ def __init__(self): # Regex patterns for common fields self.patterns = { 'invoice_number': [ r'INV[-\s]?\d+', r'Invoice\s*#?\s*:?\s*([A-Z0-9\-]+)', r'\b\d{6,}\b', # 6+ digits ], 'date': [ r'\d{1,2}[/-]\d{1,2}[/-]\d{2,4}', r'\d{4}[/-]\d{1,2}[/-]\d{1,2}', r'\d{1,2}\s+(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{2,4}', ], 'amount': [ r'[₹$€£¥]\s*[\d,]+\.?\d*', r'[\d,]+\.?\d*\s*[₹$€£¥]', r'\b\d{1,3}(?:,\d{3})*(?:\.\d{2})?\b', ], 'gst_id': [ r'\d{2}[A-Z]{5}\d{4}[A-Z]{1}[A-Z\d]{1}[Z]{1}[A-Z\d]{1}', # Indian GST r'[A-Z]{2}\d{9}', # VAT pattern ], 'email': [ r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', ], 'phone': [ r'\+?\d{1,3}[-.\s]?\(?\d{1,4}\)?[-.\s]?\d{1,4}[-.\s]?\d{1,9}', ], } def process_document( self, document_type: str, classification_confidence: float, ocr_text: str, ocr_boxes: List[Dict], ner_entities: List[Dict] ) -> Dict: """ Main post-processing pipeline Args: document_type: Predicted document type (INVOICE, RECEIPT, FORM, OTHER) classification_confidence: Confidence score from classifier ocr_text: Full OCR extracted text ocr_boxes: OCR bounding boxes with confidence ner_entities: NER predicted entities Returns: Structured document with extracted fields """ logger.info(f"Post-processing {document_type} document") # Extract fields based on document type if document_type == 'INVOICE': fields = self._extract_invoice_fields(ocr_text, ocr_boxes, ner_entities) elif document_type == 'RECEIPT': fields = self._extract_receipt_fields(ocr_text, ocr_boxes, ner_entities) elif document_type == 'FORM': fields = self._extract_form_fields(ocr_text, ocr_boxes, ner_entities) elif document_type == 'BANK_STATEMENT': fields = self._extract_bank_statement_fields(ocr_text, ocr_boxes, ner_entities) else: fields = self._extract_generic_fields(ocr_text, ocr_boxes, ner_entities) # Build structured output result = { 'document_type': document_type, 'classification_confidence': classification_confidence, 'fields': fields, 'raw_ocr_text': ocr_text[:1000], # First 1000 chars for debugging } return result def _extract_invoice_fields( self, text: str, boxes: List[Dict], entities: List[Dict] ) -> Dict: """Extract invoice-specific fields""" fields = {} # Invoice number invoice_num = self._extract_field( 'INVOICE_NUMBER', text, boxes, entities, self.patterns['invoice_number'] ) if invoice_num: fields['invoice_number'] = invoice_num # Date date = self._extract_field( 'DATE', text, boxes, entities, self.patterns['date'] ) if date: fields['date'] = self._normalize_date(date) # Total amount total = self._extract_field( 'TOTAL_AMOUNT', text, boxes, entities, self.patterns['amount'] ) if total: fields['total_amount'] = self._normalize_amount(total) # Tax amount tax = self._extract_field( 'TAX_AMOUNT', text, boxes, entities, self.patterns['amount'] ) if tax: fields['tax_amount'] = self._normalize_amount(tax) # Vendor name vendor = self._extract_field( 'VENDOR_NAME', text, boxes, entities, [] ) if vendor: fields['vendor_name'] = vendor # Customer name customer = self._extract_field( 'CUSTOMER_NAME', text, boxes, entities, [] ) if customer: fields['customer_name'] = customer # GST/Tax ID gst = self._extract_field( 'GST_ID', text, boxes, entities, self.patterns['gst_id'] ) if gst: fields['gst_id'] = gst return fields def _extract_receipt_fields( self, text: str, boxes: List[Dict], entities: List[Dict] ) -> Dict: """Extract receipt-specific fields""" # Similar to invoice but may have different priorities return self._extract_invoice_fields(text, boxes, entities) def _extract_form_fields( self, text: str, boxes: List[Dict], entities: List[Dict] ) -> Dict: """Extract form-specific fields""" fields = {} # Generic entity extraction for entity in entities: entity_type = entity['entity'].lower() if entity_type not in fields: fields[entity_type] = { 'value': entity['text'], 'confidence': entity['confidence'], 'bbox': [entity['start'], 0, entity['end'], 0] # Simplified } return fields def _extract_bank_statement_fields( self, text: str, boxes: List[Dict], entities: List[Dict] ) -> Dict: """Extract bank statement specific fields""" # Reuse invoice extraction logic for common fields like Date and Amount return self._extract_invoice_fields(text, boxes, entities) def _extract_generic_fields( self, text: str, boxes: List[Dict], entities: List[Dict] ) -> Dict: """Extract generic fields for unknown document types""" return self._extract_form_fields(text, boxes, entities) def _extract_field( self, entity_type: str, text: str, boxes: List[Dict], entities: List[Dict], fallback_patterns: List[str] ) -> Optional[Dict]: """ Extract a single field using NER + regex + heuristics Returns: Dictionary with value, confidence, and bbox, or None """ # Step 1: Check NER entities for entity in entities: if entity['entity'] == entity_type: # Get OCR confidence for this text span ocr_conf = self._get_ocr_confidence(entity, boxes) # Combined confidence combined_conf = (entity['confidence'] + ocr_conf) / 2 return { 'value': entity['text'], 'confidence': float(combined_conf), 'bbox': self._get_bbox_for_entity(entity, boxes), 'source': 'ner' } # Step 2: Fallback to regex patterns if fallback_patterns: for pattern in fallback_patterns: matches = re.finditer(pattern, text, re.IGNORECASE) for match in matches: matched_text = match.group(0) # Validate match if self._validate_field(entity_type, matched_text): # Find in OCR boxes ocr_conf = self._find_in_boxes(matched_text, boxes) return { 'value': matched_text, 'confidence': float(ocr_conf * 0.8), # Lower confidence for regex 'bbox': None, 'source': 'regex' } return None def _validate_field(self, entity_type: str, value: str) -> bool: """Validate extracted field value""" if not value or not value.strip(): return False # Type-specific validation if entity_type == 'DATE': return len(value) >= 6 # Minimum date length elif entity_type in ['TOTAL_AMOUNT', 'TAX_AMOUNT']: return any(c.isdigit() for c in value) elif entity_type == 'INVOICE_NUMBER': return len(value) >= 3 return True def _get_ocr_confidence(self, entity: Dict, boxes: List[Dict]) -> float: """Get average OCR confidence for entity text""" entity_text = entity['text'].lower() confidences = [] for box in boxes: if box['text'].lower() in entity_text or entity_text in box['text'].lower(): confidences.append(box['confidence']) return float(np.mean(confidences)) if confidences else 0.5 def _get_bbox_for_entity(self, entity: Dict, boxes: List[Dict]) -> Optional[List[float]]: """Get bounding box for entity from OCR boxes""" entity_text = entity['text'].lower() for box in boxes: if entity_text in box['text'].lower(): return box['bbox'] return None def _find_in_boxes(self, text: str, boxes: List[Dict]) -> float: """Find text in OCR boxes and return confidence""" text_lower = text.lower() for box in boxes: if text_lower in box['text'].lower(): return float(box['confidence']) return 0.5 # Default confidence def _normalize_date(self, date_field: Dict) -> Dict: """Normalize date to standard format""" date_str = date_field['value'] # Try to parse date formats = [ '%d/%m/%Y', '%d-%m-%Y', '%Y-%m-%d', '%Y/%m/%d', '%d/%m/%y', '%d-%m-%y', '%d %B %Y', '%d %b %Y', ] for fmt in formats: try: parsed = datetime.strptime(date_str, fmt) date_field['value'] = parsed.strftime('%Y-%m-%d') date_field['normalized'] = True return date_field except: continue # Could not parse, return as-is date_field['normalized'] = False return date_field def _normalize_amount(self, amount_field: Dict) -> Dict: """Normalize amount to numeric value""" amount_str = amount_field['value'] # Extract currency symbol currency_symbols = {'₹': 'INR', '$': 'USD', '€': 'EUR', '£': 'GBP', '¥': 'JPY'} currency = None for symbol, code in currency_symbols.items(): if symbol in amount_str: currency = code amount_str = amount_str.replace(symbol, '') break # Remove commas and spaces amount_str = amount_str.replace(',', '').replace(' ', '').strip() # Extract numeric value try: numeric_value = float(amount_str) amount_field['value'] = f"{numeric_value:.2f}" amount_field['numeric_value'] = numeric_value if currency: amount_field['currency'] = currency amount_field['normalized'] = True except: amount_field['normalized'] = False return amount_field if __name__ == "__main__": # Example usage processor = PostProcessor() # Dummy data ocr_text = """ INVOICE Invoice #: INV-12345 Date: 28/11/2025 Total Amount: ₹ 12,500.00 Tax: ₹ 2,250.00 GST ID: 29ABCDE1234F1Z5 """ ocr_boxes = [ {'text': 'INVOICE', 'bbox': [100, 50, 200, 80], 'confidence': 0.98}, {'text': 'INV-12345', 'bbox': [150, 100, 250, 120], 'confidence': 0.95}, {'text': '28/11/2025', 'bbox': [150, 130, 250, 150], 'confidence': 0.92}, {'text': '₹ 12,500.00', 'bbox': [200, 200, 300, 220], 'confidence': 0.96}, ] ner_entities = [ {'entity': 'INVOICE_NUMBER', 'text': 'INV-12345', 'confidence': 0.94, 'start': 30, 'end': 39}, {'entity': 'DATE', 'text': '28/11/2025', 'confidence': 0.89, 'start': 47, 'end': 57}, {'entity': 'TOTAL_AMOUNT', 'text': '₹ 12,500.00', 'confidence': 0.93, 'start': 75, 'end': 86}, ] result = processor.process_document( document_type='INVOICE', classification_confidence=0.96, ocr_text=ocr_text, ocr_boxes=ocr_boxes, ner_entities=ner_entities ) print("\nPost-processing result:") print(json.dumps(result, indent=2))