feat: Support for Pose models from Teachable Machine

Teachable Machine requires TFJS v1, where the rest of ML4K uses TFJS v2, so I'm moving the pose model support to a separate iframe so I can run it on-demand. Signed-off-by: Dale Lane <dale.lane@uk.ibm.com>
IBM · Dec 29, 2020 · 02c38b9 · 02c38b9
1 parent 3082ee4
commit 02c38b9
Show file tree

Hide file tree

Showing 7 changed files with 58 additions and 5 deletions.
diff --git a/Gruntfile.js b/Gruntfile.js
@@ -143,7 +143,8 @@ module.exports = function(grunt) {
                 expand : true,
                 cwd : 'public/scratch-components',
                 src : ['help-scratch3*',
-                       'help-scratch.css'],
+                       'help-scratch.css',
+                       'teachablemachinepose.html'],
                 dest : 'web/scratch3'
             },
             indexhtml : {

diff --git a/gulpfile.js b/gulpfile.js
@@ -165,7 +165,8 @@ gulp.task('scratch3install', gulp.series('crossdomain', function() {
     return gulp.src([
         'public/scratch3/**',
         'public/scratch-components/help-scratch3*',
-        'public/scratch-components/help-scratch.css'
+        'public/scratch-components/help-scratch.css',
+        'public/scratch-components/teachablemachinepose.html'
     ]).pipe(gulp.dest('web/scratch3'));
 }));
 

diff --git a/public/components/pretrained/tfjs.tmpl.html b/public/components/pretrained/tfjs.tmpl.html
@@ -22,7 +22,7 @@ <h2>Use a pre-trained TensorFlow model in Scratch</h2>
                 <div style="padding: 1em;">
                     <div>What type of model is it?</div>
                     <md-select name="modeltypeid" ng-model="modeltypeid" ng-change="generateScratchKey(modeljson, modeltypeid)" style="padding-left: 20px;">
-                        <md-option value="10">Teachable Machine (image)</md-option>
+                        <md-option value="10">Teachable Machine (image or pose)</md-option>
                         <md-option value="11">TensorFlow GraphDef (image)</md-option>
                     </md-select>
                 </div>

diff --git a/public/scratch-components/teachablemachinepose.html b/public/scratch-components/teachablemachinepose.html
@@ -0,0 +1,41 @@
+<html>
+    <head>
+        <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.3.1/dist/tf.min.js"></script>
+        <script src="https://cdn.jsdelivr.net/npm/@teachablemachine/pose@0.8/dist/teachablemachine-pose.min.js"></script>
+        <script type="text/javascript">
+            const posemodels = {};
+
+            function initModel(id, modelurl, metadataurl) {
+                return tmPose.load(modelurl, metadataurl)
+                    .then(function (model) {
+                        posemodels[id] = model;
+                    });
+            }
+
+            function predict(id, imageelem) {
+                return posemodels[id].estimatePose(imageelem)
+                    .then(function (resp) {
+                        return posemodels[id].predict(resp.posenetOutput);
+                    })
+                    .then(function (output) {
+                        return output.map(function (tmitem) {
+                            return {
+                                class_name  : tmitem.className,
+                                confidence  : 100 * tmitem.probability
+                            };
+                        });
+                    });
+            }
+
+            function createImage(b64imagedata, callback) {
+                const imageElement = document.createElement('img');
+                imageElement.width = 257;
+                imageElement.height = 257;
+                imageElement.onload = function () {
+                    callback(imageElement);
+                };
+                imageElement.src = 'data:image/jpeg;base64,' + b64imagedata;
+            }
+        </script>
+    </head>
+</html>
diff --git a/src/lib/scratchx/extensions.ts b/src/lib/scratchx/extensions.ts
@@ -215,6 +215,10 @@ export async function getScratchTfjsExtension(scratchkey: string): Promise<strin
     const modelinfo = await scratchtfjs.getModelInfoFromScratchKey(scratchkey);
     const metadata = await scratchtfjs.getMetadata(modelinfo);
 
+    if (metadata.packageName === '@teachablemachine/pose') {
+        modelinfo.modeltype = 'teachablemachinepose';
+    }
+
     const template: string = await fileutils.read('./resources/scratch3-tfjs-classify.js');
 
     Mustache.parse(template);

diff --git a/src/lib/scratchx/scratchtfjs.ts b/src/lib/scratchx/scratchtfjs.ts
@@ -61,6 +61,8 @@ function getModelTypeAsId(type: ScratchTypes.ScratchTfjsModelType): ScratchTypes
         return 10;
     case 'graphdefimage':
         return 11;
+    case 'teachablemachinepose':
+        return 12;
     default:
         return 99;
     }
@@ -72,6 +74,9 @@ function getModelTypeFromId(id: ScratchTypes.ScratchTfjsModelTypeId): ScratchTyp
     else if (id === 11) {
         return 'graphdefimage';
     }
+    else if (id === 12) {
+        return 'teachablemachinepose';
+    }
     else {
         return 'unknown';
     }

diff --git a/src/lib/scratchx/scratchx-types.ts b/src/lib/scratchx/scratchx-types.ts
@@ -9,9 +9,10 @@ export interface Key {
 }
 
 export type ScratchTfjsModelType = 'teachablemachineimage' |
+                                   'teachablemachinepose' |
                                    'graphdefimage' |
                                    'unknown';
-export type ScratchTfjsModelTypeId = 10 | 11 | 99;
+export type ScratchTfjsModelTypeId = 10 | 11 | 12 | 99;
 
 
 export interface ScratchTfjsExtensionEncoded {
@@ -21,7 +22,7 @@ export interface ScratchTfjsExtensionEncoded {
 
 export interface ScratchTfjsExtension {
     readonly modelurl: string;
-    readonly modeltype: ScratchTfjsModelType;
+    modeltype: ScratchTfjsModelType;
 }
 
 export interface ScratchTfjsExtensionWithId extends ScratchTfjsExtension {