patch style transfer for non-square and canvas-derived inputs (#311)

ml5js · Mar 15, 2019 · fc311e3 · fc311e3
1 parent 7d7600e
commit fc311e3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/StyleTransfer/index.js b/src/StyleTransfer/index.js
@@ -89,7 +89,9 @@ class StyleTransfer extends Video {
     let input;
     let callback = cb;
 
-    if (inputOrCallback instanceof HTMLVideoElement || inputOrCallback instanceof HTMLImageElement) {
+    if (inputOrCallback instanceof HTMLVideoElement ||
+        inputOrCallback instanceof HTMLImageElement ||
+        inputOrCallback instanceof ImageData) {
       input = inputOrCallback;
     } else if (typeof inputOrCallback === 'object' && (inputOrCallback.elt instanceof HTMLVideoElement || inputOrCallback.elt instanceof HTMLImageElement)) {
       input = inputOrCallback.elt;

diff --git a/src/utils/imageUtilities.js b/src/utils/imageUtilities.js
@@ -24,7 +24,7 @@ const processVideo = (input, size, callback = () => {}) => {
 
 // Converts a tf to DOM img
 const array3DToImage = (tensor) => {
-  const [imgWidth, imgHeight] = tensor.shape;
+  const [imgHeight, imgWidth] = tensor.shape;
   const data = tensor.dataSync();
   const canvas = document.createElement('canvas');
   canvas.width = imgWidth;