Packing TensorFlow.js models into npm packages

Currently, I made a post about how to convert a Keras model into a TensorFlow js one. But once you have the model converted, what is the easiest way to share it between different project and deploy it easily.
To solve that problem, I thought about packing the model into an npm package with a simple classification wrapper.

To achieve that we need:

• the resnet50 model pre-trained (tfjs)
• the fitting labels for the model outputs
• our model wrapper
• various config files (npm etc.)

First, we set up the npm package config by running npm init.
During that process, we need to supply the name of the package, version, GitHub and some other simple information. Afterward, already have our basic npm package which is already publishable.
Then I just use a basic babel configuration, so can to implement my code in ES6. And add a prepare command to the package.json to compile the files for publishing.

"scripts": {
  "prepare": "node_modules/@babel/cli/bin/babel.js src --out-dir lib"
}

The model converted in my previous post is now placed under the ResNet50 folder. To decode the model predictions I add a slightly simplified version of the original Keras labels file to the repo under the folder assets.
Now we can get started on the main task, to build a simple to use wrapper around our model.
First, we need install our dependencies, @tensorflow/tfjs, @tensorflow/tfjs-node, and jimp. While it's clear what we need the tfjs modules for, jimp is used to load our image into an array to make it convertible into a tensor.
Now, we build our ResNetPredictor class with a short constructor:

constructor() {
  this.model;
  this.labels = labels;
  this.modelPath = `file:///${__dirname}/../ResNet50/model.json`;
}

Because the tf.loadLayersModel() function is asynchronous we need and can't be called in the constructor, we have to use a little trick now.
We build an asynchronous factory method to initialize our object.

initialize = async () => {
  this.model = await tf.loadLayersModel(this.modelPath);
};

static create = async () => {
  const o = new ResNetPredictor();
  await o.initialize();
  return o;
};

When we want to initialize our object now in an external script we have to run:

load = async () => {
  return ResNetPredictor().create();
}

Now we need a function to load an image from a path or URL and convert it into a tensor, so we can enter it into our model. That's where we need jimp to unpack our image.

loadImg = async imgURI => {
  return Jimp.read(imgURI).then(img => {
    img.resize(224, 224);
    const p = [];
    img.scan(0, 0, img.bitmap.width, img.bitmap.height, function test(
      x,
      y,
      idx
    ) {
      p.push(this.bitmap.data[idx + 0]);
      p.push(this.bitmap.data[idx + 1]);
      p.push(this.bitmap.data[idx + 2]);
    });

    return tf.tensor4d(p, [1, img.bitmap.width, img.bitmap.height, 3]);
  });
};

This function takes any URI and loads the image from that address. Then the image is resized to 224x224 pixels, so it fits into our model and we generate a one-dimensional array from the image bitmap. This array is then loaded into a tensor with the right dimensions. We need the fourth dimension at the beginning because the predict function takes a batch of tensors to predict.

Now we can build the classify function, which is the interesting one at the end, that generates the value of the package.

classify = async imgURI => {
  const img = await this.loadImg(imgURI);
  const predictions = await this.model.predict(img);
  const prediction = predictions
    .reshape([1000])
    .argMax()
    .dataSync()[0];
  const result = this.labels[prediction];
  return result;
};

We call the function with the URI of the image, we want to have classified. Then the image gets loaded and thrown into the model to get the prediction. From the predictions we get us the id of the maximum value in the tensor and look it up in our labels object. This result is then returned and hopefully predicts the right object.

In the end, my project structure looks like following.

.
├── assets
│   └── labels.json
├── .gitignore
├── .babelrc
├── package.json
├── package-lock.json
├── README.md
├── LICENSE
├── ResNet50
│   ├── group1-shard1of25.bin
                .
                .
                .
│   ├── group1-shard25of25.bin
│   └── model.json
└── src
    └── index.js

Now we can just publish our package using npm run prepare && npm publish.

Here a short CodeSandbox example, how to use the package.

If you have any open questions, for example about my concrete babel configuration or anything else, feel free to have a look at my GitHub repo.
I would also be happy if you try out the npm package and give me feedback about usability and any ideas for improvement.

Github: https://github.com/paulsp94/tfjs_resnet_imagenet
NPM: https://www.npmjs.com/package/resnet_imagenet