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