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,110 @@
|
||||
"""Auto-rotate a document crop to right-side-up before OCR/PDF export.
|
||||
|
||||
Handles the "I placed it upside down" case that broke naming on the real T4/
|
||||
T4A test recordings: images looked fine, but OCR read garbage because the
|
||||
form was rotated 180 degrees.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytesseract
|
||||
|
||||
_ROTATIONS = {
|
||||
0: None,
|
||||
90: cv2.ROTATE_90_CLOCKWISE,
|
||||
180: cv2.ROTATE_180,
|
||||
270: cv2.ROTATE_90_COUNTERCLOCKWISE,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrientationResult:
|
||||
rotation: int # degrees clockwise applied: 0, 90, 180, or 270
|
||||
method: str # "osd" | "confidence_sweep" | "none"
|
||||
|
||||
|
||||
def apply_rotation(image_bgr: np.ndarray, rotation: int) -> np.ndarray:
|
||||
code = _ROTATIONS.get(rotation % 360)
|
||||
return image_bgr if code is None else cv2.rotate(image_bgr, code)
|
||||
|
||||
|
||||
def _osd_rotation(image_bgr: np.ndarray, lang: str) -> int | None:
|
||||
"""Tesseract's orientation-and-script-detection: fast and purpose-built.
|
||||
|
||||
Raises on pages with too little recognizable text (e.g. a mostly-blank
|
||||
or badly-cropped receipt sliver); callers should fall back in that case.
|
||||
Note: OSD's own confidence scores run on a small, model-dependent scale
|
||||
(values well under 10 are common even when correct) so we trust any
|
||||
result it returns rather than gating on a confidence threshold.
|
||||
"""
|
||||
try:
|
||||
osd = pytesseract.image_to_osd(
|
||||
image_bgr, lang=lang, output_type=pytesseract.Output.DICT
|
||||
)
|
||||
except pytesseract.TesseractError:
|
||||
return None
|
||||
return int(osd.get("rotate", 0)) % 360
|
||||
|
||||
|
||||
def _confidence_sweep_rotation(
|
||||
image_bgr: np.ndarray, lang: str, psm: int, min_words: int = 3, min_word_len: int = 3
|
||||
) -> int:
|
||||
"""Fallback: OCR all 4 rotations, keep the one with the most confident text.
|
||||
|
||||
Used when OSD can't commit (too little text). Scores by *total*
|
||||
confidence summed over recognized words at least min_word_len long.
|
||||
|
||||
Both pieces matter and were tuned against real failures, not guessed:
|
||||
- Summing (not averaging) confidence: on noisy real-world photos
|
||||
(worn thermal print, low-contrast background), mean confidence
|
||||
stays within a similar narrow band at every rotation, but the
|
||||
correct orientation recognizes many more words than a sideways one.
|
||||
Mean-confidence-only picked a 90-degrees-off rotation (9 words) over
|
||||
the correct one (23 words) on a real low-contrast receipt.
|
||||
- Filtering to words >= min_word_len: sideways/upside-down *clean*
|
||||
synthetic text tends to fragment into dozens of single-character
|
||||
glyph artifacts that Tesseract reports with deceptively high
|
||||
individual confidence, which otherwise wins the sum outright.
|
||||
Requiring 3+ characters drops almost all of that noise while barely
|
||||
touching real multi-character words.
|
||||
"""
|
||||
best_rotation = 0
|
||||
best_score = -1.0
|
||||
for rotation in (0, 90, 180, 270):
|
||||
candidate = apply_rotation(image_bgr, rotation)
|
||||
data = pytesseract.image_to_data(
|
||||
candidate, lang=lang, config=f"--psm {psm}", output_type=pytesseract.Output.DICT
|
||||
)
|
||||
confidences = [
|
||||
float(c)
|
||||
for c, t in zip(data["conf"], data["text"])
|
||||
if c not in ("-1", -1) and len(t.strip()) >= min_word_len
|
||||
]
|
||||
if len(confidences) < min_words:
|
||||
continue
|
||||
total_conf = sum(confidences)
|
||||
if total_conf > best_score:
|
||||
best_score = total_conf
|
||||
best_rotation = rotation
|
||||
return best_rotation
|
||||
|
||||
|
||||
def detect_orientation(image_bgr: np.ndarray, lang: str = "eng", psm: int = 6) -> OrientationResult:
|
||||
"""Figure out how many degrees clockwise to rotate image_bgr to be upright."""
|
||||
rotation = _osd_rotation(image_bgr, lang)
|
||||
if rotation is not None:
|
||||
return OrientationResult(rotation=rotation, method="osd")
|
||||
rotation = _confidence_sweep_rotation(image_bgr, lang, psm)
|
||||
return OrientationResult(rotation=rotation, method="confidence_sweep")
|
||||
|
||||
|
||||
def auto_rotate(
|
||||
image_bgr: np.ndarray, lang: str = "eng", psm: int = 6
|
||||
) -> tuple[np.ndarray, OrientationResult]:
|
||||
"""Detect and correct orientation. Returns (rotated_image, result)."""
|
||||
result = detect_orientation(image_bgr, lang=lang, psm=psm)
|
||||
return apply_rotation(image_bgr, result.rotation), result
|
||||
Reference in New Issue
Block a user