Lots of updates to prepare for 0.4.0

Author: Kevin Scott

README.md

@@ -64,6 +64,63 @@ When you have a trained model you're happy with, save it with:
 mlClassifier.save();
 ```
 
+## Using the saved model
+
+When you hit save, Tensorflow.js will download a weights file and a model topology file.
+
+You'll need to combine both into a single `json` file. Open up your model topology file and at the top level of the JSON file, make sure to add a `weightsManifest` key pointing to your weights, like:
+
+```
+{
+  "weightsManifest": "ml-classifier-class1-class2.weights.bin",
+  "modelTopology": {
+    ...
+  }
+}
+```
+
+When using the model in your app, there's a few things to keep in mind:
+
+1. You need to make sure you transform images into the correct dimensions, depending on the pretrained model it was trained with. (For MOBILENET, this would be 1x224x224x3).
+2. You must create a pretrained model matching the dimensions used to train. An example is below for MOBILENET.
+3. You must first run your images through the pretrained model to activate them.
+4. After getting the final prediction, you must take the arg max.
+5. You'll get back a number indicating your class.
+
+Full example for MOBILENET:
+
+```
+    const loadImage = (src) => new Promise((resolve, reject) => {
+      const image = new Image();
+      image.src = src;
+      image.crossOrigin = 'Anonymous';
+      image.onload = () => resolve(image);
+      image.onerror = (err) => reject(err);
+    });
+
+    const pretrainedModelURL = 'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json';
+
+    tf.loadModel(pretrainedModelURL).then(model => {
+      const layer = model.getLayer('conv_pw_13_relu');
+      return tf.model({
+        inputs: [model.inputs[0]],
+        outputs: layer.output,
+      });
+    }).then(pretrainedModel => {
+      return tf.loadModel('/model.json').then(model => {
+        return loadImage('/trees/tree1.png').then(loadedImage => {
+          const image = tf.reshape(tf.fromPixels(loadedImage), [1,224,224,3]);
+          const pretrainedModelPrediction = pretrainedModel.predict(image);
+          const modelPrediction = model.predict(pretrainedModelPrediction);
+          const prediction = modelPrediction.as1D().argMax().dataSync()[0];
+          console.log(prediction);
+        });
+      });
+    }).catch(err => {
+      console.error('Error', err);
+    });
+```
+
 ## API Documentation
 
 Start by instantiating a new instance of `MLClassifier` with:
@@ -76,10 +133,16 @@ This will begin loading the pretrained model and provide you with an object onto
 
 ### `constructor`
 
-`MLClassifier` accepts a number of callbacks when initialized:
+`MLClassifier` accepts a number of callbacks for beginning and end of various methods.
+
+You can provide a custom pretrained model as a `pretrainedModel`.
+
+You can provide a custom training model as a `trainingModel`.
 
 #### Parameters
 
+  * **pretrainedModel** (`string | tf.Model`) *Optional* - A string denoting which pretrained model to load from an internal config. Valid strings can be found on the exported object `PRETRAINED_MODELS`. You can also specify a preloaded pretrained model directly.
+  * **trainingModel** (`tf.Model | Function`) *Optional* - A custom model to use during training. Can be provided as a `tf.Model` or as a function that accepts `{xs: [...], ys: [...]`, number of `classes`, and `params` provided to train.
   * **onLoadStart** (`Function`) *Optional* - A callback for when `load` (loading the pre-trained model) is first called.
   * **onLoadComplete** (`Function`) *Optional* - A callback for when `load` (loading the pre-trained model) is complete.
   * **onAddDataStart** (`Function`) *Optional* - A callback for when `addData` is first called.
@@ -98,8 +161,13 @@ This will begin loading the pretrained model and provide you with an object onto
 
 #### Example
 ```
-import MLClassifier from 'ml-classifier';
+import MLClassifier, {
+  PRETRAINED_MODELS,
+} from 'ml-classifier';
+
 const mlClassifier = new MLClassifier({
+  pretrainedModel: PRETRAINED_MODELS.MOBILENET,
+
   onLoadStart: () => console.log('onLoadStart'),
   onLoadComplete: () => console.log('onLoadComplete'),
   onAddDataStart: () => console.log('onAddDataStart'),
@@ -117,6 +185,18 @@ const mlClassifier = new MLClassifier({
 });
 ```
 
+Example of specifying a preloaded pretrained model:
+
+```
+import MLClassifier from 'ml-classifier';
+
+const mlClassifier = tf.loadModel('... some pretrained model ...').then(model => {
+  return new MLClassifier({
+    pretrainedModel: model,
+  });
+});
+```
+
 ### `addData`
 
 This method takes an array of incoming images, an optional array of labels, and an optional dataType.
@@ -302,3 +382,5 @@ yarn test
 ## License
 
 This project is licensed under the MIT License - see the LICENSE file for details
+
+![](https://ga-beacon.appspot.com/UA-112845439-4/ml-classifier/readme)

package.json

@@ -1,6 +1,6 @@
 {
   "name": "ml-classifier",
-  "version": "0.3.11",
+  "version": "0.4.0",
   "description": "A machine learning engine for quickly training image classification models in your browser",
   "main": "dist/index.js",
   "typings": "dist/index.d.ts",

src/index.ts

@@ -5,45 +5,29 @@ import train from './train';
 import translateImages, {
   IImageData,
 } from './translateImages';
-import loadPretrainedModel, {
-  PRETRAINED_MODELS_KEYS,
-} from './loadPretrainedModel';
+import loadPretrainedModel from './loadPretrainedModel';
+// import log, { resetLog } from './log';
 import {
   addData,
   addLabels,
-} from './prepareTrainingData';
+} from './prepareData';
 import getDefaultDownloadHandler from './getDefaultDownloadHandler';
 
 import {
   IParams,
   // IConfigurationParams,
   IData,
   ICollectedData,
+  IArgs,
   // DataType,
 } from './types';
 
-interface IArgs {
-  onLoadStart?: Function;
-  onLoadComplete?: Function;
-  onAddDataStart?: Function;
-  onAddDataComplete?: Function;
-  onClearDataStart?: Function;
-  onClearDataComplete?: Function;
-  onTrainStart?: Function;
-  onTrainComplete?: Function;
-  onPredictComplete?: Function;
-  onPredictStart?: Function;
-  onEvaluateStart?: Function;
-  onEvaluateComplete?: Function;
-  onSaveStart?: Function;
-  onSaveComplete?: Function;
-}
-
 // export { DataType } from './types';
 
 class MLClassifier {
   // private pretrainedModel: typeof tf.model;
-  private pretrainedModel: any;
+  // private pretrainedModel: any;
+  private pretrainedModel: tf.Model;
   // private model: tf.Sequential;
   private model: any;
   private callbacks: Function[] = [];
@@ -66,11 +50,17 @@ class MLClassifier {
 
   private init = async () => {
     this.callbackFn('onLoad', 'start');
-    this.pretrainedModel = await loadPretrainedModel(PRETRAINED_MODELS_KEYS.MOBILENET);
+    this.pretrainedModel = await loadPretrainedModel(this.args.pretrainedModel);
 
     this.callbacks.map(callback => callback());
 
     this.callbackFn('onLoad', 'complete');
+
+    // Warmup the model
+    const dims = await this.getInputDims();
+    tf.tidy(() => {
+      this.pretrainedModel.predict(tf.zeros([1, ...dims, 3]));
+    });
   }
 
   private loaded = async () => new Promise(resolve => {
@@ -93,8 +83,11 @@ class MLClassifier {
       batchInputShape,
     } = inputLayers[0];
     const dims = await this.getInputDims();
+    await tf.nextFrame();
     const processedImage = await cropAndResizeImage(image, dims);
-    return this.pretrainedModel.predict(processedImage);
+    await tf.nextFrame();
+    const pred = this.pretrainedModel.predict(processedImage);
+    return pred;
   }
 
   private getInputDims = async (): Promise<[number, number]> => {
@@ -128,7 +121,7 @@ class MLClassifier {
 
   public getModel = () => this.model;
 
-  public addData = async (origImages: Array<tf.Tensor3D | IImageData | HTMLImageElement | string>, origLabels: string[], dataType: string = 'train') => {
+  public addData = async (origImages: Array<tf.Tensor | IImageData | HTMLImageElement | string>, origLabels: string[], dataType: string = 'train') => {
     this.callbackFn('onAddData', 'start', origImages, origLabels, dataType);
     if (!origImages) {
       throw new Error('You must supply images');
@@ -148,20 +141,27 @@ class MLClassifier {
     }
 
     if (dataType === 'train' || dataType === 'eval') {
-      const activatedImages = await Promise.all(images.map(async (image: tf.Tensor3D, idx: number) => {
+      const activatedImages: tf.Tensor[] = [];
+      for (let i = 0; i < images.length; i++) {
+        const image = images[i];
+        // TODO: Debug this any type
+        const activatedImage: any = await this.cropAndActivateImage(image);
+        activatedImages.push(activatedImage);
         await tf.nextFrame();
-        return await this.cropAndActivateImage(image);
-      }));
+      }
 
       this.data.classes = getClasses(labels);
       const xs = addData(activatedImages);
+      await tf.nextFrame();
       const ys = addLabels(labels, this.data.classes);
+      await tf.nextFrame();
       this.data[dataType] = {
         xs,
         ys,
       };
     }
 
+    await tf.nextFrame();
     this.callbackFn('onAddData', 'complete', origImages, labels, dataType, errors);
   }
 
@@ -196,47 +196,54 @@ class MLClassifier {
     const {
       model,
       history,
-    } = await train(data, classes, params);
+    } = await train(this.pretrainedModel, data, classes, params, this.args);
 
     this.model = model;
     this.callbackFn('onTrain', 'complete', params, history);
     return history;
   }
 
-  public predict = async (origImage: tf.Tensor3D | HTMLImageElement | string, label?: string) => {
-    this.callbackFn('onPredict', 'start', origImage);
-    await this.loaded();
-    if (!this.model) {
-      throw new Error('You must call train prior to calling predict');
+  public predict = async (origImage: tf.Tensor | HTMLImageElement | string, label?: string) => {
+    try {
+      this.callbackFn('onPredict', 'start', origImage);
+      await this.loaded();
+      if (!this.model) {
+        throw new Error('You must call train prior to calling predict');
+      }
+      const dims = await this.getInputDims();
+      const {
+        images,
+        errors,
+      } = await translateImages([origImage], dims);
+      if (errors && errors.length) {
+        throw errors[0].error;
+      }
+      const data = images[0];
+      const img = await this.cropAndActivateImage(data);
+      // TODO: Do these images need to be activated?
+      const predictedClass = tf.tidy(() => {
+        const predictions = this.model.predict(img);
+        // TODO: address this
+        return (predictions as tf.Tensor).as1D().argMax();
+      });
+
+      console.log(predictedClass.dataSync());
+
+      const classId = (await predictedClass.data())[0];
+      predictedClass.dispose();
+      const prediction = Object.entries(this.data.classes).reduce((obj, [
+        key,
+        val,
+      ]) => ({
+        ...obj,
+        [val]: key,
+      }), {})[classId];
+      this.callbackFn('onPredict', 'complete', origImage, label, prediction);
+      return prediction;
+    } catch(err) {
+      console.error(err, origImage, label);
+      throw new Error(err);
     }
-    const dims = await this.getInputDims();
-    const {
-      images,
-      errors,
-    } = await translateImages([origImage], dims);
-    if (errors && errors.length) {
-      throw errors[0].error;
-    }
-    const data = images[0];
-    const img = await this.cropAndActivateImage(data);
-    // TODO: Do these images need to be activated?
-    const predictedClass = tf.tidy(() => {
-      const predictions = this.model.predict(img);
-      // TODO: address this
-      return (predictions as tf.Tensor).as1D().argMax();
-    });
-
-    const classId = (await predictedClass.data())[0];
-    predictedClass.dispose();
-    const prediction = Object.entries(this.data.classes).reduce((obj, [
-      key,
-      val,
-    ]) => ({
-      ...obj,
-      [val]: key,
-    }), {})[classId];
-    this.callbackFn('onPredict', 'complete', origImage, label, prediction);
-    return prediction;
   }
 
   public evaluate = async (params: IParams = {}) => {
@@ -271,3 +278,5 @@ class MLClassifier {
 }
 
 export default MLClassifier;
+
+export { PRETRAINED_MODELS_KEYS as PRETRAINED_MODELS } from './loadPretrainedModel';

src/loadPretrainedModel.ts

@@ -11,7 +11,11 @@ export const PRETRAINED_MODELS = {
   },
 };
 
-const loadPretrainedModel = async (pretrainedModel: string) => {
+const loadPretrainedModel = async (pretrainedModel: string | tf.Model = PRETRAINED_MODELS_KEYS.MOBILENET) => {
+  if (pretrainedModel instanceof tf.Model) {
+    return pretrainedModel;
+  }
+
   if (!PRETRAINED_MODELS[pretrainedModel]) {
     throw new Error('You have supplied an invalid key for a pretrained model');
   }

src/prepareTrainingData.ts src/prepareData.tssrc/prepareTrainingData.ts renamed to src/prepareData.ts

@@ -16,9 +16,9 @@ const oneHot = (labelIndex: number, classLength: number) => tf.tidy(() => tf.one
 //   return newData;
 // }), undefined);
 
-export const addData = (tensors: tf.Tensor3D[]): tf.Tensor3D => {
+export const addData = (tensors: tf.Tensor[]): tf.Tensor => {
   const data = tf.keep(tensors[0]);
-  return tensors.slice(1).reduce((data: tf.Tensor3D, tensor: tf.Tensor3D) => tf.tidy(() => {
+  return tensors.slice(1).reduce((data: tf.Tensor, tensor: tf.Tensor) => tf.tidy(() => {
     const newData = tf.keep(data.concat(tensor, 0));
     data.dispose();
     return newData;