Merge pull request #8 from eternaleclipse/data-url-support

Added data url support

Author: pivanov

index.js

@@ -1,9 +1,10 @@
 global.fetch = require("node-fetch");
 
 const { Readable } = require('stream');
-const tf = require('@tensorflow/tfjs')
-const PImage = require('pureimage')
+const tf = require('@tensorflow/tfjs');
+const PImage = require('pureimage');
 const isImageUrl = require('is-image-url');
+const parseDataUrl = require('parse-data-url');
 
 const wait = ms => new Promise(r => setTimeout(r, ms));
 
@@ -27,7 +28,7 @@ const retryOperation = (operation, delay, times) => new Promise((resolve, reject
 const bufferToStream = (binary) => {
   const readableInstanceStream = new Readable({
     read() {
-      this.push(binary)
+      this.push(binary);
       this.push(null);
     }
   });
@@ -76,7 +77,7 @@ const getTopKClasses = async (logits, classes) => {
     topkIndices[i] = valuesAndIndices[i].index;
   }
 
-  const topClassesAndProbs = []
+  const topClassesAndProbs = [];
   for (let i = 0; i < topkIndices.length; i++) {
     topClassesAndProbs.push({
       class: classes[topkIndices[i]],
@@ -100,7 +101,7 @@ class SashiDoTeachableMachine {
 
     try {
       const modelURL = `${modelUrl}model.json`;
-      const response = await fetch(`${modelUrl}metadata.json`)
+      const response = await fetch(`${modelUrl}metadata.json`);
       const body = await response.text();
       this.model = await tf.loadLayersModel(modelURL);
       this.model.classes = JSON.parse(body).labels;
@@ -125,7 +126,7 @@ class SashiDoTeachableMachine {
   async classify(params) {
     const { imageUrl } = params;
 
-    if (!isImageUrl(imageUrl)) {
+    if ((!imageUrl.startsWith('data:image/')) && (!isImageUrl(imageUrl))) {
       return Promise.reject({ error: "Image URL is not valid!" });
     }
 
@@ -138,10 +139,22 @@ class SashiDoTeachableMachine {
 
   async inference({ imageUrl }) {
     try {
-      const data = await fetch(imageUrl);
-
-      const contentType = data.headers.get("Content-Type");
-      const buffer = await data.buffer();
+      let data;
+      let buffer;
+      let contentType;
+
+      if (imageUrl.startsWith('data:image/')) {
+        data = parseDataUrl(imageUrl);
+        
+        contentType = data.contentType;
+        buffer = data.toBuffer();
+      } else {
+        data = await fetch(imageUrl);
+
+        contentType = data.headers.get("Content-Type");
+        buffer = await data.buffer();
+      }
+      
       const stream = bufferToStream(buffer);
       let imageBitmap;