From 36e714cd6269b93277c974beec48cbbfb884de28 Mon Sep 17 00:00:00 2001 From: erik Date: Fri, 9 Oct 2026 07:46:03 +0000 Subject: [PATCH] Faster frame registration, page rendering and compose PageGraph.add registers consecutive frames on half-resolution high-pass images (ECC without pre-blur, phase correlation on a half-size window) and builds the high-pass with the illumination blur done on a 4x smaller copy; full-resolution high-pass images are made on demand for keyframes only. The flat field is pre-warped into undistorted coordinates, so a frame is remapped as uint8 and corrected with one multiply. add() goes from ~13-20 to ~3-7 ms per frame with the same registration. image() draws the page on one thread per row strip of the canvas, with the flat-field gain and feather weight folded into one multiply. compose crops before normalising and takes the white point from a 1/16 subsample. _agree computes its moments with OpenCV, which releases the GIL; its decisions are unchanged on all 1425 checks of a 2525-frame scan. Co-Authored-By: Claude Opus 5.5 --- myscan/pagegraph.py | 136 +++++++++++++++++++++++++++++++++----------- myscan/stitch.py | 37 +++++++----- 2 files changed, 124 insertions(+), 49 deletions(-) diff --git a/myscan/pagegraph.py b/myscan/pagegraph.py index 3e70da1..cfbdb06 100644 --- a/myscan/pagegraph.py +++ b/myscan/pagegraph.py @@ -27,7 +27,7 @@ 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, _highpass, compose, +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) @@ -45,14 +45,33 @@ def _rigid(W): return _aff(th, W[0, 2], W[1, 2]) -def _ecc(template, image, T): - """Refine T (image pixels -> template pixels). Returns (T, correlation) or (None, 0).""" - W = np.linalg.inv(T)[:2].astype(np.float32) # ECC warps template coords into image coords +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, 5) + cc, W = cv2.findTransformECC(template, image, W, cv2.MOTION_EUCLIDEAN, ECC_CRIT, None, blur) except cv2.error: return None, 0.0 - return np.linalg.inv(np.vstack([W, [0, 0, 1]]).astype(np.float64)), cc + 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): @@ -77,8 +96,19 @@ class PageGraph: 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.window = cv2.createHanningWindow((self.w, self.h), cv2.CV_32F) - self.frames, self.hp, self.registered = [], [], [] # registered: has an image-verified edge + 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 @@ -90,31 +120,41 @@ class PageGraph: return len(self.frames) def _prep(self, frame): - f = frame.astype(np.float32) if self.flat is None else self.flat(frame) - return f if self.undistort is None else self.undistort(f) + """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) - gray = np.ascontiguousarray(f.mean(axis=2) if f.ndim == 3 else f, np.float32) - hp = _highpass(gray) + 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.hp.append(hp) + 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.hp[i - 1], hp, self.window) + (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.hp[i - 1], hp, T0) + 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 @@ -124,7 +164,7 @@ class PageGraph: 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.hp[i - 2], hp, Tb0) + 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 @@ -212,7 +252,7 @@ class PageGraph: return keys def _detect(self, i): - g = cv2.normalize(self.hp[i], None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) + 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): @@ -236,7 +276,7 @@ class PageGraph: 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) + 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 @@ -245,21 +285,30 @@ class PageGraph: """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.""" - hb = cv2.warpAffine(self.hp[b], T[:2], (self.w, self.h)) + 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)) > 0 - if m.sum() < 0.05 * m.size or np.corrcoef(self.hp[a][m], hb[m])[0, 1] < min_ncc: + 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 - cut = lambda im: im[:th * tile, :tw * tile].reshape(th, tile, tw, tile).swapaxes(1, 2).reshape(th, tw, -1) # noqa: E731 - A, B, M = cut(self.hp[a]), cut(hb), cut(m.astype(np.float32)) - A = A - A.mean(-1, keepdims=True) - B = B - B.mean(-1, keepdims=True) - sa, sb = A.std(-1), B.std(-1) - use = (M.mean(-1) > 0.9) & (np.maximum(sa, sb) > 8) + # 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 = (A * B).mean(-1) / np.maximum(sa * sb, 1e-6) + 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): @@ -354,8 +403,7 @@ class PageGraph: 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 = list(self._pool.map( - lambda i: cv2.resize(self.hp[i], None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA), use)) + 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) @@ -389,12 +437,32 @@ class PageGraph: 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) - for i, M in zip(use, Ms): - f = self._prep(self.frames[i]) - _warp_add(acc, f * (feather[..., None] if f.ndim == 3 else feather), off @ M) - _warp_add(wsum, feather, off @ M) + # 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 index 7e75a29..28c7eb9 100644 --- a/myscan/stitch.py +++ b/myscan/stitch.py @@ -183,22 +183,29 @@ class Stitcher: def compose(acc, wsum, white_balance=True, background=255): """Weighted canvas (acc, wsum) -> cropped page image in page orientation.""" covered = wsum > 0.05 - ws = np.maximum(wsum, 1e-6) - out = acc / (ws[..., None] if acc.ndim == 3 else ws) - if white_balance: - out = white_patch(out, covered) - m = covered[..., None] if acc.ndim == 3 else covered - out = np.where(m, out, background).clip(0, 255).astype(np.uint8) - ys, xs = np.nonzero(covered) - if len(ys): - out = out[ys.min():ys.max() + 1, xs.min():xs.max() + 1] + 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_patch(img, mask, pct=97, white=245): - """Scale each channel so its bright percentile (the paper) maps to `white`.""" - img = img.astype(np.float32) - px = img[mask] - ref = np.percentile(px, pct, axis=0) - return img * (white / np.maximum(ref, 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)