Spaces:
Running on Zero
Running on Zero
| import os | |
| try: | |
| # Ensure local .env values (OLLAMA_*) are loaded when this module is imported | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| except Exception: | |
| pass | |
| import uuid | |
| import re | |
| import base64 | |
| from typing import Dict, Any, Optional | |
| def _mock_result(received: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: | |
| out = {"doc_id": str(uuid.uuid4()), "text": "(mock OCR output)", "metadata": {"pages": 1, "confidence": 0.0}} | |
| if received: | |
| out["received_args"] = received | |
| return out | |
| def clean_ocr_text(raw: Optional[str]) -> str: | |
| """Lightweight post-processing for OCR output to improve readability. | |
| - turns literal "\\n" sequences into real newlines | |
| - unescapes common unicode/HTML escapes when safe | |
| - collapses repeated spaces and blank lines | |
| - trims and merges lines that were broken mid-sentence | |
| """ | |
| if not raw: | |
| return "" | |
| s = raw | |
| try: | |
| # Convert literal backslash-n to actual newline | |
| s = s.replace('\\n', '\n') | |
| except Exception: | |
| pass | |
| # Try to interpret common escape sequences (conservative) | |
| try: | |
| s_decoded = s.encode('utf-8', errors='surrogatepass').decode('unicode_escape') | |
| # Accept decoded result if it reduces backslash artifacts | |
| if '\\n' not in s_decoded: | |
| s = s_decoded | |
| except Exception: | |
| pass | |
| # Normalize line endings and whitespace | |
| s = s.replace('\r\n', '\n').replace('\r', '\n') | |
| s = re.sub(r'[ \t]{2,}', ' ', s) | |
| # Collapse 3+ newlines into max 2 | |
| s = re.sub(r'\n{3,}', '\n\n', s) | |
| # Trim whitespace on each line | |
| lines = [ln.strip() for ln in s.splitlines()] | |
| # Merge lines that likely were broken mid-sentence: if a line doesn't end with | |
| # sentence punctuation and the next line begins with a lowercase letter or digit, | |
| # join them with a space. | |
| merged = [] | |
| i = 0 | |
| while i < len(lines): | |
| line = lines[i] | |
| if i < len(lines) - 1: | |
| nxt = lines[i + 1] | |
| if line and nxt and not re.search(r'[\.\!\?\:;]$', line) and re.match(r'^[a-z0-9]', nxt): | |
| line = line + ' ' + nxt | |
| i += 2 | |
| # Continue merging subsequent similar lines | |
| while i < len(lines) and lines[i] and not re.search(r'[\.\!\?\:;]$', line) and re.match(r'^[a-z0-9]', lines[i]): | |
| line = line + ' ' + lines[i] | |
| i += 1 | |
| merged.append(line) | |
| continue | |
| merged.append(line) | |
| i += 1 | |
| s = '\n'.join([ln for ln in merged if ln]) | |
| # Remove long runs of underscores or pipes | |
| s = re.sub(r'[_\|]{2,}', '', s) | |
| # Strip control chars except common whitespace | |
| s = ''.join(ch for ch in s if ord(ch) >= 9 and ord(ch) != 11 and ord(ch) != 12) | |
| try: | |
| import html | |
| s = html.unescape(s) | |
| except Exception: | |
| pass | |
| # Decode literal unicode escapes like \u0027 -> ' and \xA4 -> currency symbol | |
| try: | |
| def _u_decode(m): | |
| code = m.group(1) | |
| try: | |
| return chr(int(code, 16)) | |
| except Exception: | |
| return m.group(0) | |
| s = re.sub(r'\\u([0-9A-Fa-f]{4})', _u_decode, s) | |
| s = re.sub(r'\\x([0-9A-Fa-f]{2})', lambda m: chr(int(m.group(1), 16)), s) | |
| except Exception: | |
| pass | |
| # Remove stray leading 'n' characters that appear before uppercase words (common OCR/newline artifact) | |
| try: | |
| s = re.sub(r'(?<=\s|^)[nN](?=[A-Z])', '', s) | |
| except Exception: | |
| pass | |
| # Ensure space after commas when missing | |
| try: | |
| s = re.sub(r',(?=[^\s])', ', ', s) | |
| except Exception: | |
| pass | |
| # Remove spaces before punctuation | |
| try: | |
| s = re.sub(r'\s+([,\.\:\;\%])', r'\1', s) | |
| except Exception: | |
| pass | |
| # Attempt to fix common mojibake (double-encoded UTF-8) conservatively. | |
| # Only apply if there are visible 'Ã' sequences suggesting encoding issues. | |
| try: | |
| if 'Ã' in s or 'Â' in s: | |
| try: | |
| candidate = s.encode('latin-1').decode('utf-8') | |
| # Accept candidate if it reduces suspicious sequences | |
| if candidate.count('Ã') < s.count('Ã'): | |
| s = candidate | |
| except Exception: | |
| pass | |
| except Exception: | |
| pass | |
| return s.strip() | |
| def process_image(image_url: Optional[str] = None, image_bytes: Optional[bytes] = None, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: | |
| """Process an image using local OCR first (pytesseract), with a light transformers fallback. | |
| Accepts either `image_url` or `image_bytes` (bytes or base64 string). Returns dict: {doc_id, text, metadata} | |
| """ | |
| # Lightweight imports that may be missing in minimal environments | |
| try: | |
| from PIL import Image | |
| from io import BytesIO | |
| import requests | |
| except Exception: | |
| # Missing imaging/network deps -> return mock with echo of args | |
| return _mock_result(received={"image_url": image_url, "has_bytes": bool(image_bytes)}) | |
| # Normalize image_bytes if base64 string was passed | |
| if isinstance(image_bytes, str): | |
| try: | |
| image_bytes = base64.b64decode(image_bytes) | |
| except Exception: | |
| image_bytes = None | |
| # Acquire PIL image | |
| img = None | |
| try: | |
| if image_bytes: | |
| img = Image.open(BytesIO(image_bytes)).convert("RGB") | |
| elif image_url: | |
| resp = requests.get(image_url, timeout=15) | |
| resp.raise_for_status() | |
| img = Image.open(BytesIO(resp.content)).convert("RGB") | |
| else: | |
| return _mock_result(received={"image_url": image_url, "has_bytes": False}) | |
| except Exception: | |
| return _mock_result(received={"image_url": image_url, "has_bytes": bool(image_bytes)}) | |
| # PRIORITY: Try local Tesseract (pytesseract) first — lightweight and deterministic on Windows | |
| try: | |
| import pytesseract | |
| # Robustly find tesseract executable: prefer PATH, then common Windows install locations | |
| try: | |
| import shutil | |
| found = shutil.which('tesseract') | |
| except Exception: | |
| found = None | |
| if not found: | |
| # common install locations on Windows | |
| for candidate in (r"C:\Program Files\Tesseract-OCR\tesseract.exe", r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe"): | |
| try: | |
| if os.path.exists(candidate): | |
| found = candidate | |
| break | |
| except Exception: | |
| continue | |
| try: | |
| if found: | |
| pytesseract.pytesseract.tesseract_cmd = found | |
| print(f"[document_scanner] using tesseract executable: {found}") | |
| else: | |
| # leave default; pytesseract will raise a helpful error which we catch below | |
| print("[document_scanner] tesseract executable not found in PATH or common locations") | |
| except Exception: | |
| # ignore assignment errors and let pytesseract raise on use | |
| pass | |
| try: | |
| if img is not None: | |
| print("[document_scanner] attempting PRIMARY pytesseract OCR") | |
| ttxt = pytesseract.image_to_string(img) | |
| if ttxt and ttxt.strip(): | |
| print(f"[document_scanner] pytesseract returned text (len={len(ttxt.strip())})") | |
| return {"doc_id": str(uuid.uuid4()), "text": clean_ocr_text(ttxt), "metadata": {"pages": 1, "confidence": None}} | |
| else: | |
| print("[document_scanner] pytesseract returned no text; falling through to model-based fallbacks") | |
| except Exception as e_pyt_first: | |
| print(f"[document_scanner] primary pytesseract attempt failed: {e_pyt_first}") | |
| except Exception: | |
| # pytesseract not installed in environment — continue to available backends | |
| pass | |
| # Last resort: return a safe mock that echoes inputs | |
| # (Tesseract failed or is not installed — Q&A will still work via Granite Vision if image is saved) | |
| return _mock_result(received={"image_url": image_url, "has_bytes": bool(image_bytes)}) | |