import re from typing import Tuple, Dict, Any, List, Optional class UniversalDocumentClassifier: """ Universal Multi-Signal Document Classifier supporting 17 document types: id_card, passport, driver_license, invoice, receipt, contract, cv, bank_document, form, report, certificate, letter, screenshot, table, image_with_text, generic_document, unknown. Uses a fusion of: 1. OCR text & multilingual keywords (Arabic and English) 2. Structural signals (sections, tables, key-value ratios) 3. Document geometry & aspect ratio (receipts, cards, sheets) 4. Exact regex patterns (dates, amounts, National IDs, MRZ, IBAN) Never raises exceptions; gracefully falls back to 'generic_document' or 'unknown'. """ CATEGORIES = [ "id_card", "passport", "driver_license", "invoice", "receipt", "contract", "cv", "bank_document", "form", "report", "certificate", "letter", "screenshot", "table", "image_with_text", "generic_document", "unknown" ] @classmethod def classify_with_signals( cls, text: str, layout: Any = None, metadata: Any = None ) -> Dict[str, Any]: if not text or not text.strip(): return {"category": "unknown", "confidence": 0.0, "signals": ["empty_input"]} clean_text = text.strip() lower = clean_text.lower() word_count = len(clean_text.split()) lines = [l.strip() for l in clean_text.splitlines() if l.strip()] line_count = len(lines) signals_detected: List[str] = [] scores: Dict[str, float] = {cat: 0.0 for cat in cls.CATEGORIES} # Sparse text check if word_count <= 4 and not any(kw in lower for kw in ["جمهورية", "passport", "بطاقة", "invoice", "receipt", "عقد"]): return {"category": "image_with_text", "confidence": 0.70, "signals": ["sparse_text_word_count_le_4"]} # ------------------------------------------------------------- # 1. Structural & Layout Signals from LayoutInfo # ------------------------------------------------------------- tables_count = 0 sections_count = 0 columns_count = 1 has_signatures = False if layout: if hasattr(layout, 'tables') and layout.tables: tables_count = len(layout.tables) if hasattr(layout, 'blocks') and layout.blocks: sections_count = sum(1 for b in layout.blocks if getattr(b, 'type', '') == 'header') if hasattr(layout, 'columns'): columns_count = getattr(layout, 'columns', 1) if hasattr(layout, 'has_signatures'): has_signatures = getattr(layout, 'has_signatures', False) if tables_count > 0: scores["table"] += 2.0 * tables_count scores["invoice"] += 1.0 * tables_count scores["receipt"] += 0.8 * tables_count scores["bank_document"] += 1.2 * tables_count signals_detected.append(f"tables_detected_{tables_count}") # ------------------------------------------------------------- # 2. Geometry Signals (Aspect Ratio) # ------------------------------------------------------------- if metadata: width = getattr(metadata, 'width', 0) if not isinstance(metadata, dict) else metadata.get('width', 0) height = getattr(metadata, 'height', 0) if not isinstance(metadata, dict) else metadata.get('height', 0) if width > 0 and height > 0: aspect_ratio = height / width # Tall receipt ratio (H/W >= 1.8) if aspect_ratio >= 1.8: scores["receipt"] += 1.8 signals_detected.append("tall_aspect_ratio_receipt") # Standard ID/License landscape card (W/H >= 1.35 and relatively small) elif width / height >= 1.35 and word_count <= 50: scores["id_card"] += 1.2 scores["driver_license"] += 1.2 signals_detected.append("landscape_card_aspect_ratio") # ------------------------------------------------------------- # 3. Numeric & Specialized Pattern Signals # ------------------------------------------------------------- # Passport MRZ (e.g. P= 3: scores["id_card"] += 3.0 signals_detected.append(f"egyptian_id_field_labels_{len(matched_id_anchors)}") # High-specificity contextual combinations if ("@gmail.com" in lower or "@yahoo.com" in lower or "linkedin.com" in lower or "github.com" in lower) and ("experience" in lower or "education" in lower or "الخبرات" in lower): scores["cv"] += 2.5 signals_detected.append("cv_contact_plus_sections") if "مرور" in lower and ("رخصة" in lower or "قيادة" in lower): scores["driver_license"] += 2.5 signals_detected.append("traffic_dept_license") # Lines with pipes or tabs lines_with_pipes = sum(1 for line in lines if "|" in line or "\t" in line) if lines_with_pipes >= 3: scores["table"] += 3.0 signals_detected.append("delimiters_pipe_or_tab") # ------------------------------------------------------------- # 5. Determine Top Category & Confidence # ------------------------------------------------------------- best_cat, max_score = max(scores.items(), key=lambda item: item[1]) if max_score >= 2.0: conf = min(0.98, 0.60 + (max_score * 0.05)) category = best_cat elif max_score >= 1.2: conf = 0.55 category = best_cat elif word_count <= 4: conf = 0.60 category = "image_with_text" else: conf = 0.50 category = "generic_document" return { "category": category, "confidence": round(conf, 2), "signals": signals_detected, "scores": {k: round(v, 2) for k, v in scores.items() if v > 0} } @classmethod def classify(cls, text: str, layout: Any = None, metadata: Any = None) -> Tuple[str, float]: """Backward-compatible classification method returning (category, confidence).""" res = cls.classify_with_signals(text, layout, metadata) return res["category"], res["confidence"]