Fix clipped receipt margins, add time-of-day naming, speed up rotation detection
- vision/document.py: pad_quad() expands the detected quad outward before warping so a straight-line contour fit doesn't clip the first character of every line on a slightly curled/wrinkled receipt - ocr/extract.py + llm/vision.py: extract transaction time (HH:MM) so same-day repeat visits to a vendor get distinct filenames (2023-03-20_1525_walmart.pdf vs. a later same-day trip) - vision/orient.py: downscale before Tesseract OSD/confidence-sweep rotation detection — this was the dominant cost in the detection phase (7x speedup: 472s -> 67s on a 36s test video)
This commit is contained in:
@@ -196,3 +196,24 @@ def scale_quad(quad: np.ndarray, scale_x: float, scale_y: float) -> np.ndarray:
|
||||
quad[:, 0] *= scale_x
|
||||
quad[:, 1] *= scale_y
|
||||
return quad
|
||||
|
||||
|
||||
def pad_quad(
|
||||
quad: np.ndarray, image_shape: tuple[int, ...], margin_fraction: float = 0.04
|
||||
) -> np.ndarray:
|
||||
"""Expand a quad outward from its centroid by margin_fraction, clamped to the frame.
|
||||
|
||||
approxPolyDP fits a straight-sided polygon to a real paper edge that's
|
||||
often slightly curled or wrinkled, so the tightest-fit quad routinely
|
||||
clips a sliver of the page (a character or two off the left margin of
|
||||
every line, in the worst cases). A small outward pad recovers that
|
||||
margin; refine_crop's paper-band tightening trims any extra
|
||||
mat/background it pulls in back off afterward.
|
||||
"""
|
||||
h, w = image_shape[:2]
|
||||
quad = np.asarray(quad, dtype=np.float32)
|
||||
centroid = quad.mean(axis=0)
|
||||
padded = centroid + (quad - centroid) * (1.0 + margin_fraction)
|
||||
padded[:, 0] = np.clip(padded[:, 0], 0, w - 1)
|
||||
padded[:, 1] = np.clip(padded[:, 1], 0, h - 1)
|
||||
return padded
|
||||
|
||||
@@ -32,6 +32,19 @@ def apply_rotation(image_bgr: np.ndarray, rotation: int) -> np.ndarray:
|
||||
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
|
||||
return cv2.resize(image_bgr, (round(w * scale), round(h * scale)), 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.
|
||||
|
||||
@@ -43,7 +56,7 @@ def _osd_rotation(image_bgr: np.ndarray, lang: str) -> int | None:
|
||||
"""
|
||||
try:
|
||||
osd = pytesseract.image_to_osd(
|
||||
image_bgr, lang=lang, output_type=pytesseract.Output.DICT
|
||||
_downscale(image_bgr), lang=lang, output_type=pytesseract.Output.DICT
|
||||
)
|
||||
except pytesseract.TesseractError:
|
||||
return None
|
||||
@@ -72,10 +85,11 @@ def _confidence_sweep_rotation(
|
||||
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(image_bgr, rotation)
|
||||
candidate = apply_rotation(small, rotation)
|
||||
data = pytesseract.image_to_data(
|
||||
candidate, lang=lang, config=f"--psm {psm}", output_type=pytesseract.Output.DICT
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user