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:
parent
a43ab20df6
commit
5be9d5935d
@ -47,6 +47,11 @@ document:
|
||||
enhance: false
|
||||
# Auto-detect and correct 0/90/180/270 rotation before OCR/PDF export.
|
||||
auto_rotate: true
|
||||
# Expand the detected quad outward by this fraction before warping, to
|
||||
# recover paper margins clipped by curl/wrinkles (straight-line quad vs.
|
||||
# a slightly bowed real edge). refine.tighten below trims any extra
|
||||
# mat/background this pulls in back off.
|
||||
pad_margin: 0.04
|
||||
# Tighten crops to the actual paper region (cuts mat margins/hands) and
|
||||
# inpaint fingers pressing on the page.
|
||||
refine: true
|
||||
|
||||
@ -52,6 +52,9 @@ class DocumentConfig:
|
||||
# Fixes the "placed the document upside down" case, which otherwise
|
||||
# produces a fine-looking crop that OCR reads as garbage.
|
||||
auto_rotate: bool = True
|
||||
# Expand the detected quad outward by this fraction before warping, to
|
||||
# recover paper margins clipped by curl/wrinkles (see pad_quad).
|
||||
pad_margin: float = 0.04
|
||||
# Tighten each crop to the actual paper region (removes mat margins and
|
||||
# hands hanging off the side of receipts) and inpaint fingers pressing
|
||||
# on the page. Content *under* a finger can't be recovered — it's filled
|
||||
|
||||
@ -29,6 +29,7 @@ _PROMPT = """Look at this image of a document. Reply with ONLY a JSON object, no
|
||||
"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",
|
||||
"time": "transaction time as 24-hour HH:MM if printed on the document, else 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",
|
||||
@ -43,6 +44,7 @@ class LlmAnalysis:
|
||||
vendor: str | None
|
||||
form_code: str | None
|
||||
date: str | None # normalized YYYY-MM-DD
|
||||
time: str | None # normalized HH:MM (24h)
|
||||
tax_year: str | None
|
||||
total: str | None
|
||||
description: str
|
||||
@ -87,6 +89,15 @@ def _clean_str(value) -> str | None:
|
||||
return text if text and text.lower() not in ("null", "none", "n/a", "unknown") else None
|
||||
|
||||
|
||||
def _normalize_time(raw: str | None) -> str | None:
|
||||
if not raw:
|
||||
return None
|
||||
match = re.search(r"([01]?\d|2[0-3]):([0-5]\d)", str(raw))
|
||||
if not match:
|
||||
return None
|
||||
return f"{int(match.group(1)):02d}:{match.group(2)}"
|
||||
|
||||
|
||||
def analyze_document(
|
||||
image_bgr: np.ndarray,
|
||||
model: str = "qwen2.5vl",
|
||||
@ -128,6 +139,7 @@ def analyze_document(
|
||||
vendor=_clean_str(data.get("vendor")),
|
||||
form_code=_clean_str(data.get("form_code")),
|
||||
date=_normalize_date(data.get("date")),
|
||||
time=_normalize_time(data.get("time")),
|
||||
tax_year=tax_year,
|
||||
total=_clean_str(data.get("total")),
|
||||
description=_clean_str(data.get("description")) or "",
|
||||
|
||||
@ -27,6 +27,9 @@ class NamingResult:
|
||||
source: str # "ocr" | "none" (will include "speech" once module 4 exists)
|
||||
|
||||
|
||||
_TIME_RE = re.compile(r"^([01]\d|2[0-3]):([0-5]\d)$")
|
||||
|
||||
|
||||
def build_name(
|
||||
date: str | None,
|
||||
vendor: str | None,
|
||||
@ -35,11 +38,21 @@ def build_name(
|
||||
form_code: str | None = None,
|
||||
tax_year: str | None = None,
|
||||
org_name: str | None = None,
|
||||
time: str | None = None,
|
||||
) -> NamingResult:
|
||||
"""date must already be normalized to YYYY-MM-DD, or None."""
|
||||
"""date must already be normalized to YYYY-MM-DD, time to HH:MM, or None.
|
||||
|
||||
Time is appended to the date (e.g. 2023-03-20_1525) whenever it was
|
||||
extracted, so two same-day visits to the same vendor (Walmart in the
|
||||
morning and again in the evening) don't collide or silently overwrite.
|
||||
"""
|
||||
vendor_slug = slugify(vendor, max_vendor_len) if vendor else None
|
||||
ts = fallback_dt.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
date_part = date
|
||||
if date and time and _TIME_RE.match(time):
|
||||
date_part = f"{date}_{time.replace(':', '')}"
|
||||
|
||||
# Tax slips: prefer 2023_T4_York_University over UNSORTED_t4_york_university_<timestamp>.
|
||||
if form_code:
|
||||
org_slug = slugify(org_name, max_vendor_len) if org_name else None
|
||||
@ -58,10 +71,10 @@ def build_name(
|
||||
return NamingResult(stem=f"{code}_{org_slug}", source="ocr")
|
||||
return NamingResult(stem=code.lower(), source="ocr")
|
||||
|
||||
if date and vendor_slug:
|
||||
return NamingResult(stem=f"{date}_{vendor_slug}", source="ocr")
|
||||
if date:
|
||||
return NamingResult(stem=f"{date}_UNSORTED", source="ocr")
|
||||
if date_part and vendor_slug:
|
||||
return NamingResult(stem=f"{date_part}_{vendor_slug}", source="ocr")
|
||||
if date_part:
|
||||
return NamingResult(stem=f"{date_part}_UNSORTED", source="ocr")
|
||||
if vendor_slug:
|
||||
return NamingResult(stem=f"UNSORTED_{vendor_slug}_{ts}", source="ocr")
|
||||
return NamingResult(stem=f"UNSORTED_{ts}", source="none")
|
||||
|
||||
@ -42,6 +42,12 @@ _DATE_PATTERNS: list[tuple[re.Pattern, str]] = [
|
||||
_TOTAL_LINE_RE = re.compile(r"\btotal\b", re.IGNORECASE)
|
||||
_MONEY_RE = re.compile(r"\$?\s*(\d{1,4}\.\d{2})\b")
|
||||
|
||||
# HH:MM(:SS)? 24-hour clock, as printed near the transaction date on most
|
||||
# receipts (e.g. "03/20/23 15:25:35"). Distinct enough from other
|
||||
# colon-free receipt numbers (barcodes, reference IDs) that the first match
|
||||
# in the document is reliably the transaction time.
|
||||
_TIME_RE = re.compile(r"\b([01]?\d|2[0-3]):([0-5]\d)(?::[0-5]\d)?")
|
||||
|
||||
# Common Canadian tax-slip codes, longest/most-specific first so "T4A" and
|
||||
# "T5008" match before the more generic "T4"/"T5" would. Receipts and
|
||||
# letters won't match any of these, so this only ever affects naming for
|
||||
@ -148,6 +154,7 @@ class OcrResult:
|
||||
text: str
|
||||
vendor: str | None
|
||||
date: str | None # normalized YYYY-MM-DD
|
||||
time: str | None # normalized HH:MM (24h), e.g. "15:25"
|
||||
total: str | None # e.g. "81.52"
|
||||
mean_confidence: float # 0-100, avg Tesseract word confidence
|
||||
form_code: str | None = None # e.g. "T4", "T5008", if a known tax slip
|
||||
@ -187,6 +194,15 @@ def extract_date(text: str) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def extract_time(text: str) -> str | None:
|
||||
"""First HH:MM clock time found in OCR text, normalized to zero-padded 24h."""
|
||||
match = _TIME_RE.search(text)
|
||||
if not match:
|
||||
return None
|
||||
hour, minute = int(match.group(1)), int(match.group(2))
|
||||
return f"{hour:02d}:{minute:02d}"
|
||||
|
||||
|
||||
def extract_total(text: str) -> str | None:
|
||||
"""Dollar amount on a line containing "total" but not "subtotal"."""
|
||||
lines = text.splitlines()
|
||||
@ -295,6 +311,7 @@ def run_ocr(
|
||||
psm=psm,
|
||||
)
|
||||
date = extract_date(text)
|
||||
time = extract_time(text)
|
||||
total = extract_total(text)
|
||||
tax_year = extract_tax_year(text)
|
||||
|
||||
@ -311,6 +328,7 @@ def run_ocr(
|
||||
text=text,
|
||||
vendor=vendor,
|
||||
date=date,
|
||||
time=time,
|
||||
total=total,
|
||||
mean_confidence=mean_confidence,
|
||||
form_code=form_code,
|
||||
|
||||
@ -25,7 +25,7 @@ from paperpod.naming.filename import build_name, dedupe_stem
|
||||
from paperpod.ocr.extract import run_ocr, extract_store_name
|
||||
from paperpod.pdf.build import images_to_pdf
|
||||
from paperpod.vision.capture_events import _plausible_capture, find_document_events
|
||||
from paperpod.vision.document import detect_document, warp_document
|
||||
from paperpod.vision.document import detect_document, pad_quad, warp_document
|
||||
from paperpod.vision.enhance import enhance_for_ocr
|
||||
from paperpod.vision.motion import find_segments, motion_score, prepare_gray
|
||||
from paperpod.vision.orient import auto_rotate
|
||||
@ -132,7 +132,8 @@ def run_detection(video_path: str | Path, cfg: Config, out_dir: str | Path) -> d
|
||||
cv2.imwrite(str(frame_path), frame.image)
|
||||
record["frame_path"] = str(frame_path.relative_to(out_dir))
|
||||
|
||||
crop = warp_document(frame.image, detection.quad)
|
||||
padded_quad = pad_quad(detection.quad, frame.image.shape, cfg.document.pad_margin)
|
||||
crop = warp_document(frame.image, padded_quad)
|
||||
|
||||
refine_info: dict[str, Any] = {}
|
||||
if cfg.document.refine:
|
||||
@ -285,6 +286,7 @@ def run_export(report: dict[str, Any], cfg: Config, out_dir: str | Path) -> dict
|
||||
|
||||
vendor = (llm.vendor if llm else None) or ocr.vendor
|
||||
date = (llm.date if llm else None) or ocr.date
|
||||
time = (llm.time if llm else None) or ocr.time
|
||||
total = (llm.total if llm else None) or ocr.total
|
||||
form_code = (llm.form_code if llm else None) or ocr.form_code
|
||||
tax_year = (llm.tax_year if llm else None) or ocr.tax_year
|
||||
@ -312,6 +314,7 @@ def run_export(report: dict[str, Any], cfg: Config, out_dir: str | Path) -> dict
|
||||
form_code=form_code,
|
||||
tax_year=tax_year,
|
||||
org_name=org_name,
|
||||
time=time,
|
||||
)
|
||||
stem = dedupe_stem(naming.stem, used_stems)
|
||||
|
||||
@ -327,6 +330,7 @@ def run_export(report: dict[str, Any], cfg: Config, out_dir: str | Path) -> dict
|
||||
"ocr_vendor": vendor or "",
|
||||
"ocr_form_code": form_code or "",
|
||||
"ocr_date": date or "",
|
||||
"ocr_time": time or "",
|
||||
"ocr_total": total or "",
|
||||
"ocr_confidence": round(ocr.mean_confidence, 1),
|
||||
"llm_description": llm.description if llm else "",
|
||||
@ -341,7 +345,7 @@ def run_export(report: dict[str, Any], cfg: Config, out_dir: str | Path) -> dict
|
||||
with open(out_dir / "export_summary.csv", "w", newline="") as f:
|
||||
fieldnames = [
|
||||
"timestamp", "pod_id", "page_count", "source_of_name",
|
||||
"ocr_vendor", "ocr_form_code", "ocr_date", "ocr_total", "ocr_confidence",
|
||||
"ocr_vendor", "ocr_form_code", "ocr_date", "ocr_time", "ocr_total", "ocr_confidence",
|
||||
"llm_description", "llm_issues",
|
||||
"detection_confidence", "rotation_applied",
|
||||
"final_filename", "pdf_path",
|
||||
|
||||
@ -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
|
||||
)
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from paperpod.vision.document import detect_document, warp_document
|
||||
from paperpod.vision.document import detect_document, pad_quad, warp_document
|
||||
from paperpod.vision.sharpness import sharpness_score
|
||||
|
||||
|
||||
@ -77,3 +77,20 @@ def test_downscaled_detection_matches_full_resolution():
|
||||
det_small = detect_document(frame, detect_width=640)
|
||||
assert det_full is not None and det_small is not None
|
||||
assert abs(det_full.area_ratio - det_small.area_ratio) < 0.05
|
||||
|
||||
|
||||
def test_pad_quad_expands_outward():
|
||||
quad = np.array([[100, 100], [300, 100], [300, 300], [100, 300]], dtype=np.float32)
|
||||
padded = pad_quad(quad, image_shape=(720, 1280), margin_fraction=0.1)
|
||||
centroid = quad.mean(axis=0)
|
||||
# Every corner should move further from the centroid, not closer.
|
||||
orig_dist = np.linalg.norm(quad - centroid, axis=1)
|
||||
new_dist = np.linalg.norm(padded - centroid, axis=1)
|
||||
assert np.all(new_dist > orig_dist)
|
||||
|
||||
|
||||
def test_pad_quad_clamps_to_frame_bounds():
|
||||
quad = np.array([[5, 5], [1275, 5], [1275, 715], [5, 715]], dtype=np.float32)
|
||||
padded = pad_quad(quad, image_shape=(720, 1280), margin_fraction=0.2)
|
||||
assert padded[:, 0].min() >= 0 and padded[:, 0].max() <= 1279
|
||||
assert padded[:, 1].min() >= 0 and padded[:, 1].max() <= 719
|
||||
|
||||
@ -46,6 +46,27 @@ def test_build_name_nothing_found():
|
||||
assert result.source == "none"
|
||||
|
||||
|
||||
def test_build_name_with_time_disambiguates_same_day_visits():
|
||||
dt = datetime(2026, 1, 1, 12, 0, 0)
|
||||
morning = build_name("2023-03-20", "WALMART", dt, time="09:05")
|
||||
evening = build_name("2023-03-20", "WALMART", dt, time="18:42")
|
||||
assert morning.stem == "2023-03-20_0905_walmart"
|
||||
assert evening.stem == "2023-03-20_1842_walmart"
|
||||
assert morning.stem != evening.stem
|
||||
|
||||
|
||||
def test_build_name_invalid_time_ignored():
|
||||
dt = datetime(2026, 1, 1, 12, 0, 0)
|
||||
result = build_name("2023-03-14", "METRO", dt, time="not-a-time")
|
||||
assert result.stem == "2023-03-14_metro"
|
||||
|
||||
|
||||
def test_build_name_time_without_date_ignored():
|
||||
dt = datetime(2026, 1, 1, 12, 0, 0)
|
||||
result = build_name(None, "METRO", dt, time="09:05")
|
||||
assert result.stem == "UNSORTED_metro_20260101_120000"
|
||||
|
||||
|
||||
def test_build_name_tax_form_with_year_and_org():
|
||||
dt = datetime(2026, 1, 1, 12, 0, 0)
|
||||
result = build_name(
|
||||
|
||||
@ -5,7 +5,7 @@ deterministic and don't depend on Tesseract's OCR accuracy. End-to-end OCR
|
||||
accuracy against rendered receipts is exercised in test_ocr_integration.py.
|
||||
"""
|
||||
|
||||
from paperpod.ocr.extract import extract_date, extract_form_code, extract_org_name, extract_store_name, extract_tax_year, extract_total
|
||||
from paperpod.ocr.extract import extract_date, extract_form_code, extract_org_name, extract_store_name, extract_tax_year, extract_time, extract_total
|
||||
|
||||
|
||||
def test_extract_date_iso():
|
||||
@ -55,6 +55,26 @@ def test_extract_total_none_when_absent():
|
||||
assert extract_total("no totals here") is None
|
||||
|
||||
|
||||
def test_extract_time_hms():
|
||||
assert extract_time("03/20/23 15:25:35") == "15:25"
|
||||
|
||||
|
||||
def test_extract_time_hm():
|
||||
assert extract_time("Transaction at 09:05 today") == "09:05"
|
||||
|
||||
|
||||
def test_extract_time_zero_pads_single_digit_hour():
|
||||
assert extract_time("open at 9:05am") == "09:05"
|
||||
|
||||
|
||||
def test_extract_time_none_when_absent():
|
||||
assert extract_time("no time here, just $103.60 total") is None
|
||||
|
||||
|
||||
def test_extract_time_ignores_colonless_numbers():
|
||||
assert extract_time("TC# 3821 2498 6394 0609 7307") is None
|
||||
|
||||
|
||||
def test_extract_form_code_t4():
|
||||
text = "Employer's name\nYORK UNIVERSITY\nT4\nStatement of Remuneration Paid"
|
||||
assert extract_form_code(text) == "T4"
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user