remake posenet results so they are now grouped into one single array

ml5js · Jun 3, 2018 · 1d9148e · oveddan · Jun 5, 2018 · cvalenzuela
1 parent 9c44ab8
commit 1d9148e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/ImageAndVideo.js b/src/ImageAndVideo.js
@@ -15,7 +15,9 @@ class ImageAndVideo {
     this.videoReady = false;
     this.onVideoReady = () => {
       this.videoReady = true;
-      this.waitingPredictions.forEach(i => this.predict(i.imgToPredict, i.num, i.callback));
+      if (this.waitingPredictions) {
+        this.waitingPredictions.forEach(i => this.predict(i.imgToPredict, i.num, i.callback));
+      }
     };
 
     if (video instanceof HTMLVideoElement) {

diff --git a/src/PoseNet/index.js b/src/PoseNet/index.js
@@ -80,7 +80,7 @@ class PoseNet extends ImageAndVideo {
 
     this.net.estimateSinglePose(input, this.imageScaleFactor, this.flipHorizontal, this.outputStride)
       .then((pose) => {
-        callback([pose]);
+        callback([{ pose, skeleton: this.skeleton(pose.keypoints) }]);
         tf.nextFrame().then(() => { this.singlePose(callback); });
       });
   }
@@ -99,7 +99,8 @@ class PoseNet extends ImageAndVideo {
 
     this.net.estimateMultiplePoses(input, this.imageScaleFactor, this.flipHorizontal, this.outputStride)
       .then((poses) => {
-        callback(poses);
+        const result = poses.map(pose => ({ pose, skeleton: this.skeleton(pose.keypoints) }));
+        callback(result);
         tf.nextFrame().then(() => { this.multiPose(callback); });
       });
   }