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,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