From 64beafa83f5144215f5328afc9e64e57a329a1b1 Mon Sep 17 00:00:00 2001
From: Na Li <linazhao@google.com>
Date: Thu, 1 Apr 2021 12:58:52 -0700
Subject: [PATCH] [pose-detection]Add skeleton to demo.

---
 pose-detection/demo/src/camera.js | 36 ++++++++++++++++++++++++++-----
 pose-detection/src/constants.ts   | 11 ++++++++++
 pose-detection/src/util.ts        | 10 +++++++++
 3 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/pose-detection/demo/src/camera.js b/pose-detection/demo/src/camera.js
index 3047628718..ec1f5578c3 100644
--- a/pose-detection/demo/src/camera.js
+++ b/pose-detection/demo/src/camera.js
@@ -93,7 +93,11 @@ export class Camera {
 
   drawResult(pose, shouldScale = false) {
     if (pose.keypoints != null) {
-      this.drawKeypoints(pose.keypoints, shouldScale);
+      const scaleX = shouldScale ? this.video.videoWidth : 1;
+      const scaleY = shouldScale ? this.video.videoHeight : 1;
+
+      this.drawKeypoints(pose.keypoints, scaleY, scaleX);
+      this.drawSkeleton(pose.keypoints, scaleY, scaleX);
     }
   }
 
@@ -103,10 +107,7 @@ export class Camera {
    * @param shouldScale If the keypoints are normalized, shouldScale should be
    *     set to true.
    */
-  drawKeypoints(keypoints, shouldScale) {
-    const scaleX = shouldScale ? this.video.videoWidth : 1;
-    const scaleY = shouldScale ? this.video.videoHeight : 1;
-
+  drawKeypoints(keypoints, scaleY, scaleX) {
     const keypointInd =
         posedetection.util.getKeypointIndexBySide(params.STATE.model.model);
     this.ctx.fillStyle = 'White';
@@ -140,4 +141,29 @@ export class Camera {
       this.ctx.stroke(circle);
     }
   }
+
+  drawSkeleton(keypoints, scaleY, scaleX) {
+    this.ctx.fillStyle = 'White';
+    this.ctx.strokeStyle = 'White';
+    this.ctx.lineWidth = params.DEFAULT_LINE_WIDTH;
+
+    posedetection.util.getAdjacentPairs(params.STATE.model.model)
+        .forEach(([i, j]) => {
+          const kp1 = keypoints[i];
+          const kp2 = keypoints[j];
+
+          // If score is null, just show the keypoint.
+          const score1 = kp1.score != null ? kp1.score : 1;
+          const score2 = kp2.score != null ? kp2.score : 1;
+          const scoreThreshold =
+              params.STATE.model[params.STATE.model.model].scoreThreshold || 0;
+
+          if (score1 >= scoreThreshold && score2 >= scoreThreshold) {
+            this.ctx.beginPath();
+            this.ctx.moveTo(kp1.x * scaleX, kp1.y * scaleY);
+            this.ctx.lineTo(kp2.x * scaleX, kp2.y * scaleY);
+            this.ctx.stroke();
+          }
+        });
+  }
 }
diff --git a/pose-detection/src/constants.ts b/pose-detection/src/constants.ts
index 11e2227fbe..c406a4c783 100644
--- a/pose-detection/src/constants.ts
+++ b/pose-detection/src/constants.ts
@@ -42,3 +42,14 @@ export const COCO_KEYPOINTS_BY_SIDE = {
   right: [2, 4, 6, 8, 10, 12, 14, 16],
   middle: [0]
 };
+export const COCO_CONNECTED_KEYPOINTS_PAIRS = [
+  [0, 1], [0, 2], [1, 3], [2, 4], [5, 6], [5, 7], [5, 11], [6, 8], [6, 12],
+  [7, 9], [8, 10], [11, 12], [11, 13], [12, 14], [13, 15], [14, 16]
+];
+export const BLAZEPOSE_CONNECTED_KEYPOINTS_PAIRS = [
+  [0, 1],   [0, 4],   [1, 2],   [2, 3],   [3, 7],   [4, 5],   [5, 6],
+  [6, 8],   [9, 10],  [11, 12], [11, 13], [11, 23], [12, 14], [14, 16],
+  [12, 24], [13, 15], [15, 17], [16, 18], [15, 21], [16, 22], [17, 19],
+  [18, 20], [23, 25], [23, 24], [24, 26], [25, 27], [26, 28], [27, 29],
+  [28, 30], [27, 31], [28, 32], [29, 31], [30, 32]
+];
diff --git a/pose-detection/src/util.ts b/pose-detection/src/util.ts
index e2ef96f956..67d789738d 100644
--- a/pose-detection/src/util.ts
+++ b/pose-detection/src/util.ts
@@ -28,3 +28,13 @@ export function getKeypointIndexBySide(model: SupportedModels):
       throw new Error(`Model ${model} is not supported.`);
   }
 }
+export function getAdjacentPairs(model: SupportedModels): number[][] {
+  switch (model) {
+    case SupportedModels.MediapipeBlazepose:
+      return constants.BLAZEPOSE_CONNECTED_KEYPOINTS_PAIRS;
+    case SupportedModels.PoseNet:
+      return constants.COCO_CONNECTED_KEYPOINTS_PAIRS;
+    default:
+      throw new Error(`Model ${model} is not supported.`);
+  }
+}