added next frame to style transfer, improving performance

Author: cvalenzuela

src/StyleTransfer/index.js

@@ -9,23 +9,38 @@ Fast Style Transfer
 */
 
 import * as tf from '@tensorflow/tfjs';
+import Video from './../utils/Video';
 import CheckpointLoader from '../utils/checkpointLoader';
 import { array3DToImage } from '../utils/imageUtilities';
 
-class StyleTransfer {
-  constructor(model, callback) {
+const IMAGE_SIZE = 200;
+
+class StyleTransfer extends Video {
+  constructor(model, videoOrCallback, cb = () => {}) {
+    super(videoOrCallback, IMAGE_SIZE);
     this.ready = false;
     this.variableDictionary = {};
     this.timesScalar = tf.scalar(150);
     this.plusScalar = tf.scalar(255.0 / 2);
     this.epsilonScalar = tf.scalar(1e-3);
-    this.loadCheckpoints(model).then(() => {
-      this.ready = true;
-      if (callback) {
-        callback();
-      }
+    this.video = null;
+
+    let callback = cb;
+    if (typeof videoOrCallback === 'function') {
+      callback = videoOrCallback;
+    }
+
+    this.loadVideo().then(() => {
+      this.videoReady = true;
+      this.loadCheckpoints(model).then(() => {
+        this.ready = true;
+        if (callback) {
+          callback();
+        }
+      });
     });
   }
+
   async loadCheckpoints(path) {
     const checkpointLoader = new CheckpointLoader(path);
     this.variables = await checkpointLoader.getAllVariables();
@@ -70,7 +85,19 @@ class StyleTransfer {
     return y3;
   }
 
-  transfer(input) {
+  async transfer(inputOrCallback, cb = () => {}) {
+    let input;
+    let callback = cb;
+
+    if (inputOrCallback instanceof HTMLVideoElement || inputOrCallback instanceof HTMLImageElement) {
+      input = inputOrCallback;
+    } else if (typeof inputOrCallback === 'object' && (inputOrCallback.elt instanceof HTMLVideoElement || inputOrCallback.elt instanceof HTMLImageElement)) {
+      input = inputOrCallback;
+    } else if (typeof inputOrCallback === 'function') {
+      input = this.video;
+      callback = inputOrCallback;
+    }
+
     const image = tf.fromPixels(input);
     const result = tf.tidy(() => {
       const conv1 = this.convLayer(image, 1, true, 0);
@@ -91,7 +118,8 @@ class StyleTransfer {
       const normalized = tf.div(clamped, tf.scalar(255.0));
       return normalized;
     });
-    return array3DToImage(result);
+    await tf.nextFrame();
+    callback(array3DToImage(result));
   }
 
   // Static Methods