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?
Installing Python and TensorFlow
Follow this https://www.tensorflow.org/install/install_windows for windows installation.
- 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)
- Run pip3 install --upgrade tensorflow (for CPU only)
Or pip3 install --upgrade tensorflow-gpu (for GPU)
Convert MobileNet model to tfjs model
- Get the MobileNet model from here.
- 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
- We get the image from the image tag. Refer full code for image selection part.
let imageData = document.getElementById("imageSrc"); - 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] let pixels = tf.fromPixels(imageData) .resizeNearestNeighbor([224, 224]) .toFloat();- 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();
- Call predict on the tensor
const output = await model.predict(pixels);
- 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
