From 20cb91739d9f642e78b4f3e9214db28e7e964cfb Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Fri, 19 Nov 2021 18:30:57 -0500
Subject: [PATCH] add body 3d interpolation

---
 src/body/blazepose.ts   | 26 +++++++++++++++++---------
 src/util/interpolate.ts |  2 ++
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/body/blazepose.ts b/src/body/blazepose.ts
index e0ac35ac..8bba9f1a 100644
--- a/src/body/blazepose.ts
+++ b/src/body/blazepose.ts
@@ -64,15 +64,23 @@ function calculateBoxes(keypoints: Array<BodyKeypoint>, outputSize: [number, num
 async function prepareImage(input: Tensor): Promise<Tensor> {
   const t: Record<string, Tensor> = {};
   if (!input.shape || !input.shape[1] || !input.shape[2]) return input;
-  padding = [
-    [0, 0], // dont touch batch
-    [input.shape[2] > input.shape[1] ? Math.trunc((input.shape[2] - input.shape[1]) / 2) : 0, input.shape[2] > input.shape[1] ? Math.trunc((input.shape[2] - input.shape[1]) / 2) : 0], // height before&after
-    [input.shape[1] > input.shape[2] ? Math.trunc((input.shape[1] - input.shape[2]) / 2) : 0, input.shape[1] > input.shape[2] ? Math.trunc((input.shape[1] - input.shape[2]) / 2) : 0], // width before&after
-    [0, 0], // dont touch rbg
-  ];
-  t.pad = tf.pad(input, padding);
-  t.resize = tf.image.resizeBilinear(t.pad, [inputSize[1][0], inputSize[1][1]]);
-  const final = tf.div(t.resize, constants.tf255);
+  let final: Tensor;
+  if (input.shape[1] !== input.shape[2]) { // only pad if width different than height
+    padding = [
+      [0, 0], // dont touch batch
+      [input.shape[2] > input.shape[1] ? Math.trunc((input.shape[2] - input.shape[1]) / 2) : 0, input.shape[2] > input.shape[1] ? Math.trunc((input.shape[2] - input.shape[1]) / 2) : 0], // height before&after
+      [input.shape[1] > input.shape[2] ? Math.trunc((input.shape[1] - input.shape[2]) / 2) : 0, input.shape[1] > input.shape[2] ? Math.trunc((input.shape[1] - input.shape[2]) / 2) : 0], // width before&after
+      [0, 0], // dont touch rbg
+    ];
+    t.pad = tf.pad(input, padding);
+    t.resize = tf.image.resizeBilinear(t.pad, [inputSize[1][0], inputSize[1][1]]);
+    final = tf.div(t.resize, constants.tf255);
+  } else if (input.shape[1] !== inputSize[1][0]) { // if input needs resizing
+    t.resize = tf.image.resizeBilinear(input, [inputSize[1][0], inputSize[1][1]]);
+    final = tf.div(t.resize, constants.tf255);
+  } else { // if input is already in a correct resolution just normalize it
+    final = tf.div(input, constants.tf255);
+  }
   Object.keys(t).forEach((tensor) => tf.dispose(t[tensor]));
   return final;
 }
diff --git a/src/util/interpolate.ts b/src/util/interpolate.ts
index cddd1346..0ad5fc27 100644
--- a/src/util/interpolate.ts
+++ b/src/util/interpolate.ts
@@ -50,10 +50,12 @@ export function calc(newResult: Result, config: Config): Result {
           position: [
             bufferedResult.body[i].keypoints[j] ? ((bufferedFactor - 1) * bufferedResult.body[i].keypoints[j].position[0] + keypoint.position[0]) / bufferedFactor : keypoint.position[0],
             bufferedResult.body[i].keypoints[j] ? ((bufferedFactor - 1) * bufferedResult.body[i].keypoints[j].position[1] + keypoint.position[1]) / bufferedFactor : keypoint.position[1],
+            bufferedResult.body[i].keypoints[j] ? ((bufferedFactor - 1) * (bufferedResult.body[i].keypoints[j].position[2] || 0) + (keypoint.position[2] || 0)) / bufferedFactor : keypoint.position[2],
           ],
           positionRaw: [
             bufferedResult.body[i].keypoints[j] ? ((bufferedFactor - 1) * bufferedResult.body[i].keypoints[j].positionRaw[0] + keypoint.positionRaw[0]) / bufferedFactor : keypoint.position[0],
             bufferedResult.body[i].keypoints[j] ? ((bufferedFactor - 1) * bufferedResult.body[i].keypoints[j].positionRaw[1] + keypoint.positionRaw[1]) / bufferedFactor : keypoint.position[1],
+            bufferedResult.body[i].keypoints[j] ? ((bufferedFactor - 1) * (bufferedResult.body[i].keypoints[j].positionRaw[2] || 0) + (keypoint.positionRaw[2] || 0)) / bufferedFactor : keypoint.position[1],
           ],
         }))) as Array<{ score: number, part: string, position: [number, number, number?], positionRaw: [number, number, number?] }>;
       const annotations: Record<string, Point[][]> = {};