Run Machine Learning models in your browser with TensorFlow.js
By Omri Grossman
After reading this post, you will:
- Learn about TensorFlow.js and the ways you can use it.
- Gain the skills to create such a project by yourself
- And finally, gain more knowledge about Machine Learning.
So, how does it work?
There are several options that we can choose from:
1. Run existing models:
TensorFlow.js provided us few attractive pre-trained models that we can import into our project, provide input, and use the output to our requirements, here you can explore the models they are providing for us: TensorFlow.js Models, and they keep adding more models as time goes by.
In addition to that, you can find many attractive pre-trained models developed by the TensorFlow.js community all across the web.
2. Retrain existing models:
This option allows us to improve an existing model for our specific use-case. We can achieve that by using a method called: Transfer Learning.
Transfer learning is the improvement of learning in a new task by transferring knowledge from a related task that has already been learned.
For instance, in the real-world, the balancing logic learned while riding a bicycle can be transferred to learn driving other two-wheeled vehicles. Similarly, in machine learning, transfer learning can be used to transfer the algorithmic logic from one ML model to the other.
The third option will be used for situations where the developer wants to create a new Machine Learning model from scratch, using TensorFlow.js API, just like the regular TensorFlow version.
So let’s get started!
I built an application that demonstrates the use of a pre-trained image tag classification model created by the Tensorflow.js team. The model is called MobileNet, and you can find more information about it here: MobileNet
Let’s walk through the main parts of the application together:
After we set up our React application, we will need to install tfjs and the image classification model — mobilenet, by running the commands below:
$ npm i @tensorflow-models/mobilenet$ npm i @tensorflow/tfjs
Now, after we installed the packages, we can import them to our `App.js` file:
We imported the image classification model and the TensorFlow.js engine, which runs the machine learning models in the background every time we invoke the model.
Next up, we need to load the model into our component for future use. Please note that the model.load() function is an asynchronous function, so we need to wait for it to be completed.
The mobileNet model has a method called classify, after we loaded the model we can call this method：
This method accepts 2 arguments:
- img: A Tensor or an image element to make a classification on.
- topk: How many of the top probabilities to return. Defaults to 3.
In the next step we want to read the user input image and load the uploaded file into a canvas element of type HTMLCanvasElement
After the image is loaded into the canvas we can run the classification method.
The output of the model.classify method is an array of classified labels and their prediction score. The output looks like this:
Once we saved the predictions array in our component we can loop over the array and render them to the screen:
(I’ve chosen to only render meaningful probabilities that are above 0.000)
So that’s it, we have a living Machine Learning model in our browser, congratulations!
Please visit those links for:
- The full code of the application can be found in this repository.
- Full code + live demo: TensorFlow.js Image Classification.
You can upload your own images, get predictions and can even be more creative and try to add new features :)