[pose-detection]Add named output util. (#659)

Author: lina128

pose-detection/src/constants.ts

@@ -65,7 +65,7 @@ export const BLAZEPOSE_CONNECTED_KEYPOINTS_PAIRS_UPPER_BODY = [
   [12, 24], [13, 15], [15, 17], [16, 18], [16, 20], [15, 17], [15, 19],
   [15, 21], [16, 22], [17, 19], [18, 20], [23, 24]
 ];
-export const COCO_KEYPOINTS_NAMED_MAP: {[index: string]: number} = {
+export const COCO_KEYPOINTS_BY_NAME: {[index: string]: number} = {
   nose: 0,
   left_eye: 1,
   right_eye: 2,
@@ -84,3 +84,67 @@ export const COCO_KEYPOINTS_NAMED_MAP: {[index: string]: number} = {
   left_ankle: 15,
   right_ankle: 16
 };
+export const BLAZEPOSE_KEYPOINTS_BY_NAME_UPPER_BODY:
+    {[index: string]: number} = {
+      nose: 0,
+      left_eye_inner: 1,
+      left_eye: 2,
+      left_eye_outer: 3,
+      right_eye_inner: 4,
+      right_eye: 5,
+      right_eye_outer: 6,
+      left_ear: 7,
+      right_ear: 8,
+      mouth_left: 9,
+      mouth_right: 10,
+      left_shoulder: 11,
+      right_shoulder: 12,
+      left_elbow: 13,
+      right_elbow: 14,
+      left_wrist: 15,
+      right_wrist: 16,
+      left_pinky: 17,
+      right_pinky: 18,
+      left_index: 19,
+      right_index: 20,
+      left_thumb: 21,
+      right_thumb: 22,
+      left_hip: 23,
+      right_hip: 24
+    };
+export const BLAZEPOSE_KEYPOINTS_BY_NAME_FULL_BODY:
+    {[index: string]: number} = {
+      nose: 0,
+      left_eye_inner: 1,
+      left_eye: 2,
+      left_eye_outer: 3,
+      right_eye_inner: 4,
+      right_eye: 5,
+      right_eye_outer: 6,
+      left_ear: 7,
+      right_ear: 8,
+      mouth_left: 9,
+      mouth_right: 10,
+      left_shoulder: 11,
+      right_shoulder: 12,
+      left_elbow: 13,
+      right_elbow: 14,
+      left_wrist: 15,
+      right_wrist: 16,
+      left_pinky: 17,
+      right_pinky: 18,
+      left_index: 19,
+      right_index: 20,
+      left_thumb: 21,
+      right_thumb: 22,
+      left_hip: 23,
+      right_hip: 24,
+      left_knee: 25,
+      right_knee: 26,
+      left_ankle: 27,
+      right_ankle: 28,
+      left_heel: 29,
+      right_heel: 30,
+      left_foot_index: 31,
+      right_foot_index: 32
+    };

pose-detection/src/movenet/detector.ts

@@ -24,7 +24,7 @@ import {BoundingBox} from '../calculators/interfaces/shape_interfaces';
 import {isVideo} from '../calculators/is_video';
 import {KeypointsOneEuroFilter} from '../calculators/keypoints_one_euro_filter';
 import {LowPassFilter} from '../calculators/low_pass_filter';
-import {COCO_KEYPOINTS, COCO_KEYPOINTS_NAMED_MAP} from '../constants';
+import {COCO_KEYPOINTS, COCO_KEYPOINTS_BY_NAME} from '../constants';
 import {BasePoseDetector, PoseDetector} from '../pose_detector';
 import {InputResolution, Keypoint, Pose, PoseDetectorInput} from '../types';
 
@@ -343,13 +343,13 @@ export class MoveNetDetector extends BasePoseDetector {
 
   torsoVisible(keypoints: Keypoint[]): boolean {
     return (
-        keypoints[COCO_KEYPOINTS_NAMED_MAP['left_hip']].score >
+        keypoints[COCO_KEYPOINTS_BY_NAME['left_hip']].score >
             MIN_CROP_KEYPOINT_SCORE &&
-        keypoints[COCO_KEYPOINTS_NAMED_MAP['right_hip']].score >
+        keypoints[COCO_KEYPOINTS_BY_NAME['right_hip']].score >
             MIN_CROP_KEYPOINT_SCORE &&
-        keypoints[COCO_KEYPOINTS_NAMED_MAP['left_shoulder']].score >
+        keypoints[COCO_KEYPOINTS_BY_NAME['left_shoulder']].score >
             MIN_CROP_KEYPOINT_SCORE &&
-        keypoints[COCO_KEYPOINTS_NAMED_MAP['right_shoulder']].score >
+        keypoints[COCO_KEYPOINTS_BY_NAME['right_shoulder']].score >
             MIN_CROP_KEYPOINT_SCORE);
   }
 
@@ -381,7 +381,7 @@ export class MoveNetDetector extends BasePoseDetector {
     let maxBodyYrange = 0.0;
     let maxBodyXrange = 0.0;
     for (const key of Object.keys(targetKeypoints)) {
-      if (keypoints[COCO_KEYPOINTS_NAMED_MAP[key]].score <
+      if (keypoints[COCO_KEYPOINTS_BY_NAME[key]].score <
           MIN_CROP_KEYPOINT_SCORE) {
         continue;
       }
@@ -412,10 +412,10 @@ export class MoveNetDetector extends BasePoseDetector {
       imageWidth: number): BoundingBox {
     const targetKeypoints: {[index: string]: number[]} = {};
 
-    for (const key of Object.keys(COCO_KEYPOINTS_NAMED_MAP)) {
+    for (const key of Object.keys(COCO_KEYPOINTS_BY_NAME)) {
       targetKeypoints[key] = [
-        keypoints[COCO_KEYPOINTS_NAMED_MAP[key]].y * imageHeight,
-        keypoints[COCO_KEYPOINTS_NAMED_MAP[key]].x * imageWidth
+        keypoints[COCO_KEYPOINTS_BY_NAME[key]].y * imageHeight,
+        keypoints[COCO_KEYPOINTS_BY_NAME[key]].x * imageWidth
       ];
     }
 

pose-detection/src/util.ts

@@ -44,3 +44,18 @@ export function getAdjacentPairs(model: SupportedModels): number[][] {
       throw new Error(`Model ${model} is not supported.`);
   }
 }
+
+export function getKeypointIndexByName(model: SupportedModels):
+    {[index: string]: number} {
+  switch (model) {
+    case SupportedModels.MediapipeBlazeposeUpperBody:
+      return constants.BLAZEPOSE_KEYPOINTS_BY_NAME_UPPER_BODY;
+    case SupportedModels.MediapipeBlazeposeFullBody:
+      return constants.BLAZEPOSE_KEYPOINTS_BY_NAME_FULL_BODY;
+    case SupportedModels.PoseNet:
+    case SupportedModels.MoveNet:
+      return constants.COCO_KEYPOINTS_BY_NAME;
+    default:
+      throw new Error(`Model ${model} is not supported.`);
+  }
+}