[pose-detection]Use 3d model for BlazePose. (#762)

FEATURE

Author: lina128

pose-detection/demos/live_video/index.html

@@ -27,10 +27,9 @@
       position: relative;
       margin: 0;
     }
-    canvas {
-      position: absolute;
-      top: 0;
-      left: 0;
+    #canvas-wrapper,
+    #scatter-gl-container {
+      position: relative;
     }
   </style>
 </head>
@@ -49,6 +48,7 @@
           ">
         </video>
       </div>
+      <div id="scatter-gl-container"></div>
     </div>
   </div>
 </body>

pose-detection/demos/live_video/package-lock.json

@@ -9,12 +9,13 @@
     "node": ">=8.9.0"
   },
   "dependencies": {
+    "@mediapipe/pose": "~0.4.0",
     "@tensorflow-models/pose-detection": "file:../../dist",
-    "@tensorflow/tfjs-backend-wasm": "^3.6.0",
-    "@tensorflow/tfjs-backend-webgl": "^3.6.0",
-    "@tensorflow/tfjs-converter": "^3.6.0",
-    "@tensorflow/tfjs-core": "^3.6.0",
-    "@mediapipe/pose": "~0.3.0"
+    "@tensorflow/tfjs-backend-wasm": "^3.8.0",
+    "@tensorflow/tfjs-backend-webgl": "^3.8.0",
+    "@tensorflow/tfjs-converter": "^3.8.0",
+    "@tensorflow/tfjs-core": "^3.8.0",
+    "scatter-gl": "0.0.8"
   },
   "scripts": {
     "watch": "cross-env NODE_ENV=development parcel index.html --no-hmr --open",

pose-detection/demos/live_video/package.json

@@ -15,15 +15,26 @@
  * =============================================================================
  */
 import * as posedetection from '@tensorflow-models/pose-detection';
+import * as scatter from 'scatter-gl';
 
 import * as params from './params';
 import {isMobile} from './util';
 
+// These anchor points allow the pose pointcloud to resize according to its
+// position in the input.
+const ANCHOR_POINTS = [[0, 0, 0], [0, 1, 0], [-1, 0, 0], [-1, -1, 0]];
 export class Camera {
   constructor() {
     this.video = document.getElementById('video');
     this.canvas = document.getElementById('output');
     this.ctx = this.canvas.getContext('2d');
+    this.scatterGLEl = document.querySelector('#scatter-gl-container');
+    this.scatterGL = new scatter.ScatterGL(this.scatterGLEl, {
+      'rotateOnStart': true,
+      'selectEnabled': false,
+      'styles': {polyline: {defaultOpacity: 1, deselectedOpacity: 1}}
+    });
+    this.scatterGLHasInitialized = false;
   }
 
   /**
@@ -81,6 +92,13 @@ export class Camera {
     camera.ctx.translate(camera.video.videoWidth, 0);
     camera.ctx.scale(-1, 1);
 
+    camera.scatterGLEl.style =
+        `width: ${videoWidth}px; height: ${videoHeight}px;`;
+    camera.scatterGL.resize();
+
+    camera.scatterGLEl.style.display =
+        params.STATE.modelConfig.render3D ? 'inline-block' : 'none';
+
     return camera;
   }
 
@@ -112,6 +130,9 @@ export class Camera {
       this.drawKeypoints(pose.keypoints);
       this.drawSkeleton(pose.keypoints);
     }
+    if (pose.keypoints3D != null && params.STATE.modelConfig.render3D) {
+      this.drawKeypoints3D(pose.keypoints3D);
+    }
   }
 
   /**
@@ -121,7 +142,7 @@ export class Camera {
   drawKeypoints(keypoints) {
     const keypointInd =
         posedetection.util.getKeypointIndexBySide(params.STATE.model);
-    this.ctx.fillStyle = 'White';
+    this.ctx.fillStyle = 'Red';
     this.ctx.strokeStyle = 'White';
     this.ctx.lineWidth = params.DEFAULT_LINE_WIDTH;
 
@@ -181,4 +202,41 @@ export class Camera {
       }
     });
   }
+
+  drawKeypoints3D(keypoints) {
+    const scoreThreshold = params.STATE.modelConfig.scoreThreshold || 0;
+    const pointsData =
+        keypoints.map(keypoint => ([-keypoint.x, -keypoint.y, -keypoint.z]));
+
+    const dataset =
+        new scatter.ScatterGL.Dataset([...pointsData, ...ANCHOR_POINTS]);
+
+    const keypointInd =
+        posedetection.util.getKeypointIndexBySide(params.STATE.model);
+    this.scatterGL.setPointColorer((i) => {
+      if (keypoints[i] == null || keypoints[i].score < scoreThreshold) {
+        // hide anchor points and low-confident points.
+        return '#ffffff';
+      }
+      if (i === 0) {
+        return '#ff0000' /* Red */;
+      }
+      if (keypointInd.left.indexOf(i) > -1) {
+        return '#00ff00' /* Green */;
+      }
+      if (keypointInd.right.indexOf(i) > -1) {
+        return '#ffa500' /* Orange */;
+      }
+    });
+
+    if (!this.scatterGLHasInitialized) {
+      this.scatterGL.render(dataset);
+    } else {
+      this.scatterGL.updateDataset(dataset);
+    }
+    const connections = posedetection.util.getAdjacentPairs(params.STATE.model);
+    const sequences = connections.map(pair => ({indices: pair}));
+    this.scatterGL.setSequences(sequences);
+    this.scatterGLHasInitialized = true;
+  }
 }

pose-detection/demos/live_video/src/camera.js

@@ -53,8 +53,7 @@ async function createDetector() {
         return posedetection.createDetector(STATE.model, {
           runtime,
           modelType: STATE.modelConfig.type,
-          solutionPath:
-              'https://cdn.jsdelivr.net/npm/@mediapipe/[email protected]'
+          solutionPath: 'https://cdn.jsdelivr.net/npm/@mediapipe/[email protected]'
         });
       } else if (runtime === 'tfjs') {
         return posedetection.createDetector(

pose-detection/demos/live_video/src/index.js

@@ -189,6 +189,13 @@ function addBlazePoseControllers(modelConfigFolder, type) {
   });
 
   modelConfigFolder.add(params.STATE.modelConfig, 'scoreThreshold', 0, 1);
+
+  const render3DController =
+      modelConfigFolder.add(params.STATE.modelConfig, 'render3D');
+  render3DController.onChange(render3D => {
+    document.querySelector('#scatter-gl-container').style.display =
+        render3D ? 'inline-block' : 'none';
+  });
 }
 
 /**

pose-detection/demos/live_video/src/option_panel.js

@@ -33,7 +33,9 @@ export const STATE = {
 };
 export const BLAZEPOSE_CONFIG = {
   maxPoses: 1,
-  scoreThreshold: 0.65
+  type: 'full',
+  scoreThreshold: 0.65,
+  render3D: true
 };
 export const POSENET_CONFIG = {
   maxPoses: 1,
@@ -83,8 +85,7 @@ export const BACKEND_FLAGS_MAP = {
 export const MODEL_BACKEND_MAP = {
   [posedetection.SupportedModels.PoseNet]: ['tfjs-webgl'],
   [posedetection.SupportedModels.MoveNet]: ['tfjs-webgl', 'tfjs-wasm'],
-  [posedetection.SupportedModels.BlazePose]:
-      isiOS() ? ['tfjs-webgl'] : ['mediapipe-gpu', 'tfjs-webgl']
+  [posedetection.SupportedModels.BlazePose]: ['mediapipe-gpu', 'tfjs-webgl']
 }
 
 export const TUNABLE_FLAG_NAME_MAP = {