- Add test for YOLO

- Add test for Imageclassifier
- Resize img if necessary in utils

Author: cvalenzuela

src/ImageClassifier/index_test.js

@@ -13,7 +13,7 @@ const DEFAULTS = {
   version: 1,
 };
 
-describe('Create an image classifier', () => {
+describe('imageClassifier', () => {
   let classifier;
 
   beforeEach(async () => {
@@ -28,12 +28,14 @@ describe('Create an image classifier', () => {
     expect(classifier.ready).toBeTruthy();
   });
 
-  it('Should classify an robin, American robin, Turdus migratorius', async () => {
-    const img = new Image();
-    img.crossOrigin = '';
-    img.src = 'https://ml5js.org/docs/assets/img/bird.jpg';
-    await new Promise((resolve) => { img.onload = resolve; });
-    classifier.predict(img)
-      .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+  describe('predict', () => {
+    it('Should classify an image of a Robin', async () => {
+      const img = new Image();
+      img.crossOrigin = '';
+      img.src = 'https://ml5js.org/docs/assets/img/bird.jpg';
+      await new Promise((resolve) => { img.onload = resolve; });
+      classifier.predict(img)
+        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+    });
   });
 });

src/LSTM/index.js

@@ -30,8 +30,9 @@ class LSTM {
       length: 20,
       temperature: 0.5,
     };
+
     this.ready = callCallback(this.loadCheckpoints(model), callback);
-    this.then = this.ready.then.bind(this.ready);
+    // this.then = this.ready.then.bind(this.ready);
   }
 
   async loadCheckpoints(path) {
@@ -59,9 +60,10 @@ class LSTM {
     return this;
   }
 
-  async loadVocab(file) {
-    const json = await fetch(`${file}/vocab.json`)
-      .then(response => response.json());
+  async loadVocab(path) {
+    const json = await fetch(`${path}/vocab.json`)
+      .then(response => response.json())
+      .catch(err => console.error(err));
     this.vocab = json;
     this.vocabSize = Object.keys(json).length;
   }

src/LSTM/index_test.js

@@ -1,21 +1,23 @@
-/* eslint new-cap: 0 */
+const { LSTMGenerator } = ml5;
 
-// import LSTMGenerator from './index';
+const LSTM_MODEL_URL = 'https://raw.githubusercontent.com/ml5js/ml5-data-and-models/master/models/lstm/woolf/';
+const LSTM_MODEL_DEFAULTS = {
+  cellsAmount: 2,
+  vocabSize: 90,
+  firstChar: 61,
+};
 
-describe('LSTM', () => {
-  // let generator;
+describe('LSTMGenerator', () => {
+  let lstm;
 
-  // beforeEach(async () => {
-  //   jasmine.DEFAULT_TIMEOUT_INTERVAL = 10000;
-  //   generator = await LSTMGenerator('https://raw.githubusercontent.com/ml5js/ml5-data-and-training/master/models/lstm/dubois/');
-  // });
+  beforeEach(async () => {
+    // This never resolves.
+    // lstm = await LSTMGenerator(LSTM_MODEL_URL);
+  });
 
-  // it('instantiates a generator', () => {
-  //   expect(generator).toBeTruthy();
-  // });
-
-  // Fails with 'must be a Tensor' error that's particular to this test suite.
-  // it('generates some text', async () => {
-  //   expect(await generator.generate('Hi there')).toBeTruthy();
-  // });
+  it('instantiates a lstm generator', () => {
+    // expect(lstm.cellsAmount).toBe(LSTM_MODEL_DEFAULTS.cellsAmount);
+    // expect(lstm.vocabSize).toBe(DEFAULTS.vocabSize);
+    // expect(lstm.vocab[0]).toBe(DEFAULTS.firstChar);
+  });
 });

src/Word2vec/index_test.js

@@ -1,7 +1,7 @@
 /* eslint no-loop-func: 0 */
 const { word2vec } = ml5;
 
-const URL = 'https://raw.githubusercontent.com/ml5js/ml5-data-and-training/master/models/wordvecs/common-english/wordvecs1000.json';
+const W2V_MODEL_URL = 'https://raw.githubusercontent.com/ml5js/ml5-data-and-training/master/models/wordvecs/common-english/wordvecs1000.json';
 
 describe('word2vec', () => {
   let word2vecInstance;
@@ -10,7 +10,7 @@ describe('word2vec', () => {
   beforeAll((done) => {
     jasmine.DEFAULT_TIMEOUT_INTERVAL = 5000;
     numTensorsBeforeAll = tf.memory().numTensors;
-    word2vecInstance = word2vec(URL, done);
+    word2vecInstance = word2vec(W2V_MODEL_URL, done);
   });
 
   afterAll(() => {

src/YOLO/index.js

@@ -44,7 +44,7 @@ class YOLOBase extends Video {
     this.modelReady = false;
     this.isPredicting = false;
     this.ready = callCallback(this.loadModel(), callback);
-    this.then = this.ready.then;
+    // this.then = this.ready.then;
   }
 
   async loadModel() {
@@ -75,7 +75,7 @@ class YOLOBase extends Video {
     await tf.nextFrame();
     this.isPredicting = true;
     const [allBoxes, boxConfidence, boxClassProbs] = tf.tidy(() => {
-      const input = imgToTensor(imgToPredict);
+      const input = imgToTensor(imgToPredict, [imageSize, imageSize]);
       const activation = this.model.predict(input);
       const [boxXY, boxWH, bConfidence, bClassProbs] = head(activation, ANCHORS, 80);
       const aBoxes = boxesToCorners(boxXY, boxWH);

src/YOLO/index.test.js

@@ -0,0 +1,40 @@
+/* eslint new-cap: 0 */
+
+const { YOLO } = ml5;
+
+const YOLO_DEFAULTS = {
+  IOUThreshold: 0.4,
+  classProbThreshold: 0.4,
+  filterBoxesThreshold: 0.01,
+  size: 416,
+};
+
+describe('YOLO', () => {
+  let yolo;
+
+  async function getRobin() {
+    const img = new Image();
+    img.crossOrigin = '';
+    img.src = 'https://ml5js.org/docs/assets/img/bird.jpg';
+    await new Promise((resolve) => { img.onload = resolve; });
+    return img;
+  }
+
+  beforeEach(async () => {
+    jasmine.DEFAULT_TIMEOUT_INTERVAL = 100000;
+    yolo = await YOLO();
+  });
+
+  it('instantiates the YOLO classifier with defaults', () => {
+    expect(yolo.IOUThreshold).toBe(YOLO_DEFAULTS.IOUThreshold);
+    expect(yolo.classProbThreshold).toBe(YOLO_DEFAULTS.classProbThreshold);
+    expect(yolo.filterBoxesThreshold).toBe(YOLO_DEFAULTS.filterBoxesThreshold);
+    expect(yolo.size).toBe(YOLO_DEFAULTS.size);
+  });
+
+  it('detects a robin', async () => {
+    const robin = await getRobin();
+    const detection = await yolo.detect(robin);
+    expect(detection[0].className).toBe('bird');
+  });
+});

src/YOLO/index_test.js

@@ -62,9 +62,12 @@ const cropImage = (img) => {
 };
 
 // Static Method: image to tf tensor
-function imgToTensor(input) {
+function imgToTensor(input, size = null) {
   return tf.tidy(() => {
-    const img = tf.fromPixels(input);
+    let img = tf.fromPixels(input);
+    if (size) {
+      img = tf.image.resizeBilinear(img, size);
+    }
     const croppedImage = cropImage(img);
     const batchedImage = croppedImage.expandDims(0);
     return batchedImage.toFloat().div(tf.scalar(127)).sub(tf.scalar(1));