myscan/myscan/calib.py

"""Per-device calibration stored in the mouse's EEPROM (written at the factory).

Layout (from the Mac app, see RE_colour_stitch.md 1.4): 0x3000 u16 LE blob length N,
0x3002 4-byte checksum, 0x3006 N-byte boost::archive binary blob (i386 sizes) holding
Roi, Sensor, Undistortion, Color and Info params. Fields are located by structure, not
fully deserialised; offsets below were checked against one unit only.

Colour calibration is two sets of point samples, black and white, each
`u16 x, u16 y, u32 3, u8 b, g, r` in raw-frame coordinates, which the app fits with
NURBS surfaces and uses for a per-pixel linear correction (`slope*in + intercept`,
white -> colorDestBrightImage = 250).
"""
import os
import struct

import cv2
import numpy as np

EEPROM_BASE = 0x3000
DEST_WHITE = 250.0
CACHE = os.path.join(os.path.dirname(__file__), "..", "calib", "eeprom_3000.bin")


def read_eeprom(x, cache=CACHE):
    """Header + blob, cached on disk after the first (~15 s) read."""
    if cache and os.path.exists(cache):
        return open(cache, "rb").read()
    head = x.eeprom_read(EEPROM_BASE, 6)
    n = struct.unpack("<H", head[:2])[0]
    if not 0 < n <= 20000:
        raise ValueError(f"no calibration blob (length {n})")
    data = head + x.eeprom_read(EEPROM_BASE + 6, n)
    if cache:
        os.makedirs(os.path.dirname(cache), exist_ok=True)
        open(cache, "wb").write(data)
    return data


def _sample_runs(blob, min_len=50):
    """Find runs of 11-byte colour samples (the u32 3 marks each one)."""
    runs, i = [], 4
    while i < len(blob) - 11:
        j = i
        while blob[j:j + 4] == b"\x03\x00\x00\x00":
            j += 11
        if (j - i) // 11 >= min_len:
            runs.append((i - 4, (j - i) // 11))
            i = j
        else:
            i += 1
    return runs


def parse(data):
    blob = data[6:]
    if blob[4:26] != b"serialization::archive":
        raise ValueError("not a boost archive")
    d = lambda o: struct.unpack_from("<d", blob, o)[0]  # noqa: E731
    runs = _sample_runs(blob)
    if len(runs) != 2:
        raise ValueError(f"expected black+white sample runs, found {len(runs)}")
    samples = [np.array([struct.unpack_from("<2HI3B", blob, s + 11 * k) for k in range(n)],
                        np.float32)[:, [0, 1, 3, 4, 5]] for s, n in runs]
    end = runs[1][0] + 11 * runs[1][1]
    return {
        "roi": struct.unpack_from("<4i", blob, 0x29),            # x0, y0, x1, y1
        "meters_per_unit": (d(0x3B), d(0x43)),                   # left, right laser
        "laser_left_m": (d(0x5B), d(0x63)),
        "laser_right_m": (d(0x73), d(0x7B)),
        "focal_px": d(0x85), "principal_px": (d(0x8D), d(0x95)),
        "dist_coeffs": struct.unpack_from("<5d", blob, 0xA5),     # OpenCV k1, k2, p1, p2, k3
        "tilt_h": struct.unpack_from("<8d", blob, 0xD1) + (1.0,),  # 3x3 row-major, pixel space
        "pixel_size_m": d(0x111),
        "black": samples[0], "white": samples[1],                # columns x, y, b, g, r
        "gains_rgb": struct.unpack_from("<3i", blob, end + 4),   # sensor regs 1, 2, 3
    }


def _bspline_basis(t, n_ctrl, length):
    """Uniform cubic B-spline basis (len(t) x n_ctrl) on [0, length)."""
    seg = n_ctrl - 3
    u = np.clip(np.asarray(t, np.float64) / length * seg, 0, seg - 1e-9)
    k = u.astype(int)
    f = u - k
    w = np.stack([(1 - f) ** 3, 3 * f ** 3 - 6 * f ** 2 + 4, -3 * f ** 3 + 3 * f ** 2 + 3 * f + 1, f ** 3], 1) / 6
    B = np.zeros((len(u), n_ctrl))
    for j in range(4):
        B[np.arange(len(u)), k + j] = w[:, j]
    return B


def fit_surface(s, shape, ctrl=(6, 10), ridge=1.0):
    """Smooth surface (H x W x 3) through scattered samples (x, y, b, g, r)."""
    h, w = shape
    ny, nx = ctrl
    By, Bx = _bspline_basis(s[:, 1], ny, h), _bspline_basis(s[:, 0], nx, w)
    A = (By[:, :, None] * Bx[:, None, :]).reshape(len(s), -1)
    # Ridge towards the mean keeps sparsely sampled corners from running off.
    mean = s[:, 2:].mean(0)
    C = np.linalg.solve(A.T @ A + ridge * np.eye(A.shape[1]), A.T @ (s[:, 2:] - mean)) + mean
    Gy, Gx = _bspline_basis(np.arange(h), ny, h), _bspline_basis(np.arange(w), nx, w)
    return np.einsum("yi,ijc,xj->yxc", Gy, C.reshape(ny, nx, 3), Gx).astype(np.float32)


class FlatField:
    """Per-pixel colour correction from the EEPROM black/white surfaces.

    Calling the object divides a frame by the white surface normalised to mean 1 per
    channel: that removes the LED/lens falloff whatever the current exposure and sensor
    gains, leaving overall colour balance to the stitcher's white patch. `apply` is the
    app's full linear correction, valid only at the factory exposure and gains."""

    def __init__(self, cal, shape=(296, 640)):
        self.black = fit_surface(cal["black"], shape)
        self.white = fit_surface(cal["white"], shape)
        self.shape = self.white / self.white.reshape(-1, 3).mean(0)
        self.inv = 1 / self.shape
        self.inv_gray = 1 / self.shape.mean(2)

    def __call__(self, frame):
        return frame.astype(np.float32) * (self.inv if frame.ndim == 3 else self.inv_gray)

    def apply(self, frame, level=1.0):
        f = frame.astype(np.float32)
        b, w = self.black * level, self.white * level
        if f.ndim == 2:
            b, w = b.mean(2), w.mean(2)
        return (f - b) * (DEST_WHITE / np.maximum(w - b, 1))


class Undistort:
    """Lens + tilt correction, as the app's generateMapFromParams (EEPROM.md): for each
    output pixel p, the source pixel is K * distort(inv(K) * inv(H) * p)."""

    def __init__(self, cal, shape=(296, 640)):
        h, w = shape
        f, (cx, cy) = cal["focal_px"], cal["principal_px"]
        K = np.array([[f, 0, cx], [0, f, cy], [0, 0, 1.0]])
        M = np.linalg.inv(K) @ np.linalg.inv(np.reshape(cal["tilt_h"], (3, 3)))
        k1, k2, p1, p2, k3 = cal["dist_coeffs"]
        ys, xs = np.mgrid[0:h, 0:w].astype(np.float64)
        X, Y, W = (M @ np.stack([xs.ravel(), ys.ravel(), np.ones(xs.size)])).reshape(3, h, w)
        x, y = X / W, Y / W
        r2 = x * x + y * y
        radial = 1 + k1 * r2 + k2 * r2 ** 2 + k3 * r2 ** 3
        xd = x * radial + 2 * p1 * x * y + p2 * (r2 + 2 * x * x)
        yd = y * radial + p1 * (r2 + 2 * y * y) + 2 * p2 * x * y
        self.map = cv2.convertMaps((f * xd + cx).astype(np.float32), (f * yd + cy).astype(np.float32),
                                   cv2.CV_16SC2)

    def __call__(self, frame):
        return cv2.remap(frame, *self.map, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)


def load(x=None, cache=CACHE):
    if x is None and not (cache and os.path.exists(cache)):
        from .xu import XU
        x = XU()
    return parse(read_eeprom(x, cache))