Using MobileNet with TensorFlow.js

What?

We are going to build an image classifier with MobileNet in TensorFlow js from scratch (installing python included)

Why?

I wanted to learn how to use a pre-trained model in javascript.

How?

Find complete code here and demo

Installing Python and TensorFlow

Follow this https://www.tensorflow.org/install/install_windows for windows installation.

1. Download python 3.5 or 3.6. Version 3.7 will NOT work out as of July 2018. It doesn't download tensorflow js (henceforth referred to as tfjs)
2. Run pip3 install --upgrade tensorflow (for CPU only)
   Or pip3 install --upgrade tensorflow-gpu (for GPU)

Convert MobileNet model to tfjs model

1. Get the MobileNet model from here. 
2. This model is in tensorflow and needs to be converted to a tfjs model. Refer https://github.com/tensorflow/tfjs-converter
   
   

   tensorflowjs_converter \
       --input_format=tf_frozen_model \
       --output_node_names='MobilenetV1/Predictions/Reshape_1' \
       --saved_model_tags=serve \
       F:/image/image-
       classifier/mobilenet_v1_1.0_224/mobilenet_v1_1.0_224_frozen.pb \
       F:/image/image-classifier/modelv1
   

Load the frozen model

I use GitHub as a fast and dirty way to host static files. Upload your model to a GitHub repository and use the raw path in your code.

const MODEL_URL =
  "https://raw.githubusercontent.com/shivangidas/image-
  classifier/master/modelv1/tensorflowjs_model.pb";
const WEIGHTS_URL =
  "https://raw.githubusercontent.com/shivangidas/image-
  classifier/master/modelv1/weights_manifest.json";
let model;
(async () => {
  model = await tf.loadFrozenModel(MODEL_URL, WEIGHTS_URL);
})();

Get ImageNet classes

ImageNet Dataset has 1000 classes. You can read more here.

let IMAGENET_CLASSES = [];
$.getJSON(
    "https://raw.githubusercontent.com/shivangidas/image-
     classifier/master/mobilenet/imagenet_classes.json",
    function(data) {
      $.each(data, function(key, val) {
        IMAGENET_CLASSES.push(val);
      });
    }
  );

Make Predictions

1. We get the image from the image tag. Refer full code for image selection part.
   

   let imageData = document.getElementById("imageSrc");

2. Make tensor from the image using tfjs inbuilt function fromPixels()
   Resize the tensor to 224x224 and make the datatype float.
   The MobileNet version we are using (MobileNet_v1_1.0_224) takes an input of size [batchSize,224,224, 3]

3. let pixels = tf.fromPixels(imageData)
                  .resizeNearestNeighbor([224, 224])
                  .toFloat();

4. Normalize tensors to values between -1 and 1. Add a dimension as we have just one image to predict, no batch.
   

   let offset = tf.scalar(128);

   pixels = pixels.sub(offset).div(offset).expandDims();

5. Call predict on the tensor

   const output = await model.predict(pixels);

6. We get probabilities of 1000 classes.
   Sort and map to the classes we imported at the beginning.
   Here we are showing the top ten classes.
   
   

   const predictions = Array.from(output.dataSync())
                          .map(function(p, i) {
                             return {
                               probabilty: p,
                               classname: IMAGENET_CLASSES[i]
                             };
                          })
                          .sort((a, b) => b.probabilty - a.probabilty)
                          .slice(0, 10);
   console.log(predictions);
   

Putting it together:

let offset = tf.scalar(128);
let imageData = document.getElementById("imageSrc");
let pixels = tf.fromPixels(imageData)
               .resizeNearestNeighbor([224, 224])
               .toFloat();

pixels = pixels.sub(offset).div(offset).expandDims();
      
const output = await model.predict(pixels);
const predictions = Array.from(output.dataSync())
                         .map(function(p, i) {
                          return {
                            probabilty: p,
                            classname: IMAGENET_CLASSES[i]
                          };
                       })
                       .sort((a, b) => b.probabilty - a.probabilty)
                       .slice(0, 10);

console.log(predictions);

That's it!

Well not really. We are making 7 extra tensors that we need to dispose to prevent a memory leak. That means the tensors created will persist in GPU and take up space. I went to great lengths to dispose of them individually but if you come up with a better solution (maybe using tidy), do let me know.

References

All the links above