Files
PaperPod/paperpod/pipeline.py
T
ilia 12e5d1aeb1 Add page-marker grouping and party-aware naming
Consecutive captures whose printed PAGE X OF Y markers advance under the same total are grouped into one multi-page PDF, with duplicate pages keeping the better read. Naming gains doc-title/party components and a scan-date fallback so nothing lands in UNSORTED. Vision LLM default switches to minicpm-v (faster, fewer false blanks on dense mono pages).
2026-07-26 15:37:28 -04:00

568 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Detection pipeline: video in, cropped document candidates + JSON report out.
Two sequential passes over the video (no seeking, works on any codec):
Pass 1: sample downscaled frames, score motion, find stable/moving segments.
Pass 2: scan every sampled frame for document contours, cluster consecutive
detections into capture events (works even while hands are moving),
pick the sharpest frame per event, perspective-correct, save the crop.
"""
from __future__ import annotations
import csv
import json
import re
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any
import cv2
import numpy as np
from paperpod.capture.video import iter_frames, probe_video
from paperpod.config import Config, config_to_dict
from paperpod.llm.vision import analyze_document, ollama_available
from paperpod.naming.filename import build_name, dedupe_stem, slugify
from paperpod.ocr.extract import extract_store_name, run_ocr
from paperpod.pdf.build import images_to_pdf
from paperpod.vision.capture_events import _plausible_capture, find_document_events, score_sample
from paperpod.vision.document import detect_document, pad_quad, warp_document
from paperpod.vision.enhance import darken_print, enhance_for_ocr
from paperpod.vision.motion import find_segments, motion_score, prepare_gray
from paperpod.vision.orient import auto_rotate
from paperpod.vision.refine import refine_crop
_AMOUNT_RE = re.compile(r"\d[\d,]{2,}(?:\.\d{2})?")
def _content_overlap(a: str, b: str) -> float:
"""Shared dollar/account numbers between two OCR texts (1.0 ≈ same page)."""
fa = set(_AMOUNT_RE.findall(a or ""))
fb = set(_AMOUNT_RE.findall(b or ""))
if not fa or not fb:
return 0.0
return len(fa & fb) / min(len(fa), len(fb))
def run_detection(video_path: str | Path, cfg: Config, out_dir: str | Path) -> dict[str, Any]:
"""Run the full detect pipeline and write artifacts into out_dir.
Artifacts:
crops/window_NNN.png perspective-corrected document candidates
frames/window_NNN_full.png the full best frame per window (for debugging)
motion_scores.csv per-sample motion scores (for threshold tuning)
report.json stable windows, motion events, detections
"""
video_path = Path(video_path)
out_dir = Path(out_dir)
crops_dir = out_dir / "crops"
frames_dir = out_dir / "frames"
crops_dir.mkdir(parents=True, exist_ok=True)
frames_dir.mkdir(parents=True, exist_ok=True)
meta = probe_video(video_path)
# ---- Pass 1: motion scores on downscaled frames -------------------------
timestamps: list[float] = []
scores: list[float] = []
prev_gray: np.ndarray | None = None
small_size: tuple[int, int] | None = None # (w, h) of processing frames
for frame in iter_frames(
video_path,
sample_fps=cfg.capture.sample_fps,
resize_width=cfg.capture.processing_width,
):
gray = prepare_gray(frame.image, cfg.motion.blur_ksize)
if small_size is None:
small_size = (frame.image.shape[1], frame.image.shape[0])
# First frame has no predecessor; treat it as perfectly stable.
score = (
0.0
if prev_gray is None
else motion_score(prev_gray, gray, cfg.motion.pixel_threshold)
)
timestamps.append(frame.t)
scores.append(score)
prev_gray = gray
segments = find_segments(
timestamps, scores, cfg.motion.threshold, cfg.motion.stable_min_duration_s
)
with open(out_dir / "motion_scores.csv", "w", newline="") as f:
writer = csv.writer(f)
writer.writerow(["t", "score"])
writer.writerows(zip(timestamps, scores))
# ---- Pass 2: frame-based document capture events ------------------------
# Scan every sampled frame for a document contour, cluster consecutive
# detections into separate placements, keep the sharpest frame per cluster.
# Unlike stable-window-only detection, this still captures receipts shown
# while hands are adjusting them (motion scores stay high but the contour
# is visible frame-by-frame).
#
# Only tiny scored proxies are retained here — full-res 4K BGR would OOM
# on a long phone video. Winning frame indices are re-decoded in pass 2b.
det_kwargs = dict(
min_area_ratio=cfg.document.min_area_ratio,
max_area_ratio=cfg.document.max_area_ratio,
max_aspect=cfg.document.max_aspect,
canny_low=cfg.document.canny_low,
canny_high=cfg.document.canny_high,
detect_width=cfg.document.detect_width,
)
scored_samples = []
for frame in iter_frames(video_path, sample_fps=cfg.capture.sample_fps):
detection = detect_document(frame.image, **det_kwargs)
if detection is not None and _plausible_capture(detection):
scored_samples.append(score_sample(frame, detection))
doc_events = find_document_events(
scored_samples,
gap_s=cfg.document.event_gap_s,
min_samples=cfg.document.event_min_samples,
)
# Pass 2b: re-decode only the winning frame indices (one BGR at a time).
want = {event.best_frame_index: event for event in doc_events}
if want:
for frame in iter_frames(video_path, sample_fps=cfg.capture.sample_fps):
event = want.get(frame.index)
if event is None:
continue
event.frame = frame
del want[frame.index]
if not want:
break
window_records: list[dict[str, Any]] = []
for event in doc_events:
i = event.event_id
frame = event.frame
if frame is None:
# Winner index was skipped (sample_fps / codec edge); skip event.
continue
detection = event.detection
sharp = event.sharpness
record: dict[str, Any] = {
"window_id": i,
"t_start": round(event.t_start, 3),
"t_end": round(event.t_end, 3),
"duration_s": round(event.t_end - event.t_start, 3),
"sample_count": event.sample_count,
"document_found": False,
"best_frame_t": round(event.best_frame_t, 3),
"best_frame_index": event.best_frame_index,
"sharpness": round(sharp, 1),
"best_frame_skin_fraction": event.skin_fraction,
}
frame_path = frames_dir / f"window_{i:03d}_full.png"
cv2.imwrite(str(frame_path), frame.image)
record["frame_path"] = str(frame_path.relative_to(out_dir))
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:
refined = refine_crop(crop, remove_fingers=cfg.document.remove_fingers)
crop = refined.image
refine_info = {
"crop_tightened": refined.tightened,
"fingers_removed": refined.fingers_removed,
"skin_fraction": refined.skin_fraction,
}
rotation_applied = 0
rotation_method = "none"
if cfg.document.auto_rotate:
crop, orientation = auto_rotate(crop, lang=cfg.ocr.lang, psm=cfg.ocr.psm)
rotation_applied = orientation.rotation
rotation_method = orientation.method
if cfg.document.enhance:
crop = enhance_for_ocr(crop)
else:
# Keep the photographic look but pull pale grey print toward black.
crop = darken_print(crop)
crop_path = crops_dir / f"window_{i:03d}.png"
cv2.imwrite(str(crop_path), crop)
record.update(
document_found=True,
detection_method=detection.method,
detection_source=detection.source,
detection_confidence=round(detection.confidence, 3),
area_ratio=round(detection.area_ratio, 4),
quad=[[round(float(x), 1), round(float(y), 1)] for x, y in detection.quad],
rotation_applied=rotation_applied,
rotation_method=rotation_method,
crop_path=str(crop_path.relative_to(out_dir)),
crop_size=[crop.shape[1], crop.shape[0]],
**refine_info,
)
window_records.append(record)
report = {
"video": str(video_path),
"video_meta": {
"fps": round(meta.fps, 3),
"frame_count": meta.frame_count,
"duration_s": round(meta.duration_s, 3),
"width": meta.width,
"height": meta.height,
},
"config": config_to_dict(cfg),
"motion_events": [
{
"kind": s.kind,
"t_start": round(s.t_start, 3),
"t_end": round(s.t_end, 3),
"duration_s": round(s.duration, 3),
}
for s in segments
],
"stable_windows": window_records,
}
with open(out_dir / "report.json", "w") as f:
json.dump(report, f, indent=2)
return report
def _is_exportable(ocr) -> bool:
"""Skip obvious mat-transition junk that has no usable naming signal."""
if ocr.form_code or ocr.date or ocr.total:
return True
if extract_store_name(ocr.text) or ocr.doc_title or ocr.page_number:
return True
return False
@dataclass
class _PageCandidate:
"""One exportable capture: crop + everything extracted from it."""
window: dict[str, Any]
crop_path: Path
ocr: Any
llm: Any
vendor: str | None
date: str | None
time: str | None
total: str | None
form_code: str | None
tax_year: str | None
org_name: str | None
doc_title: str | None
party: str | None
page_number: int | None
page_total: int | None
@property
def signature(self) -> tuple:
return (self.vendor, self.date, self.total, self.form_code)
@dataclass
class _Document:
pages: list[_PageCandidate] = field(default_factory=list)
def last_marker_page(self) -> _PageCandidate | None:
"""Most recent page whose printed 'PAGE X OF Y' marker was readable."""
for page in reversed(self.pages):
if page.page_number is not None:
return page
return None
def _same_letterhead(prev: _PageCandidate, curr: _PageCandidate) -> bool:
"""Loose 'same printed letterhead' check for grouping decisions."""
if prev.doc_title and curr.doc_title:
return prev.doc_title == curr.doc_title
if prev.vendor and curr.vendor:
a, b = slugify(prev.vendor), slugify(curr.vendor)
return a.startswith(b) or b.startswith(a)
return False
def _signature_duplicate(prev: _PageCandidate, curr: _PageCandidate) -> bool:
"""Same marker-less page captured twice in a row (a contour blip).
Only trustworthy for events moments apart, and only when neither side
carries a printed page marker. Pages of one bank statement share
vendor + date, and flips are often 34s — the old 5s gate was replacing
page 3 with page 4 (higher OCR conf) and wiping the earlier page.
"""
if prev.page_number is not None or curr.page_number is not None:
return False
if curr.window["t_start"] - prev.window["t_end"] > 5.0:
return False
return bool(any(curr.signature) and curr.signature == prev.signature)
def _first(*values):
return next((v for v in values if v), None)
def _group_documents(candidates: list[_PageCandidate]) -> list[_Document]:
"""Merge consecutive captures into documents; drop duplicate captures.
Grouping is anchored on the printed "PAGE X OF Y" marker: within one
document, later captures must advance the page number under the same
page total (gaps are fine — a page whose marker OCR couldn't read must
not split one statement into several PDFs). A marker-less capture with
the same letterhead may slot into a document that is still incomplete
(e.g. the "PAGE 1 OF 2" statement's second page whose header came out
garbled); unrelated marker-less single-pagers never group.
"""
documents: list[_Document] = []
for cand in candidates:
doc = documents[-1] if documents else None
if doc is None:
documents.append(_Document(pages=[cand]))
continue
anchor = doc.last_marker_page()
last = doc.pages[-1]
if cand.page_number is not None:
if (
anchor is not None
and cand.page_total == anchor.page_total
and _same_letterhead(anchor, cand)
):
if cand.page_number == anchor.page_number:
# Same page captured twice; keep the better read.
if cand.ocr.mean_confidence > anchor.ocr.mean_confidence:
idx = next(i for i, p in enumerate(doc.pages) if p is anchor)
doc.pages[idx] = cand
continue
if cand.page_number > anchor.page_number:
doc.pages.append(cand)
continue
documents.append(_Document(pages=[cand]))
continue
if _signature_duplicate(last, cand):
if cand.ocr.mean_confidence > last.ocr.mean_confidence:
doc.pages[-1] = cand
continue
# Content-fingerprint dedupe for marker-less recaptures only when the
# amounts strongly match (≥0.55) or the new read looks garbled
# (conf < 75). A 0.250.40 overlap is normal between successive
# pages of one statement (shared account # / branch / balances);
# using that as a duplicate gate previously wiped page 3 with page 4.
overlap = _content_overlap(last.ocr.text, cand.ocr.text)
if overlap >= 0.55 or (overlap >= 0.25 and cand.ocr.mean_confidence < 75.0):
if cand.ocr.mean_confidence > last.ocr.mean_confidence:
doc.pages[-1] = cand
continue
if (
anchor is not None
and anchor.page_number < (anchor.page_total or 0)
and _same_letterhead(anchor, cand)
and cand.ocr.mean_confidence >= 70.0
and _content_overlap(anchor.ocr.text, cand.ocr.text) < 0.55
):
doc.pages.append(cand)
continue
documents.append(_Document(pages=[cand]))
return documents
def run_export(report: dict[str, Any], cfg: Config, out_dir: str | Path) -> dict[str, Any]:
"""OCR-name and PDF-export every detected document from a detect() report.
Consecutive captures that carry an incrementing printed "PAGE X OF Y"
marker are grouped into one multi-page PDF (page 1's metadata names the
file). Everything else stays one PDF per capture.
Writes:
pdf/<name>.pdf one PDF per document (possibly multi-page)
export_summary.csv timestamp, pod_id, page_count, source_of_name,
ocr_vendor, ocr_date, ocr_total, final_filename
Returns the summary rows (also embedded back for programmatic use).
"""
out_dir = Path(out_dir)
pdf_dir = out_dir / "pdf"
pdf_dir.mkdir(parents=True, exist_ok=True)
video_path = Path(report["video"])
duration_s = report["video_meta"]["duration_s"]
try:
recording_start = datetime.fromtimestamp(video_path.stat().st_mtime) - timedelta(
seconds=duration_s
)
except FileNotFoundError:
recording_start = datetime.now() - timedelta(seconds=duration_s)
use_llm = cfg.llm.enabled and ollama_available(cfg.llm.base_url)
if cfg.llm.enabled and not use_llm:
print("Ollama not reachable — falling back to Tesseract-only naming.")
# ---- Phase 1: OCR/LLM every crop, keep the exportable ones --------------
candidates: list[_PageCandidate] = []
for window in report["stable_windows"]:
if not window.get("document_found"):
continue
crop_path = out_dir / window["crop_path"]
image = cv2.imread(str(crop_path))
if image is None:
continue
# Always OCR the illumination-normalized version, even if the saved
# crop/PDF itself wasn't enhanced (document.enhance only controls the
# saved image's appearance) — this is pure naming-accuracy upside.
# Skip re-normalizing if the crop was already enhanced at detect time
# (cfg.document.enhance) to avoid dividing an already-flat image by
# its own blurred background twice.
ocr_image = image if cfg.document.enhance else enhance_for_ocr(image)
ocr = run_ocr(
ocr_image,
lang=cfg.ocr.lang,
vendor_top_fraction=cfg.ocr.vendor_top_fraction,
min_word_confidence=cfg.ocr.min_word_confidence,
psm=cfg.ocr.psm,
)
# Local vision LLM (Ollama): reads the whole document at once, so it
# names curled receipts and hand-occluded captures that defeat
# Tesseract. Its fields take priority; OCR heuristics fill gaps.
llm = None
if use_llm:
llm = analyze_document(
image,
model=cfg.llm.model,
base_url=cfg.llm.base_url,
timeout=cfg.llm.timeout_s,
)
# Drop LLM-flagged blank pages (backs of slips, empty sheets) only
# when OCR agrees there's nothing there — vision models (especially
# qwen2.5vl) have mislabeled real statement pages as "blank", and
# losing a real page is far worse than exporting a blank one.
if llm is not None and (
llm.doc_type == "blank" or llm.description.lower().startswith("blank")
):
if len(ocr.text.split()) < 8:
continue
llm = None # its other fields aren't trustworthy for this page
vendor = _first(llm.vendor if llm else None, ocr.vendor)
date = _first(llm.date if llm else None, ocr.date)
form_code = _first(llm.form_code if llm else None, ocr.form_code)
total = _first(llm.total if llm else None, ocr.total)
if llm is None and not _is_exportable(ocr):
continue
if llm is not None and not any((vendor, date, total, form_code)):
continue
candidates.append(
_PageCandidate(
window=window,
crop_path=crop_path,
ocr=ocr,
llm=llm,
vendor=vendor,
date=date,
time=_first(llm.time if llm else None, ocr.time),
total=total,
form_code=form_code,
tax_year=_first(llm.tax_year if llm else None, ocr.tax_year),
org_name=_first(
llm.vendor if llm and llm.form_code else None, ocr.org_name
),
doc_title=ocr.doc_title,
party=_first(llm.party if llm else None, ocr.party),
page_number=ocr.page_number,
page_total=ocr.page_total,
)
)
# ---- Phase 2: group consecutive captures into documents ------------------
documents = _group_documents(candidates)
# ---- Phase 3: name and write one PDF per document ------------------------
used_stems: set[str] = set()
rows: list[dict[str, Any]] = []
for doc in documents:
pages = doc.pages
first = pages[0]
# Page 1 usually carries the letterhead/date, but any page may have
# filled a field the others missed.
date = _first(*(p.date for p in pages))
doc_title = _first(*(p.doc_title for p in pages))
party = _first(*(p.party for p in pages))
# Prefer a vendor that isn't just the printed title re-read (the
# largest-font heuristic picks "Transaction History" on some pages
# and the actual "TD Canada Trust" letterhead on others).
vendors = [p.vendor for p in pages if p.vendor]
vendor = _first(
*(v for v in vendors if not doc_title or slugify(v) not in doc_title),
*vendors,
)
event_time = recording_start + timedelta(seconds=first.window["t_start"])
naming = build_name(
date,
vendor,
event_time,
cfg.naming.max_vendor_len,
form_code=_first(*(p.form_code for p in pages)),
tax_year=_first(*(p.tax_year for p in pages)),
org_name=_first(*(p.org_name for p in pages)),
time=_first(*(p.time for p in pages)),
doc_title=doc_title,
party=party,
)
stem = dedupe_stem(naming.stem, used_stems)
pdf_path = pdf_dir / f"{stem}.pdf"
images_to_pdf([p.crop_path for p in pages], pdf_path, title=stem)
llm = first.llm
rows.append(
{
"timestamp": event_time.isoformat(timespec="seconds"),
"pod_id": first.window["window_id"],
"page_count": len(pages),
"page_windows": ";".join(str(p.window["window_id"]) for p in pages),
"source_of_name": "llm" if llm else naming.source,
"ocr_vendor": vendor or "",
"ocr_form_code": _first(*(p.form_code for p in pages)) or "",
"ocr_date": date or "",
"ocr_time": _first(*(p.time for p in pages)) or "",
"ocr_total": _first(*(p.total for p in pages)) or "",
"doc_title": doc_title or "",
"party": party or "",
"ocr_confidence": round(first.ocr.mean_confidence, 1),
"llm_description": llm.description if llm else "",
"llm_issues": ";".join(llm.issues) if llm else "",
"detection_confidence": first.window.get("detection_confidence", ""),
"rotation_applied": first.window.get("rotation_applied", 0),
"final_filename": pdf_path.name,
"pdf_path": str(pdf_path.relative_to(out_dir)),
}
)
with open(out_dir / "export_summary.csv", "w", newline="") as f:
fieldnames = [
"timestamp", "pod_id", "page_count", "page_windows", "source_of_name",
"ocr_vendor", "ocr_form_code", "ocr_date", "ocr_time", "ocr_total",
"doc_title", "party", "ocr_confidence",
"llm_description", "llm_issues",
"detection_confidence", "rotation_applied",
"final_filename", "pdf_path",
]
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
return {"rows": rows, "pdf_dir": str(pdf_dir)}