Add crop refinement, finger removal, and local LLM naming
- vision/refine.py: tighten crops to the paper band (removes mat margins and hands beside receipts) and inpaint border-connected skin regions so fingers disappear from output - llm/vision.py: identify documents with a local Ollama vision model (qwen2.5vl); extracts vendor/date/total/form code and flags quality issues (fingers, blur, glare); falls back to Tesseract when down - pipeline: drop blank pages, dedupe consecutive captures of the same document, record refine/LLM fields in report and export summary
This commit is contained in:
@@ -0,0 +1,135 @@
|
||||
"""Identify documents with a local vision LLM served by Ollama.
|
||||
|
||||
Everything stays on-machine: images go to localhost:11434, never the cloud.
|
||||
A vision model (qwen2.5-vl class) reads the whole document at once, which is
|
||||
far more robust than Tesseract + regex for curled receipts, hands in frame,
|
||||
and stylized logos — it correctly reads "COSTCO WHOLESALE" where Tesseract
|
||||
produced "Cosrco". It also reports quality issues (fingers, blur, glare) so
|
||||
bad captures can be flagged for re-scanning.
|
||||
|
||||
Trade-off: ~20-40s per document on an M-series Mac, vs <2s for Tesseract.
|
||||
The pipeline uses the LLM for naming when available and falls back to the
|
||||
OCR heuristics when Ollama isn't running.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
_PROMPT = """Look at this image of a document. Reply with ONLY a JSON object, no other text:
|
||||
{"doc_type": "receipt|tax_form|letter|invoice|blank|other",
|
||||
"vendor": "business or organization name, or null",
|
||||
"form_code": "tax form code like T4/T4A/T5008 if this is a tax form, else null",
|
||||
"date": "document/transaction date as YYYY-MM-DD, or null",
|
||||
"tax_year": "tax year like 2023 if this is a tax form, else null",
|
||||
"total": "total amount like 103.60, or null",
|
||||
"description": "3-6 word summary",
|
||||
"issues": ["any of: fingers_visible, blurry, cut_off, upside_down, glare, partial_document"]}
|
||||
Use doc_type "blank" if the page has no meaningful printed content (blank back of a page, empty sheet).
|
||||
Set vendor to null unless you can actually read a real business/organization name — never output garbled text."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class LlmAnalysis:
|
||||
doc_type: str
|
||||
vendor: str | None
|
||||
form_code: str | None
|
||||
date: str | None # normalized YYYY-MM-DD
|
||||
tax_year: str | None
|
||||
total: str | None
|
||||
description: str
|
||||
issues: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def ollama_available(base_url: str = "http://localhost:11434", timeout: float = 2.0) -> bool:
|
||||
try:
|
||||
with urllib.request.urlopen(f"{base_url}/api/tags", timeout=timeout):
|
||||
return True
|
||||
except (urllib.error.URLError, OSError):
|
||||
return False
|
||||
|
||||
|
||||
def _encode_image(image_bgr: np.ndarray, max_width: int = 1200) -> str:
|
||||
"""JPEG-encode (downscaled) for the API; PNG of a 4K crop is wastefully big."""
|
||||
h, w = image_bgr.shape[:2]
|
||||
if w > max_width:
|
||||
image_bgr = cv2.resize(image_bgr, (max_width, round(h * max_width / w)))
|
||||
ok, buf = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 90])
|
||||
if not ok:
|
||||
raise ValueError("Could not JPEG-encode image")
|
||||
return base64.b64encode(buf.tobytes()).decode()
|
||||
|
||||
|
||||
def _normalize_date(raw: str | None) -> str | None:
|
||||
if not raw:
|
||||
return None
|
||||
match = re.search(r"(\d{4})[-/](\d{1,2})[-/](\d{1,2})", str(raw))
|
||||
if not match:
|
||||
return None
|
||||
y, m, d = (int(g) for g in match.groups())
|
||||
if not (1 <= m <= 12 and 1 <= d <= 31 and 1990 <= y <= 2100):
|
||||
return None
|
||||
return f"{y:04d}-{m:02d}-{d:02d}"
|
||||
|
||||
|
||||
def _clean_str(value) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value).strip()
|
||||
return text if text and text.lower() not in ("null", "none", "n/a", "unknown") else None
|
||||
|
||||
|
||||
def analyze_document(
|
||||
image_bgr: np.ndarray,
|
||||
model: str = "qwen2.5vl",
|
||||
base_url: str = "http://localhost:11434",
|
||||
timeout: float = 180.0,
|
||||
) -> LlmAnalysis | None:
|
||||
"""Ask the local vision model what this document is. None on any failure."""
|
||||
try:
|
||||
body = json.dumps(
|
||||
{
|
||||
"model": model,
|
||||
"prompt": _PROMPT,
|
||||
"images": [_encode_image(image_bgr)],
|
||||
"stream": False,
|
||||
"format": "json",
|
||||
"options": {"temperature": 0},
|
||||
}
|
||||
).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{base_url}/api/generate", data=body, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
payload = json.load(resp)
|
||||
data = json.loads(payload["response"])
|
||||
except (urllib.error.URLError, OSError, KeyError, ValueError):
|
||||
return None
|
||||
|
||||
issues = data.get("issues") or []
|
||||
if not isinstance(issues, list):
|
||||
issues = [str(issues)]
|
||||
|
||||
tax_year = _clean_str(data.get("tax_year"))
|
||||
if tax_year:
|
||||
year_match = re.search(r"(19|20)\d{2}", tax_year)
|
||||
tax_year = year_match.group(0) if year_match else None
|
||||
|
||||
return LlmAnalysis(
|
||||
doc_type=_clean_str(data.get("doc_type")) or "other",
|
||||
vendor=_clean_str(data.get("vendor")),
|
||||
form_code=_clean_str(data.get("form_code")),
|
||||
date=_normalize_date(data.get("date")),
|
||||
tax_year=tax_year,
|
||||
total=_clean_str(data.get("total")),
|
||||
description=_clean_str(data.get("description")) or "",
|
||||
issues=[str(i) for i in issues],
|
||||
)
|
||||
Reference in New Issue
Block a user