@thi.ng/pixel-flow/index.js

Naive, lightweight CPU-based dense optical flow implementation

const { abs, ceil } = Math;
class OpticalFlow {
  prev;
  flow;
  step;
  displace;
  displaceStep;
  windowSize;
  windowStep;
  amp;
  smooth;
  width;
  height;
  margin;
  threshold;
  invRange;
  mode;
  pred;
  constructor(startFrame, opts) {
    this.prev = startFrame.copy();
    this.displace = opts?.displace ?? 9;
    this.displaceStep = opts?.displaceStep ?? 3;
    this.windowSize = opts?.windowSize ?? 12;
    this.windowStep = opts?.windowStep ?? 3;
    this.amp = opts?.smooth ?? 1;
    this.smooth = opts?.smooth ?? 0.25;
    this.threshold = opts?.threshold ?? -1;
    this.invRange = 1 / (opts?.range ?? 255);
    this.mode = opts?.mode ?? "min";
    this.pred = this.mode == "min" ? (a, b) => a < b : (a, b) => a > b;
    const margin = this.margin = this.windowSize + this.displace;
    const step = this.step = opts?.step ?? 6;
    this.width = ceil((startFrame.width - 2 * margin) / step);
    this.height = ceil((startFrame.height - 2 * margin) / step);
    this.flow = new Float64Array(this.width * this.height * 2);
  }
  /**
   * Computes optical flow between given frame and (stored) previous frame,
   * according to configured options. Returns updated flowfield in a format
   * compatible with thi.ng/tensors `asTensor()`.
   *
   * @remarks
   * See {@link OpticalFlow} class docs for more details.
   *
   * @param curr
   */
  update(curr) {
    const {
      amp,
      displace,
      displaceStep,
      flow,
      invRange,
      margin,
      mode,
      pred,
      prev,
      smooth,
      step,
      threshold,
      windowSize,
      windowStep
    } = this;
    const srcA = prev.data;
    const srcB = curr.data;
    const width = prev.width;
    const w = prev.width - margin;
    const h = prev.height - margin;
    const invWindowScale = 1 / (1 + 2 * (windowSize / windowStep)) ** 2;
    const isMax = mode === "max";
    const initialDist = isMax ? threshold : Infinity;
    const dirScale = amp / displace;
    let x, y, x2, y2, maxX, maxY, i, j, idx, idxA, idxB, sum, dx, dy, candidate, meanX = 0, meanY = 0;
    for (y = margin, idx = 0; y < h; y += step) {
      maxY = y + displace;
      for (x = margin; x < w; x += step) {
        candidate = initialDist;
        dx = 0;
        dy = 0;
        maxX = x + displace;
        for (y2 = y - displace; y2 <= maxY; y2 += displaceStep) {
          for (x2 = x - displace; x2 <= maxX; x2 += displaceStep) {
            sum = 0;
            for (j = -windowSize; j <= windowSize; j += windowStep) {
              for (i = -windowSize, idxA = (y + j) * width + x, idxB = (y2 + j) * width + x2; i <= windowSize; i += windowStep) {
                sum += (abs(srcA[idxA] - srcB[idxB]) * invRange) ** 2;
                idxA += windowStep;
                idxB += windowStep;
              }
            }
            sum *= invWindowScale;
            if (sum >= threshold && pred(sum, candidate)) {
              candidate = sum;
              dx = x - x2;
              dy = y - y2;
            }
          }
        }
        dx *= dirScale;
        dy *= dirScale;
        meanX += flow[idx] += (dx - flow[idx]) * smooth;
        idx++;
        meanY += flow[idx] += (dy - flow[idx]) * smooth;
        idx++;
      }
    }
    srcA.set(srcB);
    idx >>= 1;
    return {
      type: "f64",
      data: flow,
      shape: [this.height, this.width, 2],
      stride: [this.width * 2, 2, 1],
      dir: [meanX / idx, meanY / idx]
    };
  }
}
export {
  OpticalFlow
};