myscan/tools/scan.py

#!/usr/bin/env python3
"""Scan with the scan button, like the original app: press it to start, glide the
mouse over the page, press it again to stop. Frames are registered while you scan;
after the stop press the page is assembled from all of them at once (see
myscan/pagegraph.py) and saved as <outdir>/scan-YYYYmmdd-HHMMSS.png. The mouse goes
back to mouse mode between scans (and on Ctrl-C)."""
import argparse
import csv
import os
import queue
import sys
import threading
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from myscan.session import ScanSession, quiet_libjpeg  # noqa: E402

quiet_libjpeg()
import cv2  # noqa: E402

from myscan import calib  # noqa: E402
from myscan.pagegraph import PageGraph  # noqa: E402
from myscan.xu import XU  # noqa: E402


def wait_press(x, poll=0.05):
    """Block until the scan button goes down then up again."""
    prev = 1
    while True:
        b = x.scan_button()
        if b and not prev:
            break
        prev = b
        time.sleep(poll)
    while x.scan_button():
        time.sleep(poll)


def scan_page(x, a, frames_dir=None, flat=None, undistort=None, gains=None):
    """Capture until the button is pressed again. Frames are registered (and saved) on a
    worker thread so capture never waits; returns the PageGraph, or None."""
    g, rows, work = None, [], queue.Queue()

    def worker():
        while (item := work.get()) is not None:
            n, im, v = item
            g.add(im, v)
            if frames_dir:
                cv2.imwrite(os.path.join(frames_dir, f"f{n:05d}.png"), im)
    th = None
    with ScanSession(a.video, a.target, a.exposure, xu=x, gains=gains) as s:
        down, last = 0, 0.0
        for t, n, im, rec in s.frames(color=not a.gray):
            if g is None:
                g = PageGraph(im.shape, flat=flat, undistort=undistort)
                th = threading.Thread(target=worker, daemon=True)
                th.start()
            work.put((n, im, rec.v))
            if frames_dir:
                rows.append([round(t, 4), n, rec.counter, *rec.v, rec.b14, rec.b15, rec.b16, rec.flags, s.e6,
                             round(float(im.mean()), 1), rec.raw.hex()])
            # Stop on a press: two "down" reads in a row (debounce), ignoring the first
            # 0.7 s so a bouncing start press can't end the scan.
            if n % 2 == 0:
                down = down + 1 if x.scan_button() else 0
                if down >= 2 and t > 0.7:
                    break
            if t >= a.max_time:
                print("max scan time reached", flush=True)
                break
            if t - last >= 1:
                print(f"\r  {t:5.1f}s  {n + 1} frames  ({g.n} registered)  E6={s.e6}", end="", flush=True)
                last = t
    print()
    if th:
        work.put(None)
        th.join()
    if frames_dir:   # same layout as tools/record.py, so tools/stitch.py can replay it
        with open(os.path.join(frames_dir, "records.csv"), "w", newline="") as fh:
            w = csv.writer(fh)
            w.writerow(["t", "frame", "counter", "v0", "v1", "v2", "v3", "v4", "v5",
                        "b14", "b15", "b16", "flags", "e6", "brightness", "raw"])
            w.writerows(rows)
    return g


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-o", "--outdir", default="scans")
    ap.add_argument("--once", action="store_true", help="exit after one page")
    ap.add_argument("--now", action="store_true", help="start the first scan immediately")
    ap.add_argument("--video", default="/dev/video4")
    ap.add_argument("--target", type=float, default=200, help="auto-exposure target for paper (95th pct of brightest channel)")
    ap.add_argument("--exposure", type=int, help="fixed E6 value instead of auto-exposure")
    ap.add_argument("--gray", action="store_true")
    ap.add_argument("--no-wb", action="store_true", help="skip white-patch balancing")
    ap.add_argument("--no-flat", action="store_true", help="skip the EEPROM flat-field correction")
    ap.add_argument("--no-undistort", action="store_true", help="skip the EEPROM lens correction")
    ap.add_argument("--save-frames", action="store_true", help="also keep raw frames + records.csv per scan")
    ap.add_argument("--max-time", type=float, default=180, help="stop a scan after this many seconds")
    a = ap.parse_args()
    os.makedirs(a.outdir, exist_ok=True)
    x = XU(a.video)
    print("reading calibration from the mouse's EEPROM (cached after the first time)", flush=True)
    cal = calib.load(x)
    flat = None if a.no_flat else calib.FlatField(cal)
    undistort = None if a.no_undistort else calib.Undistort(cal)
    gains = cal["gains_rgb"]
    first = True
    try:
        while True:
            if not (first and a.now):
                print("press the scan button to start (Ctrl-C to quit)", flush=True)
                wait_press(x)
            first = False
            print("scanning - press the scan button again to stop", flush=True)
            name = time.strftime("scan-%Y%m%d-%H%M%S")
            frames_dir = os.path.join(a.outdir, name) if a.save_frames else None
            if frames_dir:
                os.makedirs(frames_dir)
            g = scan_page(x, a, frames_dir, flat, undistort, gains)
            if g is None or g.n < 2:
                print("no frames captured")
            else:
                print("assembling the page...", flush=True)
                t0 = time.time()
                page = g.image(white_balance=not a.no_wb)
                path = os.path.join(a.outdir, name + ".png")
                cv2.imwrite(path, page)
                print(f"saved {path}  ({time.time() - t0:.1f} s; {g.stats})", flush=True)
            if a.once:
                break
    except KeyboardInterrupt:
        print("\nbye")


if __name__ == "__main__":
    main()