Edge TensorFlow.js tutorial

What you will build

In this tutorial you will download a TensorFlow.js Image Classification model trained and exported using AutoML Vision Edge. You will then build a web page that loads the model and makes a prediction on an image.

Objectives

You will write JavaScript code to:

  • Run a pre-trained AutoML Vision Edge Image Classification model in a web page using the TensorFlow.js library.

Before you begin

Train a model from AutoML Vision Edge

Before you can deploy a model to an Edge device you must first train and export a TensorFlow.js model from AutoML Vision Edge following the Edge device model quickstart.

In the final step, export the model to TensorFlow.js:

updated export TensorFlow.js model screenshot

After completing the quickstart you should have the following exported files on Google Cloud Storage:

  • a dict.txt file with labels
  • a model.json file
  • *.bin weight files
Tensorflow.js folder example

Download the model files

Copy the exported files from Google Cloud Storage to a local directory:

gsutil cp gs://${cloud-storage-bucket}/model-export/icn/${model-name}/* ${local-folder}

Write a small web app

After you have your TensorFlow.js model files stored locally you are ready to write your web app:

  1. Navigate to the local directory where your model files are stored if you haven't already.

  2. Create an index.html file in that same local directory with the following contents:

    index.html

    Code samples

    <script src="https://unpkg.com/@tensorflow/tfjs"></script>
<script src="https://unpkg.com/@tensorflow/tfjs-automl"></script>
<img id="daisy" crossorigin="anonymous" src="https://storage.googleapis.com/tfjs-testing/tfjs-automl/img_classification/daisy.jpg">
<script>
async function run() {
  const model = await tf.automl.loadImageClassification('model.json');
  const image = document.getElementById('daisy');
  const predictions = await model.classify(image);
  console.log(predictions);
  // Show the resulting object on the page.
  const pre = document.createElement('pre');
  pre.textContent = JSON.stringify(predictions, null, 2);
  document.body.append(pre);
}
run();
</script>

Run the web app

Start a simple HTTP server on port 8000 from the directory with the index.html file:

Python

python3 -m SimpleHTTPServer 8000

Node.js

If you have Node.js, install the http-server package globally and run:

http-server -p 8000

Open http://localhost:8000 in the browser and you will see the prediction result on the console (View > Developer > JavaScript console), as well as on the page.

TensorFlow.js prediction in browser

How does it work?

Now that you have your app running you can explore what the code is doing.

The first two script tags load the TensorFlow.js library and the AutoML library, which are available on NPM package manager.

tfjs-automl/code_snippets/img_classification.html 
<script src="https://unpkg.com/@tensorflow/tfjs"></script>
<script src="https://unpkg.com/@tensorflow/tfjs-automl"></script>

The AutoML NPM package provides a set of APIs to load and run models produced by AutoML Vision Edge. The package takes care of any pre-processing or post-processing needed to run the model such as the ability to feed an image or video element, normalizing pixel values, and returning a sorted object with labels and scores.

The image tag loads a test image from a publish Google Cloud Storage path:

tfjs-automl/code_snippets/img_classification.html 
<img id="daisy" crossorigin="anonymous" src="https://storage.googleapis.com/tfjs-testing/tfjs-automl/img_classification/daisy.jpg">

You should replace the image src with the path to your own image when testing your model.

Next, you load the model and make a prediction with your image:

tfjs-automl/code_snippets/img_classification.html 
const model = await tf.automl.loadImageClassification('model.json');
const image = document.getElementById('daisy');
const predictions = await model.classify(image);

You specify a model with tf.automl.loadImageClassification(url). This function takes an absolute or relative URL to the exported model.json file, which is this case is a relative path since the index.html file and the model files are in the same directory.

You get your predictions by calling model.classify(image). This function runs a single image through the model and returns the prediction. The input for this function is an html image element, video element, or a 3D tensor.

What's next?

You have completed a tutorial of the TensorFlow.js Image Classification web app using an Edge model. You ran the web app in a web browser and made a classification prediction using your custom Edge model and an image that you loaded from the web. You then examined parts of the sample code to understand the underlying functionality.

Next steps: