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 <noreply@anthropic.com>

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):
+def _ecc(template, image, T, scale=1.0, blur=5):
-    """Refine T (image pixels -> template pixels). Returns (T, correlation) or (None, 0)."""
+    """Refine T (image pixels -> template pixels, full-resolution units) on images given at
-    W = np.linalg.inv(T)[:2].astype(np.float32)   # ECC warps template coords into image coords
+    `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.half = (self.w // 2, self.h // 2)
-        self.frames, self.hp, self.registered = [], [], []   # registered: has an image-verified edge
+        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)
+        """Raw uint8 frame -> undistorted, flat-fielded float32 frame."""
-        return f if self.undistort is None else self.undistort(f)
+        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)
+        hs = cv2.resize(_highpass_fast(_gray(f)), self.half, interpolation=cv2.INTER_AREA)
-        hp = _highpass(gray)
         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
+        m = cv2.erode(m.astype(np.uint8), np.ones((15, 15), np.uint8))
-        if m.sum() < 0.05 * m.size or np.corrcoef(self.hp[a][m], hb[m])[0, 1] < min_ncc:
+        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
+        # Per-tile means (exact for whole tiles) of a, b, a^2, b^2, ab and the mask.
-        A, B, M = cut(self.hp[a]), cut(hb), cut(m.astype(np.float32))
+        mean = lambda im: cv2.resize(im[:th * tile, :tw * tile], (tw, th), interpolation=cv2.INTER_AREA)  # noqa: E731
-        A = A - A.mean(-1, keepdims=True)
+        A, B = ha.astype(np.float64), hb.astype(np.float64)
-        B = B - B.mean(-1, keepdims=True)
+        Ea, Eb = mean(A), mean(B)
-        sa, sb = A.std(-1), B.std(-1)
+        va = np.maximum(mean(A * A) - Ea * Ea, 0)
-        use = (M.mean(-1) > 0.9) & (np.maximum(sa, sb) > 8)
+        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(
+        small = [self.hs[i] for i in use]                   # already at half resolution
-            lambda i: cv2.resize(self.hp[i], None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA), use))
         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):
+        # Flat-field gain and feather weight in one multiply per frame.
-            f = self._prep(self.frames[i])
+        fw = self._inv.get(ch or 1, np.ones(self.frames[0].shape, np.float32)) * (feather[..., None] if ch else feather)
-            _warp_add(acc, f * (feather[..., None] if f.ndim == 3 else feather), off @ M)
+
-            _warp_add(wsum, feather, off @ M)
+        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)

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)
+    rows, cols = np.nonzero(covered.any(1))[0], np.nonzero(covered.any(0))[0]
-    out = acc / (ws[..., None] if acc.ndim == 3 else ws)
+    if len(rows):
-    if white_balance:
+        box = np.s_[rows[0]:rows[-1] + 1, cols[0]:cols[-1] + 1]
-        out = white_patch(out, covered)
+        acc, wsum, covered = acc[box], wsum[box], covered[box]
-    m = covered[..., None] if acc.ndim == 3 else covered
+    inv = 1 / np.maximum(wsum, 1e-6)
-    out = np.where(m, out, background).clip(0, 255).astype(np.uint8)
+    gain = white_point(acc, inv, covered) if white_balance else np.ones(acc.shape[2:] or 1, np.float32)
-    ys, xs = np.nonzero(covered)
+    if acc.ndim == 3:
-    if len(ys):
+        out = acc * (inv[..., None] * gain.astype(np.float32))
-        out = out[ys.min():ys.max() + 1, xs.min():xs.max() + 1]
+    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):
+def white_point(acc, inv, mask, pct=97, white=245, step=4):
-    """Scale each channel so its bright percentile (the paper) maps to `white`."""
+    """Per-channel gain mapping the bright percentile (the paper) of acc * inv to `white`,
-    img = img.astype(np.float32)
+    estimated on every `step`-th pixel in each direction (the full canvas is ~20M pixels)."""
-    px = img[mask]
+    sub = mask[::step, ::step]
-    ref = np.percentile(px, pct, axis=0)
+    if not sub.any():
-    return img * (white / np.maximum(ref, 1))
+        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)