ilia
5be9d5935d
- 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)
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
from paperpod.vision.document import detect_document, pad_quad, warp_document
|
|
from paperpod.vision.sharpness import sharpness_score
|
|
|
|
|
|
def synthetic_frame(angle: float = 8.0) -> np.ndarray:
|
|
"""Dark table with a rotated white 400x600 'document'."""
|
|
frame = np.full((720, 1280, 3), (60, 90, 120), dtype=np.uint8)
|
|
doc = np.full((600, 400, 3), 245, dtype=np.uint8)
|
|
canvas = np.zeros_like(frame)
|
|
mask = np.zeros(frame.shape[:2], dtype=np.uint8)
|
|
x0, y0 = 440, 60
|
|
canvas[y0 : y0 + 600, x0 : x0 + 400] = doc
|
|
mask[y0 : y0 + 600, x0 : x0 + 400] = 255
|
|
matrix = cv2.getRotationMatrix2D((x0 + 200, y0 + 300), angle, 1.0)
|
|
canvas = cv2.warpAffine(canvas, matrix, (1280, 720))
|
|
mask = cv2.warpAffine(mask, matrix, (1280, 720))
|
|
frame[mask > 0] = canvas[mask > 0]
|
|
return frame
|
|
|
|
|
|
def test_detects_rotated_document():
|
|
det = detect_document(synthetic_frame())
|
|
assert det is not None
|
|
assert det.quad.shape == (4, 2)
|
|
assert det.area_ratio > 0.2 # 400x600 doc in a 1280x720 frame
|
|
|
|
|
|
def test_warp_restores_aspect_ratio():
|
|
det = detect_document(synthetic_frame(angle=8.0))
|
|
crop = warp_document(synthetic_frame(angle=8.0), det.quad)
|
|
h, w = crop.shape[:2]
|
|
# Source doc is 400x600 (aspect 1.5); allow tolerance for edges/dilation.
|
|
assert 1.3 < h / w < 1.7
|
|
# Crop should be mostly white paper.
|
|
assert crop.mean() > 180
|
|
|
|
|
|
def test_no_document_on_empty_table():
|
|
frame = np.full((720, 1280, 3), (60, 90, 120), dtype=np.uint8)
|
|
assert detect_document(frame) is None
|
|
|
|
|
|
def test_sharpness_prefers_sharp_frame():
|
|
sharp = synthetic_frame()
|
|
blurred = cv2.GaussianBlur(sharp, (31, 31), 0)
|
|
assert sharpness_score(sharp) > sharpness_score(blurred)
|
|
|
|
|
|
def test_detection_reports_source_and_confidence():
|
|
det = detect_document(synthetic_frame())
|
|
assert det.source in ("edges", "threshold")
|
|
assert 0.0 < det.confidence <= 1.0
|
|
|
|
|
|
def test_rejects_thin_sliver_candidate():
|
|
"""A narrow strip (e.g. a barcode edge) shouldn't be mistaken for a page."""
|
|
frame = np.full((720, 1280, 3), (60, 90, 120), dtype=np.uint8)
|
|
# 30x600 sliver: aspect ratio 20:1, well past max_aspect.
|
|
frame[60:660, 600:630] = 245
|
|
assert detect_document(frame, max_aspect=6.0) is None
|
|
|
|
|
|
def test_rejects_near_whole_frame_candidate():
|
|
"""A candidate covering almost the entire frame is probably background/lighting."""
|
|
frame = np.full((720, 1280, 3), 245, dtype=np.uint8)
|
|
frame[10:20, 10:20] = (60, 90, 120) # tiny dark speck so it's not perfectly blank
|
|
assert detect_document(frame, max_area_ratio=0.92) is None
|
|
|
|
|
|
def test_downscaled_detection_matches_full_resolution():
|
|
"""detect_width downscaling shouldn't change which document gets found."""
|
|
frame = synthetic_frame()
|
|
det_full = detect_document(frame, detect_width=10_000) # effectively no downscale
|
|
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
|