Introducing KeypointTracker, which associated people between frames using keypoint similarity. (#769)

* Skeleton for abstract base class Tracker

* Creating updateTracks() which cleans up track store based on track age. Also, updating property names to be mixed case instead of using underscores

* Basic formatting updates.

* Adding lowerCamelCase for all variables and a couple more formatting updates

* Creating tracker_utils.ts to validate tracker config.

* Making computeSimilarity() abstract to allow for batch implementations. Also fixing some comments and validation checks

* Updating flow so that tracks are filtered by age prior to assignment with new detections

* Updating the documentation for maxTracks to indicate that this should be set higher than maxPoses

* Moving track filtering before the computation of similarity matrices

* Implementing assignTracks(), which is a greedy assigner.

* Merge remote-tracking branch 'upstream/master' into tracker, and add new tracking changes.

* Fixing linting issues.

* Updating assignTracks so that unmatched detections are kept track of

* Merge remote-tracking branch 'upstream/master' into tracker

* Introducing KeypointTracker, which tracks people from frame-to-frame based on keypoint similarity.

* Extending keypoint tracking test to show the enforcement of maxTracks.

* Fixing linting issues.

* Small surgical fixes.

* Switching condition in computeSimilarity() so that both poses and tracks must be present for similarity matrix computation.

* Making KeypointTracker area() and oks() methods private.

* Formatting update.

* Comment formatting update.

Author: ronnyvotel

pose-detection/src/calculators/interfaces/config_interfaces.ts

@@ -83,4 +83,25 @@ export interface TrackerConfig {
                    // re-identifications).
   minSimilarity: number;  // New poses will only be linked with tracks if the
                           // similarity score exceeds this threshold.
+  trackerParams: KeypointTrackerConfig;  // Config for tracker. Note that as
+                                         // more trackers are implemented, this
+                                         // should become a union of all tracker
+                                         // types.
+}
+// A tracker that links detections (i.e. poses) and tracks based on keypoint
+// similarity.
+export interface KeypointTrackerConfig {
+  keypointConfidenceThreshold: number;  // The minimum keypoint confidence
+                                        // threshold. A keypoint is only
+                                        // compared in the OKS calculation if
+                                        // both the new detected keypoint and
+                                        // the corresponding track keypoint have
+                                        // confidences above this threshold.
+  
+  keypointFalloff: number[];  // Per-keypoint falloff in OKS calculation.
+  minNumberOfKeypoints: number;  // The minimum number of keypoints that are
+                                 // necessary for computing OKS. If the number
+                                 // of confident keypoints (between a pose and
+                                 // track) are under this value, an OKS of 0.0
+                                 // will be given.
 }

pose-detection/src/calculators/keypoint_tracker.ts

@@ -0,0 +1,120 @@
+/**
+ * @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 {Keypoint, Pose} from '../types';
+import {Tracker} from './tracker';
+import {Track} from './interfaces/common_interfaces';
+import {TrackerConfig} from './interfaces/config_interfaces';
+
+/**
+ * KeypointTracker, which tracks poses based on keypoint similarity. This
+ * tracker assumes that keypoints are provided in normalized image
+ * coordinates.
+ */
+export class KeypointTracker extends Tracker {
+  private readonly keypointThreshold: number;
+  private readonly keypointFalloff: number[];
+  private readonly minNumKeyoints: number;
+
+  constructor(config: TrackerConfig) {
+    super(config);
+    //TODO(ronnyvotel): validate.
+    this.keypointThreshold = config.trackerParams.keypointConfidenceThreshold;
+    this.keypointFalloff = config.trackerParams.keypointFalloff;
+    this.minNumKeyoints = config.trackerParams.minNumberOfKeypoints;
+  }
+
+  /**
+   * Computes similarity based on Object Keypoint Similarity (OKS). It's
+   * assumed that the keypoints within each `Pose` are in normalized image
+   * coordinates. See `Tracker` for more details.
+   */
+  computeSimilarity(poses: Pose[]): number[][] {
+    if (poses.length === 0 || this.tracks.length === 0) {
+      return [[]];
+    }
+
+    const simMatrix = [];
+    for (const pose of poses) {
+      const row = [];
+      for (const track of this.tracks) {
+        row.push(this.oks(pose, track));
+      }
+      simMatrix.push(row);
+    }
+    return simMatrix;
+  }
+
+  /**
+   * Computes the Object Keypoint Similarity (OKS) between a pose and track.
+   * This is similar in spirit to the calculation used by COCO keypoint eval:
+   * https://cocodataset.org/#keypoints-eval
+   * In this case, OKS is calculated as:
+   * (1/sum_i d(c_i, c_ti)) * \sum_i exp(-d_i^2/(2*a_ti*x_i^2))*d(c_i, c_ti)
+   * where
+   *   d(x, y) is an indicator function which only produces 1 if x and y
+   *     exceed a given threshold (i.e. keypointThreshold), otherwise 0.
+   *   c_i is the confidence of keypoint i from the new pose
+   *   c_ti is the confidence of keypoint i from the track
+   *   d_i is the Euclidean distance between the pose and track keypoint
+   *   a_ti is the area of the track object (the box covering the keypoints)
+   *   x_i is a constant that controls falloff in a Gaussian distribution,
+   *    computed as 2*keypointFalloff[i].
+   * @param pose A `Pose`.
+   * @param track A `Track`.
+   * @returns The OKS score between the pose and the track. This number is
+   * between 0 and 1, and larger values indicate more keypoint similarity.
+   */
+  private oks(pose: Pose, track: Track): number {
+    const boxArea = this.area(track.keypoints) + 1e-6;
+    let oksTotal = 0;
+    let numValidKeypoints = 0;
+    for (let i = 0; i < pose.keypoints.length; ++i) {
+      const poseKpt = pose.keypoints[i];
+      const trackKpt = track.keypoints[i];
+      if (poseKpt.score < this.keypointThreshold ||
+          trackKpt.score < this.keypointThreshold) {
+        continue;
+      }
+      numValidKeypoints += 1;
+      const dSquared =
+          Math.pow(poseKpt.x - trackKpt.x, 2) +
+          Math.pow(poseKpt.y - trackKpt.y, 2);
+      const x = 2 * this.keypointFalloff[i];
+      oksTotal += Math.exp(-1 * dSquared / (2 * boxArea * x**2));
+    }
+    if (numValidKeypoints < this.minNumKeyoints) {
+      return 0.0;
+    }
+    return oksTotal / numValidKeypoints;
+  }
+
+  /**
+   * Computes the area of a bounding box that tightly covers keypoints.
+   * @param Keypoint[] An array of `Keypoint`s.
+   * @returns The area of the object.
+   */
+  private area(keypoints: Keypoint[]): number {
+    const validKeypoint = keypoints.filter(
+        kpt => kpt.score > this.keypointThreshold);
+    const minX = Math.min(1.0, ...validKeypoint.map(kpt => kpt.x));
+    const maxX = Math.max(0.0, ...validKeypoint.map(kpt => kpt.x));
+    const minY = Math.min(1.0, ...validKeypoint.map(kpt => kpt.y));
+    const maxY = Math.max(0.0, ...validKeypoint.map(kpt => kpt.y));
+    return (maxX - minX) * (maxY - minY);
+  }
+}

pose-detection/src/calculators/keypoint_tracker_test.ts

@@ -0,0 +1,242 @@
+/**
+ * @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 {Keypoint, Pose} from '../types';
+import {KeypointTracker} from './keypoint_tracker';
+import {Track} from './interfaces/common_interfaces';
+import {TrackerConfig} from './interfaces/config_interfaces';
+
+describe('Keypoint tracker', () => {
+  const trackerConfig: TrackerConfig = {
+      maxTracks: 4,
+      maxAge: 1000,
+      minSimilarity: 0.5,
+      trackerParams: {
+        keypointConfidenceThreshold: 0.2,
+        keypointFalloff: [0.1, 0.1, 0.1, 0.1],
+        minNumberOfKeypoints: 2
+      }
+    };
+
+  it('Instantiate tracker', () => {
+    const kptTracker = new KeypointTracker(trackerConfig);
+    expect(kptTracker instanceof KeypointTracker).toBe(true);
+  });
+
+  it('Compute OKS', () => {
+    const kptTracker = new KeypointTracker(trackerConfig);
+    const pose: Pose = {
+      keypoints: [
+          {x: 0.2, y: 0.2, score: 1.0},
+          {x: 0.4, y: 0.4, score: 0.8},
+          {x: 0.6, y: 0.6, score: 0.1},  // Low confidence.
+          {x: 0.8, y: 0.7, score: 0.8}
+      ]};
+    const track: Track = {
+      id: 0,
+      lastTimestamp: 1000,
+      keypoints: [
+          {x: 0.2, y: 0.2, score: 1.0},
+          {x: 0.4, y: 0.4, score: 0.8},
+          {x: 0.6, y: 0.6, score: 0.9},
+          {x: 0.8, y: 0.8, score: 0.8}
+      ]};
+    const oks = kptTracker['oks'](pose, track);
+
+    const boxArea = (0.8 - 0.2) * (0.8 - 0.2);
+    const x = 2*trackerConfig.trackerParams.keypointFalloff[3];
+    const d = 0.1;
+    const expectedOks =
+        (1 + 1 + Math.exp(-1*d**2/(2*boxArea*x**2))) / 3;
+    expect(oks).toBeCloseTo(expectedOks, 6);
+  });
+
+  it('Compute OKS returns 0.0 with less than 2 valid keypoints', () => {
+    const kptTracker = new KeypointTracker(trackerConfig);
+    const pose: Pose = {
+      keypoints: [
+          {x: 0.2, y: 0.2, score: 1.0},
+          {x: 0.4, y: 0.4, score: 0.1},  // Low confidence.
+          {x: 0.6, y: 0.6, score: 0.9},
+          {x: 0.8, y: 0.8, score: 0.8}
+      ]};
+    const track: Track = {
+      id: 0,
+      lastTimestamp: 1000,
+      keypoints: [
+          {x: 0.2, y: 0.2, score: 1.0},
+          {x: 0.4, y: 0.4, score: 0.8},
+          {x: 0.6, y: 0.6, score: 0.1},  // Low confidence.
+          {x: 0.8, y: 0.8, score: 0.0}  // Low confidence.
+      ]};
+    const oks = kptTracker['oks'](pose, track);
+    expect(oks).toBeCloseTo(0.0, 6);
+  });
+
+  it('Compute area', () => {
+    const kptTracker = new KeypointTracker(trackerConfig);
+    const keypoints: Keypoint[] = [
+          {x: 0.1, y: 0.2, score: 1.0},
+          {x: 0.3, y: 0.4, score: 0.9},
+          {x: 0.4, y: 0.6, score: 0.9},
+          {x: 0.7, y: 0.8, score: 0.1}  // Low confidence.
+      ];
+    const area = kptTracker['area'](keypoints);
+
+    const expectedArea = (0.4 - 0.1) * (0.6 - 0.2);
+    expect(area).toBeCloseTo(expectedArea, 6);
+  });
+
+  it('Apply tracker', () => {
+    // Timestamp: 0. Pose becomes the only track.
+    const kptTracker = new KeypointTracker(trackerConfig);
+    let tracks: Track[];
+    let poses: Pose[] = [
+        {keypoints: [  // Becomes id = 1.
+            {x: 0.2, y: 0.2, score: 1.0},
+            {x: 0.4, y: 0.4, score: 0.8},
+            {x: 0.6, y: 0.6, score: 0.9},
+            {x: 0.8, y: 0.8, score: 0.0}  // Low confidence.
+        ]}
+    ];
+    poses = kptTracker.apply(poses, 0);
+    tracks = kptTracker.getTracks();
+    expect(poses.length).toEqual(1);
+    expect(poses[0].id).toEqual(1);
+    expect(tracks.length).toEqual(1);
+    expect(tracks[0].id).toEqual(1);
+    expect(tracks[0].lastTimestamp).toEqual(0);
+
+    // Timestamp: 100. First pose is linked with track 1. Second pose spawns a
+    // new track (id = 2).
+    poses = [
+        {keypoints: [  // Links with id = 1.
+            {x: 0.2, y: 0.2, score: 1.0},
+            {x: 0.4, y: 0.4, score: 0.8},
+            {x: 0.6, y: 0.6, score: 0.9},
+            {x: 0.8, y: 0.8, score: 0.8}
+        ]},
+        {keypoints: [  // Becomes id = 2.
+            {x: 0.8, y: 0.8, score: 0.8},
+            {x: 0.6, y: 0.6, score: 0.3},
+            {x: 0.4, y: 0.4, score: 0.1},  // Low confidence.
+            {x: 0.2, y: 0.2, score: 0.8}
+        ]}
+    ];
+    poses = kptTracker.apply(poses, 100);
+    tracks = kptTracker.getTracks();
+    expect(poses.length).toEqual(2);
+    expect(poses[0].id).toEqual(1);
+    expect(poses[1].id).toEqual(2);
+    expect(tracks.length).toEqual(2);
+    expect(tracks[0].id).toEqual(1);
+    expect(tracks[0].lastTimestamp).toEqual(100);
+    expect(tracks[1].id).toEqual(2);
+    expect(tracks[1].lastTimestamp).toEqual(100);
+
+    // Timestamp: 900. First pose is linked with track 2. Second pose spawns a
+    // new track (id = 3).
+    poses = [
+        {keypoints: [  // Links with id = 2.
+            {x: 0.6, y: 0.7, score: 0.7},
+            {x: 0.5, y: 0.6, score: 0.7},
+            {x: 0.0, y: 0.0, score: 0.1},  // Low confidence.
+            {x: 0.2, y: 0.1, score: 1.0}
+        ]},
+        {keypoints: [  // Becomes id = 3.
+            {x: 0.5, y: 0.1, score: 0.6},
+            {x: 0.9, y: 0.3, score: 0.6},
+            {x: 0.1, y: 1.0, score: 0.9},
+            {x: 0.4, y: 0.4, score: 0.1}  // Low confidence.
+        ]},
+    ];
+    poses = kptTracker.apply(poses, 900);
+    tracks = kptTracker.getTracks();
+    expect(poses.length).toEqual(2);
+    expect(poses[0].id).toEqual(2);
+    expect(poses[1].id).toEqual(3);
+    expect(tracks.length).toEqual(3);
+    expect(tracks[0].id).toEqual(2);
+    expect(tracks[0].lastTimestamp).toEqual(900);
+    expect(tracks[1].id).toEqual(3);
+    expect(tracks[1].lastTimestamp).toEqual(900);
+    expect(tracks[2].id).toEqual(1);
+    expect(tracks[2].lastTimestamp).toEqual(100);
+
+    // Timestamp: 1200. First pose spawns a new track (id = 4), even though it
+    // has the same keypoints as track 1. This is because the age exceeds 1000
+    // msec. The second pose links with id 2. The third pose spawns a new
+    // track (id = 5).
+    poses = [
+        {keypoints: [  // Becomes id = 4.
+            {x: 0.2, y: 0.2, score: 1.0},
+            {x: 0.4, y: 0.4, score: 0.8},
+            {x: 0.6, y: 0.6, score: 0.9},
+            {x: 0.8, y: 0.8, score: 0.8}
+        ]},
+        {keypoints: [  // Links with id = 2.
+            {x: 0.55, y: 0.7, score: 0.7},
+            {x: 0.5, y: 0.6, score: 0.9},
+            {x: 1.0, y: 1.0, score: 0.1},  // Low confidence.
+            {x: 0.8, y: 0.1, score: 0.0}  // Low confidence.
+        ]},
+        {keypoints: [  // Becomes id = 5.
+            {x: 0.1, y: 0.1, score: 0.1},  // Low confidence.
+            {x: 0.2, y: 0.2, score: 0.9},
+            {x: 0.3, y: 0.3, score: 0.7},
+            {x: 0.4, y: 0.4, score: 0.8} 
+        ]},
+    ];
+    poses = kptTracker.apply(poses, 1200);
+    tracks = kptTracker.getTracks();
+    expect(poses.length).toEqual(3);
+    expect(poses[0].id).toEqual(4);
+    expect(poses[1].id).toEqual(2);
+    expect(tracks.length).toEqual(4);
+    expect(tracks[0].id).toEqual(2);
+    expect(tracks[0].lastTimestamp).toEqual(1200);
+    expect(tracks[1].id).toEqual(4);
+    expect(tracks[1].lastTimestamp).toEqual(1200);
+    expect(tracks[2].id).toEqual(5);
+    expect(tracks[2].lastTimestamp).toEqual(1200);
+    expect(tracks[3].id).toEqual(3);
+    expect(tracks[3].lastTimestamp).toEqual(900);
+
+    // Timestamp: 1300. First pose spawns a new track (id = 6). Since maxTracks
+    // is 4, the oldest track (id = 3) is removed.
+    poses = [
+        {keypoints: [  // Becomes id = 6.
+            {x: 0.1, y: 0.8, score: 1.0},
+            {x: 0.2, y: 0.9, score: 0.6},
+            {x: 0.2, y: 0.9, score: 0.5},
+            {x: 0.8, y: 0.2, score: 0.4}
+        ]},
+    ];
+    poses = kptTracker.apply(poses, 1300);
+    tracks = kptTracker.getTracks();
+    expect(poses.length).toEqual(1);
+    expect(poses[0].id).toEqual(6);
+    expect(tracks.length).toEqual(4);
+    expect(tracks[0].id).toEqual(6);
+    expect(tracks[0].lastTimestamp).toEqual(1300);
+    expect(tracks[1].id).toEqual(2);
+    expect(tracks[1].lastTimestamp).toEqual(1200);
+    expect(tracks[2].id).toEqual(4);
+    expect(tracks[2].lastTimestamp).toEqual(1200);
+    expect(tracks[3].id).toEqual(5);
+    expect(tracks[3].lastTimestamp).toEqual(1200);
+  });
+});