mtcnn - createImageBitmap alternative for nodejs env

justadudewhohacks · justadudewhohacks · commit c8b3410d6d12 · 2018-11-03T18:09:57.000+01:00
diff --git a/src/mtcnn/extractImagePatches.ts b/src/mtcnn/extractImagePatches.ts
@@ -1,5 +1,12 @@
 import * as tf from '@tensorflow/tfjs-core';
-import { Box, createCanvas, getContext2dOrThrow, IDimensions } from 'tfjs-image-recognition-base';
+import {
+  Box,
+  createCanvas,
+  createCanvasFromMedia,
+  env,
+  getContext2dOrThrow,
+  IDimensions,
+} from 'tfjs-image-recognition-base';
 
 import { normalize } from './normalize';
 
@@ -20,7 +27,7 @@ export async function extractImagePatches(
     const fromY = y - 1
     const imgData = imgCtx.getImageData(fromX, fromY, (ex - fromX), (ey - fromY))
 
-    return createImageBitmap(imgData)
+    return env.isNodejs() ? createCanvasFromMedia(imgData) : createImageBitmap(imgData)
   }))
 
   const imagePatchesDatas: number[][] = []
