How to run a pre-trained model in TensorFlow.js
BlazeFace will be my example,
as it is perhaps the smallest of those libraries.
It’s “a lightweight model that detects faces in images”.
It has a single method,
that takes an image and returns a promise of an array of “face landmarks”.
The BlazeFace library is a thin wrapper around TensorFlow.js. Here I show the essence of the BlazeFace wrapper using just the TensorFlow.js API:
This shows the basic pattern of nearly all model wrapper libraries:
- Load the pre-trained model.
In this case, the Blazeface model is in the undocumented TensorFlow.js “web-friendly format”,
so we use
- Convert the input to a tensor.
Most models work on images, so we use
tf.browser.fromPixels. It returns a 3D tensor with shape
[height, width, numChannels], with RGB values in the range
[0, 255]. (Note the surprising shape: height is the first dimension!)
- Convert the input tensor to the format expected by the model.
The Blazeface model only works on 128×128 square images,
but input images (e.g. from a webcam) can be any size and shape.
The Blazeface library chooses to squish the whole input image
into a 128×128 square.
This destroys the aspect ratio, but appears to not affect accuracy
(presumably, the model has been trained on various squished images).
The Blazeface model also expects RGB channels in the range
[-1, 1], so we have to convert from the range
- Run the model on the input tensor with
.predict. We get back a prediction tensor.
This at least requires running