ilia
17be9cef33
Adds .gitea/workflows/ci.yml (python-ci with tesseract + libgl for the OCR/opencv tests, plus gitleaks secret-scan) and ruff.toml (E4/E7/E9/F/I, line-length 120). Cleans up the remaining unused-variable findings so the lint gate is green.
219 lines
8.2 KiB
Python
219 lines
8.2 KiB
Python
"""Document/receipt contour detection and perspective correction.
|
|
|
|
Detection runs two independent candidate generators — Canny edges and Otsu
|
|
brightness thresholding — because each fails on different real-world cases
|
|
(Canny struggles with low-contrast paper edges; brightness thresholding
|
|
struggles when a bright hand or lighting glare merges with the page). Every
|
|
candidate contour from both is scored on how "paper-shaped" it is (filled,
|
|
convex, plausible aspect ratio) rather than just picking whichever contour
|
|
happens to have the largest area, which was previously fooled by fused
|
|
hand+background blobs and by fragmented outlines on high-resolution frames.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
@dataclass
|
|
class DocumentDetection:
|
|
quad: np.ndarray # 4x2 float32, corner points in source-image coordinates
|
|
area_ratio: float # contour area / frame area
|
|
method: str # "quad" (clean 4-corner polygon) | "min_rect" (fallback)
|
|
source: str # "edges" | "threshold" — which candidate generator won
|
|
confidence: float # 0-1 composite shape score (higher = more paper-like)
|
|
|
|
|
|
def _order_points(pts: np.ndarray) -> np.ndarray:
|
|
"""Order 4 points as top-left, top-right, bottom-right, bottom-left."""
|
|
pts = pts.astype(np.float32)
|
|
s = pts.sum(axis=1)
|
|
d = np.diff(pts, axis=1).ravel()
|
|
return np.array(
|
|
[
|
|
pts[np.argmin(s)], # top-left: smallest x+y
|
|
pts[np.argmin(d)], # top-right: smallest y-x
|
|
pts[np.argmax(s)], # bottom-right: largest x+y
|
|
pts[np.argmax(d)], # bottom-left: largest y-x
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
|
|
def _shape_score(contour: np.ndarray, frame_area: float) -> dict | None:
|
|
area = cv2.contourArea(contour)
|
|
area_ratio = area / frame_area
|
|
x, y, w, h = cv2.boundingRect(contour)
|
|
if w == 0 or h == 0:
|
|
return None
|
|
extent = area / (w * h)
|
|
hull = cv2.convexHull(contour)
|
|
hull_area = cv2.contourArea(hull)
|
|
solidity = area / hull_area if hull_area > 0 else 0.0
|
|
aspect = max(w, h) / max(1, min(w, h))
|
|
return {
|
|
"area_ratio": area_ratio,
|
|
"extent": extent,
|
|
"solidity": solidity,
|
|
"aspect": aspect,
|
|
}
|
|
|
|
|
|
def _composite_score(
|
|
stats: dict, min_area_ratio: float, max_area_ratio: float, max_aspect: float
|
|
) -> float:
|
|
"""Higher is more "paper-like". Negative disqualifies the candidate.
|
|
|
|
A page should be a large-ish (but not whole-frame — that's usually
|
|
background/lighting, not a document), filled (high extent), convex
|
|
(high solidity) rectangle. Weighted toward shape over raw size so a
|
|
smaller clean rectangle beats a larger fused hand+background blob.
|
|
"""
|
|
if stats["area_ratio"] < min_area_ratio or stats["area_ratio"] > max_area_ratio:
|
|
return -1.0
|
|
if stats["aspect"] > max_aspect:
|
|
return -1.0
|
|
size_term = min(stats["area_ratio"] / 0.5, 1.0)
|
|
return stats["extent"] * 0.45 + stats["solidity"] * 0.35 + size_term * 0.20
|
|
|
|
|
|
def _candidates(mask: np.ndarray, frame_area: float, top_n: int = 8) -> list[np.ndarray]:
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
contours = sorted(contours, key=cv2.contourArea, reverse=True)[:top_n]
|
|
return contours
|
|
|
|
|
|
def _edge_mask(gray: np.ndarray, canny_low: int, canny_high: int) -> np.ndarray:
|
|
blur = cv2.GaussianBlur(gray, (5, 5), 0)
|
|
edges = cv2.Canny(blur, canny_low, canny_high)
|
|
# Kernel scaled to image size so gaps in the outline close reliably
|
|
# regardless of source resolution (a fixed pixel kernel that works at
|
|
# 1000px wide leaves 4x-larger gaps unclosed at 4K).
|
|
k = max(3, round(min(gray.shape) * 0.01)) | 1 # odd
|
|
return cv2.dilate(edges, np.ones((k, k), np.uint8), iterations=2)
|
|
|
|
|
|
def _threshold_mask(gray: np.ndarray) -> np.ndarray:
|
|
blur = cv2.GaussianBlur(gray, (5, 5), 0)
|
|
_, mask = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
k = max(3, round(min(gray.shape) * 0.01)) | 1
|
|
kernel = np.ones((k, k), np.uint8)
|
|
return cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
|
|
|
|
|
|
def _quad_from_contour(contour: np.ndarray) -> tuple[np.ndarray, str]:
|
|
peri = cv2.arcLength(contour, True)
|
|
approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
|
|
if len(approx) == 4:
|
|
return _order_points(approx.reshape(4, 2)), "quad"
|
|
rect = cv2.minAreaRect(contour)
|
|
return _order_points(cv2.boxPoints(rect)), "min_rect"
|
|
|
|
|
|
def detect_document(
|
|
image_bgr: np.ndarray,
|
|
min_area_ratio: float = 0.04,
|
|
max_area_ratio: float = 0.92,
|
|
max_aspect: float = 6.0,
|
|
canny_low: int = 50,
|
|
canny_high: int = 150,
|
|
detect_width: int = 1000,
|
|
) -> DocumentDetection | None:
|
|
"""Find the best document-like contour in the frame.
|
|
|
|
Contour search runs on a copy downscaled to detect_width (edge/threshold
|
|
gap-closing kernels are sized relative to this, not the source
|
|
resolution), then the winning quad is scaled back up to source
|
|
resolution for a full-quality perspective warp.
|
|
"""
|
|
h, w = image_bgr.shape[:2]
|
|
|
|
scale = detect_width / w if w > detect_width else 1.0
|
|
small = (
|
|
cv2.resize(image_bgr, (round(w * scale), round(h * scale)), interpolation=cv2.INTER_AREA)
|
|
if scale != 1.0
|
|
else image_bgr
|
|
)
|
|
small_area = float(small.shape[0] * small.shape[1])
|
|
gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)
|
|
|
|
best: tuple[float, np.ndarray, str] | None = None # (score, contour, source)
|
|
for source, mask in (
|
|
("edges", _edge_mask(gray, canny_low, canny_high)),
|
|
("threshold", _threshold_mask(gray)),
|
|
):
|
|
for contour in _candidates(mask, small_area):
|
|
stats = _shape_score(contour, small_area)
|
|
if stats is None:
|
|
continue
|
|
score = _composite_score(stats, min_area_ratio, max_area_ratio, max_aspect)
|
|
if score < 0:
|
|
continue
|
|
if best is None or score > best[0]:
|
|
best = (score, contour, source)
|
|
|
|
if best is None:
|
|
return None
|
|
|
|
score, contour, source = best
|
|
small_quad, method = _quad_from_contour(contour)
|
|
quad = scale_quad(small_quad, 1.0 / scale, 1.0 / scale) if scale != 1.0 else small_quad
|
|
area_ratio = cv2.contourArea(contour) / small_area
|
|
|
|
return DocumentDetection(
|
|
quad=quad,
|
|
area_ratio=float(area_ratio),
|
|
method=method,
|
|
source=source,
|
|
confidence=float(score),
|
|
)
|
|
|
|
|
|
def warp_document(image_bgr: np.ndarray, quad: np.ndarray) -> np.ndarray:
|
|
"""Perspective-correct the quad region into a clean top-down crop."""
|
|
quad = _order_points(np.asarray(quad, dtype=np.float32))
|
|
tl, tr, br, bl = quad
|
|
|
|
width = int(max(np.linalg.norm(br - bl), np.linalg.norm(tr - tl)))
|
|
height = int(max(np.linalg.norm(tr - br), np.linalg.norm(tl - bl)))
|
|
width, height = max(width, 1), max(height, 1)
|
|
|
|
dst = np.array(
|
|
[[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]],
|
|
dtype=np.float32,
|
|
)
|
|
matrix = cv2.getPerspectiveTransform(quad, dst)
|
|
return cv2.warpPerspective(image_bgr, matrix, (width, height))
|
|
|
|
|
|
def scale_quad(quad: np.ndarray, scale_x: float, scale_y: float) -> np.ndarray:
|
|
"""Rescale a quad detected on a downscaled frame back to full resolution."""
|
|
quad = np.asarray(quad, dtype=np.float32).copy()
|
|
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
|