@tensorflow/tfjs-automl

AutoML Edge API

This packages provides a set of APIs to load and run models produced by AutoML Edge.

Status

Early development. This is an unpublished experimental package.

Installation

If you are using npm/yarn

npm i @tensorflow/tfjs-automl

If you are using CDN:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-automl"></script>

We support the following types of AutoML Edge models:

1. Image classification
2. Object detection

Image classification

AutoML Image classification model will output the following set of files:

• model.json, the model topology
• dict.txt, a newline-separated list of labels
• One or more of *.bin files which hold the weights

Make sure you can access those files as static assets from your web app by serving them locally or on Google Cloud Storage.

Demo

The image classification demo lives in demo/img_classification. To run it:

cd demo/img_classification
yarn
yarn watch

This will start a local HTTP server on port 1234 that serves the demo.

Loading the model

import * as automl from '@tensorflow/tfjs-automl';
const modelUrl = 'model.json'; // URL to the model.json file.
const model = await automl.loadImageClassification(modelUrl);

Making a prediction

The input img can be HTMLImageElement, HTMLCanvasElement, HTMLVideoElement, ImageData or a 3D Tensor:

<img id="img" src="PATH_TO_IMAGE" />

const img = document.getElementById('img');
const options = {centerCrop: true};
const predictions = await model.classify(img, options);

options is optional and has the following properties:

• centerCrop - Defaults to true. Since the ML model expects a square image, we need to resize. If true, the image will be cropped first to the center before resizing.

The result predictions is a sorted list of predicted labels and their probabilities:

[
  {label: "daisy", prob: 0.931},
  {label: "dandelion", prob: 0.027},
  {label: "roses", prob: 0.013},
  ...
]

Advanced usage

Advanced users can access the underlying GraphModel via model.graphModel. The GraphModel allows users to call lower level methods such as predict(), execute() and executeAsync() which return tensors.

model.dictionary gives you access to the ordered list of labels.

Object detection

AutoML Object detection model will output the following set of files:

• model.json, the model topology
• dict.txt, a newline-separated list of labels
• One or more of *.bin files which hold the weights

Make sure you can access those files as static assets from your web app by serving them locally or on Google Cloud Storage.

Demo

The object detection demo lives in demo/object_classification. To run it:

cd demo/object_detection
yarn
yarn watch

This will start a local HTTP server on port 1234 that serves the demo.

Loading the model

import * as automl from '@tensorflow/tfjs-automl';
const modelUrl = 'model.json'; // URL to the model.json file.
const model = await automl.loadObjectDetection(modelUrl);

Making a prediction

The input img can be HTMLImageElement, HTMLCanvasElement, HTMLVideoElement, ImageData or a 3D Tensor:

<img id="img" src="PATH_TO_IMAGE" />

const img = document.getElementById('img');
const options = {score: 0.5, iou: 0.5, topk: 20};
const predictions = await model.detect(img, options);

options is optional and has the following properties:

• score - Probability score between 0 and 1. Defaults to 0.5. Boxes with score lower than this threshold will be ignored.
• topk - Only the topk most likely objects are returned. The actual number of objects might be less than this number.
• iou - Intersection over union threshold. IoU is a metric between 0 and 1 used to measure the overlap of two boxes. The predicted boxes will not overlap more than the specified threshold.

The result predictions is a sorted list of predicted objects:

[
  {
    box: {
      left: 105.1,
      top: 22.2,
      width: 70.6,
      height: 55.7
    },
    label: "Tomato",
    score: 0.972
  },
  ...
]

Advanced usage

Advanced users can access the underlying GraphModel via model.graphModel. The GraphModel allows users to call lower level methods such as predict(), execute() and executeAsync() which return tensors.

model.dictionary gives you access to the ordered list of labels.