From 62ccbe41bd9a1cd65a21d6a78bb13a272feeecd2 Mon Sep 17 00:00:00 2001
From: Moe Alam <github@m03.ca>
Date: Mon, 18 Jan 2021 12:17:47 -0800
Subject: [PATCH] update tensorflow module test

---
 plugins/tensorflow/test.js | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/plugins/tensorflow/test.js b/plugins/tensorflow/test.js
index 2efa84c4..c3920e12 100644
--- a/plugins/tensorflow/test.js
+++ b/plugins/tensorflow/test.js
@@ -19,20 +19,47 @@
 // npm rebuild @tensorflow/tfjs-node-gpu@1.7.3 build-addon-from-source --unsafe-perm
 // ==============================================================
 
+console.log('############################################')
+console.log('@tensorflow/tfjs-node(-gpu) module test for Object Detection')
+console.log('############################################')
 
 // Base Init >>
 var fs = require('fs');
 const fetch = require('node-fetch');
 // Base Init />>
-
 var tf
+const tfjsBuild = process.argv[2]
 try{
-    tf = require('@tensorflow/tfjs-node-gpu')
+    switch(tfjsBuild){
+        case'gpu':
+            tf = require('@tensorflow/tfjs-node-gpu')
+            console.log('GPU Test for Tensorflow Module')
+        break;
+        case'cpu':
+            tf = require('@tensorflow/tfjs-node')
+            console.log('CPU Test for Tensorflow Module')
+        break;
+        default:
+            console.log('Nothing selected, using CPU Module for test.')
+            console.log(`Hint : Run the script like one of the following to specify cpu or gpu.`)
+            console.log(`node test.js cpu`)
+            console.log(`node test.js gpu`)
+            tf = require('@tensorflow/tfjs-node')
+        break;
+    }
 }catch(err){
+    console.log(`Selection Failed. Could not load desired module. ${tfjsBuild}`)
+    console.log(err)
+}
+if(!tf){
     try{
-        tf = require('@tensorflow/tfjs-node')
+        tf = require('@tensorflow/tfjs-node-gpu')
     }catch(err){
-        return console.log('tfjs-node could not be loaded')
+        try{
+            tf = require('@tensorflow/tfjs-node')
+        }catch(err){
+            return console.log('tfjs-node could not be loaded')
+        }
     }
 }