[wasm] Optimization of tf.browser.fromPixels.

I noticed that tf.browser.fromPixels can be very slow when it spends time in this loop:

    const numPixels = width * height;
    values = new Int32Array(numPixels * numChannels);
    for (let i = 0; i < numPixels; i++) {
      for (let channel = 0; channel < numChannels; ++channel) {
        values[i * numChannels + channel] = vals[i * 4 + channel];
      }
    }

(Also reported by #4218)

I was able to get significantly better performance with my application by calling this function instead:

function convertArrayBufferRGBAToUInt8RGB(buffer) {
  // create view of ArrayBuffer
  const x = new Uint8Array(buffer);
  const n = x.length;
  let i = 0;
  let j = 0;
  // loop through all pixels and shift them over to remove alpha
  while (i < n) {
    x[j] = x[i]; i++; j++;
    x[j] = x[i]; i++; j++;
    x[j] = x[i]; i++; j++;
    i++;
  }
  // return a subset of the original
  return x.slice(0, j);
}

And then calling it like sort of like this:

const fromPixels2DContext = document.createElement('canvas').getContext('2d'); // only call this once
function fromPixelsFast(inputElt, width, height) {
  return tf.tidy(() => {
    fromPixels2DContext.canvas.width = width;
    fromPixels2DContext.canvas.height = height;
    fromPixels2DContext.drawImage(inputElt, 0, 0, width, height);
    const buffer = fromPixels2DContext.getImageData(0, 0, width, height).data.buffer;
    const UInt8RGB = convertArrayBufferRGBAToUInt8RGB(inputArray);
    return tf.tensor3d(UInt8RGB, [height, width, 3]);
  });
}

I think the thing that makes it faster is connected to:

1. Using Uint8Array instead of UInt32Array (small effect)
2. The unrolled loop is better optimized by the browser. (medium effect)
3. It is more efficient to shuffle data in one array than between two arrays due to cache effects. (big effect)

(But I have not done rigorous testing)

It might even be possible to get a little more speed from .subarray instead of .slice to avoid making the copy.

A complete example using this function with blazeface and a custom model on a webcam inside a webworker is available in this repo.

[vladmandic]

@kylemcdonald This is very much inline in what I’m suggesting in #4218

Two suggestions that can enhance your code:

Try OffscreenCanvas before DOM, it’s slightly faster and works inside web workers. But it’s not supported on Firefox yet.

const fromPixelsCanvas (typeof OffscreenCanvas !== 'undefined') ? new OffscreenCanvas(width, height) : document.createElement('canvas');

Try GL readPixels method first as its quite a bit faster.
And fallback to 2D context and getImageData if cannot get GL context.

const fromPixelsGLContext  = fromPixelsCanvas.getContext('webgl');
const buffer = new Uint8Array(fromPixelsCanvas.width * fromPixelsCanvas.height * 4);
fromPixelsGLContext.readPixels(0, 0, fromPixelsCanvas.width, fromPixelsCanvas.height, gl.RGBA, gl.UNSIGNED_BYTE, buffer);

But…one issue here would be that getImageData returns data starting top/left and readPixels returns data starting bottom/left, so processing loop needs to be inverted.

Of course, best method would be to implement fromPixels directly in WASM backend just like it’s implemented in WebGL backend.

[annxingyuan]

Thanks for reporting this @kylemcdonald ! I’m going to close this as a duplicate of https://github.com/tensorflow/tfjs/issues/4218 - which was created 1 day earlier 😃

Thanks for the greaat discussion on this issue - this will be a really helpful reference for us!