ilia
9af29dc614
OCR enhancement chain (illumination normalize, CLAHE, gamma darkening) now always runs before Tesseract, with a conditional 1.5x upscale for narrow crops. Adds printed doc-title and PAGE X OF Y marker extraction.
127 lines
5.0 KiB
Python
127 lines
5.0 KiB
Python
"""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 _downscale(image_bgr: np.ndarray, max_dim: int = 1000) -> np.ndarray:
|
|
"""Shrink before feeding Tesseract: orientation only needs word shapes,
|
|
not per-character resolution, and _confidence_sweep_rotation runs this
|
|
up to 4x per document — at full 4K-crop resolution that made rotation
|
|
detection by far the slowest part of the whole pipeline."""
|
|
h, w = image_bgr.shape[:2]
|
|
longest = max(h, w)
|
|
if longest <= max_dim:
|
|
return image_bgr
|
|
scale = max_dim / longest
|
|
new_w = max(1, round(w * scale))
|
|
new_h = max(1, round(h * scale))
|
|
return cv2.resize(image_bgr, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
|
|
|
|
|
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(
|
|
_downscale(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.
|
|
"""
|
|
small = _downscale(image_bgr)
|
|
best_rotation = 0
|
|
best_score = -1.0
|
|
for rotation in (0, 90, 180, 270):
|
|
candidate = apply_rotation(small, 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
|