Aniruddha7/QueryLens-Text2SQL_DocVQA-V2

1import os2try:3    # Ensure local .env values (OLLAMA_*) are loaded when this module is imported4    from dotenv import load_dotenv5    load_dotenv()6except Exception:7    pass8import uuid9import re10import base6411from typing import Dict, Any, Optional12 13def _mock_result(received: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:14    out = {"doc_id": str(uuid.uuid4()), "text": "(mock OCR output)", "metadata": {"pages": 1, "confidence": 0.0}}15    if received:16        out["received_args"] = received17    return out18 19 20def clean_ocr_text(raw: Optional[str]) -> str:21    """Lightweight post-processing for OCR output to improve readability.22 23    - turns literal "\\n" sequences into real newlines24    - unescapes common unicode/HTML escapes when safe25    - collapses repeated spaces and blank lines26    - trims and merges lines that were broken mid-sentence27    """28    if not raw:29        return ""30    s = raw31    try:32        # Convert literal backslash-n to actual newline33        s = s.replace('\\n', '\n')34    except Exception:35        pass36 37    # Try to interpret common escape sequences (conservative)38    try:39        s_decoded = s.encode('utf-8', errors='surrogatepass').decode('unicode_escape')40        # Accept decoded result if it reduces backslash artifacts41        if '\\n' not in s_decoded:42            s = s_decoded43    except Exception:44        pass45 46    # Normalize line endings and whitespace47    s = s.replace('\r\n', '\n').replace('\r', '\n')48    s = re.sub(r'[ \t]{2,}', ' ', s)49    # Collapse 3+ newlines into max 250    s = re.sub(r'\n{3,}', '\n\n', s)51 52    # Trim whitespace on each line53    lines = [ln.strip() for ln in s.splitlines()]54 55    # Merge lines that likely were broken mid-sentence: if a line doesn't end with56    # sentence punctuation and the next line begins with a lowercase letter or digit,57    # join them with a space.58    merged = []59    i = 060    while i < len(lines):61        line = lines[i]62        if i < len(lines) - 1:63            nxt = lines[i + 1]64            if line and nxt and not re.search(r'[\.\!\?\:;]$', line) and re.match(r'^[a-z0-9]', nxt):65                line = line + ' ' + nxt66                i += 267                # Continue merging subsequent similar lines68                while i < len(lines) and lines[i] and not re.search(r'[\.\!\?\:;]$', line) and re.match(r'^[a-z0-9]', lines[i]):69                    line = line + ' ' + lines[i]70                    i += 171                merged.append(line)72                continue73        merged.append(line)74        i += 175 76    s = '\n'.join([ln for ln in merged if ln])77 78    # Remove long runs of underscores or pipes79    s = re.sub(r'[_\|]{2,}', '', s)80 81    # Strip control chars except common whitespace82    s = ''.join(ch for ch in s if ord(ch) >= 9 and ord(ch) != 11 and ord(ch) != 12)83 84    try:85        import html86 87        s = html.unescape(s)88    except Exception:89        pass90 91    # Decode literal unicode escapes like \u0027 -> ' and \xA4 -> currency symbol92    try:93        def _u_decode(m):94            code = m.group(1)95            try:96                return chr(int(code, 16))97            except Exception:98                return m.group(0)99 100        s = re.sub(r'\\u([0-9A-Fa-f]{4})', _u_decode, s)101        s = re.sub(r'\\x([0-9A-Fa-f]{2})', lambda m: chr(int(m.group(1), 16)), s)102    except Exception:103        pass104 105    # Remove stray leading 'n' characters that appear before uppercase words (common OCR/newline artifact)106    try:107        s = re.sub(r'(?<=\s|^)[nN](?=[A-Z])', '', s)108    except Exception:109        pass110 111    # Ensure space after commas when missing112    try:113        s = re.sub(r',(?=[^\s])', ', ', s)114    except Exception:115        pass116 117    # Remove spaces before punctuation118    try:119        s = re.sub(r'\s+([,\.\:\;\%])', r'\1', s)120    except Exception:121        pass122 123    # Attempt to fix common mojibake (double-encoded UTF-8) conservatively.124    # Only apply if there are visible 'Ã' sequences suggesting encoding issues.125    try:126        if 'Ã' in s or 'Â' in s:127            try:128                candidate = s.encode('latin-1').decode('utf-8')129                # Accept candidate if it reduces suspicious sequences130                if candidate.count('Ã') < s.count('Ã'):131                    s = candidate132            except Exception:133                pass134    except Exception:135        pass136 137    return s.strip()138 139 140def process_image(image_url: Optional[str] = None, image_bytes: Optional[bytes] = None, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:141    """Process an image using local OCR first (pytesseract), with a light transformers fallback.142 143    Accepts either `image_url` or `image_bytes` (bytes or base64 string). Returns dict: {doc_id, text, metadata}144    """145 146    # Lightweight imports that may be missing in minimal environments147    try:148        from PIL import Image149        from io import BytesIO150        import requests151    except Exception:152        # Missing imaging/network deps -> return mock with echo of args153        return _mock_result(received={"image_url": image_url, "has_bytes": bool(image_bytes)})154 155    # Normalize image_bytes if base64 string was passed156    if isinstance(image_bytes, str):157        try:158            image_bytes = base64.b64decode(image_bytes)159        except Exception:160            image_bytes = None161 162    # Acquire PIL image163    img = None164    try:165        if image_bytes:166            img = Image.open(BytesIO(image_bytes)).convert("RGB")167        elif image_url:168            resp = requests.get(image_url, timeout=15)169            resp.raise_for_status()170            img = Image.open(BytesIO(resp.content)).convert("RGB")171        else:172            return _mock_result(received={"image_url": image_url, "has_bytes": False})173    except Exception:174        return _mock_result(received={"image_url": image_url, "has_bytes": bool(image_bytes)})175 176    # PRIORITY: Try local Tesseract (pytesseract) first — lightweight and deterministic on Windows177 178    try:179        import pytesseract180        # Robustly find tesseract executable: prefer PATH, then common Windows install locations181        try:182            import shutil183            found = shutil.which('tesseract')184        except Exception:185            found = None186        if not found:187            # common install locations on Windows188            for candidate in (r"C:\Program Files\Tesseract-OCR\tesseract.exe", r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe"):189                try:190                    if os.path.exists(candidate):191                        found = candidate192                        break193                except Exception:194                    continue195        try:196            if found:197                pytesseract.pytesseract.tesseract_cmd = found198                print(f"[document_scanner] using tesseract executable: {found}")199            else:200                # leave default; pytesseract will raise a helpful error which we catch below201                print("[document_scanner] tesseract executable not found in PATH or common locations")202        except Exception:203            # ignore assignment errors and let pytesseract raise on use204            pass205 206        try:207            if img is not None:208                print("[document_scanner] attempting PRIMARY pytesseract OCR")209                ttxt = pytesseract.image_to_string(img)210                if ttxt and ttxt.strip():211                    print(f"[document_scanner] pytesseract returned text (len={len(ttxt.strip())})")212                    return {"doc_id": str(uuid.uuid4()), "text": clean_ocr_text(ttxt), "metadata": {"pages": 1, "confidence": None}}213                else:214                    print("[document_scanner] pytesseract returned no text; falling through to model-based fallbacks")215        except Exception as e_pyt_first:216            print(f"[document_scanner] primary pytesseract attempt failed: {e_pyt_first}")217    except Exception:218        # pytesseract not installed in environment — continue to available backends219        pass220 221    # Last resort: return a safe mock that echoes inputs222    # (Tesseract failed or is not installed — Q&A will still work via Granite Vision if image is saved)223    return _mock_result(received={"image_url": image_url, "has_bytes": bool(image_bytes)})224