""" InvoiceForge AI — utils/confidence.py Multi-engine OCR confidence fusion. Implements weighted voting at the token level: - PaddleOCR weight: 0.50 - EasyOCR weight: 0.30 - TrOCR weight: 0.20 (only for handwritten documents) Tokens from multiple engines are spatially matched using IoU overlap of bounding boxes, then the highest weighted-confidence text is selected. """ from __future__ import annotations import logging from dataclasses import dataclass, field import numpy as np logger = logging.getLogger(__name__) # ───────────────────────────────────────────────────────────────────────────── # CONSTANTS # ───────────────────────────────────────────────────────────────────────────── ENGINE_WEIGHTS: dict[str, float] = { "paddleocr": 0.50, "easyocr": 0.30, "trocr": 0.20, } IOU_MATCH_THRESHOLD: float = 0.40 # Minimum IoU to consider tokens overlapping X_SNAP_PX: float = 30.0 # Fallback positional match if no bbox Y_SNAP_PX: float = 15.0 @dataclass class FusedToken: """Token produced after multi-engine confidence fusion.""" text: str confidence: float x: float y: float bbox: list = field(default_factory=list) engine: str = "fused" candidates: list[dict] = field(default_factory=list) # ───────────────────────────────────────────────────────────────────────────── # IOu HELPERS # ───────────────────────────────────────────────────────────────────────────── def _bbox_to_xyxy(bbox: list) -> tuple[float, float, float, float]: """Convert [[x1,y1],[x2,y1],[x2,y2],[x1,y2]] to (x1,y1,x2,y2).""" if not bbox: return (0.0, 0.0, 0.0, 0.0) xs = [p[0] for p in bbox] ys = [p[1] for p in bbox] return (min(xs), min(ys), max(xs), max(ys)) def _iou(a: list, b: list) -> float: """Compute intersection-over-union of two bboxes.""" ax1, ay1, ax2, ay2 = _bbox_to_xyxy(a) bx1, by1, bx2, by2 = _bbox_to_xyxy(b) ix1 = max(ax1, bx1) iy1 = max(ay1, by1) ix2 = min(ax2, bx2) iy2 = min(ay2, by2) if ix2 <= ix1 or iy2 <= iy1: return 0.0 inter = (ix2 - ix1) * (iy2 - iy1) area_a = max((ax2 - ax1) * (ay2 - ay1), 1e-6) area_b = max((bx2 - bx1) * (by2 - by1), 1e-6) return inter / (area_a + area_b - inter) def _pos_match(a: dict, b: dict) -> bool: """Positional proximity match when bounding boxes are unavailable.""" return abs(a["x"] - b["x"]) < X_SNAP_PX and abs(a["y"] - b["y"]) < Y_SNAP_PX # ───────────────────────────────────────────────────────────────────────────── # CONFIDENCE FUSION # ───────────────────────────────────────────────────────────────────────────── class ConfidenceFusion: """ Fuses OCR outputs from multiple engines using confidence-weighted voting. Usage: fusion = ConfidenceFusion() fusion.add_engine_result("paddleocr", paddle_tokens) fusion.add_engine_result("easyocr", easy_tokens) fused = fusion.fuse() """ def __init__(self) -> None: self._engine_results: dict[str, list[dict]] = {} def add_engine_result(self, engine_name: str, tokens: list) -> None: """ Add OCR results from one engine. Args: engine_name: One of "paddleocr", "easyocr", "trocr". tokens: List of token dicts or OCRToken objects. """ normalised: list[dict] = [] for tok in tokens: if hasattr(tok, "__dict__"): # OCRToken dataclass normalised.append( { "text": tok.text, "confidence": tok.confidence, "bbox": tok.bbox, "x": tok.x, "y": tok.y, "engine": engine_name, } ) else: d = dict(tok) d["engine"] = engine_name normalised.append(d) self._engine_results[engine_name] = normalised logger.debug( "Fusion: added %d tokens from %s.", len(normalised), engine_name ) def fuse(self) -> list[FusedToken]: """ Perform confidence-weighted fusion. Algorithm: 1. Build a list of all tokens across engines. 2. For each unmatched token, find all overlapping tokens from other engines using IoU or positional proximity. 3. Compute weighted confidence for each text candidate. 4. Select the text with the highest weighted score. Returns: Sorted list of FusedToken objects. """ all_tokens: list[dict] = [] for engine, tokens in self._engine_results.items(): for tok in tokens: tok = dict(tok) tok["engine"] = engine tok["_used"] = False all_tokens.append(tok) fused_tokens: list[FusedToken] = [] used_indices: set[int] = set() for i, anchor in enumerate(all_tokens): if i in used_indices: continue group: list[dict] = [anchor] used_indices.add(i) # Find matching tokens from other engines for j, candidate in enumerate(all_tokens): if j in used_indices: continue if candidate["engine"] == anchor["engine"]: continue matched = False if anchor["bbox"] and candidate["bbox"]: matched = _iou(anchor["bbox"], candidate["bbox"]) >= IOU_MATCH_THRESHOLD else: matched = _pos_match(anchor, candidate) if matched: group.append(candidate) used_indices.add(j) # Weighted voting across group fused = _vote(group) fused_tokens.append(fused) # Sort by reading order (top to bottom, left to right) fused_tokens.sort(key=lambda t: (t.y, t.x)) logger.debug("Fusion: %d fused tokens produced.", len(fused_tokens)) return fused_tokens def reset(self) -> None: """Clear all engine results for reuse.""" self._engine_results.clear() def _vote(group: list[dict]) -> FusedToken: """ Select the best text from a group of candidate tokens using confidence-weighted voting. """ text_scores: dict[str, float] = {} best_bbox: list = [] best_x: float = 0.0 best_y: float = 0.0 for tok in group: engine = tok.get("engine", "paddleocr") weight = ENGINE_WEIGHTS.get(engine, 0.30) text = tok.get("text", "").strip() conf = float(tok.get("confidence", 0.5)) weighted = conf * weight text_scores[text] = text_scores.get(text, 0.0) + weighted if tok.get("bbox"): best_bbox = tok["bbox"] best_x = tok.get("x", 0.0) best_y = tok.get("y", 0.0) # Pick text with highest aggregate weighted score if not text_scores: return FusedToken(text="", confidence=0.0, x=best_x, y=best_y) best_text = max(text_scores, key=lambda t: text_scores[t]) total_weight = sum(ENGINE_WEIGHTS.values()) normalised_conf = min(text_scores[best_text] / total_weight, 1.0) return FusedToken( text=best_text, confidence=round(normalised_conf, 4), x=best_x, y=best_y, bbox=best_bbox, engine="fused", candidates=group, ) # ───────────────────────────────────────────────────────────────────────────── # DOCUMENT-LEVEL CONFIDENCE SCORING # ───────────────────────────────────────────────────────────────────────────── def compute_document_confidence( fused_tokens: list[FusedToken], validation_errors: list[str], ) -> float: """ Compute an aggregate document-level confidence score. Factors: - Mean per-token confidence from fusion - Penalty for each validation error (−0.05 per error, max −0.30) Returns: Float 0.0 – 1.0. """ if not fused_tokens: return 0.0 mean_conf = float(np.mean([t.confidence for t in fused_tokens])) error_penalty = min(len(validation_errors) * 0.05, 0.30) score = max(0.0, mean_conf - error_penalty) return round(score, 4)