[In Progress] Offline support for ml5 (#553)

* added modelLoader util

* added proof of concept for offline sentiment

* adds offline support for uNet

* updates sketchRnn for offline support

* adds offline support for yolo

* adds modelUrl to config

* adds package lock

* fix imageSize ref

* adds support for external URL defs

* updates packagelock

* fixes unet tests

Author: joeyklee

package-lock.json

@@ -74,8 +74,17 @@ class Mobilenet {
       learningRate: options.learningRate || DEFAULTS.learningRate,
       batchSize: options.batchSize || DEFAULTS.batchSize,
       layer: options.layer || DEFAULTS.layer,
-      alpha: options.alpha || DEFAULTS.alpha
+      alpha: options.alpha || DEFAULTS.alpha,
     }
+
+    // for graph model
+    this.model = null;
+    this.url = MODEL_INFO[this.config.version][this.config.alpha];
+    this.normalizationOffset = tf.scalar(127.5);
+
+    // check if a mobilenet URL is given
+    this.mobilenetURL = options.mobilenetURL || `${BASE_URL}${this.config.version}_${this.config.alpha}_${IMAGE_SIZE}/model.json`; 
+    this.graphModelURL = options.graphModelURL || this.url;
     /**
      * Boolean value to check if the model is predicting.
      * @public
@@ -92,15 +101,17 @@ class Mobilenet {
     this.usageType = null;
     this.ready = callCallback(this.loadModel(), callback);
 
-    // for graph model
-    this.model = null;
-    this.url = MODEL_INFO[this.config.version][this.config.alpha];
-    this.normalizationOffset = tf.scalar(127.5);
+    
   }
 
   async loadModel() {
-    this.mobilenet = await tf.loadLayersModel(`${BASE_URL}${this.config.version}_${this.config.alpha}_${IMAGE_SIZE}/model.json`);
-    this.model = await tf.loadGraphModel(this.url, {fromTFHub: true});
+    this.mobilenet = await tf.loadLayersModel(this.mobilenetURL);
+    if(this.graphModelURL.includes('https://tfhub.dev/')){
+      this.model = await tf.loadGraphModel(this.graphModelURL, {fromTFHub: true});
+    } else {
+      this.model = await tf.loadGraphModel(this.graphModelURL, {fromTFHub: false});
+    }
+    
 
 
     const layer = this.mobilenet.getLayer(this.config.layer);

src/FeatureExtractor/Mobilenet.js

@@ -61,6 +61,7 @@ class PoseNet extends EventEmitter {
      * @type {String}
      * @public
      */
+    this.modelUrl = options.modelUrl || null;
     this.architecture = options.architecture || DEFAULTS.architecture;
     this.detectionType = detectionType || options.detectionType || DEFAULTS.detectionType;
     this.imageScaleFactor = options.imageScaleFactor || DEFAULTS.imageScaleFactor;
@@ -85,7 +86,8 @@ class PoseNet extends EventEmitter {
         outputStride: this.outputStride,
         inputResolution: this.inputResolution,
         multiplier: this.multiplier,
-        quantBytes: this.quantBytes
+        quantBytes: this.quantBytes,
+        modelUrl: this.modelUrl
       }
     } else {
       modelJson = {

src/PoseNet/index.js

@@ -1,8 +1,10 @@
 import * as tf from '@tensorflow/tfjs';
 import callCallback from '../utils/callcallback';
+import modelLoader from '../utils/modelLoader';
+
 /**
-* Initializes the Sentiment demo.
-*/
+ * Initializes the Sentiment demo.
+ */
 
 const OOV_CHAR = 2;
 const PAD_CHAR = 0;
@@ -51,19 +53,32 @@ class Sentiment {
   }
 
   /**
-  * Initializes the Sentiment demo.
-  */
+   * Initializes the Sentiment demo.
+   */
 
   async loadModel(modelName) {
 
+    const movieReviews = {
+      model: null,
+      metadata: null,
+    }
+    
     if (modelName.toLowerCase() === 'moviereviews') {
+        
+      movieReviews.model = 'https://storage.googleapis.com/tfjs-models/tfjs/sentiment_cnn_v1/model.json';
+      movieReviews.metadata = 'https://storage.googleapis.com/tfjs-models/tfjs/sentiment_cnn_v1/metadata.json';
+
+    } else if(modelLoader.isAbsoluteURL(modelName) === true ) {
+      const modelPath = modelLoader.getModelPath(modelName);
+
+      movieReviews.model = `${modelPath}/model.json`;
+      movieReviews.metadata = `${modelPath}/metadata.json`;
+
+    } else {
+      console.error('problem loading model');
+      return this;
+    }
 
-    const movieReviews = {
-      model:
-        'https://storage.googleapis.com/tfjs-models/tfjs/sentiment_cnn_v1/model.json',
-      metadata:
-        'https://storage.googleapis.com/tfjs-models/tfjs/sentiment_cnn_v1/metadata.json',
-    };
 
     /**
      * The model being used.
@@ -80,9 +95,6 @@ class Sentiment {
     this.wordIndex = sentimentMetadata.word_index;
     this.vocabularySize = sentimentMetadata.vocabulary_size;
 
-    } else {
-      console.error('problem loading model')
-    }
     return this;
   }
 
@@ -112,15 +124,12 @@ class Sentiment {
     const score = predictOut.dataSync()[0];
     predictOut.dispose();
 
-    return { score };
+    return {
+      score
+    };
   }
 }
 
-const sentiment = (modelName, callback) => new Sentiment( modelName, callback ) ;
-
-export default sentiment;
-
-
-
-
+const sentiment = (modelName, callback) => new Sentiment(modelName, callback);
 
+export default sentiment;

src/Sentiment/index.js

@@ -12,10 +12,21 @@ SketchRNN
 import * as ms from '@magenta/sketch';
 import callCallback from '../utils/callcallback';
 import modelPaths from './models';
+import modelLoader from '../utils/modelLoader';
 
-const PATH_START_LARGE = 'https://storage.googleapis.com/quickdraw-models/sketchRNN/large_models/';
-const PATH_START_SMALL = 'https://storage.googleapis.com/quickdraw-models/sketchRNN/models/';
-const PATH_END = '.gen.json';
+// const PATH_START_LARGE = 'https://storage.googleapis.com/quickdraw-models/sketchRNN/large_models/';
+// const PATH_START_SMALL = 'https://storage.googleapis.com/quickdraw-models/sketchRNN/models/';
+// const PATH_END = '.gen.json';
+
+
+const DEFAULTS = {
+  modelPath: 'https://storage.googleapis.com/quickdraw-models/sketchRNN/large_models/',
+  modelPath_large: 'https://storage.googleapis.com/quickdraw-models/sketchRNN/models/',
+  modelPath_small: 'https://storage.googleapis.com/quickdraw-models/sketchRNN/models/',
+  PATH_END: '.gen.json',
+  temperature: 0.65,
+  pixelFactor: 3.0,
+}
 
 class SketchRNN {
   /**
@@ -27,21 +38,37 @@ class SketchRNN {
    */
   constructor(model, callback, large = true) {
     let checkpointUrl = model;
-    if (modelPaths.has(checkpointUrl)) {
-      checkpointUrl = (large ? PATH_START_LARGE : PATH_START_SMALL) + checkpointUrl + PATH_END;
-    }
-    this.defaults = {
+
+    this.config = {
       temperature: 0.65,
       pixelFactor: 3.0,
+      modelPath: DEFAULTS.modelPath,
+      modelPath_small: DEFAULTS.modelPath_small,
+      modelPath_large: DEFAULTS.modelPath_large,
+      PATH_END: DEFAULTS.PATH_END,
     };
-    this.model = new ms.SketchRNN(checkpointUrl);
+    
+    
+    if(modelLoader.isAbsoluteURL(checkpointUrl) === true){
+      const modelPath = modelLoader.getModelPath(checkpointUrl);
+      this.config.modelPath = modelPath;
+
+    } else if(modelPaths.has(checkpointUrl)) {
+        checkpointUrl = (large ? this.config.modelPath : this.config.modelPath_small) + checkpointUrl + this.config.PATH_END;
+        this.config.modelPath = checkpointUrl;
+    } else {
+      console.log('no model found!');
+      return this;
+    }
+
+    this.model = new ms.SketchRNN(this.config.modelPath);
     this.penState = this.model.zeroInput();
     this.ready = callCallback(this.model.initialize(), callback);
   }
 
   async generateInternal(options, strokes) {
-    const temperature = +options.temperature || this.defaults.temperature;
-    const pixelFactor = +options.pixelFactor || this.defaults.pixelFactor;
+    const temperature = +options.temperature || this.config.temperature;
+    const pixelFactor = +options.pixelFactor || this.config.pixelFactor;
 
     await this.ready;
     if (!this.rnnState) {

src/SketchRNN/index.js

@@ -12,8 +12,10 @@ import callCallback from '../utils/callcallback';
 import { array3DToImage } from '../utils/imageUtilities';
 import p5Utils from '../utils/p5Utils';
 
-const URL = 'https://raw.githubusercontent.com/zaidalyafeai/HostedModels/master/unet-128/model.json';
-const imageSize = 128;
+const DEFAULTS = {
+  modelPath: 'https://raw.githubusercontent.com/zaidalyafeai/HostedModels/master/unet-128/model.json',
+  imageSize: 128
+}
 
 class UNET {
   /**
@@ -26,11 +28,15 @@ class UNET {
   constructor(video, options, callback) {
     this.modelReady = false;
     this.isPredicting = false;
+    this.config = {
+      modelPath: typeof options.modelPath !== 'undefined' ? options.modelPath : DEFAULTS.modelPath,
+      imageSize: typeof options.imageSize !== 'undefined' ? options.imageSize : DEFAULTS.imageSize
+    };
     this.ready = callCallback(this.loadModel(), callback);
   }
 
   async loadModel() {
-    this.model = await tf.loadLayersModel(URL);
+    this.model = await tf.loadLayersModel(this.config.modelPath);
     this.modelReady = true;
     return this;
   }
@@ -80,7 +86,7 @@ class UNET {
     const tensor = tf.tidy(() => {
       // preprocess
       const tfImage = tf.browser.fromPixels(imgToPredict).toFloat();
-      const resizedImg = tf.image.resizeBilinear(tfImage, [imageSize, imageSize]);
+      const resizedImg = tf.image.resizeBilinear(tfImage, [this.config.imageSize, this.config.imageSize]);
       const normTensor = resizedImg.div(tf.scalar(255));
 
       const batchedImage = normTensor.expandDims(0);
@@ -102,7 +108,7 @@ class UNET {
     let image;
 
     if (p5Utils.checkP5()) {
-        const blob1 = await p5Utils.rawToBlob(raw, imageSize, imageSize);
+        const blob1 = await p5Utils.rawToBlob(raw, this.config.imageSize, this.config.imageSize);
         const p5Image1 = await p5Utils.blobToP5Image(blob1);
         image = p5Image1;
     }
@@ -129,7 +135,7 @@ const uNet = (videoOr, optionsOr, cb) => {
     callback = videoOr;
   } else if (typeof videoOr === 'object') {
     options = videoOr;
-  }
+  } 
 
   if (typeof optionsOr === 'object') {
     options = optionsOr;

src/UNET/index.js

@@ -14,6 +14,7 @@ import Video from '../utils/Video';
 import { imgToTensor } from '../utils/imageUtilities';
 import callCallback from '../utils/callcallback';
 import CLASS_NAMES from './../utils/COCO_CLASSES';
+import modelLoader from '../utils/modelLoader';
 
 import {
   nonMaxSuppression,
@@ -23,9 +24,8 @@ import {
   ANCHORS,
 } from './postprocess';
 
-const URL = 'https://raw.githubusercontent.com/ml5js/ml5-data-and-training/master/models/YOLO/model.json';
-
 const DEFAULTS = {
+  modelUrl: 'https://raw.githubusercontent.com/ml5js/ml5-data-and-training/master/models/YOLO/model.json',
   filterBoxesThreshold: 0.01,
   IOUThreshold: 0.4,
   classProbThreshold: 0.4,
@@ -51,6 +51,7 @@ class YOLOBase extends Video {
   constructor(video, options, callback) {
     super(video, imageSize);
 
+    this.modelUrl = options.modelUrl || DEFAULTS.modelUrl;
     this.filterBoxesThreshold = options.filterBoxesThreshold || DEFAULTS.filterBoxesThreshold;
     this.IOUThreshold = options.IOUThreshold || DEFAULTS.IOUThreshold;
     this.classProbThreshold = options.classProbThreshold || DEFAULTS.classProbThreshold;
@@ -64,7 +65,15 @@ class YOLOBase extends Video {
     if (this.videoElt && !this.video) {
       this.video = await this.loadVideo();
     }
-    this.model = await tf.loadLayersModel(URL);
+
+    if(modelLoader.isAbsoluteURL(this.modelUrl) === true){
+      this.model = await tf.loadLayersModel(this.modelUrl);
+    } else {
+      const modelPath = modelLoader.getModelPath(this.modelUrl);
+      this.modelUrl = `${modelPath}/model.json`;
+      this.model = await tf.loadLayersModel(this.modelUrl);
+    }
+    
     this.modelReady = true;
     return this;
   }

src/YOLO/index.js

@@ -0,0 +1,14 @@
+function isAbsoluteURL(str) {
+    const pattern = new RegExp('^(?:[a-z]+:)?//', 'i');
+    return !!pattern.test(str);
+}
+
+function getModelPath(absoluteOrRelativeUrl) {
+    const modelJsonPath = isAbsoluteURL(absoluteOrRelativeUrl) ? absoluteOrRelativeUrl : window.location.pathname + absoluteOrRelativeUrl
+    return modelJsonPath;
+}
+
+export default {
+    isAbsoluteURL,
+    getModelPath
+}