[pose-detection]Initial blazepose checkin. (#616)

FEATURE

Author: lina128

pose-detection/demo/src/camera.js

@@ -89,17 +89,25 @@ export class Camera {
     this.ctx.clearRect(0, 0, this.video.videoWidth, this.video.videoHeight);
   }
 
-  drawResult(pose) {
-    this.drawKeypoints(pose.keypoints);
+  drawResult(pose, shouldScale = false) {
+    this.drawKeypoints(pose.keypoints, shouldScale);
   }
 
-  drawKeypoints(keypoints) {
+  /**
+   * Draw the keypoints on the video.
+   * @param keypoints A list of keypoints, may be normalized.
+   * @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;
     this.ctx.fillStyle = 'red';
     this.ctx.strokeStyle = 'white';
     this.ctx.lineWidth = 4;
     keypoints.forEach(keypoint => {
       const circle = new Path2D();
-      circle.arc(keypoint.x, keypoint.y, 4, 0, 2 * Math.PI);
+      circle.arc(keypoint.x * scaleX, keypoint.y * scaleY, 4, 0, 2 * Math.PI);
       this.ctx.fill(circle);
       this.ctx.stroke(circle);
     });

pose-detection/demo/src/index.js

@@ -27,6 +27,22 @@ import {setupStats} from './stats_panel';
 
 let detector, camera, stats;
 
+async function createDetector(model) {
+  switch (model) {
+    case posedetection.SupportedModels.PoseNet:
+      return posedetection.createDetector(STATE.model.model, {
+        quantBytes: 4,
+        architecture: 'MobileNetV1',
+        outputStride: 16,
+        inputResolution: {width: 500, height: 500},
+        multiplier: 0.75
+      });
+    case posedetection.SupportedModels.MediapipeBlazepose:
+      return posedetection.createDetector(
+          STATE.model.model, {quantBytes: 4, upperBodyOnly: false});
+  }
+}
+
 async function checkGuiUpdate() {
   if (STATE.changeToTargetFPS || STATE.changeToSizeOption) {
     if (STATE.changeToTargetFPS) {
@@ -42,6 +58,14 @@ async function checkGuiUpdate() {
     camera = await Camera.setupCamera(STATE.camera);
   }
 
+  if (STATE.changeToModel) {
+    STATE.model.model = STATE.changeToModel;
+    STATE.changeToModel = null;
+
+    detector.dispose();
+    detector = await createDetector(STATE.model.model);
+  }
+
   await tf.nextFrame();
 }
 
@@ -50,12 +74,15 @@ async function renderResult() {
     camera.lastVideoTime = camera.video.currentTime;
 
     const poses = await detector.estimatePoses(
-        video, {maxPoses: 1, flipHorizontal: false});
+        camera.video, {maxPoses: 1, flipHorizontal: false});
 
     camera.drawCtx();
 
     if (poses.length > 0) {
-      camera.drawResult(poses[0]);
+      const shouldScale = STATE.model.model ===
+          posedetection.SupportedModels.MediapipeBlazepose;
+
+      camera.drawResult(poses[0], shouldScale);
     }
   }
 }
@@ -78,14 +105,7 @@ async function app() {
   stats = setupStats();
   camera = await Camera.setupCamera(STATE.camera);
 
-  detector = await posedetection.createDetector(
-      posedetection.SupportedModels.PoseNet, {
-        quantBytes: 4,
-        architecture: 'MobileNetV1',
-        outputStride: 16,
-        inputResolution: {width: 500, height: 500},
-        multiplier: 0.75
-      });
+  detector = await createDetector(STATE.model.model);
 
   renderPrediction();
 };

pose-detection/demo/src/option_panel.js

@@ -14,6 +14,8 @@
  * limitations under the License.
  * =============================================================================
  */
+import * as posedetection from '@tensorflow-models/posedetection';
+
 import {STATE, VIDEO_SIZE} from './params';
 
 export function setupDatGui() {
@@ -32,5 +34,14 @@ export function setupDatGui() {
   });
   cameraFolder.open();
 
+  // The model folder contains options for model settings.
+  const modelFolder = gui.addFolder('Model');
+  const modelController = modelFolder.add(
+      STATE.model, 'model', Object.values(posedetection.SupportedModels));
+  modelController.onChange(model => {
+    STATE.changeToModel = model;
+  });
+  modelFolder.open();
+
   return gui;
 }

pose-detection/demo/src/params.js

@@ -14,12 +14,15 @@
  * limitations under the License.
  * =============================================================================
  */
+import * as posedetection from '@tensorflow-models/posedetection';
+
 export const DEFAULT_LINE_WIDTH = 4;
 
 export const VIDEO_SIZE = {
   '640 X 480': {width: 640, height: 480},
   '640 X 360': {width: 640, height: 360}
 };
 export const STATE = {
-  camera: {targetFPS: 60, sizeOption: '640 X 480'}
+  camera: {targetFPS: 60, sizeOption: '640 X 480'},
+  model: {model: posedetection.SupportedModels.PoseNet}
 };

pose-detection/src/blazepose/blazepose_test.ts

@@ -0,0 +1,54 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tf from '@tensorflow/tfjs-core';
+// tslint:disable-next-line: no-imports-from-dist
+import {ALL_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
+
+import * as poseDetection from '../index';
+
+describeWithFlags('Blazepose', ALL_ENVS, () => {
+  let detector: poseDetection.PoseDetector;
+  let startTensors: number;
+
+  beforeEach(async () => {
+    startTensors = tf.memory().numTensors;
+
+    // Note: this makes a network request for model assets.
+    const modelConfig: poseDetection.BlazeposeModelConfig = {
+      quantBytes: 4,
+      upperBodyOnly: false
+    };
+    detector = await poseDetection.createDetector(
+        poseDetection.SupportedModels.MediapipeBlazepose, modelConfig);
+  });
+
+  it('estimatePoses does not leak memory', async () => {
+    const input: tf.Tensor3D = tf.zeros([128, 128, 3]);
+
+    const beforeTensors = tf.memory().numTensors;
+
+    await detector.estimatePoses(input);
+
+    expect(tf.memory().numTensors).toEqual(beforeTensors);
+
+    detector.dispose();
+    input.dispose();
+
+    expect(tf.memory().numTensors).toEqual(startTensors);
+  });
+});

pose-detection/src/blazepose/calculators/calculate_alignment_points_rects.ts

@@ -0,0 +1,60 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {ImageSize} from '../../calculators/interfaces/common_interfaces';
+import {Rect} from '../../calculators/interfaces/shape_interfaces';
+import {computeRotation} from './detection_to_rect';
+import {DetectionToRectConfig} from './interfaces/config_interfaces';
+import {Detection} from './interfaces/shape_interfaces';
+
+// ref:
+// https://github.com/google/mediapipe/blob/master/mediapipe/calculators/util/alignment_points_to_rects_calculator.cc
+export function calculateAlignmentPointsRects(
+    detection: Detection, imageSize: ImageSize,
+    config: DetectionToRectConfig): Rect {
+  const startKeypoint = config.rotationVectorStartKeypointIndex;
+  const endKeypoint = config.rotationVectorEndKeypointIndex;
+
+
+  const locationData = detection.locationData;
+  const xCenter =
+      locationData.relativeKeypoints[startKeypoint].x * imageSize.width;
+  const yCenter =
+      locationData.relativeKeypoints[startKeypoint].y * imageSize.height;
+
+  const xScale =
+      locationData.relativeKeypoints[endKeypoint].x * imageSize.width;
+  const yScale =
+      locationData.relativeKeypoints[endKeypoint].y * imageSize.height;
+
+  // Bounding box size as double distance from center to scale point.
+  const boxSize = Math.sqrt(
+                      (xScale - xCenter) * (xScale - xCenter) +
+                      (yScale - yCenter) * (yScale - yCenter)) *
+      2;
+
+  const rotation = computeRotation(detection, imageSize, config);
+
+  // Set resulting bounding box.
+  return {
+    xCenter: xCenter / imageSize.width,
+    yCenter: yCenter / imageSize.height,
+    width: boxSize / imageSize.width,
+    height: boxSize / imageSize.height,
+    rotation
+  };
+}

pose-detection/src/blazepose/calculators/calculate_landmark_projection.ts

@@ -0,0 +1,59 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {Rect} from '../../calculators/interfaces/shape_interfaces';
+import {Keypoint} from '../../types';
+
+/**
+ * Projects normalized landmarks in a rectangle to its original coordinates. The
+ * rectangle must also be in normalized coordinates.
+ * @param landmarks A normalized Landmark list representing landmarks in a
+ *     normalized rectangle.
+ * @param inputRect A normalized rectangle.
+ * @param config Config object has one field ignoreRotation, default to false.
+ */
+// ref:
+// https://github.com/google/mediapipe/blob/master/mediapipe/calculators/util/landmark_projection_calculator.cc
+export function calculateLandmarkProjection(
+    landmarks: Keypoint[], inputRect: Rect,
+    config: {ignoreRotation: boolean} = {
+      ignoreRotation: false
+    }) {
+  const outputLandmarks = [];
+  for (const landmark of landmarks) {
+    const x = landmark.x - 0.5;
+    const y = landmark.y - 0.5;
+    const angle = config.ignoreRotation ? 0 : inputRect.rotation;
+    let newX = Math.cos(angle) * x - Math.sin(angle) * y;
+    let newY = Math.sin(angle) * x + Math.cos(angle) * y;
+
+    newX = newX * inputRect.width + inputRect.xCenter;
+    newY = newY * inputRect.height + inputRect.yCenter;
+
+    const newZ = landmark.z * inputRect.width;  // Scale Z coordinate as x.
+
+    const newLandmark = {...landmark};
+
+    newLandmark.x = newX;
+    newLandmark.y = newY;
+    newLandmark.z = newZ;
+
+    outputLandmarks.push(newLandmark);
+  }
+
+  return outputLandmarks;
+}