Image Prediction on tfjs-node (with model made by Teachable Machine Image)

image-predict-on-tfjs-node.js

const tf = require('@tensorflow/tfjs-node');
const Jimp = require('jimp');

// Directory path for model files (model.json, metadata.json, weights.bin)
// NOTE: It can be obtained from [Export Model] -> [Tensorflow.js] -> [Download my model]
//        on https://teachablemachine.withgoogle.com/train/image
const MODEL_DIR_PATH = `${__dirname}`;

// Path for image file to predict class
const IMAGE_FILE_PATH = `${__dirname}/example.jpg`;

(async () => {

  const labels = require(`${MODEL_DIR_PATH}/metadata.json`).labels;

  const model = await tf.loadLayersModel(`file://${MODEL_DIR_PATH}/model.json`);
  model.summary();

  const image = await Jimp.read(IMAGE_FILE_PATH);
  image.cover(224, 224, Jimp.HORIZONTAL_ALIGN_CENTER | Jimp.VERTICAL_ALIGN_MIDDLE);

  const NUM_OF_CHANNELS = 3;
  let values = new Float32Array(224 * 224 * NUM_OF_CHANNELS);

  let i = 0;
  image.scan(0, 0, image.bitmap.width, image.bitmap.height, (x, y, idx) => {
    const pixel = Jimp.intToRGBA(image.getPixelColor(x, y));
    pixel.r = pixel.r / 127.0 - 1;
    pixel.g = pixel.g / 127.0 - 1;
    pixel.b = pixel.b / 127.0 - 1;
    pixel.a = pixel.a / 127.0 - 1;
    values[i * NUM_OF_CHANNELS + 0] = pixel.r;
    values[i * NUM_OF_CHANNELS + 1] = pixel.g;
    values[i * NUM_OF_CHANNELS + 2] = pixel.b;
    i++;
  });

  const outShape = [224, 224, NUM_OF_CHANNELS];
  let img_tensor = tf.tensor3d(values, outShape, 'float32');
  img_tensor = img_tensor.expandDims(0);

  const predictions = await model.predict(img_tensor).dataSync();

  for (let i = 0; i < predictions.length; i++) {
    const label = labels[i];
    const probability = predictions[i];
    console.log(`${label}: ${probability}`);
  }

})();

package.json

{
  "name": "image-predict-on-tfjs-node",
  "description": "Prediction using tfjs-node (with model made by Teachable Machine Image)",
  "scripts": {
    "start": "node image-predict-on-tfjs-node.js"
  },
  "engines": {
    "node": "14"
  },
  "dependencies": {
    "jimp": "^0.12.1",
    "@tensorflow/tfjs-node": "^1.3.1",
  },
  "devDependencies": {
  },
  "private": true
}