PaperPod/paperpod/pipeline.py

"""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: within each stable window, pick the sharpest full-resolution frame,
          detect the document contour, perspective-correct, save the crop.
"""

from __future__ import annotations

import csv
import json
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.vision.document import detect_document, warp_document
from paperpod.vision.enhance import enhance_for_ocr
from paperpod.vision.motion import Segment, find_segments, motion_score, prepare_gray
from paperpod.vision.sharpness import sharpness_score


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
    )
    stable_windows = [s for s in segments if s.kind == "stable"]

    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: best frame per stable window, document detection -----------
    # Map each sampled timestamp range to its stable window for quick lookup.
    window_records: list[dict[str, Any]] = []
    best_frames: dict[int, tuple[float, Any]] = {}  # window_idx -> (sharpness, frame)

    def window_for(t: float) -> int | None:
        for i, w in enumerate(stable_windows):
            if w.t_start <= t <= w.t_end:
                return i
        return None

    if stable_windows:
        for frame in iter_frames(video_path, sample_fps=cfg.capture.sample_fps):
            idx = window_for(frame.t)
            if idx is None:
                continue
            score = sharpness_score(frame.image)
            if idx not in best_frames or score > best_frames[idx][0]:
                best_frames[idx] = (score, frame)

    for i, window in enumerate(stable_windows):
        record: dict[str, Any] = {
            "window_id": i,
            "t_start": round(window.t_start, 3),
            "t_end": round(window.t_end, 3),
            "duration_s": round(window.duration, 3),
            "document_found": False,
        }
        if i in best_frames:
            sharp, frame = best_frames[i]
            record["best_frame_t"] = round(frame.t, 3)
            record["best_frame_index"] = frame.index
            record["sharpness"] = round(sharp, 1)

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

            detection = detect_document(
                frame.image,
                min_area_ratio=cfg.document.min_area_ratio,
                canny_low=cfg.document.canny_low,
                canny_high=cfg.document.canny_high,
            )
            if detection is not None:
                crop = warp_document(frame.image, detection.quad)
                if cfg.document.enhance:
                    crop = enhance_for_ocr(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,
                    area_ratio=round(detection.area_ratio, 4),
                    quad=[[round(float(x), 1), round(float(y), 1)] for x, y in detection.quad],
                    crop_path=str(crop_path.relative_to(out_dir)),
                    crop_size=[crop.shape[1], crop.shape[0]],
                )
        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