[tflite] Add task library wrapper client (tensorflow#4979)

* exports stuff from tflite web api client better

* save

* save

* update

* update

* update

* address comments

* update

* fix

* update version

Author: jinjingforever

tfjs-tflite/demo/models/cartoongan_fp16.tflite

@@ -20,7 +20,6 @@
   },
   "staticFiles": {
     "staticPath": [
-      "models",
       "../dist"
     ]
   },

tfjs-tflite/demo/package.json

@@ -29,9 +29,10 @@ const PHYSICAL_CORES = navigator.hardwareConcurrency / 2;
 async function start() {
   // Load model runner with the cartoonizer tflite model.
   const start = Date.now();
-  const tfliteModel = await loadTFLiteModel('cartoongan_fp16.tflite', {
-    numThreads: PHYSICAL_CORES,
-  });
+  const tfliteModel = await loadTFLiteModel(
+      'https://tfhub.dev/sayakpaul/lite-model/cartoongan/fp16/1', {
+        numThreads: PHYSICAL_CORES,
+      });
   ele('.loading-msg').innerHTML = `Loaded WASM module and <a href='${
       CARTOONIZER_LINK}' target='blank'>TFLite model</a> in ${
       Date.now() - start}ms`;

tfjs-tflite/demo/src/script.ts

@@ -1,6 +1,6 @@
 {
   "name": "@tensorflow/tfjs-tflite",
-  "version": "0.0.1-alpha.0",
+  "version": "0.0.1-alpha.2",
   "description": "TFLite support for TensorFlow.js",
   "main": "dist/tf-tflite.node.js",
   "jsnext:main": "dist/index.js",

tfjs-tflite/package.json

@@ -85,7 +85,7 @@ function config({
 module.exports = cmdOptions => {
   const bundles = [];
 
-  const name = 'tf';
+  const name = 'tflite';
   const extend = true;
   const fileName = 'tf-tflite';
 

tfjs-tflite/rollup.config.js

@@ -21,7 +21,7 @@ set -e
 cd "$(dirname "$0")"
 
 # The default version.
-CURRENT_VERSION=0.0.2
+CURRENT_VERSION=0.0.3
 
 # Get the version from the first parameter.
 # Default to the value in CURRENT_VERSION.
@@ -37,3 +37,7 @@ fi
 mkdir -p ../deps
 GCP_DIR="gs://tfweb/${VERSION}/dist"
 gsutil -m cp "${GCP_DIR}/*" ../deps/
+
+# Append module exports to the JS client to make it a valid CommonJS module.
+# This is needed to help bundler correctly initialize the tfweb namespace.
+echo "var tfweb = (window && window['tfweb']) || this['tfweb']; exports.tfweb = tfweb;" >> ../deps/tflite_web_api_client.js

tfjs-tflite/scripts/download-tflite-web-api.sh

@@ -17,6 +17,9 @@
 
 export * from './tflite_model';
 export * from './types/tflite_web_model_runner';
-export * from './types/nl_classifier';
-export * from './types/image_classifier';
-export * from './types/tflite_web_api';
+export * from './tflite_task_library_client/image_classifier';
+export * from './tflite_task_library_client/image_segmenter';
+export * from './tflite_task_library_client/object_detector';
+export * from './tflite_task_library_client/nl_classifier';
+export * from './tflite_task_library_client/bert_nl_classifier';
+export * from './tflite_task_library_client/bert_qa';

tfjs-tflite/src/index.ts

@@ -269,24 +269,24 @@ export class TFLiteModel implements InferenceModel {
 /**
  * Loads a TFLiteModel from the given model url.
  *
- * @param modelUrl The path to the model.
+ * @param model The path to the model (string), or the model content in memory
+ *     (ArrayBuffer).
  * @param options Options related to model inference.
  *
  * @doc {heading: 'Models', subheading: 'TFLiteModel'}
  */
 export async function loadTFLiteModel(
-    modelUrl: string,
+    model: string|ArrayBuffer,
     options: TFLiteWebModelRunnerOptions =
         DEFAULT_TFLITE_MODEL_RUNNER_OPTIONS): Promise<TFLiteModel> {
   // Handle tfhub links.
-  if (modelUrl.includes('tfhub.dev')) {
-    if (!modelUrl.endsWith(TFHUB_SEARCH_PARAM)) {
-      modelUrl = `${modelUrl}${TFHUB_SEARCH_PARAM}`;
-    }
+  if (typeof model === 'string' && model.includes('tfhub.dev') &&
+      model.includes('lite-model') && !model.endsWith(TFHUB_SEARCH_PARAM)) {
+    model = `${model}${TFHUB_SEARCH_PARAM}`;
   }
   const tfliteModelRunner =
       await tfliteWebAPIClient.tfweb.TFLiteWebModelRunner.create(
-          modelUrl, options);
+          model, options);
   return new TFLiteModel(tfliteModelRunner);
 }
 

tfjs-tflite/src/tflite_model.ts

@@ -0,0 +1,67 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tfliteWebAPIClient from '../tflite_web_api_client';
+import {BertNLClassifier as TaskLibraryBertNLClassifier} from '../types/bert_nl_classifier';
+import {Class} from './common';
+
+export interface BertNLClassifierOptions {
+  /**
+   * Max number of tokens to pass to the model.
+   *
+   * Default to 128.
+   */
+  maxSeqLen?: number;
+}
+
+/**
+ * Client for BertNLClassifier TFLite Task Library.
+ *
+ * It is a wrapper around the underlying javascript API to make it more
+ * convenient to use. See comments in the corresponding type declaration file in
+ * src/types for more info.
+ */
+export class BertNLClassifier {
+  constructor(private instance: TaskLibraryBertNLClassifier) {}
+
+  static async create(
+      model: string|ArrayBuffer,
+      options?: BertNLClassifierOptions): Promise<BertNLClassifier> {
+    const protoOptions = new tfliteWebAPIClient.tfweb.BertNLClassifierOptions();
+    if (options) {
+      if (options.maxSeqLen) {
+        protoOptions.setMaxSeqLen(options.maxSeqLen);
+      }
+    }
+    const instance = await tfliteWebAPIClient.tfweb.BertNLClassifier.create(
+        model, protoOptions);
+    return new BertNLClassifier(instance);
+  }
+
+  classify(input: string): Class[] {
+    return this.instance.classify(input).map(category => {
+      return {
+        probability: category.score,
+        className: category.className,
+      };
+    });
+  }
+
+  cleanUp() {
+    this.instance.cleanUp();
+  }
+}

tfjs-tflite/src/tflite_task_library_client/bert_nl_classifier.ts

@@ -0,0 +1,73 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tfliteWebAPIClient from '../tflite_web_api_client';
+import {BertQuestionAnswerer as TaskLibraryBertQuestionAnswerer} from '../types/bert_qa';
+
+/** A single answer. */
+export interface QaAnswer {
+  /** The text of the answer. */
+  text: string;
+  /** The position and logit of the answer. */
+  pos: Pos;
+}
+
+/** Answer position. */
+export interface Pos {
+  /** The start position. */
+  start: number;
+  /** The end position. */
+  end: number;
+  /** The logit. */
+  logit: number;
+}
+
+/**
+ * Client for BertQA TFLite Task Library.
+ *
+ * It is a wrapper around the underlying javascript API to make it more
+ * convenient to use. See comments in the corresponding type declaration file in
+ * src/types for more info.
+ */
+export class BertQuestionAnswerer {
+  constructor(private instance: TaskLibraryBertQuestionAnswerer) {}
+
+  static async create(model: string|
+                      ArrayBuffer): Promise<BertQuestionAnswerer> {
+    const instance =
+        await tfliteWebAPIClient.tfweb.BertQuestionAnswerer.create(model);
+    return new BertQuestionAnswerer(instance);
+  }
+
+  answer(context: string, question: string): QaAnswer[] {
+    const result = this.instance.answer(context, question);
+    if (!result) {
+      return [];
+    }
+
+    return result.map(answer => {
+      return {
+        text: answer.text,
+        pos: answer.pos,
+      };
+    });
+  }
+
+  cleanUp() {
+    this.instance.cleanUp();
+  }
+}

tfjs-tflite/src/tflite_task_library_client/bert_qa.ts

@@ -0,0 +1,53 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {Class as ProtoClass} from '../types/common';
+
+/** Common options for all task library tasks. */
+export interface CommonTaskLibraryOptions {
+  /**
+   * The number of threads to be used for TFLite ops that support
+   * multi-threading when running inference with CPU. num_threads should be
+   * greater than 0 or equal to -1. Setting num_threads to -1 has the effect to
+   * let TFLite runtime set the value.
+   *
+   * Default to -1.
+   */
+  numThreads?: number;
+}
+
+/** A single class in the classification result. */
+export interface Class {
+  /** The name of the class. */
+  className: string;
+
+  /** The probability/score of the class. */
+  probability: number;
+}
+
+/** Convert proto Class array to our own Class array. */
+export function convertProtoClassesToClasses(protoClasses: ProtoClass[]):
+    Class[] {
+  const classes: Class[] = [];
+  protoClasses.forEach(cls => {
+    classes.push({
+      className: cls.getDisplayName() || cls.getClassName(),
+      probability: cls.getScore(),
+    });
+  });
+  return classes;
+}

tfjs-tflite/src/tflite_task_library_client/common.ts

@@ -0,0 +1,86 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tfliteWebAPIClient from '../tflite_web_api_client';
+import {ImageClassifier as TaskLibraryImageClassifier} from '../types/image_classifier';
+
+import {Class, CommonTaskLibraryOptions, convertProtoClassesToClasses} from './common';
+
+/** ImageClassifier options. */
+export interface ImageClassifierOptions extends CommonTaskLibraryOptions {
+  /**
+   * Maximum number of top scored results to return. If < 0, all results will
+   * be returned. If 0, an invalid argument error is returned.
+   */
+  maxResults?: number;
+
+  /**
+   * Score threshold in [0,1), overrides the ones provided in the model metadata
+   * (if any). Results below this value are rejected.
+   */
+  scoreThreshold?: number;
+}
+
+/**
+ * Client for ImageClassifier TFLite Task Library.
+ *
+ * It is a wrapper around the underlying javascript API to make it more
+ * convenient to use. See comments in the corresponding type declaration file in
+ * src/types for more info.
+ */
+export class ImageClassifier {
+  constructor(private instance: TaskLibraryImageClassifier) {}
+
+  static async create(
+      model: string|ArrayBuffer,
+      options?: ImageClassifierOptions): Promise<ImageClassifier> {
+    const optionsProto = new tfliteWebAPIClient.tfweb.ImageClassifierOptions();
+    if (options) {
+      if (options.maxResults !== undefined) {
+        optionsProto.setMaxResults(options.maxResults);
+      }
+      if (options.scoreThreshold !== undefined) {
+        optionsProto.setScoreThreshold(options.scoreThreshold);
+      }
+      if (options.numThreads !== undefined) {
+        optionsProto.setNumThreads(options.numThreads);
+      }
+    }
+    const instance = await tfliteWebAPIClient.tfweb.ImageClassifier.create(
+        model, optionsProto);
+    return new ImageClassifier(instance);
+  }
+
+  classify(input: ImageData|HTMLImageElement|HTMLCanvasElement|
+           HTMLVideoElement): Class[] {
+    const result = this.instance.classify(input);
+    if (!result) {
+      return [];
+    }
+
+    let classes: Class[] = [];
+    if (result.getClassificationsList().length > 0) {
+      classes = convertProtoClassesToClasses(
+          result.getClassificationsList()[0].getClassesList());
+    }
+    return classes;
+  }
+
+  cleanUp() {
+    this.instance.cleanUp();
+  }
+}