Class ImageClassifier

electron/image_classifier.js:33–187

 * dataset and can be seen in `./imagenet_classes.js`.
 */
export class ImageClassifier {
  constructor() {
    this.model = null;
  }

  /**
   * Perform classification on a batch of image tensors.
   *
   * @param {tf.Tensor} images Batch image tensor of shape
   *   `[numExamples, height, width, channels]`. The values of `height`,
   *   `width` and `channel` must match the underlying MobileNetV2 model
   *   (default: 224, 224, 3).
   * @param {number} topK How many results with top probability / logit values
   *   to return for each example.
   * @return {Array<{className: string, prob: number}>} An array of classes
   *   with the highest `topK` probability scores, sorted in the descending
   *   order of the probability scores. Each element of the array corresponds
   *   to one example in `images`. The order of the elements matches that
   *   of `images`.
   */
  async classify(images, topK = 5) {
    await this.ensureModelLoaded();
    return tf.tidy(() => {
      const probs = this.model.predict(images);
      const sorted = true;
      const {values, indices} = tf.topk(probs, topK, sorted);

      const classProbs = values.arraySync();
      const classIndices = indices.arraySync();

      const results = [];
      classIndices.forEach((indices, i) => {
        const classesAndProbs = [];
        indices.forEach((index, j) => {
          classesAndProbs.push({
            className: IMAGENET_CLASSES[index],
            prob: classProbs[i][j]
          });
        });
        results.push(classesAndProbs);
      })

      return results;
    });
  }

  /**
   * If the underlying model is not loaded, load it.
   *
   * @param {() => any} loadingCallback An optional callback function that will
   *   be invoked when the model is being loaded.
   */
  async ensureModelLoaded(loadingCallback) {
    if (this.model == null) {
      console.log('Loading image classifier model...');
      if (loadingCallback != null) {
        loadingCallback();
      }