Add crop refinement, finger removal, and local LLM naming
- vision/refine.py: tighten crops to the paper band (removes mat margins and hands beside receipts) and inpaint border-connected skin regions so fingers disappear from output - llm/vision.py: identify documents with a local Ollama vision model (qwen2.5vl); extracts vendor/date/total/form code and flags quality issues (fingers, blur, glare); falls back to Tesseract when down - pipeline: drop blank pages, dedupe consecutive captures of the same document, record refine/LLM fields in report and export summary
This commit is contained in:
@@ -0,0 +1,124 @@
|
||||
"""Group per-frame document detections into capture events.
|
||||
|
||||
Stable-window-only detection misses documents shown while hands are still
|
||||
adjusting them — motion scores stay above threshold even though a receipt
|
||||
is clearly visible and detectable frame-by-frame. This clusters consecutive
|
||||
frames where a document contour was found (with a time-gap to split separate
|
||||
placements) and keeps the sharpest frame from each cluster.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
from paperpod.capture.video import Frame
|
||||
from paperpod.vision.document import DocumentDetection
|
||||
from paperpod.vision.sharpness import sharpness_score
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentEvent:
|
||||
event_id: int
|
||||
t_start: float
|
||||
t_end: float
|
||||
sample_count: int
|
||||
best_frame_t: float
|
||||
best_frame_index: int
|
||||
sharpness: float
|
||||
frame: Frame
|
||||
detection: DocumentDetection
|
||||
|
||||
|
||||
def _same_document(
|
||||
prev: DocumentDetection,
|
||||
curr: DocumentDetection,
|
||||
frame_shape: tuple[int, ...],
|
||||
area_tolerance: float = 0.12,
|
||||
centroid_tolerance: float = 0.12,
|
||||
) -> bool:
|
||||
"""True if two detections likely belong to the same physical document.
|
||||
|
||||
On a black mat, the empty table is often detected as a large contour on
|
||||
almost every frame, which bridges separate receipts/forms into one giant
|
||||
cluster if we only use time gaps. A sudden change in area or position
|
||||
means a new document was placed (Costco receipt area ~0.26 vs mat ~0.65).
|
||||
"""
|
||||
if abs(prev.area_ratio - curr.area_ratio) > area_tolerance:
|
||||
return False
|
||||
h, w = frame_shape[:2]
|
||||
prev_c = prev.quad.mean(axis=0)
|
||||
curr_c = curr.quad.mean(axis=0)
|
||||
if float(np.linalg.norm(prev_c - curr_c)) > centroid_tolerance * max(w, h):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _plausible_capture(det: DocumentDetection) -> bool:
|
||||
"""Skip mat/table false positives that sit between receipt and full-page sizes.
|
||||
|
||||
On a black mat, empty-table contours usually land around area_ratio
|
||||
0.45-0.57. Real receipts cluster below 0.40; full tax forms above 0.58.
|
||||
"""
|
||||
ar = det.area_ratio
|
||||
return ar <= 0.40 or ar >= 0.58
|
||||
|
||||
|
||||
def find_document_events(
|
||||
samples: list[tuple[Frame, DocumentDetection]],
|
||||
gap_s: float = 2.0,
|
||||
min_samples: int = 2,
|
||||
) -> list[DocumentEvent]:
|
||||
"""Cluster consecutive detections; split when gap_s passes with no detection.
|
||||
|
||||
samples must be in time order. Each sample is a (frame, detection) pair
|
||||
from a frame where detect_document() succeeded.
|
||||
"""
|
||||
if not samples:
|
||||
return []
|
||||
|
||||
clusters: list[tuple[float, float, int, float, Frame, DocumentDetection]] = []
|
||||
active: tuple[float, float, int, float, Frame, DocumentDetection] | None = None
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal active
|
||||
if active is None:
|
||||
return
|
||||
t_start, t_end, count, sharp, frame, det = active
|
||||
if count >= min_samples:
|
||||
clusters.append(active)
|
||||
active = None
|
||||
|
||||
for frame, det in samples:
|
||||
if active is None:
|
||||
active = (frame.t, frame.t, 1, sharpness_score(frame.image), frame, det)
|
||||
continue
|
||||
|
||||
t_start, t_end, count, best_sharp, best_frame, best_det = active
|
||||
if frame.t - t_end > gap_s or not _same_document(best_det, det, frame.image.shape):
|
||||
flush()
|
||||
active = (frame.t, frame.t, 1, sharpness_score(frame.image), frame, det)
|
||||
continue
|
||||
|
||||
sharp = sharpness_score(frame.image)
|
||||
if sharp > best_sharp:
|
||||
best_sharp, best_frame, best_det = sharp, frame, det
|
||||
active = (t_start, frame.t, count + 1, best_sharp, best_frame, best_det)
|
||||
|
||||
flush()
|
||||
|
||||
return [
|
||||
DocumentEvent(
|
||||
event_id=i,
|
||||
t_start=t_start,
|
||||
t_end=t_end,
|
||||
sample_count=count,
|
||||
best_frame_t=frame.t,
|
||||
best_frame_index=frame.index,
|
||||
sharpness=sharp,
|
||||
frame=frame,
|
||||
detection=det,
|
||||
)
|
||||
for i, (t_start, t_end, count, sharp, frame, det) in enumerate(clusters)
|
||||
]
|
||||
+126
-27
@@ -1,4 +1,14 @@
|
||||
"""Document/receipt contour detection and perspective correction."""
|
||||
"""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
|
||||
|
||||
@@ -13,6 +23,8 @@ 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:
|
||||
@@ -31,47 +43,134 @@ def _order_points(pts: np.ndarray) -> np.ndarray:
|
||||
)
|
||||
|
||||
|
||||
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 largest document-like contour in the frame.
|
||||
"""Find the best document-like contour in the frame.
|
||||
|
||||
Pipeline: grayscale -> blur -> Canny -> dilate -> contours -> largest
|
||||
contour above min_area_ratio. If its polygon approximation has 4 corners
|
||||
we use it directly; otherwise we fall back to the minimum-area rectangle.
|
||||
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]
|
||||
frame_area = float(h * w)
|
||||
|
||||
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
|
||||
gray = cv2.GaussianBlur(gray, (5, 5), 0)
|
||||
edges = cv2.Canny(gray, canny_low, canny_high)
|
||||
# Close gaps in the document outline so it forms one contour.
|
||||
edges = cv2.dilate(edges, np.ones((5, 5), np.uint8), iterations=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)
|
||||
|
||||
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
if not contours:
|
||||
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
|
||||
|
||||
best = max(contours, key=cv2.contourArea)
|
||||
area_ratio = cv2.contourArea(best) / frame_area
|
||||
if area_ratio < min_area_ratio:
|
||||
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
|
||||
|
||||
peri = cv2.arcLength(best, True)
|
||||
approx = cv2.approxPolyDP(best, 0.02 * peri, True)
|
||||
if len(approx) == 4:
|
||||
quad = _order_points(approx.reshape(4, 2))
|
||||
method = "quad"
|
||||
else:
|
||||
rect = cv2.minAreaRect(best)
|
||||
quad = _order_points(cv2.boxPoints(rect))
|
||||
method = "min_rect"
|
||||
|
||||
return DocumentDetection(quad=quad, area_ratio=float(area_ratio), method=method)
|
||||
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:
|
||||
|
||||
+28
-10
@@ -1,4 +1,12 @@
|
||||
"""Optional image enhancement for OCR-readiness."""
|
||||
"""Scan-style image enhancement, for OCR-readiness and optionally the saved crop.
|
||||
|
||||
`normalize_illumination` is the "flatten it, make it easier to read" trick
|
||||
scanner apps like CamScanner use: divide the image by a heavily-blurred copy
|
||||
of itself (an estimate of the local lighting/shadow), which cancels out
|
||||
gradients, glare, and wood-grain/shadow texture far better than a plain
|
||||
contrast boost. Confirmed on a real low-contrast receipt-on-wood-grain photo:
|
||||
raw OCR read almost nothing, normalized OCR recovered full lines of text.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -6,15 +14,25 @@ import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def enhance_for_ocr(image_bgr: np.ndarray) -> np.ndarray:
|
||||
"""Boost local contrast (CLAHE on the luminance channel).
|
||||
def normalize_illumination(image_bgr: np.ndarray, blur_fraction: float = 1 / 15) -> np.ndarray:
|
||||
"""Cancel out uneven lighting/shadow/background texture. Returns grayscale."""
|
||||
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
|
||||
sigma = max(image_bgr.shape[1] * blur_fraction, 10.0)
|
||||
background = cv2.GaussianBlur(gray, (0, 0), sigmaX=sigma)
|
||||
# Background approaches white; dividing pulls the true page background
|
||||
# up to ~255 wherever lighting made it dim, while text (locally much
|
||||
# darker than its surroundings) stays dark.
|
||||
return cv2.divide(gray, background, scale=255)
|
||||
|
||||
Deliberately conservative: no binarization, so the crop stays pleasant to
|
||||
read in the PDF while giving OCR more to work with. Thresholding can be
|
||||
added behind a config flag later if OCR needs it.
|
||||
|
||||
def enhance_for_ocr(image_bgr: np.ndarray) -> np.ndarray:
|
||||
"""Illumination-normalize + local contrast boost. Returns a BGR image.
|
||||
|
||||
Always grayscale-looking (3 identical channels) since normalization
|
||||
operates on luminance. Intended primarily as an OCR preprocessing step;
|
||||
only baked into the saved crop/PDF when document.enhance is enabled.
|
||||
"""
|
||||
lab = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2LAB)
|
||||
l_chan, a_chan, b_chan = cv2.split(lab)
|
||||
normalized = normalize_illumination(image_bgr)
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
lab = cv2.merge((clahe.apply(l_chan), a_chan, b_chan))
|
||||
return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
|
||||
boosted = clahe.apply(normalized)
|
||||
return cv2.cvtColor(boosted, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""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 _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(
|
||||
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.
|
||||
"""
|
||||
best_rotation = 0
|
||||
best_score = -1.0
|
||||
for rotation in (0, 90, 180, 270):
|
||||
candidate = apply_rotation(image_bgr, 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
|
||||
@@ -0,0 +1,169 @@
|
||||
"""Post-warp crop refinement: tighten to the paper region and remove fingers.
|
||||
|
||||
The initial contour crop often includes a margin of mat/table and any hand
|
||||
holding the receipt (thermal receipts curl, so people press them flat).
|
||||
This module:
|
||||
1. builds a paper mask (bright pixels that are NOT skin-colored),
|
||||
2. tightens the crop to the rows/columns that actually contain paper,
|
||||
3. optionally inpaints skin regions so fingers disappear from the output.
|
||||
|
||||
Skin detection uses the classic YCrCb range, which cleanly separates skin
|
||||
from both white paper (low Cr) and the black mat.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
_SKIN_LOW = (0, 133, 77)
|
||||
_SKIN_HIGH = (255, 178, 127)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefineResult:
|
||||
image: np.ndarray
|
||||
tightened: bool # crop bounds were shrunk to the paper band
|
||||
fingers_removed: bool # skin regions were inpainted
|
||||
skin_fraction: float # fraction of the crop that looked like skin
|
||||
|
||||
|
||||
def skin_mask(image_bgr: np.ndarray) -> np.ndarray:
|
||||
"""Binary mask (255 = skin-colored) with small speckles removed."""
|
||||
ycrcb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2YCrCb)
|
||||
mask = cv2.inRange(ycrcb, _SKIN_LOW, _SKIN_HIGH)
|
||||
return cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((9, 9), np.uint8))
|
||||
|
||||
|
||||
def _border_connected(mask: np.ndarray) -> np.ndarray:
|
||||
"""Keep only mask components touching the image border.
|
||||
|
||||
Fingers always reach in from an edge; skin-colored false positives in
|
||||
the middle of the page (beige logos, tinted paper) do not, and
|
||||
inpainting those would smear real content.
|
||||
"""
|
||||
n, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
|
||||
h, w = mask.shape
|
||||
keep = np.zeros_like(mask)
|
||||
for i in range(1, n):
|
||||
x, y, bw, bh, _area = stats[i]
|
||||
if x == 0 or y == 0 or x + bw >= w or y + bh >= h:
|
||||
keep[labels == i] = 255
|
||||
return keep
|
||||
|
||||
|
||||
def _paper_mask(image_bgr: np.ndarray, skin: np.ndarray) -> np.ndarray:
|
||||
gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
|
||||
blur = cv2.GaussianBlur(gray, (7, 7), 0)
|
||||
_, bright = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
||||
skin_dilated = cv2.dilate(skin, np.ones((15, 15), np.uint8), iterations=1)
|
||||
paper = cv2.bitwise_and(bright, cv2.bitwise_not(skin_dilated))
|
||||
return cv2.morphologyEx(paper, cv2.MORPH_OPEN, np.ones((11, 11), np.uint8), iterations=2)
|
||||
|
||||
|
||||
def _density_band(
|
||||
density: np.ndarray, threshold: float, gap_fraction: float = 0.15
|
||||
) -> tuple[int, int] | None:
|
||||
"""Longest contiguous run of rows/columns above threshold.
|
||||
|
||||
First-to-last-above-threshold would bridge across gaps: a bright strip
|
||||
of wood grain at the far edge of the frame can extend the band across
|
||||
the entire (dark) mat, keeping all the background in the crop.
|
||||
|
||||
Gaps up to gap_fraction of the axis length are closed first, so dark
|
||||
printed regions *on* the paper (a promo banner on a receipt, a filled
|
||||
table header) don't split the paper band in two.
|
||||
"""
|
||||
above = density > threshold
|
||||
if not above.any():
|
||||
return None
|
||||
# 1-D morphological closing: dilate then erode with the gap window.
|
||||
gap = max(1, int(len(density) * gap_fraction))
|
||||
kernel = np.ones(gap, dtype=bool)
|
||||
closed = np.convolve(above, kernel, mode="same") > 0 # dilate
|
||||
eroded = np.convolve(~closed, kernel, mode="same") == 0 # erode
|
||||
above = eroded if eroded.any() else above
|
||||
|
||||
best_start, best_len = 0, 0
|
||||
start = None
|
||||
for i, flag in enumerate(np.append(above, False)):
|
||||
if flag and start is None:
|
||||
start = i
|
||||
elif not flag and start is not None:
|
||||
if i - start > best_len:
|
||||
best_start, best_len = start, i - start
|
||||
start = None
|
||||
return best_start, best_start + best_len - 1
|
||||
|
||||
|
||||
def refine_crop(
|
||||
image_bgr: np.ndarray,
|
||||
tighten: bool = True,
|
||||
remove_fingers: bool = True,
|
||||
density_threshold: float = 0.25,
|
||||
min_skin_fraction: float = 0.01,
|
||||
inpaint_radius: int = 15,
|
||||
inpaint_max_width: int = 1200,
|
||||
) -> RefineResult:
|
||||
"""Tighten a warped document crop to its paper band and erase fingers.
|
||||
|
||||
tighten: cut rows/columns whose paper coverage is below
|
||||
density_threshold (removes mat margins and the hand hanging
|
||||
off the side of a receipt).
|
||||
remove_fingers: inpaint border-connected skin blobs. Content under a
|
||||
finger is unrecoverable — inpainting fills it with plausible
|
||||
paper texture — but margins/blank areas clean up completely.
|
||||
"""
|
||||
skin = skin_mask(image_bgr)
|
||||
skin_fraction = float((skin > 0).mean())
|
||||
|
||||
result = image_bgr
|
||||
tightened = False
|
||||
|
||||
if tighten:
|
||||
paper = _paper_mask(image_bgr, skin)
|
||||
h, w = paper.shape
|
||||
col_band = _density_band(paper.sum(axis=0) / (255.0 * h), density_threshold)
|
||||
row_band = _density_band(paper.sum(axis=1) / (255.0 * w), density_threshold)
|
||||
if col_band and row_band:
|
||||
x0, x1 = col_band
|
||||
y0, y1 = row_band
|
||||
# Only shrink, never grow; skip degenerate bands.
|
||||
if (x1 - x0) > w * 0.2 and (y1 - y0) > h * 0.2:
|
||||
if x0 > 0 or x1 < w - 1 or y0 > 0 or y1 < h - 1:
|
||||
result = result[y0 : y1 + 1, x0 : x1 + 1]
|
||||
skin = skin[y0 : y1 + 1, x0 : x1 + 1]
|
||||
tightened = True
|
||||
|
||||
fingers_removed = False
|
||||
if remove_fingers and skin_fraction >= min_skin_fraction:
|
||||
fingers = _border_connected(skin)
|
||||
if (fingers > 0).any():
|
||||
fingers = cv2.dilate(fingers, np.ones((25, 25), np.uint8), iterations=1)
|
||||
# Inpaint at reduced resolution: cv2.inpaint is O(radius * area)
|
||||
# and full-res 4K crops take tens of seconds for no visible gain.
|
||||
h, w = result.shape[:2]
|
||||
if w > inpaint_max_width:
|
||||
scale = inpaint_max_width / w
|
||||
small = cv2.resize(result, (inpaint_max_width, round(h * scale)))
|
||||
small_mask = cv2.resize(
|
||||
fingers, (small.shape[1], small.shape[0]), interpolation=cv2.INTER_NEAREST
|
||||
)
|
||||
inpainted = cv2.inpaint(small, small_mask, inpaint_radius, cv2.INPAINT_TELEA)
|
||||
inpainted = cv2.resize(inpainted, (w, h))
|
||||
# Blend: keep original pixels outside the finger mask so
|
||||
# only the finger region loses resolution.
|
||||
mask3 = cv2.cvtColor(fingers, cv2.COLOR_GRAY2BGR) > 0
|
||||
result = np.where(mask3, inpainted, result)
|
||||
else:
|
||||
result = cv2.inpaint(result, fingers, inpaint_radius, cv2.INPAINT_TELEA)
|
||||
fingers_removed = True
|
||||
|
||||
return RefineResult(
|
||||
image=result,
|
||||
tightened=tightened,
|
||||
fingers_removed=fingers_removed,
|
||||
skin_fraction=round(skin_fraction, 4),
|
||||
)
|
||||
Reference in New Issue
Block a user