From c8b9f610bf954a1ddfea13308b7c32b93fcb7995 Mon Sep 17 00:00:00 2001 From: erik Date: Fri, 9 Oct 2026 07:55:22 +0000 Subject: [PATCH] Split stitching into a hardware-independent myscan.stitching package stitch.py and pagegraph.py mixed laser odometry, image registration, frame correction, canvas handling, pose-graph maths and the two stitchers. They are now modules of myscan/stitching, which imports nothing from the device code: geometry rigid transforms, pose vectors, warp_add filters grayscale, high-pass, feather weights correct Corrector: undistortion + flat field (and weighted variant) odometry LaserOdometry: the laser model and the phase-correlation override, shared by both stitchers register ECC, phase correlation, SIFT matching, overlap agreement canvas growable Canvas, layout, parallel render, normalize/compose posegraph robust Gauss-Newton, largest connected component loops keyframes and loop-closure candidates pagegraph PageGraph, now orchestration only incremental Stitcher, on the shared Canvas The stitchers take a Corrector and an odometry object (anything with predict()) instead of building them, so another device can supply its own motion prior. live.py drops its own canvas and normalisation for the shared ones; PageGraph's mosaic check uses the parallel renderer. Output on scan_text is unchanged (global stitcher identical, incremental within 1 grey level). stitch.py --live now applies lens correction too (--no-undistort restores the old output). Co-Authored-By: Claude Opus 5.5 --- FINDINGS.md | 4 +- README.md | 12 +- myscan/live.py | 83 ++---- myscan/pagegraph.py | 468 -------------------------------- myscan/stitch.py | 211 -------------- myscan/stitching/__init__.py | 30 ++ myscan/stitching/canvas.py | 149 ++++++++++ myscan/stitching/correct.py | 34 +++ myscan/stitching/filters.py | 35 +++ myscan/stitching/geometry.py | 51 ++++ myscan/stitching/incremental.py | 70 +++++ myscan/stitching/loops.py | 39 +++ myscan/stitching/odometry.py | 73 +++++ myscan/stitching/pagegraph.py | 221 +++++++++++++++ myscan/stitching/posegraph.py | 73 +++++ myscan/stitching/register.py | 92 +++++++ tools/scan.py | 14 +- tools/stitch.py | 13 +- 18 files changed, 907 insertions(+), 765 deletions(-) delete mode 100644 myscan/pagegraph.py delete mode 100644 myscan/stitch.py create mode 100644 myscan/stitching/__init__.py create mode 100644 myscan/stitching/canvas.py create mode 100644 myscan/stitching/correct.py create mode 100644 myscan/stitching/filters.py create mode 100644 myscan/stitching/geometry.py create mode 100644 myscan/stitching/incremental.py create mode 100644 myscan/stitching/loops.py create mode 100644 myscan/stitching/odometry.py create mode 100644 myscan/stitching/pagegraph.py create mode 100644 myscan/stitching/posegraph.py create mode 100644 myscan/stitching/register.py diff --git a/FINDINGS.md b/FINDINGS.md index 75bfec8..d4aa156 100644 --- a/FINDINGS.md +++ b/FINDINGS.md @@ -339,8 +339,8 @@ handling**. There is nothing in it to build a Linux driver from. - Record v0/v1 = laser 1 cumulative (a, b) counts, v3/v4 = laser 2; signed 16-bit, wrapping. ~3 counts per frame pixel. - Frame content shift: dx = -0.022*a - 0.355*b, dy = 0.324*a + 0.002*b (a, b = mean of both lasers). - Rotation: dtheta ≈ -4.4e-4 rad * (b1 - b2) (fit against ORB rotation over 10-30 frame windows; corr ~0.88). -- At fast motion (~25 px/frame) the lasers disagree (one over-counts by ~40%), so stitch.py checks each step against - frame-to-frame phase correlation, then ECC-refines every frame against the canvas. +- At fast motion (~25 px/frame) the lasers disagree (one over-counts by ~40%), so the odometry (myscan/stitching/odometry.py) + checks each step against frame-to-frame phase correlation, and the stitchers then ECC-refine it. - Frames are transposed + mirrored relative to the page; output = flip(transpose(canvas), -1). - `tools/stitch.py ` -> /page.png. Full 751-frame session stitches cleanly. diff --git a/README.md b/README.md index 8873762..b2448a8 100644 --- a/README.md +++ b/README.md @@ -6,9 +6,15 @@ Linux driver tools for the Hama myScan (00052343) scanner mouse, a Dacuda design - `myscan/xu.py` — UVC extension-unit access: switch between mouse and scan mode, LEDs, exposure - `myscan/frame.py` — fix the lying JPEG header and decode the per-frame sensor record (laser dx/dy) - `myscan/calib.py` — per-device calibration from the mouse's EEPROM: colour gains, flat field, lens undistortion -- `myscan/pagegraph.py` — assemble a page from all frames at once: frame-to-frame registration, loop closures - between passes, pose-graph optimisation (used by `scan.py` and `stitch.py`) -- `myscan/stitch.py` — the older live, frame-by-frame stitcher (`stitch.py --live`) and shared helpers +- `myscan/stitching/` — page stitching, independent of the scanner hardware (no device imports), so it can + be used elsewhere: `from myscan.stitching import Corrector, PageGraph`. Its modules, lowest level first: + `geometry` (rigid transforms, poses), `filters` (grayscale, high-pass, feather weights), `correct` + (`Corrector`: lens undistortion + flat field), `odometry` (`LaserOdometry`, the myScan's laser motion + prior; another device can pass its own object with `predict()`), `register` (ECC, phase correlation, SIFT + matching, overlap check), `canvas` (weighted canvases, page composition, parallel rendering), `posegraph` + (robust pose-graph optimisation), `loops` (loop-closure candidates), and the two stitchers built from them: + `pagegraph` (`PageGraph`: assembles the page from all frames at once with loop closures between passes; + used by `scan.py` and `stitch.py`) and `incremental` (`Stitcher`: frame by frame, `stitch.py --live`) - `tools/scan.py` — scan with the button: `python tools/scan.py` (press to start, press again to stop). `--view [PORT]` serves a live view at `http://:8000/`: the camera, the page as it is stitched (re-solved every 1.5 s, current position outlined), Start/Stop buttons, PNG/JPEG download of the finished diff --git a/myscan/live.py b/myscan/live.py index 046cc3f..f64a495 100644 --- a/myscan/live.py +++ b/myscan/live.py @@ -13,53 +13,9 @@ import traceback import numpy as np import cv2 -from .stitch import _aff, _feather - - -def page_orient(img): - """Frames and canvas are transposed and mirrored relative to the page (see stitch.compose).""" - return cv2.flip(cv2.transpose(img), -1) - - -class _Canvas: - """Growable weighted canvas in scaled page coordinates.""" - - def __init__(self, channels): - self.ch = channels - self.acc = self.ws = None - self.off = np.eye(3) # scaled page coords -> canvas pixels - - def paste(self, img, fe, A): - """Add img (with feather weights fe) placed by A (img px -> scaled page coords).""" - h, w = img.shape[:2] - q = (A @ np.array([[0, 0, 1], [w, 0, 1], [0, h, 1], [w, h, 1.0]]).T)[:2] - self._fit(q.min(1), q.max(1)) - x0, y0 = np.maximum(np.floor((self.off @ np.r_[q.min(1), 1])[:2]).astype(int) - 1, 0) - x1, y1 = np.ceil((self.off @ np.r_[q.max(1), 1])[:2]).astype(int) + 2 - x1, y1 = min(x1, self.ws.shape[1]), min(y1, self.ws.shape[0]) - L = (_aff(0, -x0, -y0) @ self.off @ A)[:2] - few = fe[..., None] if img.ndim == 3 else fe - self.acc[y0:y1, x0:x1] += cv2.warpAffine(img.astype(np.float32) * few, L, (x1 - x0, y1 - y0)) - self.ws[y0:y1, x0:x1] += cv2.warpAffine(fe, L, (x1 - x0, y1 - y0)) - - def _fit(self, lo, hi, pad=150): - """Grow the canvas so scaled page coords lo..hi are inside it.""" - if self.acc is None: - self.off = _aff(0, pad - lo[0], pad - lo[1]) - W, H = int(hi[0] - lo[0]) + 2 * pad, int(hi[1] - lo[1]) + 2 * pad - self.acc = np.zeros((H, W) + self.ch, np.float32) - self.ws = np.zeros((H, W), np.float32) - return - a = (self.off @ np.r_[lo, 1])[:2] - b = (self.off @ np.r_[hi, 1])[:2] - H, W = self.ws.shape - left, top = int(max(0, pad - a[0])) if a[0] < 10 else 0, int(max(0, pad - a[1])) if a[1] < 10 else 0 - right = int(max(0, b[0] + pad - W)) if b[0] > W - 10 else 0 - bottom = int(max(0, b[1] + pad - H)) if b[1] > H - 10 else 0 - if left or top or right or bottom: - self.acc = cv2.copyMakeBorder(self.acc, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0) - self.ws = cv2.copyMakeBorder(self.ws, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0) - self.off = _aff(0, left, top) @ self.off +from .stitching.canvas import Canvas, bbox, normalize, page_orient +from .stitching.filters import feather +from .stitching.geometry import corners, pose_matrix class LivePreview: @@ -69,7 +25,7 @@ class LivePreview: self.Si = np.diag([1 / scale, 1 / scale, 1.0]) self.small = [] # downscaled corrected frames, uint8 self.lock = threading.Lock() - self.c = None # current _Canvas + self.c = None # current Canvas self.pasted = 0 # frames < pasted are on the canvas self.cur = None # pose of the last pasted frame self.seen = None # solution the canvas was drawn from @@ -81,8 +37,8 @@ class LivePreview: small = cv2.resize(corrected, None, fx=self.s, fy=self.s, interpolation=cv2.INTER_AREA) self.small.append(small.clip(0, 255).astype(np.uint8)) if self.c is None: - self.fe = _feather(*small.shape[:2]) - self.c = _Canvas(small.shape[2:]) + self.fe = feather(*small.shape[:2]) + self.c = Canvas(small.shape[2] if small.ndim == 3 else 0) self.refresh() def refresh(self): @@ -94,7 +50,7 @@ class LivePreview: while self.pasted < len(self.small): i = self.pasted M = self.g.pose(i) if self.cur is None else self.cur @ self.g.odo[i - 1][2] - self.c.paste(self.small[i], self.fe, self.S @ M @ self.Si) + self.c.add(self.small[i], self.fe, self.S @ M @ self.Si) self.cur, self.pasted = M, i + 1 self.version += 1 @@ -106,13 +62,13 @@ class LivePreview: return x, use = sol k = min(len(x), len(self.small)) - c = _Canvas(self.c.ch) + c = Canvas(self.c.ch) for i in use: if i < k: - c.paste(self.small[i], self.fe, self.S @ _aff(x[i, 2], x[i, 0], x[i, 1]) @ self.Si) + c.add(self.small[i], self.fe, self.S @ pose_matrix(x[i]) @ self.Si) with self.lock: self.c, self.seen = c, sol - self.cur, self.pasted = _aff(x[k - 1, 2], x[k - 1, 0], x[k - 1, 1]), k + self.cur, self.pasted = pose_matrix(x[k - 1]), k self._catch_up() def render(self, background=40): @@ -121,22 +77,15 @@ class LivePreview: if self.c is None or self.c.acc is None: return None acc, ws, off, cur = self.c.acc.copy(), self.c.ws.copy(), self.c.off.copy(), self.cur - covered = ws > 0.05 - out = acc / np.maximum(ws, 1e-6)[..., None] if acc.ndim == 3 else acc / np.maximum(ws, 1e-6) - px = out[covered][::7] - if len(px): - out = out * (245 / np.maximum(np.percentile(px, 97, axis=0), 1)) - m = covered[..., None] if out.ndim == 3 else covered - out = np.where(m, out, background).clip(0, 255).astype(np.uint8) + out, covered = normalize(acc, ws, background=background) if out.ndim == 2: out = cv2.cvtColor(out, cv2.COLOR_GRAY2BGR) if cur is not None: - h, w = self.g.h, self.g.w - q = off @ self.S @ cur @ np.array([[0, 0, 1], [w, 0, 1], [w, h, 1], [0, h, 1.0]]).T - cv2.polylines(out, [np.round(q[:2].T).astype(np.int32)], True, (40, 40, 230), 2, cv2.LINE_AA) - ys, xs = np.nonzero(covered) - if len(ys): - out = out[max(ys.min() - 20, 0):ys.max() + 21, max(xs.min() - 20, 0):xs.max() + 21] + q = corners(off @ self.S @ cur, self.g.w, self.g.h)[:, [0, 1, 3, 2]] + cv2.polylines(out, [np.round(q.T).astype(np.int32)], True, (40, 40, 230), 2, cv2.LINE_AA) + box = bbox(covered, 20) + if box is not None: + out = out[box] return page_orient(out) diff --git a/myscan/pagegraph.py b/myscan/pagegraph.py deleted file mode 100644 index cfbdb06..0000000 --- a/myscan/pagegraph.py +++ /dev/null @@ -1,468 +0,0 @@ -"""Global page assembly: place all frames of a scan at once (pose-graph optimisation). - -The live Stitcher places each frame relative to the canvas built so far, so its error -accumulates and a later pass that comes back over earlier content can lock on a wrong -copy. Here every frame keeps its own pose and the poses are solved together: - -1. Odometry edges between consecutive frames: laser prediction (with the phase-correlation - override, as in Stitcher), refined by ECC between the two frames when that converges. - After a frame that ECC could not register, the next frame is also registered against - the one before it, which bridges single corrupt frames. -2. Loop-closure edges between non-consecutive keyframes whose current poses overlap: - SIFT matches + RANSAC (rigid), so large drift is recovered, then ECC for sub-pixel. -3. Robust Gauss-Newton over (x, y, theta) per frame. Loop edges that still disagree with - the solution are dropped (repeated content gives false matches) and the solve repeats; - rounds of 2-3 continue with the improved poses so passes far apart find each other. -4. Frames linked to the main page by image-verified edges are blended at their final - pose; the rest (off the page, lifted, corrupt) only had the lasers and are left out, - as are frames that disagree with the mosaic of all the others. - -A pose M (3x3) maps frame pixels to canvas pixels; an edge (i, j, T) says M_j ~ M_i @ T. -""" -import os -from concurrent.futures import ThreadPoolExecutor - -import numpy as np -import cv2 -from scipy.sparse import coo_matrix, diags -from scipy.sparse.linalg import spsolve - -from .stitch import (MIN_PHASE_RESP, ROT_GAIN, _aff, _feather, compose, - laser_deltas, laser_shift, relative_motion) - -ECC_CRIT = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 50, 1e-4) -SIGMA_ECC = 1.0 # px, edge registered by ECC -SIGMA_LASER = 6.0 # px, odometry edge from lasers / phase correlation only -SIGMA_LOOP = 1.5 # px, loop closure -HUBER = 3.0 # robust kernel width, in sigmas -PRUNE = 5.0 # drop loop edges whose residual exceeds this many sigmas -MAX_STEP_ROT = np.radians(2.0) # larger laser rotation between two frames is a miscount - - -def _rigid(W): - """2x3 (possibly with a little scale) -> 3x3 rigid transform.""" - th = np.arctan2(W[1, 0] - W[0, 1], W[0, 0] + W[1, 1]) - return _aff(th, W[0, 2], W[1, 2]) - - -def _ecc(template, image, T, scale=1.0, blur=5): - """Refine T (image pixels -> template pixels, full-resolution units) on images given at - `scale`. Returns (T, correlation) or (None, 0). Most of ECC's cost is per-call setup on - the whole image, so half resolution is ~4x faster; it then needs no pre-blur (blur=1).""" - S = np.diag([scale, scale, 1.0]) - Si = np.diag([1 / scale, 1 / scale, 1.0]) - W = np.linalg.inv(S @ T @ Si)[:2].astype(np.float32) # ECC warps template coords into image coords - try: - cc, W = cv2.findTransformECC(template, image, W, cv2.MOTION_EUCLIDEAN, ECC_CRIT, None, blur) - except cv2.error: - return None, 0.0 - return Si @ np.linalg.inv(np.vstack([W, [0, 0, 1]]).astype(np.float64)) @ S, cc - - -_GRAY = np.full((1, 3), 1 / 3, np.float32) - - -def _gray(f): - return cv2.transform(f, _GRAY) if f.ndim == 3 else f - - -def _highpass_fast(img, sigma=12, k=4): - """img minus its Gaussian-blurred self, the blur done on a k-times smaller copy (the - removed illumination is far smoother than k pixels): ~7x faster than at full size.""" - small = cv2.resize(img, None, fx=1 / k, fy=1 / k, interpolation=cv2.INTER_AREA) - low = cv2.GaussianBlur(small, (0, 0), sigma / k) - return cv2.subtract(img, cv2.resize(low, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LINEAR)) - - -def _warp_add(dst, img, M): - """dst += img warped by M (3x3, img pixels -> dst pixels), touching only the covered box.""" - h, w = img.shape[:2] - q = (M @ np.array([[0, 0, 1], [w, 0, 1], [0, h, 1], [w, h, 1.0]]).T)[:2] - x0, y0 = np.maximum(np.floor(q.min(1)).astype(int) - 1, 0) - x1, y1 = np.minimum(np.ceil(q.max(1)).astype(int) + 2, (dst.shape[1], dst.shape[0])) - if x1 > x0 and y1 > y0: - dst[y0:y1, x0:x1] += cv2.warpAffine(img, (_aff(0, -x0, -y0) @ M)[:2], (x1 - x0, y1 - y0), - flags=cv2.INTER_LINEAR) - - -def _close(A, B, max_deg, max_px): - D = np.linalg.solve(B, A) - return abs(np.degrees(np.arctan2(D[1, 0], D[0, 0]))) <= max_deg and np.hypot(D[0, 2], D[1, 2]) <= max_px - - -class PageGraph: - def __init__(self, frame_shape, flat=None, undistort=None, rot_gain=ROT_GAIN, min_cc=0.6): - """flat, undistort: calib.FlatField / calib.Undistort (raw frame -> corrected), or None.""" - self.h, self.w = frame_shape[:2] - self.flat, self.undistort, self.rot_gain, self.min_cc = flat, undistort, rot_gain, min_cc - self.center = (self.w / 2, self.h / 2) - self.half = (self.w // 2, self.h // 2) - self.window = cv2.createHanningWindow(self.half, cv2.CV_32F) - # Flat field pre-warped into undistorted coordinates, so a frame is remapped as uint8 - # and corrected with one multiply (cheaper than remapping float32). - self._inv = {} - if flat is not None: - self._inv = {3: flat.inv.astype(np.float32), 1: flat.inv_gray.astype(np.float32)} - if undistort is not None: - self._inv = {k: undistort(v) for k, v in self._inv.items()} - # registered: frame has an image-verified edge. hs: half-resolution high-pass images - # (odometry, verification); full-resolution ones are made on demand for keyframes. - self.frames, self.hs, self.registered = [], [], [] - self._hpf = {} - self.odo, self.bridges, self.loops = [], [], [] # edges: (i, j, T, sigma) - self.prev_v = None - self._feat, self._tried = {}, set() # kept across solves - self.solution = None # (poses (n, 3), indices of frames to draw) - self.stats = {"frames": 0, "odo_ecc": 0, "odo_laser": 0, "bridges": 0, "loop_tries": 0, "pruned": 0} - - @property - def n(self): - return len(self.frames) - - def _prep(self, frame): - """Raw uint8 frame -> undistorted, flat-fielded float32 frame.""" - if self.undistort is not None: - frame = self.undistort(frame) - if not self._inv: - return frame.astype(np.float32) - return cv2.multiply(frame, self._inv[3 if frame.ndim == 3 else 1], dtype=cv2.CV_32F) - - def _hp(self, i): - """Full-resolution high-pass image of frame i (cached; used for keyframes).""" - if i not in self._hpf: - self._hpf[i] = _highpass_fast(_gray(self._prep(self.frames[i]))) - return self._hpf[i] - - def add(self, frame, v): - """frame: HxW or HxWx3 uint8 raw frame; v: its record's six raw values. - Returns the corrected (flat-fielded, undistorted) float frame. - Safe to call while solve() runs on another thread.""" - f = self._prep(frame) - hs = cv2.resize(_highpass_fast(_gray(f)), self.half, interpolation=cv2.INTER_AREA) - v = np.asarray(v) - i = self.n - self.frames.append(frame) - self.hs.append(hs) - self.registered.append(False) - self.stats["frames"] += 1 - if i: - d = laser_deltas(np.stack([self.prev_v, v]))[1] - shift = None - (dx, dy), resp = cv2.phaseCorrelate(self.hs[i - 1], hs, self.window) - dx, dy = 2 * dx, 2 * dy - if resp > MIN_PHASE_RESP and np.hypot(*(laser_shift(d) - (dx, dy))) > 2: - shift = (dx, dy) - gain = self.rot_gain if abs(self.rot_gain * (d[1] - d[3])) < MAX_STEP_ROT else 0.0 - T0 = relative_motion(d, self.center, gain, shift) - T, cc = _ecc(self.hs[i - 1], hs, T0, 0.5, 1) - if T is not None and cc > self.min_cc and _close(T, T0, 3.0, 15): - self.odo.append((i - 1, i, T, SIGMA_ECC)) - self.registered[i - 1] = self.registered[i] = True - self.stats["odo_ecc"] += 1 - else: - self.odo.append((i - 1, i, T0, SIGMA_LASER)) - self.stats["odo_laser"] += 1 - if i >= 2 and self.odo[-2][3] != SIGMA_ECC: - Tb0 = self.odo[-2][2] @ self.odo[-1][2] - Tb, cc = _ecc(self.hs[i - 2], hs, Tb0, 0.5, 1) - if Tb is not None and cc > self.min_cc and _close(Tb, Tb0, 6.0, 60): - self.bridges.append((i - 2, i, Tb, SIGMA_ECC)) - self.registered[i - 2] = self.registered[i] = True - self.stats["bridges"] += 1 - self.prev_v = v - return f - - # --- poses ------------------------------------------------------------------- - - def pose(self, i): - """Current best pose (3x3) of frame i: the last solution, extended by odometry.""" - x = self.solution[0] if self.solution is not None else np.zeros((1, 3)) - k = min(i, len(x) - 1) - M = _aff(x[k, 2], x[k, 0], x[k, 1]) - for _, _, T, _ in self.odo[k:i]: - M = M @ T - return M - - def _chain(self, n): - """Initial poses for frames < n: the last solution, then odometry for newer frames.""" - x0 = self.solution[0] if self.solution is not None else np.zeros((1, 3)) - x0 = x0[:n] - P = [_aff(x0[-1, 2], x0[-1, 0], x0[-1, 1])] - for _, _, T, _ in self.odo[len(x0) - 1:n - 1]: - P.append(P[-1] @ T) - ext = np.array([[M[0, 2], M[1, 2], np.arctan2(M[1, 0], M[0, 0])] for M in P[1:]]).reshape(-1, 3) - return np.vstack([x0, ext]) - - def _optimise(self, x, edges, iters=20): - """Robust Gauss-Newton with frame 0 held fixed. x: (N, 3) -> (x, residual per edge in sigmas).""" - n, E = len(x), len(edges) - a = 150.0 # lever arm of the two residual points - P2 = np.array([[self.center[0] - a, self.center[1]], [self.center[0] + a, self.center[1]]]) - I = np.array([e[0] for e in edges]) - J = np.array([e[1] for e in edges]) - sig = np.array([e[3] for e in edges], np.float64) - Tp = np.einsum("eab,pb->epa", np.array([e[2] for e in edges])[:, :2], - np.hstack([P2, np.ones((2, 1))])) # (E, 2 points, 2): in frame i - Pj = np.broadcast_to(P2, (E, 2, 2)) - - def rot(c, s, q): - return np.stack([c[:, None] * q[..., 0] - s[:, None] * q[..., 1], - s[:, None] * q[..., 0] + c[:, None] * q[..., 1]], -1) - - def drot(c, s, q): - return np.stack([-s[:, None] * q[..., 0] - c[:, None] * q[..., 1], - c[:, None] * q[..., 0] - s[:, None] * q[..., 1]], -1) - - pin = np.zeros(n * 3) - pin[:3] = 1e9 - rows = np.arange(E * 4).reshape(E, 2, 2) - for _ in range(iters): - ci, si, cj, sj = np.cos(x[I, 2]), np.sin(x[I, 2]), np.cos(x[J, 2]), np.sin(x[J, 2]) - r = (rot(cj, sj, Pj) + x[J, None, :2] - rot(ci, si, Tp) - x[I, None, :2]) / sig[:, None, None] - nrm = np.sqrt((r ** 2).sum((1, 2)) / 2) - w = np.where(nrm <= HUBER, 1.0, HUBER / np.maximum(nrm, 1e-9)) - dj, di = drot(cj, sj, Pj) / sig[:, None, None], -drot(ci, si, Tp) / sig[:, None, None] - R, C, V = [], [], [] - for c in range(2): - for frame, sign, dth in ((J, 1.0, dj), (I, -1.0, di)): - for p in range(2): - R += [rows[:, p, c], rows[:, p, c]] - C += [frame * 3 + c, frame * 3 + 2] - V += [sign / sig, dth[:, p, c]] - Jm = coo_matrix((np.concatenate(V), (np.concatenate(R), np.concatenate(C))), - shape=(E * 4, n * 3)).tocsr() - Wd = diags(np.repeat(w, 4)) - H = (Jm.T @ Wd @ Jm + diags(pin + 1e-6)).tocsc() - dx = spsolve(H, -(Jm.T @ (Wd @ r.reshape(-1)))) - x = x + dx.reshape(n, 3) - if np.abs(dx).max() < 1e-3: - break - return x, nrm - - # --- loop closures --------------------------------------------------------------- - - def _keyframes(self, x, step_px=60.0): - keys, last = [], None - for i in range(len(x)): - if not self.registered[i]: - continue - if last is None or np.hypot(*(x[i, :2] - x[last, :2])) > step_px or i - last > 15: - keys.append(i) - last = i - return keys - - def _detect(self, i): - g = cv2.normalize(self._hp(i), None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U) - return cv2.SIFT_create(nfeatures=600).detectAndCompute(g, None) - - def _features(self, i): - if i not in self._feat: - self._feat[i] = self._detect(i) - return self._feat[i] - - def _match(self, a, b): - """Rigid transform b -> a from features + ECC, or None.""" - ka, da = self._features(a) - kb, db = self._features(b) - if da is None or db is None or len(ka) < 15 or len(kb) < 15: - return None - good = [m for m, n2 in (p for p in cv2.BFMatcher(cv2.NORM_L2).knnMatch(db, da, k=2) if len(p) == 2) - if m.distance < 0.75 * n2.distance] - if len(good) < 12: - return None - pb = np.float32([kb[m.queryIdx].pt for m in good]) - pa = np.float32([ka[m.trainIdx].pt for m in good]) - A, inl = cv2.estimateAffinePartial2D(pb, pa, method=cv2.RANSAC, ransacReprojThreshold=3.0) - if A is None or inl.sum() < 12 or abs(np.hypot(A[0, 0], A[1, 0]) - 1) > 0.05: - return None - T = _rigid(A) - T2, cc = _ecc(self._hp(a), self._hp(b), T) - if T2 is not None and cc > 0.5 and _close(T2, T, 1.0, 4) and self._agree(a, b, T2): - return T2 - return None - - def _agree(self, a, b, T, min_ncc=0.85, min_tiles=0.8, tile=32): - """Whole-overlap check. Repeated text (two lines ending in the same words) gives - matches that ECC confirms on the shared words while the rest of the overlap - disagrees: require high NCC over the overlap and agreement in most textured tiles.""" - ha = self._hp(a) - hb = cv2.warpAffine(self._hp(b), T[:2], (self.w, self.h)) - m = cv2.warpAffine(np.ones((self.h, self.w), np.float32), T[:2], (self.w, self.h)) > 0.99 - m = cv2.erode(m.astype(np.uint8), np.ones((15, 15), np.uint8)) - if cv2.countNonZero(m) < 0.05 * m.size: - return False - # Moments with OpenCV rather than numpy: this runs on many threads at once, and - # OpenCV releases the GIL. - (ma,), (sa,) = (v.ravel() for v in cv2.meanStdDev(ha, mask=m)) - (mb,), (sb,) = (v.ravel() for v in cv2.meanStdDev(hb, mask=m)) - if (cv2.mean(cv2.multiply(ha, hb), m)[0] - ma * mb) / max(sa * sb, 1e-6) < min_ncc: - return False - th, tw = self.h // tile, self.w // tile - # Per-tile means (exact for whole tiles) of a, b, a^2, b^2, ab and the mask. - mean = lambda im: cv2.resize(im[:th * tile, :tw * tile], (tw, th), interpolation=cv2.INTER_AREA) # noqa: E731 - A, B = ha.astype(np.float64), hb.astype(np.float64) - Ea, Eb = mean(A), mean(B) - va = np.maximum(mean(A * A) - Ea * Ea, 0) - vb = np.maximum(mean(B * B) - Eb * Eb, 0) - sa, sb = np.sqrt(va), np.sqrt(vb) - use = (mean(m.astype(np.float32)) > 0.9) & (np.maximum(sa, sb) > 8) - if not use.any(): - return True - ncc = (mean(A * B) - Ea * Eb) / np.maximum(sa * sb, 1e-6) - return (ncc[use] > 0.5).mean() >= min_tiles - - def _close_loops(self, x, keys, radius, min_gap=12, tries_per_run=3): - """Match each keyframe against overlapping later keyframes. Candidates are grouped - into runs of consecutive keyframes (one pass over the spot); each run is tried - nearest first and stops at its first match. Runs are matched in parallel.""" - P = x[keys, :2] - jobs = [] - for ia, a in enumerate(keys): - d = np.hypot(*(P[ia + 1:] - P[ia]).T) - near = [ib for ib in np.nonzero(d < radius)[0] + ia + 1 if keys[ib] - a >= min_gap] - runs, cur = [], [] - for ib in near: - if cur and ib - cur[-1] > 2: - runs.append(cur) - cur = [] - cur.append(ib) - runs += [cur] if cur else [] - for run in runs: - bs = [keys[ib] for ib in sorted(run, key=lambda k: d[k - ia - 1])[:tries_per_run]] - bs = [b for b in bs if (a, b) not in self._tried] - if bs: - jobs.append((a, bs)) - - def run_job(job): - a, bs = job - for k, b in enumerate(bs): - T = self._match(a, b) - if T is not None: - return (a, b, T, SIGMA_LOOP), k + 1 - return None, len(bs) - new = 0 - for (a, bs), (edge, tries) in zip(jobs, self._pool.map(run_job, jobs)): - self._tried.update((a, b) for b in bs[:tries]) - self.stats["loop_tries"] += tries - if edge is not None: - self.loops.append(edge) - new += 1 - return new - - def solve(self, n=None, rounds=4): - """Solve the poses of frames < n (default: all). Can run on another thread while - add() continues; earlier loop closures and features are reused, so repeated solves - during a scan only try the new pairs. Sets self.solution and returns the poses.""" - # add() appends a frame before its odometry edge; only take frames that have one. - n = min(self.n if n is None else n, len(self.odo) + 1) - x = self._chain(n) - keys = self._keyframes(x) - self._pool = ThreadPoolExecutor(os.cpu_count()) - todo = [k for k in keys if k not in self._feat] - self._feat.update(zip(todo, self._pool.map(self._detect, todo))) - base = self.odo[:n - 1] + [e for e in self.bridges if e[1] < n] - for r in range(rounds): - # Wide search first (drift can be hundreds of px), tighter once poses are good. - new = self._close_loops(x, keys, radius=700 if r == 0 else 500) - for _ in range(5): - x, nrm = self._optimise(x, base + self.loops) - bad = nrm[len(base):] > PRUNE - if not bad.any(): - break - self.loops = [e for e, b in zip(self.loops, bad) if not b] - self.stats["pruned"] += int(bad.sum()) - if r and not new: - break - use = self._verify(x, self._main_component(n, base)) - self._pool.shutdown() - self._pool = None - self.stats["loops"] = len(self.loops) - self.stats["rendered"] = len(use) - self.solution = (x, use) - return x - - def _main_component(self, n, base): - """Frames linked to the largest group by image-verified edges. Off-page stretches - register frame to frame but connect to the page only through the lasers.""" - parent = list(range(n)) - - def find(k): - while parent[k] != k: - parent[k] = parent[parent[k]] - k = parent[k] - return k - for i, j, _, s in base + self.loops: - if s != SIGMA_LASER: - parent[find(i)] = find(j) - roots = np.array([find(k) for k in range(n)]) - return list(np.nonzero(roots == np.bincount(roots).argmax())[0]) - - def _verify(self, x, use, min_ncc=0.5, scale=0.5): - """Frames of `use` that agree with the mosaic of all of them (drops a stretch held - only by a wrong registration). Compared on high-passed images at half resolution.""" - S = np.diag([scale, scale, 1.0]) - Ms, off, W, H = self._layout(x, use, S) - fe = cv2.resize(_feather(self.h, self.w), None, fx=scale, fy=scale) - small = [self.hs[i] for i in use] # already at half resolution - A = [off @ M @ np.linalg.inv(S) for M in Ms] - acc = np.zeros((H, W), np.float32) - ws = np.zeros((H, W), np.float32) - for im, M in zip(small, A): - _warp_add(acc, im * fe, M) - _warp_add(ws, fe, M) - mos = acc / np.maximum(ws, 1e-6) - h, w = fe.shape - - def score(k): - back = cv2.warpAffine(mos, np.linalg.inv(A[k])[:2], (w, h)) - return np.corrcoef(small[k][10:-10, 10:-10].ravel(), back[10:-10, 10:-10].ravel())[0, 1] - ncc = np.array(list(self._pool.map(score, range(len(use))))) - self.stats["inconsistent"] = int(np.sum(~(ncc >= min_ncc))) - return [i for i, c in zip(use, ncc) if c >= min_ncc] - - def _layout(self, x, use, S=np.eye(3)): - """Poses of frames `use` scaled by S, the offset that puts them all on a canvas, and its size.""" - Ms = [S @ _aff(x[i, 2], x[i, 0], x[i, 1]) for i in use] - c = np.array([[0, 0, 1], [self.w, 0, 1], [0, self.h, 1], [self.w, self.h, 1.0]]).T - q = np.concatenate([(M @ c)[:2] for M in Ms], 1) - off = _aff(0, 2 - q[0].min(), 2 - q[1].min()) - W, H = int(np.ceil(np.ptp(q[0]))) + 4, int(np.ceil(np.ptp(q[1]))) + 4 - return Ms, off, W, H - - # --- rendering --------------------------------------------------------------- - - def image(self, white_balance=True, background=255): - """Full-resolution page from the latest solution (solving first if frames were added since).""" - if self.solution is None or len(self.solution[0]) < self.n: - self.solve() - x, use = self.solution - Ms, off, W, H = self._layout(x, use) - Ms = [off @ M for M in Ms] - ch = self.frames[0].shape[2] if self.frames[0].ndim == 3 else 0 - acc = np.zeros((H, W, ch) if ch else (H, W), np.float32) - wsum = np.zeros((H, W), np.float32) - feather = _feather(self.h, self.w) - # Flat-field gain and feather weight in one multiply per frame. - fw = self._inv.get(ch or 1, np.ones(self.frames[0].shape, np.float32)) * (feather[..., None] if ch else feather) - - def weighted(i): - f = self.frames[i] if self.undistort is None else self.undistort(self.frames[i]) - return cv2.multiply(f, fw, dtype=cv2.CV_32F) - c = np.array([[0, 0, 1], [self.w, 0, 1], [0, self.h, 1], [self.w, self.h, 1.0]]).T - ys = np.array([(M @ c)[1] for M in Ms]) - top, bottom = ys.min(1) - 2, ys.max(1) + 2 - - # Threads draw disjoint row strips of the canvas (more strips only adds per-call - # overhead: drawing is memory-bound); a frame spanning several strips is weighted - # once per strip, which still beats one thread. - def strip(rows): - y0, y1 = rows - sh = _aff(0, 0, -y0) - for k in np.nonzero((bottom > y0) & (top < y1))[0]: - _warp_add(acc[y0:y1], weighted(use[k]), sh @ Ms[k]) - _warp_add(wsum[y0:y1], feather, sh @ Ms[k]) - n = os.cpu_count() or 1 - edges = np.linspace(0, H, n + 1).astype(int) - with ThreadPoolExecutor(n) as pool: - list(pool.map(strip, zip(edges[:-1], edges[1:]))) - return compose(acc, wsum, white_balance, background) diff --git a/myscan/stitch.py b/myscan/stitch.py deleted file mode 100644 index 28c7eb9..0000000 --- a/myscan/stitch.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Assemble scan frames into a page image. - -Poses come from dead reckoning on the two laser sensors embedded in each -frame's record (see frame.py), then each frame is refined against the canvas -built so far with ECC (Euclidean model) before it is blended in. - -Laser calibration (from scan_text, 2026-10-02; fit against image registration): - v0/v1 = laser 1 cumulative (a, b) counts, v3/v4 = laser 2; 16-bit wrapping. - content shift between consecutive frames, in frame pixels: - dx = -0.022*a - 0.355*b, dy = 0.324*a + 0.002*b (a, b = mean of both lasers) - image rotation (rad) = ROT_GAIN * (db1 - db2) -""" -import numpy as np -import cv2 - -SHIFT = np.array([[-0.0222, -0.3551], # dx from (a, b) - [0.3239, 0.0020]]) # dy from (a, b) -ROT_GAIN = -4.4e-4 -MIN_PHASE_RESP = 0.2 # phase-correlation peak needed to override the lasers - - -def _s16(v): - v = np.asarray(v, np.int64) & 0xFFFF - return np.where(v >= 0x8000, v - 0x10000, v) - - -def laser_deltas(v): - """v: (N, 6) record values -> (N, 4) per-frame deltas (a1, b1, a2, b2); first row zero.""" - p = _s16(np.asarray(v)[:, [0, 1, 3, 4]]) - d = np.diff(p, axis=0, prepend=p[:1]) - return (d + 0x8000) % 0x10000 - 0x8000 - - -def _aff(theta, tx, ty): - c, s = np.cos(theta), np.sin(theta) - return np.array([[c, -s, tx], [s, c, ty], [0, 0, 1.0]]) - - -def laser_shift(d): - """Content shift (dx, dy) in frame pixels predicted from one frame's laser deltas.""" - return SHIFT @ ((d[0] + d[2]) / 2, (d[1] + d[3]) / 2) - - -def relative_motion(d, center, rot_gain=ROT_GAIN, shift=None): - """3x3 transform mapping current-frame pixels to previous-frame pixels. - shift overrides the laser-predicted content shift.""" - dx, dy = laser_shift(d) if shift is None else shift - th = rot_gain * (d[1] - d[3]) - cx, cy = center - # content moved by (dx, dy) and rotated by th about the center: p_prev = R(-th)(p - c - d) + c - return _aff(0, cx, cy) @ _aff(-th, 0, 0) @ _aff(0, -cx - dx, -cy - dy) - - -def _highpass(img, mask=None, sigma=12): - """Remove illumination (vignetting, LED falloff) so registration sees only page texture. - With a mask, the local mean is taken over masked pixels only.""" - if mask is None: - return img - cv2.GaussianBlur(img, (0, 0), sigma) - m = mask.astype(np.float32) - mean = cv2.GaussianBlur(img * m, (0, 0), sigma) / np.maximum(cv2.GaussianBlur(m, (0, 0), sigma), 1e-3) - return (img - mean) * m - - -def _feather(h, w): - y = np.minimum(np.arange(h), np.arange(h)[::-1]) + 1 - x = np.minimum(np.arange(w), np.arange(w)[::-1]) + 1 - return np.minimum.outer(y, x).astype(np.float32).clip(max=40) / 40 - - -class Stitcher: - """Feed frames (grayscale or BGR) with their raw record values; registration runs on - grayscale, blending keeps all channels. Dropped frames are fine: the laser - positions are cumulative, so deltas are taken against the last frame added.""" - - def __init__(self, frame_shape, margin=600, refine=True, rot_gain=ROT_GAIN, flat=None): - self.h, self.w = frame_shape[:2] - self.ch = frame_shape[2] if len(frame_shape) > 2 else 1 - self.refine, self.rot_gain, self.flat = refine, rot_gain, flat - self.size = (self.w + 2 * margin, self.h + 2 * margin) # grows as needed - self.acc = np.zeros(self.size[::-1] + ((self.ch,) if self.ch > 1 else ()), np.float32) - self.wsum = np.zeros(self.size[::-1], np.float32) - self.pose = _aff(0, margin, margin) # frame -> canvas - self.feather = _feather(self.h, self.w) - self.window = cv2.createHanningWindow((self.w, self.h), cv2.CV_32F) - self.prev = None # previous frame, high-passed - self.prev_v = None # previous raw record values - self.n = 0 - self.stats = {"refined": 0, "rejected": 0, "skipped": 0, "image_shift": 0} - - def _corners(self, M): - c = np.array([[0, 0, 1], [self.w, 0, 1], [0, self.h, 1], [self.w, self.h, 1.0]]).T - return (M @ c)[:2] - - def _grow(self, M): - q = self._corners(M) - pad = 400 - left = int(max(0, pad - q[0].min())) if q[0].min() < 50 else 0 - top = int(max(0, pad - q[1].min())) if q[1].min() < 50 else 0 - right = int(max(0, q[0].max() + pad - self.size[0])) if q[0].max() > self.size[0] - 50 else 0 - bottom = int(max(0, q[1].max() + pad - self.size[1])) if q[1].max() > self.size[1] - 50 else 0 - if left or top or right or bottom: - self.acc = cv2.copyMakeBorder(self.acc, top, bottom, left, right, cv2.BORDER_CONSTANT, 0) - self.wsum = cv2.copyMakeBorder(self.wsum, top, bottom, left, right, cv2.BORDER_CONSTANT, 0) - self.size = self.wsum.shape[::-1] - shift = _aff(0, left, top) - self.pose = shift @ self.pose - return shift @ M - return M - - def _roi(self, M): - """Canvas rectangle (x0, y0, x1, y1) covered by a frame at pose M.""" - q = self._corners(M) - x0, y0 = np.floor(q.min(1)).astype(int) - 1 - x1, y1 = np.ceil(q.max(1)).astype(int) + 2 - return max(x0, 0), max(y0, 0), min(x1, self.size[0]), min(y1, self.size[1]) - - def _gray_canvas(self, x0, y0, x1, y1): - acc = self.acc[y0:y1, x0:x1] - return acc.mean(axis=2) if self.ch > 1 else acc - - def _refine(self, gray, M): - # Canvas content seen through the predicted pose, in frame coordinates. - x0, y0, x1, y1 = self._roi(M) - Mi = np.linalg.inv(_aff(0, -x0, -y0) @ M)[:2] - acc = cv2.warpAffine(self._gray_canvas(x0, y0, x1, y1), Mi, (self.w, self.h), flags=cv2.INTER_LINEAR) - ws = cv2.warpAffine(self.wsum[y0:y1, x0:x1], Mi, (self.w, self.h), flags=cv2.INTER_LINEAR) - mask = ws > 0.3 - if mask.mean() < 0.3: - self.stats["skipped"] += 1 - return M - patch = _highpass(np.where(mask, acc / np.maximum(ws, 1e-6), 0).astype(np.float32), mask) - W = np.eye(2, 3, dtype=np.float32) - try: - _, W = cv2.findTransformECC(patch, gray, W, cv2.MOTION_EUCLIDEAN, - (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 60, 1e-4), - mask.astype(np.uint8), 5) - except cv2.error: - self.stats["rejected"] += 1 - return M - # patch(x) ~ frame(W x), patch pixel x sits at canvas M x, so frame pixel y sits at M W^-1 y. - W3 = np.vstack([W, [0, 0, 1]]).astype(np.float64) - if abs(np.arctan2(W3[1, 0], W3[0, 0])) > np.radians(3) or np.hypot(W3[0, 2], W3[1, 2]) > 25: - self.stats["rejected"] += 1 - return M - self.stats["refined"] += 1 - return M @ np.linalg.inv(W3) - - def add(self, frame, v): - """frame: HxW or HxWx3 uint8; v: the record's six raw values (cumulative laser counts).""" - frame = frame.astype(np.float32) if self.flat is None else self.flat(frame) - gray = frame.mean(axis=2) if frame.ndim == 3 else frame - gray = np.ascontiguousarray(gray, np.float32) - hp = _highpass(gray) - v = np.asarray(v) - if self.n: - deltas = laser_deltas(np.stack([self.prev_v, v]))[1] - # The lasers sometimes miscount at speed (one or both), so check their translation - # against frame-to-frame phase correlation. High-passed frames keep the fixed - # illumination pattern from reading as "no motion" on blank paper. - shift = None - (dx, dy), resp = cv2.phaseCorrelate(self.prev, hp, self.window) - if resp > MIN_PHASE_RESP and np.hypot(*(laser_shift(deltas) - (dx, dy))) > 2: - shift = (dx, dy) - self.stats["image_shift"] += 1 - M = self.pose @ relative_motion(deltas, (self.w / 2, self.h / 2), self.rot_gain, shift) - M = self._grow(M) - if self.refine: - M = self._refine(hp, M) - else: - M = self.pose - self.pose, self.prev, self.prev_v = M, hp, v - x0, y0, x1, y1 = self._roi(M) - Ml = (_aff(0, -x0, -y0) @ M)[:2] - wf = self.feather[..., None] if frame.ndim == 3 else self.feather - self.acc[y0:y1, x0:x1] += cv2.warpAffine(frame * wf, Ml, (x1 - x0, y1 - y0), flags=cv2.INTER_LINEAR) - self.wsum[y0:y1, x0:x1] += cv2.warpAffine(self.feather, Ml, (x1 - x0, y1 - y0), flags=cv2.INTER_LINEAR) - self.n += 1 - - def image(self, white_balance=True, background=255): - return compose(self.acc, self.wsum, white_balance, background) - - -def compose(acc, wsum, white_balance=True, background=255): - """Weighted canvas (acc, wsum) -> cropped page image in page orientation.""" - covered = wsum > 0.05 - rows, cols = np.nonzero(covered.any(1))[0], np.nonzero(covered.any(0))[0] - if len(rows): - box = np.s_[rows[0]:rows[-1] + 1, cols[0]:cols[-1] + 1] - acc, wsum, covered = acc[box], wsum[box], covered[box] - inv = 1 / np.maximum(wsum, 1e-6) - gain = white_point(acc, inv, covered) if white_balance else np.ones(acc.shape[2:] or 1, np.float32) - if acc.ndim == 3: - out = acc * (inv[..., None] * gain.astype(np.float32)) - else: - out = acc * (inv * np.float32(gain[0])) - np.clip(out, 0, 255, out=out) - out = out.astype(np.uint8) - out[~covered] = background - # Frames arrive transposed and mirrored relative to the page. - return cv2.flip(cv2.transpose(out), -1) - - -def white_point(acc, inv, mask, pct=97, white=245, step=4): - """Per-channel gain mapping the bright percentile (the paper) of acc * inv to `white`, - estimated on every `step`-th pixel in each direction (the full canvas is ~20M pixels).""" - sub = mask[::step, ::step] - if not sub.any(): - return np.ones(acc.shape[2:] or 1, np.float32) - px = acc[::step, ::step][sub] * (inv[::step, ::step][sub][:, None] if acc.ndim == 3 else inv[::step, ::step][sub]) - ref = np.atleast_1d(np.percentile(px, pct, axis=0)) - return white / np.maximum(ref, 1) diff --git a/myscan/stitching/__init__.py b/myscan/stitching/__init__.py new file mode 100644 index 0000000..78d91ec --- /dev/null +++ b/myscan/stitching/__init__.py @@ -0,0 +1,30 @@ +"""Page stitching from overlapping camera frames. + +Independent of the scanner hardware: frames come in as uint8 images with whatever +sensor values the odometry needs. Typical use: + + from myscan.stitching import Corrector, PageGraph + g = PageGraph(frame.shape, Corrector(flat, undistort)) + for frame, v in frames: + g.add(frame, v) + page = g.image() + +Modules, lowest level first: + geometry rigid 2D transforms, pose vectors, warp-and-accumulate + filters grayscale, illumination high-pass, feather weights + correct Corrector: lens undistortion + flat field for raw frames + odometry motion prior between consecutive frames (the myScan's lasers) + register pairwise registration: ECC, phase correlation, SIFT matching, overlap check + canvas weighted canvases, page composition, parallel rendering + posegraph robust pose-graph optimisation + loops choosing loop-closure candidates + pagegraph PageGraph: global assembly of all frames (best quality) + incremental Stitcher: frame-by-frame stitching onto a growing canvas +""" +from .canvas import Canvas, compose, page_orient +from .correct import Corrector +from .incremental import Stitcher +from .odometry import ROT_GAIN, LaserOdometry +from .pagegraph import PageGraph + +__all__ = ["Canvas", "Corrector", "LaserOdometry", "PageGraph", "ROT_GAIN", "Stitcher", "compose", "page_orient"] diff --git a/myscan/stitching/canvas.py b/myscan/stitching/canvas.py new file mode 100644 index 0000000..0a777de --- /dev/null +++ b/myscan/stitching/canvas.py @@ -0,0 +1,149 @@ +"""Weighted canvases: frames are accumulated as sum(w * frame) and sum(w), and the page is +their ratio. Frames arrive transposed and mirrored relative to the page, so finished +images are turned with page_orient(). +""" +import os +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import cv2 + +from .geometry import affine, corners, warp_add + +COVERED = 0.05 # canvas weight below which a pixel counts as not scanned + + +class Canvas: + """Growable weighted canvas. Frames are placed in page coordinates; `off` maps page + coordinates to canvas pixels and changes when the canvas grows by `pad` pixels on a + side that a frame comes within `edge` pixels of.""" + + def __init__(self, channels=0, pad=150, edge=10): + self.ch, self.pad, self.edge = channels, pad, edge + self.acc = self.ws = None + self.off = np.eye(3) + + def add(self, img, weight, A): + """Add img (float, HxW or HxWxC) with weights `weight` (HxW), placed by A (img px -> page).""" + h, w = img.shape[:2] + q = corners(A, w, h) + self._fit(q.min(1), q.max(1)) + M = self.off @ A + warp_add(self.acc, img * (weight[..., None] if img.ndim == 3 else weight), M) + warp_add(self.ws, weight, M) + + def view(self, A, w, h): + """(acc, ws) of the canvas as seen through a w x h frame placed by A, in frame pixels.""" + if self.acc is None: + return np.zeros((h, w) + ((self.ch,) if self.ch else ()), np.float32), np.zeros((h, w), np.float32) + Mi = np.linalg.inv(self.off @ A)[:2] + return (cv2.warpAffine(self.acc, Mi, (w, h), flags=cv2.INTER_LINEAR), + cv2.warpAffine(self.ws, Mi, (w, h), flags=cv2.INTER_LINEAR)) + + def image(self, white_balance=True, background=255): + return compose(self.acc, self.ws, white_balance, background) + + def _fit(self, lo, hi): + """Grow the canvas so page coordinates lo..hi are inside it.""" + pad, edge = self.pad, self.edge + if self.acc is None: + self.off = affine(0, pad - lo[0], pad - lo[1]) + W, H = int(hi[0] - lo[0]) + 2 * pad, int(hi[1] - lo[1]) + 2 * pad + self.acc = np.zeros((H, W) + ((self.ch,) if self.ch else ()), np.float32) + self.ws = np.zeros((H, W), np.float32) + return + a = (self.off @ np.r_[lo, 1])[:2] + b = (self.off @ np.r_[hi, 1])[:2] + H, W = self.ws.shape + left = int(max(0, pad - a[0])) if a[0] < edge else 0 + top = int(max(0, pad - a[1])) if a[1] < edge else 0 + right = int(max(0, b[0] + pad - W)) if b[0] > W - edge else 0 + bottom = int(max(0, b[1] + pad - H)) if b[1] > H - edge else 0 + if left or top or right or bottom: + self.acc = cv2.copyMakeBorder(self.acc, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0) + self.ws = cv2.copyMakeBorder(self.ws, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0) + self.off = affine(0, left, top) @ self.off + + +def layout(poses, w, h, margin=2): + """Offset (3x3) that puts w x h frames at `poses` on a canvas with `margin` pixels to + spare, and the canvas size (W, H).""" + q = np.concatenate([corners(M, w, h) for M in poses], 1) + off = affine(0, margin - q[0].min(), margin - q[1].min()) + return off, int(np.ceil(np.ptp(q[0]))) + 2 * margin, int(np.ceil(np.ptp(q[1]))) + 2 * margin + + +def render(shape, poses, frame_size, weighted, weight, workers=None): + """Accumulate frames on a fixed canvas of `shape` ((H, W) or (H, W, C)): frame k, + weighted(k) (already multiplied by `weight`), placed by poses[k] (frame px -> canvas px). + Returns (acc, wsum). + + Threads draw disjoint row strips (more strips only adds per-call overhead: drawing is + memory-bound); a frame spanning several strips is weighted once per strip, which still + beats one thread.""" + H = shape[0] + acc = np.zeros(shape, np.float32) + wsum = np.zeros(shape[:2], np.float32) + w, h = frame_size + ys = np.array([corners(M, w, h)[1] for M in poses]).reshape(-1, 4) + top, bottom = ys.min(1) - 2, ys.max(1) + 2 + + def strip(rows): + y0, y1 = rows + sh = affine(0, 0, -y0) + for k in np.nonzero((bottom > y0) & (top < y1))[0]: + warp_add(acc[y0:y1], weighted(k), sh @ poses[k]) + warp_add(wsum[y0:y1], weight, sh @ poses[k]) + n = workers or os.cpu_count() or 1 + edges = np.linspace(0, H, n + 1).astype(int) + with ThreadPoolExecutor(n) as pool: + list(pool.map(strip, zip(edges[:-1], edges[1:]))) + return acc, wsum + + +def page_orient(img): + """Canvas or frame -> page orientation (frames are transposed and mirrored).""" + return cv2.flip(cv2.transpose(img), -1) + + +def bbox(covered, margin=0): + """Slice of the rows and columns that contain covered pixels (plus margin), or None.""" + rows, cols = np.nonzero(covered.any(1))[0], np.nonzero(covered.any(0))[0] + if not len(rows): + return None + return np.s_[max(rows[0] - margin, 0):rows[-1] + 1 + margin, max(cols[0] - margin, 0):cols[-1] + 1 + margin] + + +def normalize(acc, wsum, white_balance=True, background=255): + """Weighted canvas -> (uint8 image, covered mask); uncovered pixels get `background`.""" + covered = wsum > COVERED + inv = 1 / np.maximum(wsum, 1e-6) + gain = white_point(acc, inv, covered) if white_balance else np.ones(acc.shape[2:] or 1, np.float32) + if acc.ndim == 3: + out = acc * (inv[..., None] * gain.astype(np.float32)) + else: + out = acc * (inv * np.float32(gain[0])) + np.clip(out, 0, 255, out=out) + out = out.astype(np.uint8) + out[~covered] = background + return out, covered + + +def compose(acc, wsum, white_balance=True, background=255): + """Weighted canvas (acc, wsum) -> page image, cropped to the scanned area and in page + orientation.""" + box = bbox(wsum > COVERED) + if box is not None: + acc, wsum = acc[box], wsum[box] + return page_orient(normalize(acc, wsum, white_balance, background)[0]) + + +def white_point(acc, inv, mask, pct=97, white=245, step=4): + """Per-channel gain mapping the bright percentile (the paper) of acc * inv to `white`, + estimated on every `step`-th pixel in each direction (the full canvas is ~20M pixels).""" + sub = mask[::step, ::step] + if not sub.any(): + return np.ones(acc.shape[2:] or 1, np.float32) + px = acc[::step, ::step][sub] * (inv[::step, ::step][sub][:, None] if acc.ndim == 3 else inv[::step, ::step][sub]) + ref = np.atleast_1d(np.percentile(px, pct, axis=0)) + return white / np.maximum(ref, 1) diff --git a/myscan/stitching/correct.py b/myscan/stitching/correct.py new file mode 100644 index 0000000..abc72bb --- /dev/null +++ b/myscan/stitching/correct.py @@ -0,0 +1,34 @@ +"""Raw frame -> corrected float frame: lens undistortion, then flat-field gain.""" +import numpy as np +import cv2 + + +class Corrector: + """flat: an object with per-pixel gains `inv` (HxWx3) and `inv_gray` (HxW), such as + calib.FlatField; undistort: a callable remapping a frame, such as calib.Undistort. + Either may be None. The gain is pre-warped into undistorted coordinates, so a frame is + remapped as uint8 and corrected with one multiply (cheaper than remapping float32).""" + + def __init__(self, flat=None, undistort=None): + self.undistort = undistort + self.gain = {} # by channel count + if flat is not None: + self.gain = {3: flat.inv.astype(np.float32), 1: flat.inv_gray.astype(np.float32)} + if undistort is not None: + self.gain = {k: undistort(v) for k, v in self.gain.items()} + + def _remap(self, frame): + return frame if self.undistort is None else self.undistort(frame) + + def __call__(self, frame): + """HxW or HxWx3 uint8 -> corrected float32.""" + frame = self._remap(frame) + g = self.gain.get(frame.shape[2] if frame.ndim == 3 else 1) + return frame.astype(np.float32) if g is None else cv2.multiply(frame, g, dtype=cv2.CV_32F) + + def weighted(self, weight, channels): + """A function frame -> corrected frame * weight (HxW), with the gain and the weight + folded into one multiply per frame.""" + g = self.gain.get(channels or 1, np.float32(1)) + gw = g * (cv2.merge([weight] * channels) if channels else weight) + return lambda frame: cv2.multiply(self._remap(frame), gw, dtype=cv2.CV_32F) diff --git a/myscan/stitching/filters.py b/myscan/stitching/filters.py new file mode 100644 index 0000000..8f7c71d --- /dev/null +++ b/myscan/stitching/filters.py @@ -0,0 +1,35 @@ +"""Per-image filters: grayscale for registration, illumination removal, blend weights.""" +import numpy as np +import cv2 + +_GRAY = np.full((1, 3), 1 / 3, np.float32) + + +def gray(f): + """Channel mean of a float32 BGR image (unchanged if already single-channel).""" + return cv2.transform(f, _GRAY) if f.ndim == 3 else f + + +def highpass(img, mask=None, sigma=12): + """Remove illumination (vignetting, LED falloff) so registration sees only page texture. + With a mask, the local mean is taken over masked pixels only.""" + if mask is None: + return img - cv2.GaussianBlur(img, (0, 0), sigma) + m = mask.astype(np.float32) + mean = cv2.GaussianBlur(img * m, (0, 0), sigma) / np.maximum(cv2.GaussianBlur(m, (0, 0), sigma), 1e-3) + return (img - mean) * m + + +def highpass_fast(img, sigma=12, k=4): + """highpass() without a mask, the blur done on a k-times smaller copy (the removed + illumination is far smoother than k pixels): ~7x faster than at full size.""" + small = cv2.resize(img, None, fx=1 / k, fy=1 / k, interpolation=cv2.INTER_AREA) + low = cv2.GaussianBlur(small, (0, 0), sigma / k) + return cv2.subtract(img, cv2.resize(low, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LINEAR)) + + +def feather(h, w, ramp=40): + """Blend weight: 1 inside, falling linearly to 0 over `ramp` pixels at the borders.""" + y = np.minimum(np.arange(h), np.arange(h)[::-1]) + 1 + x = np.minimum(np.arange(w), np.arange(w)[::-1]) + 1 + return np.minimum.outer(y, x).astype(np.float32).clip(max=ramp) / ramp diff --git a/myscan/stitching/geometry.py b/myscan/stitching/geometry.py new file mode 100644 index 0000000..580c3ad --- /dev/null +++ b/myscan/stitching/geometry.py @@ -0,0 +1,51 @@ +"""Rigid 2D transforms as 3x3 matrices, and pose vectors (x, y, theta). + +A pose maps frame pixels to page (canvas) pixels; an edge transform T between frames +i and j says M_j ~ M_i @ T. +""" +import numpy as np +import cv2 + + +def affine(theta, tx, ty): + """Rotation by theta, then translation by (tx, ty).""" + c, s = np.cos(theta), np.sin(theta) + return np.array([[c, -s, tx], [s, c, ty], [0, 0, 1.0]]) + + +def pose_matrix(p): + """Pose vector (x, y, theta) -> 3x3.""" + return affine(p[2], p[0], p[1]) + + +def pose_vector(M): + """3x3 rigid transform -> pose vector (x, y, theta).""" + return np.array([M[0, 2], M[1, 2], np.arctan2(M[1, 0], M[0, 0])]) + + +def rigid(W): + """2x3 (possibly with a little scale) -> 3x3 rigid transform.""" + th = np.arctan2(W[1, 0] - W[0, 1], W[0, 0] + W[1, 1]) + return affine(th, W[0, 2], W[1, 2]) + + +def is_close(A, B, max_deg, max_px): + """A and B differ by at most max_deg of rotation and max_px of translation.""" + D = np.linalg.solve(B, A) + return abs(np.degrees(np.arctan2(D[1, 0], D[0, 0]))) <= max_deg and np.hypot(D[0, 2], D[1, 2]) <= max_px + + +def corners(M, w, h): + """Corners (2, 4) of a w x h image placed by M.""" + return (M @ np.array([[0, 0, 1], [w, 0, 1], [0, h, 1], [w, h, 1.0]]).T)[:2] + + +def warp_add(dst, img, M): + """dst += img warped by M (3x3, img pixels -> dst pixels), touching only the covered box.""" + h, w = img.shape[:2] + q = corners(M, w, h) + x0, y0 = np.maximum(np.floor(q.min(1)).astype(int) - 1, 0) + x1, y1 = np.minimum(np.ceil(q.max(1)).astype(int) + 2, (dst.shape[1], dst.shape[0])) + if x1 > x0 and y1 > y0: + dst[y0:y1, x0:x1] += cv2.warpAffine(img, (affine(0, -x0, -y0) @ M)[:2], (x1 - x0, y1 - y0), + flags=cv2.INTER_LINEAR) diff --git a/myscan/stitching/incremental.py b/myscan/stitching/incremental.py new file mode 100644 index 0000000..c674a93 --- /dev/null +++ b/myscan/stitching/incremental.py @@ -0,0 +1,70 @@ +"""Incremental stitching: each frame is placed by odometry from the previous one, refined +against the canvas built so far with ECC, and blended in at once. Cheap and immediate, +but its error accumulates; see pagegraph.py for the global assembly.""" +import numpy as np +import cv2 + +from .canvas import Canvas +from .correct import Corrector +from .filters import feather, gray, highpass +from .geometry import is_close +from .odometry import LaserOdometry +from .register import ecc, phase_shift + +REFINE_CRIT = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 60, 1e-4) + + +class Stitcher: + """Feed frames (grayscale or BGR) with their sensor values; registration runs on + grayscale, blending keeps all channels. Dropped frames are fine with the myScan lasers: + their positions are cumulative, so deltas are taken against the last frame added.""" + + def __init__(self, frame_shape, corrector=None, odometry=None, refine=True): + self.h, self.w = frame_shape[:2] + self.correct = corrector or Corrector() + self.odometry = odometry or LaserOdometry(max_step_rot=None) + self.refine = refine + self.canvas = Canvas(frame_shape[2] if len(frame_shape) > 2 else 0, pad=400, edge=50) + self.pose = np.eye(3) # frame -> page + self.feather = feather(self.h, self.w) + self.window = cv2.createHanningWindow((self.w, self.h), cv2.CV_32F) + self.prev = None # previous frame, high-passed + self.prev_v = None # previous sensor values + self.n = 0 + self.stats = {"refined": 0, "rejected": 0, "skipped": 0, "image_shift": 0} + + def _refine(self, hp, M): + """Pose M refined by registering the frame against the canvas seen through it.""" + acc, ws = self.canvas.view(M, self.w, self.h) + mask = ws > 0.3 + if mask.mean() < 0.3: + self.stats["skipped"] += 1 + return M + patch = highpass(np.where(mask, gray(acc) / np.maximum(ws, 1e-6), 0).astype(np.float32), mask) + T, _ = ecc(patch, hp, np.eye(3), crit=REFINE_CRIT, mask=mask.astype(np.uint8)) + if T is None or not is_close(T, np.eye(3), 3.0, 25): + self.stats["rejected"] += 1 + return M + self.stats["refined"] += 1 + return M @ T + + def add(self, frame, v): + """frame: HxW or HxWx3 uint8; v: its sensor values for the odometry (the myScan + record's six raw values).""" + f = self.correct(frame) + hp = highpass(np.ascontiguousarray(gray(f), np.float32)) + v = np.asarray(v) + M = self.pose + if self.n: + shift = phase_shift(self.prev, hp, self.window) + T, used = self.odometry.predict(self.prev_v, v, (self.w / 2, self.h / 2), shift) + self.stats["image_shift"] += used + M = self.pose @ T + if self.refine: + M = self._refine(hp, M) + self.pose, self.prev, self.prev_v = M, hp, v + self.canvas.add(f, self.feather, M) + self.n += 1 + + def image(self, white_balance=True, background=255): + return self.canvas.image(white_balance, background) diff --git a/myscan/stitching/loops.py b/myscan/stitching/loops.py new file mode 100644 index 0000000..ec00054 --- /dev/null +++ b/myscan/stitching/loops.py @@ -0,0 +1,39 @@ +"""Choosing which frames to try matching for loop closures.""" +import numpy as np + + +def keyframes(x, usable, step_px=60.0, max_gap=15): + """Indices of usable poses (rows of x) spaced at least step_px apart, or max_gap frames.""" + keys, last = [], None + for i in range(len(x)): + if not usable[i]: + continue + if last is None or np.hypot(*(x[i, :2] - x[last, :2])) > step_px or i - last > max_gap: + keys.append(i) + last = i + return keys + + +def candidates(x, keys, radius, tried=(), min_gap=12, tries_per_run=3): + """Jobs (a, [b, ...]): keyframe a and later keyframes b whose current poses lie within + `radius` px, at least min_gap frames on, not yet in `tried`. Candidates are grouped into + runs of consecutive keyframes (one pass over the spot); each run is one job, nearest + first, to be tried until its first match.""" + P = x[keys, :2] + jobs = [] + for ia, a in enumerate(keys): + d = np.hypot(*(P[ia + 1:] - P[ia]).T) + near = [ib for ib in np.nonzero(d < radius)[0] + ia + 1 if keys[ib] - a >= min_gap] + runs, cur = [], [] + for ib in near: + if cur and ib - cur[-1] > 2: + runs.append(cur) + cur = [] + cur.append(ib) + runs += [cur] if cur else [] + for run in runs: + bs = [keys[ib] for ib in sorted(run, key=lambda k: d[k - ia - 1])[:tries_per_run]] + bs = [b for b in bs if (a, b) not in tried] + if bs: + jobs.append((a, bs)) + return jobs diff --git a/myscan/stitching/odometry.py b/myscan/stitching/odometry.py new file mode 100644 index 0000000..be00ffd --- /dev/null +++ b/myscan/stitching/odometry.py @@ -0,0 +1,73 @@ +"""Motion prior from the myScan's two laser sensors (cumulative counts in each frame's record). + +Laser calibration (from scan_text, 2026-10-02; fit against image registration): + v0/v1 = laser 1 cumulative (a, b) counts, v3/v4 = laser 2; 16-bit wrapping. + content shift between consecutive frames, in frame pixels: + dx = -0.022*a - 0.355*b, dy = 0.324*a + 0.002*b (a, b = mean of both lasers) + image rotation (rad) = ROT_GAIN * (db1 - db2) + +The stitchers only use an odometry object's predict(); another device can supply its own. +""" +import numpy as np + +from .geometry import affine + +SHIFT = np.array([[-0.0222, -0.3551], # dx from (a, b) + [0.3239, 0.0020]]) # dy from (a, b) +ROT_GAIN = -4.4e-4 +MIN_PHASE_RESP = 0.2 # phase-correlation peak needed to override the lasers +MAX_STEP_ROT = np.radians(2.0) # larger laser rotation between two frames is a miscount + + +def _s16(v): + v = np.asarray(v, np.int64) & 0xFFFF + return np.where(v >= 0x8000, v - 0x10000, v) + + +def laser_deltas(v): + """v: (N, 6) record values -> (N, 4) per-frame deltas (a1, b1, a2, b2); first row zero.""" + p = _s16(np.asarray(v)[:, [0, 1, 3, 4]]) + d = np.diff(p, axis=0, prepend=p[:1]) + return (d + 0x8000) % 0x10000 - 0x8000 + + +def laser_shift(d): + """Content shift (dx, dy) in frame pixels predicted from one frame's laser deltas.""" + return SHIFT @ ((d[0] + d[2]) / 2, (d[1] + d[3]) / 2) + + +def relative_motion(d, center, rot_gain=ROT_GAIN, shift=None): + """3x3 transform mapping current-frame pixels to previous-frame pixels. + shift overrides the laser-predicted content shift.""" + dx, dy = laser_shift(d) if shift is None else shift + th = rot_gain * (d[1] - d[3]) + cx, cy = center + # content moved by (dx, dy) and rotated by th about the center: p_prev = R(-th)(p - c - d) + c + return affine(0, cx, cy) @ affine(-th, 0, 0) @ affine(0, -cx - dx, -cy - dy) + + +class LaserOdometry: + """predict(prev_v, v, center, image_shift) -> (T, used_image_shift): the transform from + the current frame to the previous one. + + The lasers sometimes miscount at speed (one or both), so their translation is checked + against the image's own frame-to-frame shift (phase correlation of high-passed frames, + which keeps the fixed illumination pattern from reading as "no motion" on blank paper) + and replaced by it when they disagree. With max_step_rot, a laser rotation larger than + that is taken as a miscount and ignored.""" + + def __init__(self, rot_gain=ROT_GAIN, max_step_rot=MAX_STEP_ROT, min_phase_resp=MIN_PHASE_RESP): + self.rot_gain, self.max_step_rot, self.min_phase_resp = rot_gain, max_step_rot, min_phase_resp + + def predict(self, prev_v, v, center, image_shift=None): + """image_shift: ((dx, dy), response) of the content between the two frames, or None.""" + d = laser_deltas(np.stack([prev_v, v]))[1] + shift = None + if image_shift is not None: + (dx, dy), resp = image_shift + if resp > self.min_phase_resp and np.hypot(*(laser_shift(d) - (dx, dy))) > 2: + shift = (dx, dy) + gain = self.rot_gain + if self.max_step_rot is not None and abs(gain * (d[1] - d[3])) >= self.max_step_rot: + gain = 0.0 + return relative_motion(d, center, gain, shift), shift is not None diff --git a/myscan/stitching/pagegraph.py b/myscan/stitching/pagegraph.py new file mode 100644 index 0000000..b363ff1 --- /dev/null +++ b/myscan/stitching/pagegraph.py @@ -0,0 +1,221 @@ +"""Global page assembly: place all frames of a scan at once (pose-graph optimisation). + +The incremental Stitcher places each frame relative to the canvas built so far, so its +error accumulates and a later pass that comes back over earlier content can lock on a +wrong copy. Here every frame keeps its own pose and the poses are solved together: + +1. Odometry edges between consecutive frames: the odometry prediction (lasers, checked + against phase correlation), refined by ECC between the two frames when that converges. + After a frame that ECC could not register, the next frame is also registered against + the one before it, which bridges single corrupt frames. +2. Loop-closure edges between non-consecutive keyframes whose current poses overlap: + SIFT matches + RANSAC (rigid), so large drift is recovered, then ECC for sub-pixel. +3. Robust Gauss-Newton over (x, y, theta) per frame. Loop edges that still disagree with + the solution are dropped (repeated content gives false matches) and the solve repeats; + rounds of 2-3 continue with the improved poses so passes far apart find each other. +4. Frames linked to the main page by image-verified edges are blended at their final + pose; the rest (off the page, lifted, corrupt) only had the odometry and are left out, + as are frames that disagree with the mosaic of all the others. + +A pose M (3x3) maps frame pixels to canvas pixels; an edge (i, j, T) says M_j ~ M_i @ T. +""" +import os +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import cv2 + +from . import loops, register +from .canvas import compose, layout, render +from .correct import Corrector +from .filters import feather, gray, highpass_fast +from .geometry import is_close, pose_matrix, pose_vector +from .odometry import LaserOdometry +from .posegraph import largest_component, optimise + +SIGMA_ECC = 1.0 # px, edge registered by ECC +SIGMA_ODO = 6.0 # px, odometry edge from the motion prior only +SIGMA_LOOP = 1.5 # px, loop closure +PRUNE = 5.0 # drop loop edges whose residual exceeds this many sigmas + + +class PageGraph: + def __init__(self, frame_shape, corrector=None, odometry=None, min_cc=0.6): + """corrector: Corrector for raw frames (default: none); odometry: motion prior with + predict(prev_v, v, center, image_shift) (default: the myScan lasers).""" + self.h, self.w = frame_shape[:2] + self.correct = corrector or Corrector() + self.odometry = odometry or LaserOdometry() + self.min_cc = min_cc + self.center = (self.w / 2, self.h / 2) + self.half = (self.w // 2, self.h // 2) + self.window = cv2.createHanningWindow(self.half, cv2.CV_32F) + # registered: frame has an image-verified edge. hs: half-resolution high-pass images + # (odometry, verification); full-resolution ones are made on demand for keyframes. + self.frames, self.hs, self.registered = [], [], [] + self._hpf = {} + self.odo, self.bridges, self.loops = [], [], [] # edges: (i, j, T, sigma) + self.prev_v = None + self._feat, self._tried = {}, set() # kept across solves + self._pool = None + self.solution = None # (poses (n, 3), indices of frames to draw) + self.stats = {"frames": 0, "odo_ecc": 0, "odo_laser": 0, "bridges": 0, "loop_tries": 0, "pruned": 0} + + @property + def n(self): + return len(self.frames) + + def _hp(self, i): + """Full-resolution high-pass image of frame i (cached; used for keyframes).""" + if i not in self._hpf: + self._hpf[i] = highpass_fast(gray(self.correct(self.frames[i]))) + return self._hpf[i] + + def _features(self, i): + if i not in self._feat: + self._feat[i] = register.detect(self._hp(i)) + return self._feat[i] + + def add(self, frame, v): + """frame: HxW or HxWx3 uint8 raw frame; v: its sensor values for the odometry (the + myScan record's six raw values). Returns the corrected float frame. + Safe to call while solve() runs on another thread.""" + f = self.correct(frame) + hs = cv2.resize(highpass_fast(gray(f)), self.half, interpolation=cv2.INTER_AREA) + v = np.asarray(v) + i = self.n + self.frames.append(frame) + self.hs.append(hs) + self.registered.append(False) + self.stats["frames"] += 1 + if i: + shift = register.phase_shift(self.hs[i - 1], hs, self.window, 0.5) + T0, _ = self.odometry.predict(self.prev_v, v, self.center, shift) + T, cc = register.ecc(self.hs[i - 1], hs, T0, 0.5, 1) + if T is not None and cc > self.min_cc and is_close(T, T0, 3.0, 15): + self.odo.append((i - 1, i, T, SIGMA_ECC)) + self.registered[i - 1] = self.registered[i] = True + self.stats["odo_ecc"] += 1 + else: + self.odo.append((i - 1, i, T0, SIGMA_ODO)) + self.stats["odo_laser"] += 1 + if i >= 2 and self.odo[-2][3] != SIGMA_ECC: + Tb0 = self.odo[-2][2] @ self.odo[-1][2] + Tb, cc = register.ecc(self.hs[i - 2], hs, Tb0, 0.5, 1) + if Tb is not None and cc > self.min_cc and is_close(Tb, Tb0, 6.0, 60): + self.bridges.append((i - 2, i, Tb, SIGMA_ECC)) + self.registered[i - 2] = self.registered[i] = True + self.stats["bridges"] += 1 + self.prev_v = v + return f + + # --- poses ------------------------------------------------------------------- + + def pose(self, i): + """Current best pose (3x3) of frame i: the last solution, extended by odometry.""" + x = self.solution[0] if self.solution is not None else np.zeros((1, 3)) + k = min(i, len(x) - 1) + M = pose_matrix(x[k]) + for _, _, T, _ in self.odo[k:i]: + M = M @ T + return M + + def _chain(self, n): + """Initial poses for frames < n: the last solution, then odometry for newer frames.""" + x0 = self.solution[0] if self.solution is not None else np.zeros((1, 3)) + x0 = x0[:n] + P = [pose_matrix(x0[-1])] + for _, _, T, _ in self.odo[len(x0) - 1:n - 1]: + P.append(P[-1] @ T) + return np.vstack([x0, np.array([pose_vector(M) for M in P[1:]]).reshape(-1, 3)]) + + def solve(self, n=None, rounds=4): + """Solve the poses of frames < n (default: all). Can run on another thread while + add() continues; earlier loop closures and features are reused, so repeated solves + during a scan only try the new pairs. Sets self.solution and returns the poses.""" + # add() appends a frame before its odometry edge; only take frames that have one. + n = min(self.n if n is None else n, len(self.odo) + 1) + x = self._chain(n) + keys = loops.keyframes(x, self.registered) + self._pool = ThreadPoolExecutor(os.cpu_count()) + list(self._pool.map(self._features, keys)) + base = self.odo[:n - 1] + [e for e in self.bridges if e[1] < n] + for r in range(rounds): + # Wide search first (drift can be hundreds of px), tighter once poses are good. + new = self._close_loops(x, keys, radius=700 if r == 0 else 500) + for _ in range(5): + x, nrm = optimise(x, base + self.loops, self.center) + bad = nrm[len(base):] > PRUNE + if not bad.any(): + break + self.loops = [e for e, b in zip(self.loops, bad) if not b] + self.stats["pruned"] += int(bad.sum()) + if r and not new: + break + # Frames linked to the largest group by image-verified edges. Off-page stretches + # register frame to frame but connect to the page only through the odometry. + main = largest_component(n, [e for e in base + self.loops if e[3] != SIGMA_ODO]) + use = self._verify(x, main) + self._pool.shutdown() + self._pool = None + self.stats["loops"] = len(self.loops) + self.stats["rendered"] = len(use) + self.solution = (x, use) + return x + + def _close_loops(self, x, keys, radius): + """Try the loop-closure candidates (loops.candidates) in parallel; adds the new loop + edges and returns how many there are.""" + jobs = loops.candidates(x, keys, radius, self._tried) + + def run_job(job): + a, bs = job + for k, b in enumerate(bs): + T = register.match(self._features(a), self._features(b), self._hp(a), self._hp(b)) + if T is not None: + return (a, b, T, SIGMA_LOOP), k + 1 + return None, len(bs) + new = 0 + for (a, bs), (edge, tries) in zip(jobs, self._pool.map(run_job, jobs)): + self._tried.update((a, b) for b in bs[:tries]) + self.stats["loop_tries"] += tries + if edge is not None: + self.loops.append(edge) + new += 1 + return new + + def _verify(self, x, use, min_ncc=0.5, scale=0.5): + """Frames of `use` that agree with the mosaic of all of them (drops a stretch held + only by a wrong registration). Compared on high-passed images at half resolution.""" + S = np.diag([scale, scale, 1.0]) + Ms = [S @ pose_matrix(x[i]) for i in use] + off, W, H = layout(Ms, self.w, self.h) + fe = cv2.resize(feather(self.h, self.w), None, fx=scale, fy=scale) + small = [self.hs[i] for i in use] # already at half resolution + A = [off @ M @ np.linalg.inv(S) for M in Ms] + h, w = fe.shape + acc, ws = render((H, W), A, (w, h), lambda k: small[k] * fe, fe) + mos = acc / np.maximum(ws, 1e-6) + + def score(k): + back = cv2.warpAffine(mos, np.linalg.inv(A[k])[:2], (w, h)) + return np.corrcoef(small[k][10:-10, 10:-10].ravel(), back[10:-10, 10:-10].ravel())[0, 1] + ncc = np.array(list(self._pool.map(score, range(len(use))))) + self.stats["inconsistent"] = int(np.sum(~(ncc >= min_ncc))) + return [i for i, c in zip(use, ncc) if c >= min_ncc] + + # --- rendering --------------------------------------------------------------- + + def image(self, white_balance=True, background=255): + """Full-resolution page from the latest solution (solving first if frames were added since).""" + if self.solution is None or len(self.solution[0]) < self.n: + self.solve() + x, use = self.solution + Ms = [pose_matrix(x[i]) for i in use] + off, W, H = layout(Ms, self.w, self.h) + ch = self.frames[0].shape[2] if self.frames[0].ndim == 3 else 0 + weight = feather(self.h, self.w) + weighted = self.correct.weighted(weight, ch) + acc, wsum = render((H, W, ch) if ch else (H, W), [off @ M for M in Ms], (self.w, self.h), + lambda k: weighted(self.frames[use[k]]), weight) + return compose(acc, wsum, white_balance, background) diff --git a/myscan/stitching/posegraph.py b/myscan/stitching/posegraph.py new file mode 100644 index 0000000..e0d827d --- /dev/null +++ b/myscan/stitching/posegraph.py @@ -0,0 +1,73 @@ +"""Pose-graph optimisation over rigid 2D poses. + +Poses are rows (x, y, theta) of an (N, 3) array; an edge (i, j, T, sigma) says pose j ~ +pose i @ T, with sigma its uncertainty in pixels. +""" +import numpy as np +from scipy.sparse import coo_matrix, diags +from scipy.sparse.linalg import spsolve + +HUBER = 3.0 # robust kernel width, in sigmas + + +def optimise(x, edges, center, lever=150.0, huber=HUBER, iters=20): + """Robust Gauss-Newton with pose 0 held fixed. Each edge is measured at two points + `lever` px either side of `center` (frame pixels), so its residual weighs rotation + and translation in pixels. Returns (x, residual per edge in sigmas).""" + n, E = len(x), len(edges) + P2 = np.array([[center[0] - lever, center[1]], [center[0] + lever, center[1]]]) + I = np.array([e[0] for e in edges]) + J = np.array([e[1] for e in edges]) + sig = np.array([e[3] for e in edges], np.float64) + Tp = np.einsum("eab,pb->epa", np.array([e[2] for e in edges])[:, :2], + np.hstack([P2, np.ones((2, 1))])) # (E, 2 points, 2): in frame i + Pj = np.broadcast_to(P2, (E, 2, 2)) + + def rot(c, s, q): + return np.stack([c[:, None] * q[..., 0] - s[:, None] * q[..., 1], + s[:, None] * q[..., 0] + c[:, None] * q[..., 1]], -1) + + def drot(c, s, q): + return np.stack([-s[:, None] * q[..., 0] - c[:, None] * q[..., 1], + c[:, None] * q[..., 0] - s[:, None] * q[..., 1]], -1) + + pin = np.zeros(n * 3) + pin[:3] = 1e9 + rows = np.arange(E * 4).reshape(E, 2, 2) + for _ in range(iters): + ci, si, cj, sj = np.cos(x[I, 2]), np.sin(x[I, 2]), np.cos(x[J, 2]), np.sin(x[J, 2]) + r = (rot(cj, sj, Pj) + x[J, None, :2] - rot(ci, si, Tp) - x[I, None, :2]) / sig[:, None, None] + nrm = np.sqrt((r ** 2).sum((1, 2)) / 2) + w = np.where(nrm <= huber, 1.0, huber / np.maximum(nrm, 1e-9)) + dj, di = drot(cj, sj, Pj) / sig[:, None, None], -drot(ci, si, Tp) / sig[:, None, None] + R, C, V = [], [], [] + for c in range(2): + for frame, sign, dth in ((J, 1.0, dj), (I, -1.0, di)): + for p in range(2): + R += [rows[:, p, c], rows[:, p, c]] + C += [frame * 3 + c, frame * 3 + 2] + V += [sign / sig, dth[:, p, c]] + Jm = coo_matrix((np.concatenate(V), (np.concatenate(R), np.concatenate(C))), + shape=(E * 4, n * 3)).tocsr() + Wd = diags(np.repeat(w, 4)) + H = (Jm.T @ Wd @ Jm + diags(pin + 1e-6)).tocsc() + dx = spsolve(H, -(Jm.T @ (Wd @ r.reshape(-1)))) + x = x + dx.reshape(n, 3) + if np.abs(dx).max() < 1e-3: + break + return x, nrm + + +def largest_component(n, edges): + """Indices of the largest set of poses < n connected by `edges`.""" + parent = list(range(n)) + + def find(k): + while parent[k] != k: + parent[k] = parent[parent[k]] + k = parent[k] + return k + for e in edges: + parent[find(e[0])] = find(e[1]) + roots = np.array([find(k) for k in range(n)]) + return list(np.nonzero(roots == np.bincount(roots).argmax())[0]) diff --git a/myscan/stitching/register.py b/myscan/stitching/register.py new file mode 100644 index 0000000..216d8cc --- /dev/null +++ b/myscan/stitching/register.py @@ -0,0 +1,92 @@ +"""Pairwise registration of high-passed grayscale images. + +All transforms map pixels of the second image (b) to pixels of the first (a), as 3x3 +rigid matrices in full-resolution units. +""" +import numpy as np +import cv2 + +from .geometry import is_close, rigid + +ECC_CRIT = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 50, 1e-4) + + +def ecc(template, image, T, scale=1.0, blur=5, crit=ECC_CRIT, mask=None): + """Refine T (image pixels -> template pixels, full-resolution units) on images given at + `scale`, using only template pixels where mask (uint8) is set. Returns (T, correlation) + or (None, 0). Most of ECC's cost is per-call setup on the whole image, so half + resolution is ~4x faster; it then needs no pre-blur (blur=1).""" + S = np.diag([scale, scale, 1.0]) + Si = np.diag([1 / scale, 1 / scale, 1.0]) + W = np.linalg.inv(S @ T @ Si)[:2].astype(np.float32) # ECC warps template coords into image coords + try: + cc, W = cv2.findTransformECC(template, image, W, cv2.MOTION_EUCLIDEAN, crit, mask, blur) + except cv2.error: + return None, 0.0 + return Si @ np.linalg.inv(np.vstack([W, [0, 0, 1]]).astype(np.float64)) @ S, cc + + +def phase_shift(a, b, window, scale=1.0): + """Translation of b's content relative to a by phase correlation, on images given at + `scale`: ((dx, dy) in full-resolution pixels, peak response).""" + (dx, dy), resp = cv2.phaseCorrelate(a, b, window) + return (dx / scale, dy / scale), resp + + +def detect(hp, nfeatures=600): + """SIFT keypoints and descriptors of a high-passed image.""" + g = cv2.normalize(hp, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U) + return cv2.SIFT_create(nfeatures=nfeatures).detectAndCompute(g, None) + + +def match(feat_a, feat_b, ha, hb): + """Rigid transform b -> a from features (any offset or rotation), refined by ECC and + checked over the whole overlap with agree(); None if the pair does not register.""" + ka, da = feat_a + kb, db = feat_b + if da is None or db is None or len(ka) < 15 or len(kb) < 15: + return None + good = [m for m, n2 in (p for p in cv2.BFMatcher(cv2.NORM_L2).knnMatch(db, da, k=2) if len(p) == 2) + if m.distance < 0.75 * n2.distance] + if len(good) < 12: + return None + pb = np.float32([kb[m.queryIdx].pt for m in good]) + pa = np.float32([ka[m.trainIdx].pt for m in good]) + A, inl = cv2.estimateAffinePartial2D(pb, pa, method=cv2.RANSAC, ransacReprojThreshold=3.0) + if A is None or inl.sum() < 12 or abs(np.hypot(A[0, 0], A[1, 0]) - 1) > 0.05: + return None + T = rigid(A) + T2, cc = ecc(ha, hb, T) + if T2 is not None and cc > 0.5 and is_close(T2, T, 1.0, 4) and agree(ha, hb, T2): + return T2 + return None + + +def agree(ha, hb, T, min_ncc=0.85, min_tiles=0.8, tile=32): + """Whole-overlap check. Repeated text (two lines ending in the same words) gives + matches that ECC confirms on the shared words while the rest of the overlap + disagrees: require high NCC over the overlap and agreement in most textured tiles.""" + h, w = ha.shape + hb = cv2.warpAffine(hb, T[:2], (w, h)) + m = cv2.warpAffine(np.ones((h, w), np.float32), T[:2], (w, h)) > 0.99 + m = cv2.erode(m.astype(np.uint8), np.ones((15, 15), np.uint8)) + if cv2.countNonZero(m) < 0.05 * m.size: + return False + # Moments with OpenCV rather than numpy: this runs on many threads at once, and + # OpenCV releases the GIL. + (ma,), (sa,) = (v.ravel() for v in cv2.meanStdDev(ha, mask=m)) + (mb,), (sb,) = (v.ravel() for v in cv2.meanStdDev(hb, mask=m)) + if (cv2.mean(cv2.multiply(ha, hb), m)[0] - ma * mb) / max(sa * sb, 1e-6) < min_ncc: + return False + th, tw = h // tile, w // tile + # Per-tile means (exact for whole tiles) of a, b, a^2, b^2, ab and the mask. + mean = lambda im: cv2.resize(im[:th * tile, :tw * tile], (tw, th), interpolation=cv2.INTER_AREA) # noqa: E731 + A, B = ha.astype(np.float64), hb.astype(np.float64) + Ea, Eb = mean(A), mean(B) + sa = np.sqrt(np.maximum(mean(A * A) - Ea * Ea, 0)) + sb = np.sqrt(np.maximum(mean(B * B) - Eb * Eb, 0)) + use = (mean(m.astype(np.float32)) > 0.9) & (np.maximum(sa, sb) > 8) + if not use.any(): + return True + ncc = (mean(A * B) - Ea * Eb) / np.maximum(sa * sb, 1e-6) + return (ncc[use] > 0.5).mean() >= min_tiles diff --git a/tools/scan.py b/tools/scan.py index 943d0ad..9930ff2 100644 --- a/tools/scan.py +++ b/tools/scan.py @@ -2,7 +2,7 @@ """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 /scan-YYYYmmdd-HHMMSS.png. The mouse goes +myscan/stitching/pagegraph.py) and saved as /scan-YYYYmmdd-HHMMSS.png. The mouse goes back to mouse mode between scans (and on Ctrl-C). With --view, a browser viewer (myscan/webview.py) shows the camera and the page as it @@ -24,7 +24,7 @@ import cv2 # noqa: E402 from myscan import calib # noqa: E402 from myscan.live import LiveSession # noqa: E402 -from myscan.pagegraph import PageGraph # noqa: E402 +from myscan.stitching import Corrector, PageGraph # noqa: E402 from myscan.xu import XU # noqa: E402 @@ -43,7 +43,7 @@ def wait_press(x, viewer=None, poll=0.05): time.sleep(poll) -def scan_page(x, a, frames_dir=None, flat=None, undistort=None, gains=None, viewer=None): +def scan_page(x, a, frames_dir=None, corrector=None, gains=None, viewer=None): """Capture until the button is pressed again. Frames are registered (and saved) on a worker thread so capture never waits; returns (PageGraph or None, LiveSession or None).""" g, live, rows, work, saves = None, None, [], queue.Queue(), queue.Queue() @@ -68,7 +68,7 @@ def scan_page(x, a, frames_dir=None, flat=None, undistort=None, gains=None, view 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) + g = PageGraph(im.shape, corrector) if viewer: live = LiveSession(g, viewer) th = threading.Thread(target=worker, daemon=True) @@ -141,8 +141,8 @@ def main(): 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) + corrector = Corrector(None if a.no_flat else calib.FlatField(cal), + None if a.no_undistort else calib.Undistort(cal)) gains = cal["gains_rgb"] first = True try: @@ -160,7 +160,7 @@ def main(): frames_dir = os.path.join(a.outdir, name) if a.save_frames else None if frames_dir: os.makedirs(frames_dir) - g, live = scan_page(x, a, frames_dir, flat, undistort, gains, viewer) + g, live = scan_page(x, a, frames_dir, corrector, gains, viewer) if live: live.stop_solving() live.close() diff --git a/tools/stitch.py b/tools/stitch.py index 3af9193..4e86ca5 100644 --- a/tools/stitch.py +++ b/tools/stitch.py @@ -11,15 +11,14 @@ import cv2 # noqa: E402 import numpy as np # noqa: E402 from myscan import calib # noqa: E402 -from myscan.pagegraph import PageGraph # noqa: E402 -from myscan.stitch import ROT_GAIN, Stitcher # noqa: E402 +from myscan.stitching import ROT_GAIN, Corrector, LaserOdometry, PageGraph, Stitcher # noqa: E402 def main(): ap = argparse.ArgumentParser() ap.add_argument("dir") ap.add_argument("-o", "--out", help="output image (default /page.png)") - ap.add_argument("--live", action="store_true", help="use the live frame-by-frame stitcher instead of the global one") + ap.add_argument("--live", action="store_true", help="use the incremental frame-by-frame stitcher") ap.add_argument("--no-refine", action="store_true", help="with --live: lasers only, no image registration") ap.add_argument("--rot-gain", type=float, default=ROT_GAIN) ap.add_argument("--no-flat", action="store_true", help="skip the EEPROM flat-field correction") @@ -35,14 +34,14 @@ def main(): rng = range(a.start, a.end if a.end is not None else len(rows)) cal = calib.load() - flat = None if a.no_flat else calib.FlatField(cal) - undistort = None if a.no_undistort else calib.Undistort(cal) + corrector = Corrector(None if a.no_flat else calib.FlatField(cal), + None if a.no_undistort else calib.Undistort(cal)) st = None for i in rng: im = cv2.imread(os.path.join(a.dir, f"f{frames[i]:05d}.png"), cv2.IMREAD_UNCHANGED) if st is None: - st = (Stitcher(im.shape, refine=not a.no_refine, rot_gain=a.rot_gain, flat=flat) if a.live - else PageGraph(im.shape, flat=flat, undistort=undistort, rot_gain=a.rot_gain)) + st = (Stitcher(im.shape, corrector, LaserOdometry(a.rot_gain, max_step_rot=None), refine=not a.no_refine) + if a.live else PageGraph(im.shape, corrector, LaserOdometry(a.rot_gain))) st.add(im, v[i]) out = a.out or os.path.join(a.dir, "page.png") cv2.imwrite(out, st.image(white_balance=not a.no_wb))