Getting Started - AI Image Classification With TensorFlow.js

• Download source - 10.6 MB

The software industry has evolved greatly over the past few years. With improved software architecture, rapid increases in computing power, and embedded decision-making abilities in machines, users now interact with more intelligent systems. From Facebook’s People You May Know feature to smart cities and self-driving cars, companies are looking for new ways to leverage the "silver bullet" of AI in their products and services.

One of the key areas of AI that has been getting a lot of attention is image recognition, whose popularity continues to grow as more sophisticated tools appear. With advanced tools and techniques in AI and computer vision, it has become possible to do image recognition inside web browsers, which means we no longer need to send input images to the server. Instead, only the detected result is sent. This means your users won’t have to worry about privacy.

What You Will Learn in This Series

You'll learn how to implement this: a real-time gender, age and (most importantly) grumpyness detector using TensorFlow.js that runs in a browser. Click the "Launch Demo" button below. (Note: this only works on modern browsers, and you'll have to give your browser permission to use your webcam).

Launch Demo

The Benefits of Using TensorFlow.js

TensorFlow.js is a JavaScript library developed by Google that allows developers to train and use machine learning (ML) models in the browser. This also allows developers familiar with JavaScript to get into building and using machine learning models without learning a new programming language. It’s a companion library to TensorFlow, which is a popular ML library for Python. So why opt for TensorFlow.js when there’s TensorFlow? Let’s quickly analyze what TensorFlow.js has to offer.

It runs anywhere.

Being JavaScript, it runs anywhere without any pre-requisite installations. Previously, to run an ML model, one needed to either host the model and provide the processing power to users, or the users would need to download and install all the required packages to run the model, which is not a feasible setup for commercial products. Instead of hosting static sites, TensorFlow.js allows us to actually use our users’ processing power to make intelligent decisions for our business.

You’ll get good performance.

TensorFlow.js offers surprisingly good performance because it uses WebGL (a JavaScript graphics API) and thus is hardware-accelerated. To get even more improved performance, you can use tfjs-node (the Node.js version of TensorFlow).

You can load existing models.

TensorFlow.js allows you to load pre-trained models right in your browser. If you have trained a TensorFlow or Keras model offline in Python, you can save it to a location available to the web and load it into the browser for inference. You can also use different libraries to include features such as image classification and pose detection without having to train your model from scratch. We will see an example of such a scenario later in the series.

You can retrain an imported model.

TensorFlow.js also allows you to use transfer learning to elevate an existing model using a small amount of data collected in the browser. This allows us to train an accurate model quickly and efficiently. We will also see an example of transfer learning in action later in the series.

You can create models directly in the browser.

At the most basic level, you can use TensorFlow.js to define, train, and run models entirely in the browser.

As mentioned earlier, using TensorFlow.js means that you can create and run AI models in a static HTML document with no installation required. At the time of writing, TensorFlow.js covers almost 90% of TensorFlow’s functionality. Depending on the problem you’re trying to solve, there may very well already exist a pre-trained model for you to import into your code.

Predicting the Breed of Dogs

Let’s say we’re interested in predicting the breed of a dog (image classification). One of the most popular image classification models we can use is available as a pre-trained model with TensorFlow.js, known as MobileNet. MobileNets are small, low-latency, low-power models parameterized to meet the resource constraints of a variety of use cases. These models have been trained on millions of images to recognize around 1000 different categories including the breed of the dogs that we’re going to see as an example. You can easily build upon these models to classify, detect, and segment images.

Importing the Models and Image

To follow along, create a static HTML document using your preferred text editor or IDE. We then start off by importing the required models and the MobileNet model in the head tag of our HTML document:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js"> </script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet"></script>

To make the image recognizable to the browser, we will add an image element inside the body of our HTML document that will be used for predictions:

<img id="image" alt="dog breed" src="https://encrypted-tbn0.gstatic.com/images?q=tbn%3AANd9GcQ05Ityl0QQfz-ZHTHyJcPLuyKcxMZKyVdrPg&usqp=CAU" crossorigin='anonymous'/>

We included the attribute crossorigin='anonymous' so that we can open the HTML document in the browser without setting up a server for the image. If you want to use a local image, one way to do it is to serve the image by a server allowing cors. You can use http-server with the command line as follows:

http-server -c1 --cors

Using the Model

We need to create a function to load the pre-trained MobileNet model and classify the image specified in the image tag.

const predictBreed = async () => {
        console.log("Model is loading...");
        const model = await mobilenet.load();
        console.log("Model loaded Successfully!")
 
        const predictions = await model.classify(image);
        console.log('Predictions: ', predictions);
      }

Putting it All Together

Once everything’s in place, your code should look like this:

<html lang="en">
  <head>
    <meta charset="UTF-8">
    <title>Predict dog breed</title>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js"> 
	</script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet">
	</script>
  </head>
  <body>
    <h3>Predicting the breed of the dog</h3>
    <img id="image" alt="dog breed" src="https://encrypted-tbn0.gstatic.com/images?q=tbn%3AANd9GcQ05Ityl0QQfz-ZHTHyJcPLuyKcxMZKyVdrPg&usqp=CAU" crossorigin='anonymous'/>
    <script>
      const image = document.getElementById('image');
 
      const predictBreed = async () => {
        console.log("Model is loading...");
        const model = await mobilenet.load();
        console.log("Model loaded Successfully!")
        const predictions = await model.classify(image);
        console.log('The dog in the picture could be: ', predictions);
      }
      predictBreed();
    </script>
  </body>
</html>

This code will return an array of three predictions ordered by probability score. If you open the console of the developer’s tool in your browser, the output should be similar to the following:

What’s Next?

In this article, we loaded a pre-trained model using TensorFlow.js to predict the breed of the dog inside our web browser. What could be next? What if we want to recognize objects outside the scope of defined labels?

In the next article of this series, we will create a web app that trains the classifier on the go and recognizes grumpy facial expressions.