Spaces:
Runtime error
Runtime error
| """ | |
| DocuLens - Core Invoice/Document Extraction Engine | |
| Supports multiple extraction backends: | |
| 1. Hugging Face Inference API (Qwen2.5-VL) - Primary, highest accuracy | |
| 2. Tesseract OCR + regex patterns - Fallback, CPU-only, zero cost | |
| """ | |
| import os | |
| import re | |
| import json | |
| import base64 | |
| import logging | |
| from io import BytesIO | |
| from typing import Optional | |
| from dataclasses import dataclass, field, asdict | |
| from datetime import datetime | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| import requests | |
| from PIL import Image | |
| from huggingface_hub import InferenceClient | |
| from hf_client import create_inference_client, robust_chat_completion | |
| from gemini_client import extract_with_gemini, GEMINI_API_KEY | |
| logger = logging.getLogger(__name__) | |
| # --------------------------------------------------------------------------- | |
| # Data Models | |
| # --------------------------------------------------------------------------- | |
| class LineItem: | |
| description: str = "" | |
| quantity: Optional[float] = None | |
| unit_price: Optional[float] = None | |
| amount: Optional[float] = None | |
| class BoundingBox: | |
| field: str = "" | |
| x1: int = 0 | |
| y1: int = 0 | |
| x2: int = 0 | |
| y2: int = 0 | |
| page: int = 0 | |
| class ExtractionResult: | |
| # Legacy named fields (populated for invoice use case, empty for others) | |
| vendor_name: str = "" | |
| vendor_address: str = "" | |
| invoice_number: str = "" | |
| invoice_date: str = "" | |
| due_date: str = "" | |
| currency: str = "" | |
| subtotal: Optional[float] = None | |
| tax_amount: Optional[float] = None | |
| tax_rate: Optional[str] = None | |
| total_amount: Optional[float] = None | |
| line_items: list = field(default_factory=list) | |
| payment_terms: str = "" | |
| customer_name: str = "" | |
| customer_address: str = "" | |
| raw_text: str = "" | |
| extraction_method: str = "" | |
| confidence: float = 0.0 | |
| processing_time_ms: int = 0 | |
| bounding_boxes: list = field(default_factory=list) | |
| page_image: Optional[str] = None # data-URI of rendered page (for PDFs) | |
| # Dynamic fields — primary data carrier for all use cases | |
| document_type: str = "" | |
| extracted_fields: dict = field(default_factory=dict) | |
| def to_dict(self): | |
| d = asdict(self) | |
| d["line_items"] = [asdict(li) if isinstance(li, LineItem) else li for li in self.line_items] | |
| d["bounding_boxes"] = [asdict(bb) if isinstance(bb, BoundingBox) else bb for bb in self.bounding_boxes] | |
| return d | |
| # --------------------------------------------------------------------------- | |
| # Bounding Box Generation - matches extracted values to OCR word positions | |
| # --------------------------------------------------------------------------- | |
| def _get_ocr_word_boxes(image: Image.Image) -> list[dict]: | |
| """Run pytesseract image_to_data to get word-level bounding boxes.""" | |
| try: | |
| import pytesseract | |
| data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT, config="--psm 3") | |
| words = [] | |
| for i in range(len(data["text"])): | |
| text = data["text"][i].strip() | |
| conf = int(data["conf"][i]) | |
| if text and conf > 10: | |
| words.append({ | |
| "text": text, | |
| "x": data["left"][i], | |
| "y": data["top"][i], | |
| "w": data["width"][i], | |
| "h": data["height"][i], | |
| }) | |
| return words | |
| except Exception as e: | |
| logger.warning(f"OCR word boxes failed: {e}") | |
| return [] | |
| def _date_search_variants(date_str: str) -> list[str]: | |
| """Generate alternative date representations for OCR matching. | |
| VLM returns YYYY-MM-DD but documents may show DD-Mon-YYYY, DD/MM/YYYY, etc. | |
| """ | |
| variants = [date_str] | |
| # Try parsing YYYY-MM-DD | |
| m = re.match(r"(\d{4})-(\d{2})-(\d{2})", date_str) | |
| if m: | |
| y, mo, d = m.group(1), m.group(2), m.group(3) | |
| d_int = str(int(d)) # strip leading zero | |
| months = ["Jan","Feb","Mar","Apr","May","Jun","Jul","Aug","Sep","Oct","Nov","Dec"] | |
| month_name = months[int(mo) - 1] if 1 <= int(mo) <= 12 else mo | |
| month_full = ["January","February","March","April","May","June", | |
| "July","August","September","October","November","December"] | |
| full_name = month_full[int(mo) - 1] if 1 <= int(mo) <= 12 else mo | |
| variants += [ | |
| f"{d}-{month_name}-{y}", # 15-Aug-2026 | |
| f"{d_int}-{month_name}-{y}", # 5-Aug-2026 | |
| f"{d}/{mo}/{y}", # 15/08/2026 | |
| f"{d_int}/{mo}/{y}", # 5/08/2026 | |
| f"{mo}/{d}/{y}", # 08/15/2026 | |
| f"{d}.{mo}.{y}", # 15.08.2026 | |
| f"{d}-{mo}-{y}", # 15-08-2026 | |
| f"{d_int} {month_name} {y}", # 5 Aug 2026 | |
| f"{d} {month_name} {y}", # 15 Aug 2026 | |
| f"{month_name} {d_int}, {y}", # Aug 5, 2026 | |
| f"{d_int} {full_name} {y}", # 5 August 2026 | |
| f"{d} {full_name} {y}", # 15 August 2026 | |
| f"{full_name} {d_int}, {y}", # August 5, 2026 | |
| ] | |
| return list(dict.fromkeys(variants)) | |
| def _currency_search_variants(currency: str) -> list[str]: | |
| """Generate alternative currency representations for OCR matching.""" | |
| symbol_map = { | |
| "INR": ["₹", "Rs", "Rs.", "INR"], | |
| "USD": ["$", "USD", "US$"], | |
| "EUR": ["€", "EUR"], | |
| "GBP": ["£", "GBP"], | |
| "TRY": ["₺", "TL", "TRY"], | |
| "JPY": ["¥", "JPY"], | |
| } | |
| return symbol_map.get(currency.upper(), [currency]) | |
| def _word_matches_token(word_text: str, token: str) -> bool: | |
| """Check if an OCR word matches a search token. | |
| Uses full-word matching for alphabetic tokens (avoiding "tax" matching | |
| "taxation") and substring matching for numeric tokens (since OCR often | |
| merges numbers with punctuation like "$1,234.56"). | |
| """ | |
| wt = word_text.lower() | |
| # Numeric tokens: allow substring (OCR may merge currency symbols, commas) | |
| # Also strip common OCR artifacts before checking | |
| if any(c.isdigit() for c in token): | |
| # Direct substring | |
| if token in wt: | |
| return True | |
| # Strip non-digit non-dot non-comma chars and retry | |
| wt_digits = re.sub(r"[^0-9.,]", "", wt) | |
| tk_digits = re.sub(r"[^0-9.,]", "", token) | |
| return tk_digits != "" and tk_digits in wt_digits | |
| # Short tokens (<=2 chars): exact match only to avoid noise | |
| if len(token) <= 2: | |
| return wt == token | |
| # Alphabetic tokens: check word boundaries — strip common punctuation | |
| # from OCR word and compare | |
| wt_clean = re.sub(r"[^a-z0-9]", "", wt) | |
| tk_clean = re.sub(r"[^a-z0-9]", "", token) | |
| if not tk_clean: | |
| return False | |
| return wt_clean == tk_clean or wt_clean.startswith(tk_clean) or wt_clean.endswith(tk_clean) | |
| def _find_all_text_bboxes(words: list[dict], search_text: str) -> list[BoundingBox]: | |
| """Find ALL bounding boxes for a text value by matching against OCR words. | |
| Returns every match location so callers can disambiguate using context | |
| (e.g. proximity to a field label). | |
| """ | |
| if not search_text or not words: | |
| return [] | |
| search_lower = search_text.lower().strip() | |
| search_tokens = search_lower.split() | |
| if not search_tokens: | |
| return [] | |
| results = [] | |
| if len(search_tokens) == 1: | |
| token = search_tokens[0] | |
| for w in words: | |
| if _word_matches_token(w["text"], token): | |
| results.append(BoundingBox( | |
| x1=w["x"], y1=w["y"], | |
| x2=w["x"] + w["w"], y2=w["y"] + w["h"], | |
| )) | |
| # Fallback: try joining the search text without spaces (OCR may merge) | |
| if not results: | |
| merged = search_lower.replace(" ", "") | |
| if len(merged) >= 3: | |
| for w in words: | |
| wt_clean = re.sub(r"[^a-z0-9]", "", w["text"].lower()) | |
| if wt_clean == merged: | |
| results.append(BoundingBox( | |
| x1=w["x"], y1=w["y"], | |
| x2=w["x"] + w["w"], y2=w["y"] + w["h"], | |
| )) | |
| else: | |
| # Multi-word: try anchoring on EACH token position, not just the first. | |
| # OCR may miss or garble the first word but read later words fine. | |
| for anchor_idx in range(len(search_tokens)): | |
| anchor_token = search_tokens[anchor_idx] | |
| for i, w in enumerate(words): | |
| if not _word_matches_token(w["text"], anchor_token): | |
| continue | |
| # Expand backward to find earlier tokens | |
| x1, y1 = w["x"], w["y"] | |
| x2, y2 = w["x"] + w["w"], w["y"] + w["h"] | |
| matched = 1 | |
| # Look backward for tokens before the anchor | |
| if anchor_idx > 0: | |
| back_j = anchor_idx - 1 | |
| for bk in range(1, min(anchor_idx * 2 + 1, i + 1)): | |
| if back_j < 0: | |
| break | |
| prev_w = words[i - bk] | |
| if abs(prev_w["y"] - w["y"]) > 80: | |
| break | |
| if _word_matches_token(prev_w["text"], search_tokens[back_j]): | |
| x1 = min(x1, prev_w["x"]) | |
| y1 = min(y1, prev_w["y"]) | |
| x2 = max(x2, prev_w["x"] + prev_w["w"]) | |
| y2 = max(y2, prev_w["y"] + prev_w["h"]) | |
| matched += 1 | |
| back_j -= 1 | |
| # Look forward for tokens after the anchor | |
| search_j = anchor_idx + 1 | |
| scan_limit = min((len(search_tokens) - anchor_idx) * 2, len(words) - i) | |
| for k in range(1, scan_limit): | |
| if search_j >= len(search_tokens): | |
| break | |
| next_w = words[i + k] | |
| if abs(next_w["y"] - w["y"]) > 80: | |
| break | |
| if _word_matches_token(next_w["text"], search_tokens[search_j]): | |
| x2 = max(x2, next_w["x"] + next_w["w"]) | |
| y2 = max(y2, next_w["y"] + next_w["h"]) | |
| x1 = min(x1, next_w["x"]) | |
| y1 = min(y1, next_w["y"]) | |
| matched += 1 | |
| search_j += 1 | |
| # Require at least half the tokens to match | |
| threshold = max(1, (len(search_tokens) + 1) // 2) | |
| if matched >= threshold: | |
| results.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2)) | |
| # If we found matches anchoring on this token, don't try other anchors | |
| # (avoids duplicate/overlapping results) | |
| if results: | |
| break | |
| return results | |
| def _find_text_bbox(words: list[dict], search_text: str) -> Optional[BoundingBox]: | |
| """Find bounding box for a text value (first match). Legacy convenience wrapper.""" | |
| matches = _find_all_text_bboxes(words, search_text) | |
| return matches[0] if matches else None | |
| def _bbox_center(bb: BoundingBox) -> tuple[float, float]: | |
| """Return (cx, cy) center of a bounding box.""" | |
| return ((bb.x1 + bb.x2) / 2, (bb.y1 + bb.y2) / 2) | |
| def _bbox_distance(a: BoundingBox, b: BoundingBox) -> float: | |
| """Euclidean distance between bounding box centers.""" | |
| ax, ay = _bbox_center(a) | |
| bx, by = _bbox_center(b) | |
| return ((ax - bx) ** 2 + (ay - by) ** 2) ** 0.5 | |
| def _bbox_overlaps(a: BoundingBox, b: BoundingBox, margin: int = 5) -> bool: | |
| """Check if two bounding boxes overlap (with optional margin).""" | |
| return not ( | |
| a.x2 + margin < b.x1 or b.x2 + margin < a.x1 or | |
| a.y2 + margin < b.y1 or b.y2 + margin < a.y1 | |
| ) | |
| def _get_surrounding_text( | |
| candidate: BoundingBox, | |
| words: list[dict], | |
| radius: int = 5, | |
| ) -> list[str]: | |
| """Return the OCR word texts surrounding a candidate bbox. | |
| Finds the OCR word(s) that overlap the candidate, then collects up to | |
| *radius* words before and after in reading order. Returns lowercased | |
| word texts — callers use these to check whether a field's label tokens | |
| appear near the candidate value. | |
| """ | |
| # Find word indices that overlap the candidate | |
| overlap_indices: list[int] = [] | |
| for idx, w in enumerate(words): | |
| w_bb = BoundingBox(x1=w["x"], y1=w["y"], | |
| x2=w["x"] + w["w"], y2=w["y"] + w["h"]) | |
| if _bbox_overlaps(candidate, w_bb, margin=3): | |
| overlap_indices.append(idx) | |
| if not overlap_indices: | |
| # Fallback: find the nearest word by center distance | |
| cx, cy = _bbox_center(candidate) | |
| best_idx, best_dist = 0, float("inf") | |
| for idx, w in enumerate(words): | |
| wx = w["x"] + w["w"] / 2 | |
| wy = w["y"] + w["h"] / 2 | |
| d = ((cx - wx) ** 2 + (cy - wy) ** 2) ** 0.5 | |
| if d < best_dist: | |
| best_dist = d | |
| best_idx = idx | |
| overlap_indices = [best_idx] | |
| lo = max(0, min(overlap_indices) - radius) | |
| hi = min(len(words), max(overlap_indices) + 1 + radius) | |
| # Exclude the candidate's own words | |
| surrounding = [] | |
| for idx in range(lo, hi): | |
| if idx not in overlap_indices: | |
| surrounding.append(words[idx]["text"].lower()) | |
| return surrounding | |
| def _surrounding_text_score( | |
| candidate: BoundingBox, | |
| words: list[dict], | |
| label_tokens: list[str], | |
| radius: int = 5, | |
| ) -> float: | |
| """Score how well the surrounding OCR words match a field's label tokens. | |
| Returns a value between 0.0 (no label tokens found nearby) and 1.0 | |
| (all label tokens present in the surrounding words). This helps | |
| disambiguate when the same value (e.g. "0" or "2790") appears in | |
| multiple places on the page. | |
| """ | |
| if not label_tokens: | |
| return 0.0 | |
| surrounding = _get_surrounding_text(candidate, words, radius) | |
| if not surrounding: | |
| return 0.0 | |
| # Check how many label tokens appear in the surrounding text | |
| joined_surrounding = " ".join(surrounding) | |
| matched = 0 | |
| for lt in label_tokens: | |
| # Allow partial matches for long tokens (e.g. "statement" matching "statements") | |
| if any(lt in sw or sw in lt for sw in surrounding if len(sw) > 2): | |
| matched += 1 | |
| elif lt in joined_surrounding: | |
| matched += 1 | |
| return matched / len(label_tokens) | |
| def _field_name_to_label_tokens(field_name: str) -> list[str]: | |
| """Derive likely label keywords from a snake_case field name. | |
| e.g. 'invoice_date' → ['invoice', 'date'] | |
| 'total_amount' → ['total', 'amount'] | |
| 'pan' → ['pan'] | |
| """ | |
| return [t.lower() for t in field_name.split("_") if len(t) > 1] | |
| def _find_label_bboxes(words: list[dict], field_name: str) -> list[BoundingBox]: | |
| """Find ALL positions where a field's label appears in the document OCR. | |
| Returns all matches so the caller can pick the one closest to the value. | |
| Searches for the label keywords derived from the field name. Looks for | |
| multi-word labels first (e.g. "Invoice Date"), then single keywords. | |
| """ | |
| tokens = _field_name_to_label_tokens(field_name) | |
| if not tokens: | |
| return [] | |
| # Try the full label as a phrase first | |
| full_label = " ".join(tokens) | |
| matches = _find_all_text_bboxes(words, full_label) | |
| if matches: | |
| return matches | |
| # Fall back to the most distinctive single token (longest word) | |
| tokens_by_len = sorted(tokens, key=len, reverse=True) | |
| for token in tokens_by_len: | |
| # Skip very common short words that appear everywhere | |
| if token in ("of", "no", "in", "to", "at", "by", "id", "or"): | |
| continue | |
| matches = _find_all_text_bboxes(words, token) | |
| if matches: | |
| return matches | |
| return [] | |
| def _find_label_bbox(words: list[dict], field_name: str) -> Optional[BoundingBox]: | |
| """Find where a field's label appears — convenience wrapper returning first match.""" | |
| matches = _find_label_bboxes(words, field_name) | |
| return matches[0] if matches else None | |
| def _pick_best_match( | |
| candidates: list[BoundingBox], | |
| label_bbox: Optional[BoundingBox], | |
| claimed: list[BoundingBox], | |
| label_bboxes: Optional[list[BoundingBox]] = None, | |
| words: Optional[list[dict]] = None, | |
| field_name: Optional[str] = None, | |
| ) -> Optional[BoundingBox]: | |
| """Pick the best bbox from candidates, preferring proximity to label and avoiding claimed regions. | |
| Scoring: lower is better. | |
| - Distance to label (if known) is the primary signal. | |
| - **Surrounding-text context** — when ``words`` and ``field_name`` are | |
| provided, candidates whose neighbouring OCR words contain the | |
| field's label tokens receive a significant bonus. This is the | |
| strongest disambiguation signal for repeated short values. | |
| - Values to the right of or below their label get a bonus (documents | |
| typically place values right-of or below labels). | |
| - Values that overlap the label itself are penalised (that's the label, | |
| not the value). | |
| - Overlapping a claimed region adds a large penalty. | |
| - When multiple label positions are given (label_bboxes), the scoring | |
| picks the best (label, value) pair. | |
| """ | |
| if not candidates: | |
| return None | |
| CLAIMED_PENALTY = 100_000 | |
| LABEL_OVERLAP_PENALTY = 50_000 | |
| WRONG_DIRECTION_PENALTY = 200 # values above or far-left of label | |
| CONTEXT_BONUS_MAX = 500 # max bonus for surrounding-text match | |
| all_labels = label_bboxes or ([label_bbox] if label_bbox else []) | |
| # Pre-compute label tokens for surrounding-text scoring | |
| label_tokens = _field_name_to_label_tokens(field_name) if field_name else [] | |
| # Filter out very common/short tokens that add noise | |
| label_tokens = [t for t in label_tokens if t not in ("of", "no", "in", "to", "at", "by", "id", "or", "is")] | |
| # Determine if the value is ambiguous (short/numeric → more likely repeated) | |
| is_ambiguous = len(candidates) > 1 | |
| def score_with_label(bb: BoundingBox, lbl: BoundingBox) -> float: | |
| s = _bbox_distance(bb, lbl) | |
| # Penalise if the candidate overlaps the label itself (likely IS the label) | |
| if _bbox_overlaps(bb, lbl, margin=2): | |
| s += LABEL_OVERLAP_PENALTY | |
| # Bonus for value being to the right of label (same row) or below | |
| bb_cx, bb_cy = _bbox_center(bb) | |
| lbl_cx, lbl_cy = _bbox_center(lbl) | |
| if bb_cx > lbl_cx or bb_cy > lbl_cy + 10: | |
| # Good position — slight bonus (reduce score) | |
| s *= 0.8 | |
| elif bb_cy < lbl_cy - 20 and bb_cx < lbl_cx: | |
| # Value is above-left of label — unlikely | |
| s += WRONG_DIRECTION_PENALTY | |
| return s | |
| def score(bb: BoundingBox) -> float: | |
| s = 0.0 | |
| if all_labels: | |
| # Try each label position and take the best score | |
| s = min(score_with_label(bb, lbl) for lbl in all_labels) | |
| # Surrounding-text context bonus — strongly favours candidates | |
| # whose nearby OCR words contain the field's label tokens. | |
| if words and label_tokens and is_ambiguous: | |
| ctx_score = _surrounding_text_score(bb, words, label_tokens, radius=6) | |
| # Reduce the score (= make it better) proportionally | |
| # A perfect context match (1.0) can override 500px of distance | |
| s -= ctx_score * CONTEXT_BONUS_MAX | |
| # Penalty for overlapping already-claimed regions | |
| for c in claimed: | |
| if _bbox_overlaps(bb, c): | |
| s += CLAIMED_PENALTY | |
| return s | |
| scored = sorted(candidates, key=score) | |
| return scored[0] if scored else None | |
| def generate_bounding_boxes(image: Image.Image, result: "ExtractionResult") -> list[BoundingBox]: | |
| """Generate bounding boxes by matching extracted field values to OCR word positions. | |
| Only returns bounding boxes when OCR can reliably match at least 30% of | |
| extracted fields. For handwritten or complex documents where OCR can't read | |
| the text (but VLM could), returns an empty list — no boxes is better than | |
| wrong boxes. | |
| """ | |
| # Count how many fields actually have values | |
| field_values = _get_field_values(result) | |
| numeric_fields = _get_numeric_fields(result) | |
| total_fields_with_values = sum(1 for v in field_values.values() if v) + len(numeric_fields) | |
| if total_fields_with_values == 0: | |
| return [] | |
| # Try OCR-based matching | |
| words = _get_ocr_word_boxes(image) | |
| if not words: | |
| logger.info("No OCR words detected — skipping bounding boxes") | |
| return [] | |
| bboxes = [] | |
| claimed: list[BoundingBox] = [] # track assigned regions to avoid duplicates | |
| for field_name, value in field_values.items(): | |
| if not value: | |
| continue | |
| # Try alternative representations for dates and currency | |
| if "date" in field_name: | |
| search_variants = _date_search_variants(value) | |
| elif field_name == "currency": | |
| search_variants = _currency_search_variants(value) | |
| else: | |
| search_variants = [value] | |
| # Gather ALL matches across all variants | |
| all_matches: list[BoundingBox] = [] | |
| for variant in search_variants: | |
| all_matches.extend(_find_all_text_bboxes(words, variant)) | |
| # Fallback: for multi-word values with no matches, try progressively | |
| # shorter sub-phrases, then individual distinctive words | |
| if not all_matches and len(value.split()) > 1: | |
| tokens = value.split() | |
| # Try progressively shorter windows, longest first | |
| for window in range(len(tokens) - 1, 0, -1): | |
| for start in range(len(tokens) - window + 1): | |
| sub = " ".join(tokens[start:start + window]) | |
| # For single words, require length >= 5 to be distinctive | |
| min_len = 5 if window == 1 else 4 | |
| if len(sub) >= min_len: | |
| sub_matches = _find_all_text_bboxes(words, sub) | |
| all_matches.extend(sub_matches) | |
| if all_matches: | |
| break # found matches at this window size | |
| if all_matches: | |
| # Find ALL label positions for context-aware disambiguation | |
| label_bboxes = _find_label_bboxes(words, field_name) | |
| label_bbox = label_bboxes[0] if label_bboxes else None | |
| best = _pick_best_match(all_matches, label_bbox, claimed, | |
| label_bboxes=label_bboxes, words=words, | |
| field_name=field_name) | |
| if best: | |
| best.field = field_name | |
| bboxes.append(best) | |
| claimed.append(best) | |
| for field_name, num_value in numeric_fields.items(): | |
| search_variants = _numeric_search_variants(num_value) | |
| all_matches = [] | |
| for variant in search_variants: | |
| all_matches.extend(_find_all_text_bboxes(words, variant)) | |
| if all_matches: | |
| label_bboxes = _find_label_bboxes(words, field_name) | |
| label_bbox = label_bboxes[0] if label_bboxes else None | |
| best = _pick_best_match(all_matches, label_bbox, claimed, | |
| label_bboxes=label_bboxes, words=words, | |
| field_name=field_name) | |
| if best: | |
| best.field = field_name | |
| bboxes.append(best) | |
| claimed.append(best) | |
| # Line items — try to locate each row by its description and/or amount. | |
| # Uses claimed list so successive line items with the same amount | |
| # (e.g. two items at $50) each get their own bbox. | |
| for idx, li in enumerate(result.line_items): | |
| li_field = f"line_item_{idx}" | |
| # Support both LineItem objects and raw dicts | |
| if isinstance(li, dict): | |
| li_desc = str(li.get("description", "") or "") | |
| # For non-standard schemas, pick the longest string value as | |
| # the "description" for bbox matching (e.g. transaction_description) | |
| if not li_desc: | |
| str_vals = [(k, str(v)) for k, v in li.items() | |
| if isinstance(v, str) and len(str(v)) > 3 | |
| and k not in ("amount", "quantity", "unit_price")] | |
| if str_vals: | |
| li_desc = max(str_vals, key=lambda x: len(x[1]))[1] | |
| li_amount = li.get("amount") | |
| else: | |
| li_desc = li.description | |
| li_amount = li.amount | |
| # Find best unclaimed description match | |
| desc_bbox = None | |
| if li_desc: | |
| desc_matches = _find_all_text_bboxes(words, li_desc) | |
| desc_bbox = _pick_best_match(desc_matches, None, claimed) if desc_matches else None | |
| # Find best unclaimed amount match | |
| amt_bbox = None | |
| if li_amount is not None and li_amount != 0: | |
| amt_matches = [] | |
| for variant in _numeric_search_variants(li_amount): | |
| amt_matches.extend(_find_all_text_bboxes(words, variant)) | |
| if amt_matches: | |
| # Prefer an amount on the same line as the description | |
| amt_bbox = _pick_best_match(amt_matches, desc_bbox, claimed) | |
| if desc_bbox and amt_bbox: | |
| # Only merge if both are on roughly the same line (within 60px at 200 DPI) | |
| desc_mid_y = (desc_bbox.y1 + desc_bbox.y2) / 2 | |
| amt_mid_y = (amt_bbox.y1 + amt_bbox.y2) / 2 | |
| if abs(desc_mid_y - amt_mid_y) < 60: | |
| merged = BoundingBox( | |
| field=li_field, | |
| x1=min(desc_bbox.x1, amt_bbox.x1), | |
| y1=min(desc_bbox.y1, amt_bbox.y1), | |
| x2=max(desc_bbox.x2, amt_bbox.x2), | |
| y2=max(desc_bbox.y2, amt_bbox.y2), | |
| ) | |
| bboxes.append(merged) | |
| claimed.append(merged) | |
| else: | |
| desc_bbox.field = li_field | |
| bboxes.append(desc_bbox) | |
| claimed.append(desc_bbox) | |
| elif desc_bbox: | |
| desc_bbox.field = li_field | |
| bboxes.append(desc_bbox) | |
| claimed.append(desc_bbox) | |
| elif amt_bbox: | |
| amt_bbox.field = li_field | |
| bboxes.append(amt_bbox) | |
| claimed.append(amt_bbox) | |
| # Only return OCR matches if we matched a meaningful portion of fields. | |
| # Sparse matches (1-2 out of 10) are usually noise on handwritten docs. | |
| match_ratio = len(bboxes) / total_fields_with_values if total_fields_with_values > 0 else 0 | |
| # Log which fields were NOT matched for diagnostics | |
| matched_fields = {bb.field for bb in bboxes} | |
| all_field_names = set(k for k, v in field_values.items() if v) | set(numeric_fields.keys()) | |
| unmatched = all_field_names - matched_fields | |
| if unmatched: | |
| logger.info(f"Unmatched fields: {sorted(unmatched)}") | |
| if match_ratio < 0.2: | |
| logger.info( | |
| f"OCR matched only {len(bboxes)}/{total_fields_with_values} fields " | |
| f"({match_ratio:.0%}) — too few, skipping bounding boxes" | |
| ) | |
| return [] | |
| logger.info(f"OCR matched {len(bboxes)}/{total_fields_with_values} fields ({match_ratio:.0%})") | |
| return bboxes | |
| def _get_field_values(result: "ExtractionResult") -> dict: | |
| """Get string field name→value mapping for bbox matching. | |
| Prefers extracted_fields (works for any document type). | |
| Falls back to legacy named fields for older responses. | |
| """ | |
| # Primary: extracted_fields dict (dynamic, works for all use cases) | |
| if result.extracted_fields: | |
| return { | |
| k: str(v) for k, v in result.extracted_fields.items() | |
| if v is not None and not isinstance(v, (int, float)) | |
| } | |
| # Fallback: legacy named fields (invoice-specific) | |
| return { | |
| "vendor_name": result.vendor_name, | |
| "vendor_address": result.vendor_address, | |
| "invoice_number": result.invoice_number, | |
| "invoice_date": result.invoice_date, | |
| "due_date": result.due_date, | |
| "currency": result.currency, | |
| "customer_name": result.customer_name, | |
| "customer_address": result.customer_address, | |
| "payment_terms": result.payment_terms, | |
| } | |
| def _get_numeric_fields(result: "ExtractionResult") -> dict: | |
| """Get numeric field name→value mapping for bbox matching. | |
| Prefers extracted_fields (works for any document type). | |
| Falls back to legacy named fields for older responses. | |
| """ | |
| # Primary: extracted_fields dict (dynamic, works for all use cases) | |
| if result.extracted_fields: | |
| return { | |
| k: v for k, v in result.extracted_fields.items() | |
| if isinstance(v, (int, float)) and v is not None | |
| } | |
| # Fallback: legacy named fields | |
| fields = {} | |
| if result.subtotal is not None: | |
| fields["subtotal"] = result.subtotal | |
| if result.tax_amount is not None: | |
| fields["tax_amount"] = result.tax_amount | |
| if result.total_amount is not None: | |
| fields["total_amount"] = result.total_amount | |
| return fields | |
| def _numeric_search_variants(value: float) -> list[str]: | |
| """Generate multiple string representations of a number for OCR matching.""" | |
| variants = [] | |
| # Integer form if it's a whole number: "353" | |
| if value == int(value): | |
| variants.append(str(int(value))) | |
| # Standard float: "353.0" | |
| variants.append(str(value)) | |
| # Two-decimal: "353.00" | |
| variants.append(f"{value:.2f}") | |
| # Comma-separated thousands (Western): "1,250.00" | |
| if value >= 1000: | |
| variants.append(f"{value:,.2f}") | |
| variants.append(f"{value:,.0f}" if value == int(value) else f"{value:,.2f}") | |
| # Indian lakh/crore grouping: "1,23,456.00" | |
| if value >= 1000: | |
| int_part = int(value) | |
| dec_part = f"{value:.2f}".split(".")[1] | |
| s = str(int_part) | |
| if len(s) > 3: | |
| last3 = s[-3:] | |
| rest = s[:-3] | |
| # Group remaining digits in pairs from right | |
| groups = [] | |
| while rest: | |
| groups.append(rest[-2:] if len(rest) >= 2 else rest) | |
| rest = rest[:-2] | |
| groups.reverse() | |
| indian = ",".join(groups) + "," + last3 | |
| else: | |
| indian = s | |
| variants.append(f"{indian}.{dec_part}") | |
| if value == int(value): | |
| variants.append(indian) | |
| return list(dict.fromkeys(variants)) # deduplicate, preserve order | |
| # --------------------------------------------------------------------------- | |
| # Prompt Library — per use-case extraction prompts | |
| # --------------------------------------------------------------------------- | |
| # The library maps (stage, use_case) → prompt text. | |
| # If a use_case-specific prompt doesn't exist, falls back to "_default". | |
| EXTRACTION_PROMPTS: dict[str, str] = {} | |
| EXTRACTION_PROMPTS["invoice"] = """You are a document extraction AI. Extract ALL structured data from this invoice, receipt, or financial document. | |
| Return ONLY valid JSON with this schema (use null for missing fields): | |
| { | |
| "document_type": "invoice", | |
| "confidence": 0.0 to 1.0, | |
| "fields": { | |
| "vendor_name": "string", | |
| "vendor_address": "string", | |
| "invoice_number": "string or null", | |
| "invoice_date": "string (YYYY-MM-DD)", | |
| "due_date": "string (YYYY-MM-DD) or null", | |
| "currency": "string (ISO 4217) or null", | |
| "subtotal": number or null, | |
| "tax_amount": number or null, | |
| "tax_rate": "string (e.g. 18% GST) or null", | |
| "total_amount": number, | |
| "payment_terms": "string or null", | |
| "customer_name": "string or null", | |
| "customer_address": "string or null", | |
| "...any_other_field": "extract ALL other visible fields as additional flat key-value pairs here" | |
| }, | |
| "line_items": [{"description": "string", "quantity": number or null, "unit_price": number or null, "amount": number or null}] | |
| } | |
| IMPORTANT RULES: | |
| - confidence: your self-assessed confidence in the extraction accuracy (0.0 = pure guess, 1.0 = perfectly clear). | |
| - Be precise with numbers. Parse dates into YYYY-MM-DD format. | |
| - currency: ONLY set this if a currency symbol or code is VISIBLE. Do NOT infer. | |
| - line_items: Extract ONLY what is explicitly written. Do NOT invent quantity or unit_price. Never default quantity to 1. | |
| - Subtotal: only set if the document explicitly labels a subtotal. Do not copy total_amount into subtotal. | |
| - CRITICAL: You MUST include ALL other visible fields directly inside "fields" as additional key-value pairs using snake_case names. Examples: "gstin", "hsn_code", "po_number", "email", "phone", "account_number", "credit_card_no", "statement_date", "minimum_due", "reward_points", "disbursed", etc. Do NOT limit yourself to only the fields listed in the schema above. Every piece of text data visible in the document should appear as a field. The schema above is a MINIMUM — add every other field you can find.""" | |
| EXTRACTION_PROMPTS["tax_form"] = """You are a document extraction AI specializing in tax forms, challans, and government tax documents. | |
| Extract ALL structured data from this tax document. | |
| Return ONLY valid JSON: | |
| { | |
| "document_type": "tax_form", | |
| "confidence": 0.0 to 1.0, | |
| "fields": { | |
| "form_type": "string (e.g. Form 26QB, Challan 280, Form 16, TDS Certificate)", | |
| "pan": "string or null", | |
| "tan": "string or null", | |
| "assessment_year": "string (e.g. 2026-27) or null", | |
| "financial_year": "string or null", | |
| "nature_of_payment": "string or null", | |
| "section": "string (e.g. 194IA, 234E) or null", | |
| "deductor_name": "string or null", | |
| "deductee_name": "string or null", | |
| "amount_paid": number or null, | |
| "tds_amount": number or null, | |
| "surcharge": number or null, | |
| "education_cess": number or null, | |
| "total_tax_deposited": number or null, | |
| "date_of_payment": "string (YYYY-MM-DD) or null", | |
| "date_of_deduction": "string (YYYY-MM-DD) or null", | |
| "acknowledgement_number": "string or null", | |
| "challan_number": "string or null", | |
| "bsr_code": "string or null", | |
| "demand_reference_number": "string or null" | |
| }, | |
| "line_items": [] | |
| } | |
| IMPORTANT RULES: | |
| - Include ALL visible fields, even those not listed above — add them to "fields". | |
| - Use snake_case for any extra field names, derived from the label in the document. | |
| - Parse dates into YYYY-MM-DD format. | |
| - Numbers should be numeric types, not strings. | |
| - PAN and TAN should be uppercase strings. | |
| - confidence: your self-assessed confidence in the extraction accuracy (0.0 = pure guess, 1.0 = perfectly clear).""" | |
| EXTRACTION_PROMPTS["receipt"] = """You are a document extraction AI. Extract ALL structured data from this receipt or purchase record. | |
| Return ONLY valid JSON: | |
| { | |
| "document_type": "receipt", | |
| "confidence": 0.0 to 1.0, | |
| "fields": { | |
| "store_name": "string", | |
| "store_address": "string or null", | |
| "receipt_number": "string or null", | |
| "date": "string (YYYY-MM-DD)", | |
| "currency": "string (ISO 4217) or null", | |
| "subtotal": number or null, | |
| "tax_amount": number or null, | |
| "tax_rate": "string or null", | |
| "total_amount": number, | |
| "payment_method": "string or null", | |
| "card_last_four": "string or null" | |
| }, | |
| "line_items": [{"description": "string", "quantity": number or null, "unit_price": number or null, "amount": number or null}] | |
| } | |
| IMPORTANT RULES: | |
| - Include ALL visible fields, even those not listed — add them to "fields". | |
| - currency: ONLY set if a symbol or code is VISIBLE. Do NOT infer. | |
| - line_items: Extract ONLY what is explicitly written. Never invent quantity. | |
| - Parse dates into YYYY-MM-DD format. | |
| - confidence: your self-assessed confidence in the extraction accuracy (0.0 = pure guess, 1.0 = perfectly clear).""" | |
| EXTRACTION_PROMPTS["_default"] = """You are a document extraction AI. Analyze this document and extract ALL visible structured data. | |
| Steps: | |
| 1. Identify the document type (e.g. invoice, receipt, tax_form, contract, bank_statement, id_card, letter, application, certificate, etc.) | |
| 2. Extract every labeled field visible in the document as key-value pairs. | |
| 3. Extract any tabular data as line_items. | |
| Return ONLY valid JSON: | |
| { | |
| "document_type": "string", | |
| "confidence": 0.0 to 1.0, | |
| "fields": { | |
| "field_name": "value", | |
| ... | |
| }, | |
| "line_items": [ | |
| {"column_name": "value", ...}, | |
| ... | |
| ] | |
| } | |
| IMPORTANT RULES: | |
| - confidence: your self-assessed confidence in the extraction accuracy (0.0 = pure guess, 1.0 = perfectly clear). | |
| - Use snake_case for field names, derived from the label text in the document. | |
| - Parse dates into YYYY-MM-DD format. | |
| - Numbers should be numeric types, not strings. | |
| - Include ALL visible fields — headers, form fields, metadata, identifiers, dates, amounts. | |
| - For tables, use column headers as keys in each row object. | |
| - Do NOT infer or guess values. Only extract what is explicitly visible.""" | |
| # --------------------------------------------------------------------------- | |
| # Validation Prompts | |
| # --------------------------------------------------------------------------- | |
| VALIDATION_PROMPTS: dict[str, str] = {} | |
| VALIDATION_PROMPTS["invoice"] = """You are a document validation AI. You will be given: | |
| 1. An image of the original document. | |
| 2. A JSON object of extracted data. | |
| Verify EVERY extracted field against the document image. For each field, check: | |
| - Is the value present in the document? | |
| - Is it accurately transcribed (exact text, numbers, dates)? | |
| - Are numeric fields consistent (subtotal + tax = total, qty × price = amount)? | |
| Return ONLY valid JSON: | |
| { | |
| "checks": [ | |
| { | |
| "field": "field_name", | |
| "status": "pass" | "fail" | "warn", | |
| "message": "short explanation", | |
| "expected": "what the document shows (if different)", | |
| "actual": "what was extracted" | |
| } | |
| ], | |
| "summary": "short overall assessment" | |
| } | |
| Focus on: amounts, dates, vendor/customer names, invoice numbers, tax calculations. | |
| Flag any field where the extracted value doesn't match what's visible in the image.""" | |
| VALIDATION_PROMPTS["tax_form"] = """You are a document validation AI specializing in tax forms and government documents. | |
| You will be given: | |
| 1. An image of the original tax document. | |
| 2. A JSON object of extracted data. | |
| Verify EVERY extracted field against the document image. For tax forms, pay special attention to: | |
| - PAN/TAN format validity (10 chars: 5 letters, 4 digits, 1 letter) | |
| - Assessment year / financial year consistency | |
| - TDS amount + surcharge + cess = total tax deposited | |
| - Date fields match what's printed | |
| - Section numbers and challan/acknowledgement numbers are exact | |
| Return ONLY valid JSON: | |
| { | |
| "checks": [ | |
| { | |
| "field": "field_name", | |
| "status": "pass" | "fail" | "warn", | |
| "message": "short explanation", | |
| "expected": "what the document shows (if different)", | |
| "actual": "what was extracted" | |
| } | |
| ], | |
| "summary": "short overall assessment" | |
| }""" | |
| VALIDATION_PROMPTS["receipt"] = """You are a document validation AI. You will be given: | |
| 1. An image of the original receipt. | |
| 2. A JSON object of extracted data. | |
| Verify EVERY extracted field against the receipt image. For receipts, check: | |
| - Store name and details match the header | |
| - Line item amounts are correct | |
| - Subtotal/tax/total arithmetic | |
| - Date and receipt number accuracy | |
| - Payment method matches what's printed | |
| Return ONLY valid JSON: | |
| { | |
| "checks": [ | |
| { | |
| "field": "field_name", | |
| "status": "pass" | "fail" | "warn", | |
| "message": "short explanation", | |
| "expected": "what the document shows (if different)", | |
| "actual": "what was extracted" | |
| } | |
| ], | |
| "summary": "short overall assessment" | |
| }""" | |
| VALIDATION_PROMPTS["_default"] = """You are a document validation AI. You will be given: | |
| 1. An image of the original document. | |
| 2. A JSON object of extracted data. | |
| Verify EVERY extracted field against the document image: | |
| - Is each field value present and accurately transcribed? | |
| - Are numeric fields consistent with each other? | |
| - Are dates in the correct format and matching the document? | |
| - Are identifiers (numbers, codes, references) exact? | |
| Return ONLY valid JSON: | |
| { | |
| "checks": [ | |
| { | |
| "field": "field_name", | |
| "status": "pass" | "fail" | "warn", | |
| "message": "short explanation", | |
| "expected": "what the document shows (if different)", | |
| "actual": "what was extracted" | |
| } | |
| ], | |
| "summary": "short overall assessment" | |
| }""" | |
| # Stage-level prompt library: maps (stage, use_case) → prompt. | |
| PROMPT_LIBRARY: dict[str, dict[str, str]] = { | |
| "extraction": EXTRACTION_PROMPTS, | |
| "validation": VALIDATION_PROMPTS, | |
| } | |
| def get_prompt(stage: str, use_case: str) -> str: | |
| """Look up the prompt for a (stage, use_case) pair. | |
| Priority: | |
| 1. Active prompt from Supabase prompt_versions table (if configured) | |
| 2. Hardcoded PROMPT_LIBRARY dict (in-code defaults) | |
| 3. _default prompt for the stage | |
| """ | |
| # Try DB first — gracefully returns None when Supabase is not configured | |
| try: | |
| from db.supabase import get_active_prompt | |
| db_prompt = get_active_prompt(use_case, stage) | |
| if db_prompt: | |
| return db_prompt | |
| except Exception: | |
| pass # Fall through to hardcoded prompts | |
| stage_prompts = PROMPT_LIBRARY.get(stage, {}) | |
| return stage_prompts.get(use_case, stage_prompts.get("_default", "")) | |
| # Legacy alias kept so callers that reference EXTRACTION_PROMPT still work | |
| EXTRACTION_PROMPT = EXTRACTION_PROMPTS["invoice"] | |
| def image_to_base64(image: Image.Image, max_size: int = 1024) -> str: | |
| """Convert PIL Image to base64 string, resizing if needed. | |
| IMPORTANT: Works on a copy so the original image is never mutated. | |
| This preserves full resolution for downstream OCR bounding-box generation. | |
| """ | |
| img = image.copy() | |
| if max(img.size) > max_size: | |
| img.thumbnail((max_size, max_size), Image.LANCZOS) | |
| if img.mode == "RGBA": | |
| img = img.convert("RGB") | |
| buffer = BytesIO() | |
| img.save(buffer, format="JPEG", quality=85) | |
| return base64.b64encode(buffer.getvalue()).decode("utf-8") | |
| def _parse_vlm_json(raw: str) -> dict: | |
| """Extract JSON from a VLM response, handling markdown fences.""" | |
| json_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", raw) | |
| if json_match: | |
| return json.loads(json_match.group(1).strip()) | |
| json_match = re.search(r"\{[\s\S]*\}", raw) | |
| if json_match: | |
| return json.loads(json_match.group(0)) | |
| return json.loads(raw) | |
| def _compute_confidence(fields: dict, line_items: list, model_id: str) -> float: | |
| """Compute an extraction confidence score based on field completeness. | |
| Heuristic factors: | |
| - Ratio of non-null/non-empty fields → higher is better | |
| - Having line items when fields suggest a tabular document | |
| - Model tier bonus (larger models are more reliable) | |
| """ | |
| if not fields: | |
| return 0.3 | |
| # Count meaningful fields (non-null, non-empty) | |
| total = len(fields) | |
| filled = 0 | |
| for v in fields.values(): | |
| if v is None: | |
| continue | |
| if isinstance(v, str) and v.strip() in ("", "-", "N/A", "n/a", "None", "none"): | |
| continue | |
| filled += 1 | |
| if total == 0: | |
| return 0.3 | |
| field_ratio = filled / total | |
| # Line-item bonus: having structured rows is a strong signal | |
| li_bonus = min(0.05, len(line_items) * 0.01) if line_items else 0.0 | |
| # Model tier adjustment | |
| model_lower = model_id.lower() | |
| if "72b" in model_lower: | |
| tier_bonus = 0.05 | |
| elif "gemini" in model_lower: | |
| tier_bonus = 0.03 # Gemini Flash: between 72B and 7B quality | |
| elif "7b" in model_lower: | |
| tier_bonus = 0.0 | |
| else: | |
| tier_bonus = -0.05 | |
| # Base confidence: scale field ratio into 0.55–0.95 range | |
| confidence = 0.55 + (field_ratio * 0.40) + li_bonus + tier_bonus | |
| return round(max(0.1, min(0.99, confidence)), 2) | |
| def _vlm_data_to_result(data: dict, model_id: str, use_case: str) -> ExtractionResult: | |
| """Convert parsed VLM JSON into an ExtractionResult. | |
| Handles two response shapes: | |
| 1. New format: {"document_type": "...", "fields": {...}, "line_items": [...]} | |
| 2. Legacy flat format: {"vendor_name": "...", "total_amount": 100, ...} | |
| """ | |
| short_model = model_id.split("/")[-1] | |
| # Detect response shape | |
| fields = data.get("fields", {}) | |
| doc_type = data.get("document_type", use_case) | |
| if not fields: | |
| # Legacy flat format — all keys are fields (except line_items) | |
| fields = {k: v for k, v in data.items() | |
| if k not in ("line_items", "document_type", "fields", "confidence")} | |
| line_items = data.get("line_items", []) | |
| # Use model-reported confidence if present, otherwise compute heuristic | |
| reported_confidence = data.get("confidence") | |
| if isinstance(reported_confidence, (int, float)) and 0 < reported_confidence <= 1: | |
| confidence = round(float(reported_confidence), 2) | |
| else: | |
| confidence = _compute_confidence(fields, line_items, model_id) | |
| result = ExtractionResult( | |
| extraction_method=f"vlm:{short_model}", | |
| confidence=confidence, | |
| document_type=doc_type, | |
| extracted_fields=fields, | |
| ) | |
| # Populate legacy named fields from fields dict (backward compat for invoice) | |
| result.vendor_name = str(fields.get("vendor_name", "") or "") | |
| result.vendor_address = str(fields.get("vendor_address", "") or "") | |
| result.invoice_number = str(fields.get("invoice_number", "") or "") | |
| result.invoice_date = str(fields.get("invoice_date", "") or "") | |
| result.due_date = str(fields.get("due_date", "") or "") | |
| result.currency = str(fields.get("currency", "") or "") | |
| result.subtotal = fields.get("subtotal") | |
| result.tax_amount = fields.get("tax_amount") | |
| result.tax_rate = str(fields.get("tax_rate", "") or "") | |
| result.total_amount = fields.get("total_amount") | |
| result.payment_terms = str(fields.get("payment_terms", "") or "") | |
| result.customer_name = str(fields.get("customer_name", "") or "") | |
| result.customer_address = str(fields.get("customer_address", "") or "") | |
| # Parse line items — preserve the raw dict so the frontend can display | |
| # whatever columns the VLM returned (e.g. date, transaction_description | |
| # for credit-card statements). Only coerce into the rigid LineItem | |
| # dataclass when the item carries exactly the invoice-schema keys. | |
| _LINEITEM_KEYS = {"description", "quantity", "unit_price", "amount"} | |
| for item in data.get("line_items", []): | |
| if isinstance(item, dict): | |
| item_keys = set(item.keys()) | |
| if item_keys <= _LINEITEM_KEYS: | |
| # Standard invoice row — use LineItem for backward compat | |
| result.line_items.append(LineItem( | |
| description=item.get("description", ""), | |
| quantity=item.get("quantity"), | |
| unit_price=item.get("unit_price"), | |
| amount=item.get("amount"), | |
| )) | |
| else: | |
| # Non-standard columns (credit-card statement, etc.) | |
| # Keep the raw dict so the frontend's dynamic column | |
| # builder can show all fields. | |
| result.line_items.append(item) | |
| return result | |
| def _is_gemini_model(model_id: str) -> bool: | |
| """Check whether a model id refers to a Gemini model.""" | |
| return model_id.lower().startswith("gemini") | |
| def _extract_with_gemini( | |
| image: Image.Image, | |
| model_id: str, | |
| use_case: str, | |
| ) -> Optional[ExtractionResult]: | |
| """Extract using Google Gemini Vision API.""" | |
| prompt = get_prompt("extraction", use_case) | |
| if not prompt: | |
| logger.warning(f"No extraction prompt for use_case={use_case}, using default") | |
| prompt = get_prompt("extraction", "_default") | |
| img_b64 = image_to_base64(image) | |
| try: | |
| result_text = extract_with_gemini(img_b64, prompt, model=model_id) | |
| data = _parse_vlm_json(result_text) | |
| return _vlm_data_to_result(data, model_id, use_case) | |
| except Exception as e: | |
| error_msg = str(e)[:200] | |
| logger.warning(f"Gemini extraction failed with {model_id}: {e}") | |
| return ExtractionResult( | |
| extraction_method=f"failed:gemini({model_id})", | |
| confidence=0.0, | |
| raw_text=f"{model_id}: {type(e).__name__}({error_msg})", | |
| ) | |
| def _extract_with_hf( | |
| image: Image.Image, | |
| hf_token: str, | |
| model_id: str, | |
| use_case: str, | |
| ) -> Optional[ExtractionResult]: | |
| """Extract using HF Inference API. When model_id is a specific Qwen model, | |
| tries only that model. When empty or unrecognised, cascades through all. | |
| """ | |
| prompt = get_prompt("extraction", use_case) | |
| if not prompt: | |
| logger.warning(f"No extraction prompt for use_case={use_case}, using default") | |
| prompt = get_prompt("extraction", "_default") | |
| # Determine which models to try | |
| all_models = [ | |
| "Qwen/Qwen2.5-VL-72B-Instruct", | |
| "Qwen/Qwen2.5-VL-7B-Instruct", | |
| "Qwen/Qwen2.5-VL-3B-Instruct", | |
| ] | |
| if model_id and model_id in all_models: | |
| # User picked a specific model — try it first, then cascade | |
| models = [model_id] + [m for m in all_models if m != model_id] | |
| else: | |
| models = all_models | |
| img_b64 = image_to_base64(image) | |
| data_url = f"data:image/jpeg;base64,{img_b64}" | |
| errors = [] | |
| for mid in models: | |
| try: | |
| client = create_inference_client(api_key=hf_token) | |
| response = robust_chat_completion( | |
| client, | |
| model=mid, | |
| messages=[ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": prompt}, | |
| { | |
| "type": "image_url", | |
| "image_url": {"url": data_url}, | |
| }, | |
| ], | |
| } | |
| ], | |
| max_tokens=2048, | |
| temperature=0.1, | |
| ) | |
| result_text = response.choices[0].message.content | |
| data = _parse_vlm_json(result_text) | |
| return _vlm_data_to_result(data, mid, use_case) | |
| except Exception as e: | |
| short_model = mid.split("/")[-1] | |
| error_msg = str(e)[:150] | |
| errors.append(f"{short_model}: {type(e).__name__}({error_msg})") | |
| logger.warning(f"VLM extraction failed with {mid}: {e}") | |
| continue | |
| return ExtractionResult( | |
| extraction_method="failed:vlm_all_models", | |
| confidence=0.0, | |
| raw_text="; ".join(errors) if errors else "all models failed silently", | |
| ) | |
| def extract_with_vlm( | |
| image: Image.Image, | |
| hf_token: str, | |
| use_case: str = "invoice", | |
| model_id: str = "", | |
| ) -> Optional[ExtractionResult]: | |
| """Extract document data using a Vision-Language Model. | |
| Routes to the appropriate provider based on model_id: | |
| - gemini-* → Google Gemini API (requires GEMINI_API_KEY) | |
| - Qwen/* → HF Inference API (requires hf_token) | |
| - empty → Gemini if GEMINI_API_KEY is set, else HF cascade | |
| Selects the prompt from the prompt library based on use_case. | |
| Falls back to the generic prompt when no use_case-specific prompt exists. | |
| """ | |
| # Route: Gemini models | |
| if _is_gemini_model(model_id): | |
| if not GEMINI_API_KEY: | |
| logger.warning("Gemini model requested but GEMINI_API_KEY not set, falling back to HF") | |
| return _extract_with_hf(image, hf_token, "", use_case) | |
| return _extract_with_gemini(image, model_id, use_case) | |
| # Route: HF models (explicit Qwen selection or fallback) | |
| if model_id and model_id.startswith("Qwen/"): | |
| return _extract_with_hf(image, hf_token, model_id, use_case) | |
| # Route: No specific model — prefer Gemini if configured, else HF cascade | |
| if GEMINI_API_KEY: | |
| logger.info("No model specified, defaulting to Gemini Flash") | |
| return _extract_with_gemini(image, "gemini-3.5-flash-lite", use_case) | |
| return _extract_with_hf(image, hf_token, "", use_case) | |
| # --------------------------------------------------------------------------- | |
| # OCR + Regex Extraction (Fallback) - CPU only, zero cost | |
| # --------------------------------------------------------------------------- | |
| def extract_text_ocr(image: Image.Image) -> str: | |
| """Extract text from image using Tesseract OCR.""" | |
| try: | |
| import pytesseract | |
| text = pytesseract.image_to_string(image, config="--psm 3") | |
| return text | |
| except ImportError: | |
| logger.warning("pytesseract not available") | |
| return "" | |
| except Exception as e: | |
| logger.warning(f"OCR failed: {e}") | |
| return "" | |
| def parse_amount(text: str) -> Optional[float]: | |
| """Parse a monetary amount from text.""" | |
| text = text.replace(",", "").replace(" ", "") | |
| match = re.search(r"[\d]+\.?\d*", text) | |
| if match: | |
| try: | |
| return float(match.group()) | |
| except ValueError: | |
| pass | |
| return None | |
| def detect_currency(text: str) -> str: | |
| """Detect currency from text symbols and keywords.""" | |
| currency_patterns = { | |
| "INR": [r"₹", r"Rs\.?", r"INR", r"rupee"], | |
| "USD": [r"\$", r"USD", r"dollar"], | |
| "EUR": [r"€", r"EUR", r"euro"], | |
| "GBP": [r"£", r"GBP", r"pound"], | |
| } | |
| for currency, patterns in currency_patterns.items(): | |
| for p in patterns: | |
| if re.search(p, text, re.IGNORECASE): | |
| return currency | |
| return "USD" # default | |
| def extract_with_regex(text: str) -> ExtractionResult: | |
| """Extract structured data from OCR text using regex patterns.""" | |
| result = ExtractionResult( | |
| raw_text=text, | |
| extraction_method="ocr:tesseract+regex", | |
| confidence=0.65, | |
| ) | |
| # Invoice number — require a separator (No, #, :) between label and value | |
| inv_patterns = [ | |
| r"(?:invoice|inv|bill)\s*(?:#|no\.?|number)\s*:?\s*([A-Z0-9\-/]+)", | |
| r"(?:invoice|inv)\s*:\s*([A-Z0-9\-/]+)", | |
| r"(?:#|no\.?)\s*:?\s*([A-Z0-9\-/]+)", | |
| ] | |
| for p in inv_patterns: | |
| m = re.search(p, text, re.IGNORECASE) | |
| if m: | |
| result.invoice_number = m.group(1).strip() | |
| break | |
| # Dates | |
| date_patterns = [ | |
| r"(\d{1,2}[/\-\.]\d{1,2}[/\-\.]\d{2,4})", | |
| r"(\d{4}[/\-\.]\d{1,2}[/\-\.]\d{1,2})", | |
| r"(\w+\s+\d{1,2},?\s+\d{4})", | |
| ] | |
| dates_found = [] | |
| for p in date_patterns: | |
| dates_found.extend(re.findall(p, text)) | |
| if dates_found: | |
| result.invoice_date = dates_found[0] | |
| if len(dates_found) > 1: | |
| result.due_date = dates_found[1] | |
| # Total amount — try grand total / total due first, then plain "total" | |
| # but skip lines that say "subtotal" or specific tax labels | |
| total_patterns = [ | |
| r"(?:grand\s*total|amount\s*due|balance\s*due|total\s*amount)\s*:?\s*(?:INR|USD|EUR|GBP|[₹$€£])?\s*([\d,]+\.?\d*)", | |
| r"(?<!sub)total\s*:?\s*(?:INR|USD|EUR|GBP|[₹$€£])?\s*([\d,]+\.?\d*)", | |
| ] | |
| for p in total_patterns: | |
| matches = list(re.finditer(p, text, re.IGNORECASE | re.MULTILINE)) | |
| if matches: | |
| # Take the last match (usually the grand total at the bottom) | |
| result.total_amount = parse_amount(matches[-1].group(1)) | |
| break | |
| # Subtotal | |
| m = re.search(r"(?:subtotal|sub\s*total)\s*:?\s*[₹$€£]?\s*([\d,]+\.?\d*)", text, re.IGNORECASE) | |
| if m: | |
| result.subtotal = parse_amount(m.group(1)) | |
| # Tax | |
| tax_patterns = [ | |
| r"(?:tax|gst|vat|igst|cgst|sgst)\s*(?:\(?\d+%?\)?)?\s*:?\s*[₹$€£]?\s*([\d,]+\.?\d*)", | |
| ] | |
| for p in tax_patterns: | |
| m = re.search(p, text, re.IGNORECASE) | |
| if m: | |
| result.tax_amount = parse_amount(m.group(1)) | |
| break | |
| # Tax rate — also match "CGST (9%)" and "GST @18%" patterns | |
| tax_rate_patterns = [ | |
| r"(\d+(?:\.\d+)?)\s*%\s*(?:tax|gst|vat|igst)", | |
| r"(?:tax|gst|vat|cgst|sgst|igst)\s*[@(]?\s*(\d+(?:\.\d+)?)\s*%", | |
| ] | |
| for p in tax_rate_patterns: | |
| m = re.search(p, text, re.IGNORECASE) | |
| if m: | |
| result.tax_rate = f"{m.group(1)}%" | |
| break | |
| # Currency | |
| result.currency = detect_currency(text) | |
| # Vendor name (heuristic: first line or line before address) | |
| lines = [l.strip() for l in text.split("\n") if l.strip()] | |
| if lines: | |
| result.vendor_name = lines[0] | |
| return result | |
| # --------------------------------------------------------------------------- | |
| # Main Extraction Pipeline | |
| # --------------------------------------------------------------------------- | |
| def extract_invoice( | |
| image: Image.Image, | |
| hf_token: Optional[str] = None, | |
| force_ocr: bool = False, | |
| use_case: str = "invoice", | |
| model_id: str = "", | |
| ) -> ExtractionResult: | |
| """ | |
| Main extraction pipeline. Tries VLM first, falls back to OCR+regex. | |
| Args: | |
| image: PIL Image of the document | |
| hf_token: Hugging Face API token (for VLM extraction) | |
| force_ocr: Skip VLM and use OCR only | |
| use_case: Document type for prompt selection (invoice, tax_form, receipt, etc.) | |
| model_id: Model to use (e.g. "gemini-3.5-flash-lite", "Qwen/Qwen2.5-VL-72B-Instruct"). | |
| When empty, defaults based on available API keys. | |
| Returns: | |
| ExtractionResult with all extracted fields | |
| """ | |
| start = datetime.now() | |
| result = None | |
| vlm_errors = "" | |
| # Gemini models don't need an HF token — allow VLM when Gemini is selected | |
| is_gemini = _is_gemini_model(model_id) if model_id else False | |
| has_vlm_access = (hf_token or is_gemini or GEMINI_API_KEY) and not force_ocr | |
| # Try VLM extraction first | |
| if has_vlm_access: | |
| vlm_result = extract_with_vlm(image, hf_token or "", use_case=use_case, model_id=model_id) | |
| if vlm_result and not vlm_result.extraction_method.startswith("failed"): | |
| result = vlm_result | |
| else: | |
| # VLM failed — capture diagnostics, fall through to OCR | |
| vlm_errors = vlm_result.raw_text if vlm_result else "unknown" | |
| logger.warning(f"VLM failed ({vlm_errors}), trying OCR fallback") | |
| elif not force_ocr: | |
| vlm_errors = "no_token" | |
| logger.warning("No API keys set (HF_TOKEN / GEMINI_API_KEY) — using OCR fallback only") | |
| # Fallback to OCR + regex | |
| if result is None: | |
| text = extract_text_ocr(image) | |
| if text.strip(): | |
| result = extract_with_regex(text) | |
| if vlm_errors: | |
| # Note that this was a fallback result | |
| result.extraction_method += f" (vlm failed: {vlm_errors})" | |
| else: | |
| method = f"failed:vlm({vlm_errors})+ocr_empty" if vlm_errors else "failed:ocr_empty" | |
| result = ExtractionResult( | |
| extraction_method=method, | |
| confidence=0.0, | |
| ) | |
| # Generate bounding boxes by matching extracted values to OCR word positions. | |
| # Falls back to layout-based estimation for handwritten documents. | |
| try: | |
| result.bounding_boxes = generate_bounding_boxes(image, result) | |
| except Exception as e: | |
| logger.warning(f"Bounding box generation failed: {e}") | |
| elapsed = (datetime.now() - start).total_seconds() * 1000 | |
| result.processing_time_ms = int(elapsed) | |
| return result | |
| def extract_from_pdf( | |
| pdf_bytes: bytes, | |
| hf_token: Optional[str] = None, | |
| force_ocr: bool = False, | |
| use_case: str = "invoice", | |
| max_seconds: int = 240, | |
| model_id: str = "", | |
| ) -> list[ExtractionResult]: | |
| """Extract from PDF - converts each page to image and extracts. | |
| The first page's rendered image is attached as ``page_image`` (a | |
| data-URI) so the frontend can display the PDF and overlay bounding | |
| boxes without needing a separate PDF renderer. | |
| All pages are rendered first (fast), then extraction runs per page | |
| with a time budget. If the budget is exceeded the pages already | |
| extracted are returned so the user still gets partial results and | |
| can see all page images in the viewer. | |
| """ | |
| import time as _time | |
| try: | |
| import fitz # PyMuPDF | |
| except ImportError: | |
| raise ImportError("PyMuPDF (fitz) is required for PDF processing. Install with: pip install PyMuPDF") | |
| deadline = _time.monotonic() + max_seconds | |
| doc = fitz.open(stream=pdf_bytes, filetype="pdf") | |
| num_pages = min(len(doc), 10) # Max 10 pages | |
| # Phase 1 — render all page images up-front (fast, < 1s per page) | |
| page_imgs: list[tuple[Image.Image, str]] = [] | |
| for page_num in range(num_pages): | |
| page = doc[page_num] | |
| pix = page.get_pixmap(dpi=200) | |
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) | |
| buf = BytesIO() | |
| img.save(buf, format="PNG", optimize=True) | |
| b64 = base64.b64encode(buf.getvalue()).decode() | |
| page_imgs.append((img, f"data:image/png;base64,{b64}")) | |
| doc.close() | |
| # Phase 2 — extract pages in parallel (I/O-bound API calls benefit | |
| # from concurrency). Cap at 2 workers to stay within Render free-tier | |
| # memory limits (512 MB) and Gemini rate limits (~15 RPM). | |
| max_workers = min(num_pages, 2) | |
| def _extract_page(page_idx: int) -> tuple[int, ExtractionResult]: | |
| img, data_uri = page_imgs[page_idx] | |
| result = extract_invoice( | |
| img, hf_token=hf_token, force_ocr=force_ocr, | |
| use_case=use_case, model_id=model_id, | |
| ) | |
| result.page_image = data_uri | |
| return page_idx, result | |
| results: list[Optional[ExtractionResult]] = [None] * num_pages | |
| extracted_count = 0 | |
| with ThreadPoolExecutor(max_workers=max_workers) as executor: | |
| futures = { | |
| executor.submit(_extract_page, idx): idx | |
| for idx in range(num_pages) | |
| } | |
| for future in as_completed(futures): | |
| idx = futures[future] | |
| try: | |
| page_idx, result = future.result() | |
| results[page_idx] = result | |
| extracted_count += 1 | |
| except Exception as e: | |
| logger.warning("PDF page %d extraction failed: %s", idx, e) | |
| results[idx] = ExtractionResult( | |
| extraction_method=f"failed:parallel_error", | |
| confidence=0.0, | |
| raw_text=str(e)[:200], | |
| ) | |
| results[idx].page_image = page_imgs[idx][1] # type: ignore[union-attr] | |
| extracted_count += 1 | |
| # Check time budget — cancel remaining futures if exceeded | |
| if _time.monotonic() > deadline and extracted_count < num_pages: | |
| logger.warning( | |
| "PDF extraction time budget exceeded after %d/%d pages", | |
| extracted_count, num_pages, | |
| ) | |
| for f in futures: | |
| f.cancel() | |
| break | |
| # Return only completed pages in order, dropping trailing Nones | |
| return [r for r in results if r is not None] | |
| # --------------------------------------------------------------------------- | |
| # VLM-based Semantic Validation | |
| # --------------------------------------------------------------------------- | |
| def validate_with_vlm( | |
| image: Image.Image, | |
| extraction_data: dict, | |
| hf_token: str, | |
| model_id: str = "Qwen/Qwen2.5-VL-72B-Instruct", | |
| use_case: str = "_default", | |
| ) -> list[dict]: | |
| """Use a VLM to cross-check extracted data against the source image. | |
| Returns a list of check dicts: {field, status, message, expected, actual}. | |
| """ | |
| if not hf_token: | |
| return [{ | |
| "name": "VLM validation", | |
| "status": "skipped", | |
| "message": "HF_TOKEN not set — cannot run VLM validation", | |
| }] | |
| prompt = get_prompt("validation", use_case) | |
| if not prompt: | |
| prompt = get_prompt("validation", "_default") | |
| # Build the extraction summary for the prompt | |
| fields = extraction_data.get("extracted_fields", {}) | |
| if not fields: | |
| # Legacy format | |
| fields = {k: v for k, v in extraction_data.items() | |
| if k not in ("line_items", "bounding_boxes", "raw_text", | |
| "extraction_method", "confidence", "processing_time_ms", | |
| "page_image", "page_images", "run_id", "error", | |
| "document_type", "extracted_fields")} | |
| extraction_json = json.dumps({ | |
| "fields": fields, | |
| "line_items": extraction_data.get("line_items", []), | |
| }, indent=2) | |
| full_prompt = f"{prompt}\n\nExtracted data to validate:\n```json\n{extraction_json}\n```" | |
| try: | |
| b64 = image_to_base64(image) | |
| client = create_inference_client(api_key=hf_token) | |
| response = robust_chat_completion( | |
| client, | |
| model=model_id, | |
| messages=[{ | |
| "role": "user", | |
| "content": [ | |
| {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}, | |
| {"type": "text", "text": full_prompt}, | |
| ], | |
| }], | |
| max_tokens=2048, | |
| ) | |
| raw = response.choices[0].message.content | |
| data = _parse_vlm_json(raw) | |
| checks = data.get("checks", []) | |
| # Normalize check format | |
| result = [] | |
| for c in checks: | |
| result.append({ | |
| "name": f"VLM: {c.get('field', 'unknown')}", | |
| "status": c.get("status", "warn"), | |
| "message": c.get("message", ""), | |
| "expected": c.get("expected"), | |
| "actual": c.get("actual"), | |
| }) | |
| return result if result else [{ | |
| "name": "VLM validation", | |
| "status": "pass", | |
| "message": data.get("summary", "All fields appear correct"), | |
| }] | |
| except Exception as e: | |
| logger.warning(f"VLM validation failed: {e}") | |
| return [{ | |
| "name": "VLM validation", | |
| "status": "skipped", | |
| "message": f"VLM validation unavailable: {str(e)[:100]}", | |
| }] | |
| def run_dynamic_math_checks(data: dict) -> list[dict]: | |
| """Run math/consistency checks dynamically based on extracted_fields. | |
| Works for any document type by detecting numeric relationships. | |
| """ | |
| checks = [] | |
| fields = data.get("extracted_fields", {}) | |
| if not fields: | |
| # Fall back to top-level keys | |
| fields = data | |
| line_items = data.get("line_items", []) | |
| # ── Collect numeric fields ── | |
| numeric = {} | |
| for k, v in fields.items(): | |
| if v is not None: | |
| try: | |
| numeric[k] = float(v) | |
| except (ValueError, TypeError): | |
| pass | |
| # ── Rule 1: subtotal + tax-like = total-like ── | |
| subtotal_key = next((k for k in numeric if "subtotal" in k), None) | |
| total_key = next((k for k in numeric if "total" in k and "sub" not in k), None) | |
| tax_keys = [k for k in numeric if any(t in k for t in ("tax", "cess", "surcharge", "duty", "vat", "gst"))] | |
| if subtotal_key and total_key and tax_keys: | |
| tax_sum = sum(numeric[k] for k in tax_keys) | |
| expected = numeric[subtotal_key] + tax_sum | |
| diff = abs(expected - numeric[total_key]) | |
| checks.append({ | |
| "name": "Subtotal + taxes = Total", | |
| "status": "pass" if diff < 0.02 else "warn" if diff < 1.0 else "fail", | |
| "message": f"{subtotal_key} + {' + '.join(tax_keys)} = {total_key}", | |
| "expected": numeric[total_key], | |
| "actual": round(expected, 2), | |
| }) | |
| # ── Rule 2: Line item amounts sum to subtotal or total ── | |
| if line_items: | |
| item_sum = sum( | |
| float(item.get("amount", 0) or 0) for item in line_items | |
| ) | |
| compare_key = subtotal_key or total_key | |
| if compare_key and item_sum > 0: | |
| diff = abs(item_sum - numeric[compare_key]) | |
| checks.append({ | |
| "name": "Line items sum", | |
| "status": "pass" if diff < 0.02 else "warn" if diff < 1.0 else "fail", | |
| "message": f"Sum of line item amounts vs {compare_key}", | |
| "expected": numeric[compare_key], | |
| "actual": round(item_sum, 2), | |
| }) | |
| # ── Rule 3: Each line item qty × unit_price = amount ── | |
| for i, item in enumerate(line_items): | |
| qty = item.get("quantity") | |
| price = item.get("unit_price") | |
| amount = item.get("amount") | |
| if qty is not None and price is not None and amount is not None: | |
| try: | |
| expected = round(float(qty) * float(price), 2) | |
| diff = abs(expected - float(amount)) | |
| checks.append({ | |
| "name": f"Line {i+1} math", | |
| "status": "pass" if diff < 0.02 else "fail", | |
| "message": f"qty ({qty}) × unit_price ({price}) = amount", | |
| "expected": float(amount), | |
| "actual": expected, | |
| }) | |
| except (ValueError, TypeError): | |
| pass | |
| # ── Rule 4: Total positive ── | |
| if total_key: | |
| checks.append({ | |
| "name": "Total > 0", | |
| "status": "pass" if numeric[total_key] > 0 else "fail", | |
| "message": f"{total_key} should be positive", | |
| "expected": "> 0", | |
| "actual": numeric[total_key], | |
| }) | |
| # ── Rule 5: PAN/TAN format (tax forms) ── | |
| for k, v in fields.items(): | |
| if isinstance(v, str): | |
| if k == "pan" and v: | |
| valid = bool(re.match(r'^[A-Z]{5}[0-9]{4}[A-Z]$', v)) | |
| checks.append({ | |
| "name": "PAN format", | |
| "status": "pass" if valid else "fail", | |
| "message": "PAN should be 5 letters + 4 digits + 1 letter", | |
| "expected": "XXXXX9999X", | |
| "actual": v, | |
| }) | |
| elif k == "tan" and v: | |
| valid = bool(re.match(r'^[A-Z]{4}[0-9]{5}[A-Z]$', v)) | |
| checks.append({ | |
| "name": "TAN format", | |
| "status": "pass" if valid else "fail", | |
| "message": "TAN should be 4 letters + 5 digits + 1 letter", | |
| "expected": "XXXX99999X", | |
| "actual": v, | |
| }) | |
| if not checks: | |
| checks.append({ | |
| "name": "Insufficient data", | |
| "status": "skipped", | |
| "message": "Not enough numeric or validatable fields to run checks", | |
| }) | |
| return checks | |