Support canvas elements and canvases from p5.js in ImageClassifier (#206)

Author: meiamsome

src/ImageClassifier/index.js

@@ -70,6 +70,12 @@ class ImageClassifier {
       imgToPredict = inputNumOrCallback;
     } else if (typeof inputNumOrCallback === 'object' && inputNumOrCallback.elt instanceof HTMLImageElement) {
       imgToPredict = inputNumOrCallback.elt; // Handle p5.js image
+    } else if (inputNumOrCallback instanceof HTMLCanvasElement) {
+      imgToPredict = inputNumOrCallback;
+    } else if (typeof inputNumOrCallback === 'object' && inputNumOrCallback.elt instanceof HTMLCanvasElement) {
+      imgToPredict = inputNumOrCallback.elt; // Handle p5.js image
+    } else if (typeof inputNumOrCallback === 'object' && inputNumOrCallback.canvas instanceof HTMLCanvasElement) {
+      imgToPredict = inputNumOrCallback.canvas; // Handle p5.js image
     } else if (!(this.video instanceof HTMLVideoElement)) {
       // Handle unsupported input
       throw new Error('No input image provided. If you want to classify a video, pass the video element in the constructor. ');

src/ImageClassifier/index_test.js

@@ -27,6 +27,15 @@ describe('imageClassifier', () => {
     return img;
   }
 
+  async function getCanvas() {
+    const img = await getImage();
+    const canvas = document.createElement('canvas');
+    canvas.width = img.width;
+    canvas.height = img.height;
+    canvas.getContext('2d').drawImage(img, 0, 0);
+    return canvas;
+  }
+
   beforeEach(async () => {
     jasmine.DEFAULT_TIMEOUT_INTERVAL = 5000;
     classifier = await imageClassifier('MobileNet', undefined, {});
@@ -45,6 +54,29 @@ describe('imageClassifier', () => {
       await classifier.predict(img)
         .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
     });
+
+    it('Should support p5 elements with an image on .elt', async () => {
+      const img = await getImage();
+      await classifier.predict({ elt: img })
+        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+    });
+
+    it('Should support HTMLCanvasElement', async () => {
+      const canvas = await getCanvas();
+      await classifier.predict(canvas)
+        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+    });
+
+    it('Should support p5 elements with canvas on .canvas', async () => {
+      const canvas = await getCanvas();
+      await classifier.predict({ canvas })
+        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+    });
+
+    it('Should support p5 elements with canvas on .elt', async () => {
+      const canvas = await getCanvas();
+      await classifier.predict({ elt: canvas })
+        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+    });
   });
 });
-