DEV Community


Bring Machine Learning to the Browser With TensorFlow.js — Part III

vabarbosa profile image va barbosa Originally published at Medium on ・6 min read

cross-post from

Edited 2019 Mar 11 to include changes introduced in TensorFlow.js 1.0. Additional information about some of these TensorFlow.js 1.0 updates can be found here.

How to go from a web friendly format to a web application

“assorted-color textile leaves hanging decor” by DESIGNECOLOGIST on Unsplash

Welcome to part three of a series of posts where I walk you through how TensorFlow.js makes it possible to bring machine learning to the browser. First, there’s an overview of how to bring a pre-trained model into a browser application. Then you’ll find greater detail on how to convert your pre-trained model to a web friendly format. Now in this post, we step through using that web friendly model in a web application.

We continue with the Image Segmenter from the Model Asset Exchange (MAX) converted in part two of this series. The goal here is to give you a greater understanding of TensorFlow.js and how to utilize the model we made. We will create a basic web application without much style or additional libraries. To keep this article from getting too long and keep focus on TensorFlow.js, we will skip over the HTML and non-TensorFlow.js specific JavaScript code. But you can review the complete application on GitHub.

image-segmenter application output

Importing the model

The first step in importing the model to a browser readable format is to include the TensorFlow.js library in your HTML via script tag.

This will load the latest version of TensorFlow.js but you can also target a specific version or load it via NPM.

With the library loaded, a global tf variable becomes available for accessing its API. You can load the Image Segmenter model using the loadGraphModel API. 

NOTE: With TensorFlow.js version before 1.0 use the loadFrozenModel for API.

Pass the URL to the dataflow graph using the appropriate API.

Using Tensorflow.js version 0.x.x:

Using TensorFlow.js version 1.x.x:

Depending on the model size, loading may take some time. Once loaded, the model is ready to accept inputs and return a prediction.

Pre-processing the input

Models will need the inputs to be of a specific type and/or size. In most cases, the input needs some sort of pre-processing before sending it to the model. For example, some models may require a one-dimensional array of a certain size while others may require more complex multi-dimensional inputs. So the input (e.g., image, sentence, etc.) would need to be pre-processed to the expected format.

For the Image Segmenter, recall, when inspecting the model graph, the input was an ImageTensor. It was of type and shape uint8[1,?,?,3].

This is a four-dimensional array of 8-bit unsigned integer values. The ? s are placeholders and can represent any length. They would correspond to the length and width of the image. The 1 corresponds to the batch size and the 3 corresponds to the length of the RGB value for a given pixel.

For 8-bit unsigned integer valid values are from 0 to 255. This corresponds with an image’s pixel RGB value which also ranges from 0 to 255. So, we should be able to take an image convert it to a multi-dimension array of RGB values and send that to the model.

To get a Tensor with the pixel values, you can use the tf.browser.fromPixels (or tf.fromPixels for TensorFlow.js versions before 1.0) function provided by TensorFlow.js. This will produce three-dimensional array with the shape [?, ?, 3] from the given HTMLImageElement. But, the Image Segmenter is expecting a four-dimensional array. To insert an extra dimension and get the shape needed, you also will need to call the expandDims function.

Using Tensorflow.js version 0.x.x:

Using Tensorflow.js version 1.x.x:

You should now have the required input data to run the model.

Running the model

Run the model by calling predict with the input data. The function takes the input Tensor(s) and some optional configuration parameters. It returns a prediction.

Computations are in batches. If needed you can run prediction on a single batch with the predictOnBatch function.

Depending on the model complexity, the prediction may take some time.

Processing the output

The type and shape of the output returned depends on the model. To do something meaningful extra processing of the prediction is most likely required.

For the Image Segmenter, the output is a segmentation map with integers between 0 and 20. The integers corresponds to one of the pre-defined labels for each pixel in the input image.

In our web application, we are going to overlay the original image with the segments found. And each segment color coded. For example, RGB (192, 0, 0) for chairs and RGB (0, 64, 0) for potted plants.

To achieve this, start with the dataSync Tensor function. The function will download the output tensor into a TypedArray. Next convert the TypedArray into a regular array with Array.from.

With a color map, go through the converted array and assign the appropriate color to each segment. Then take this data to create the desired overlay image.

You can now add the resulting image to your HTML page.

Completing the web application

To complete the application, add buttons to load the model, upload an image, and run the model. Also, add the code to overlay the input image and output prediction.

running image-segmenter application

You can find the completed project here. The repository contains the demo web application. It also includes the web friendly format from the tensorflowjs_converter. You will also find a Jupyter notebook to play with the Image Segmenter in Python.

Machine Learning in JavaScript

The ability to use machine learning technology on the web is often limited. Creating and training some models involve massive data and intense computations. The browser may not be the ideal environment. But, an exciting use case is to take models trained elsewhere then import and run them in the browser.

With TensorFlow.js you can convert some of these models into a web friendly format. Then bring them into your web application. Making machine learning in JavaScript that much easier.

To check out even more interesting applications take a look at Nick Kasten’s magicat or the TensorFlow.js version of his Magic Cropping Tool.

Discussion (0)

Editor guide