Implement ModelAnimator.Update

Assets/Scripts/Animator/Joint.cs

@@ -31,6 +31,14 @@ public class Joint {
   public Joint Child;
   public Joint Parent;
 
+  // <summary>
+  // The prediction by ML models. It is used by the animator to determine the rotations of
+  // the joints during Update.
+  //
+  // Different from position vector, this is in the same scale but different coordination.
+  // </summary>
+  public Vector3 Prediction { get; set; }
+
   // <summary>
   // The initial placement in the scene.
   // </summary>

Assets/Scripts/Animator/ModelAnimator.cs

@@ -17,10 +17,16 @@
 using UnityEngine.Assertions;
 
 public class ModelAnimator : MonoBehaviour {
+  public float Width = 150;
+  public float Height = 150;
   private Animator _anim;
   private Joint[] _joints = new Joint[Landmarks.Total];
   [SerializeField] private Transform _nose;
+  [SerializeField] private float _zScale = 0.3f;
 
+  // <summary>
+  // Creates the joints and set Child, Parent and Forward of each joint.
+  // </summary>
   public void Start() {
     _anim = GetComponent<Animator>();
 
@@ -106,10 +112,72 @@ public void Start() {
     head.Forward = Quaternion.LookRotation(gaze);  
   }
 
+  // <summary>
+  // Calls SetRotation on each joint by their Predictions.
+  // </summary>
   public void Update() {
+    Vector3 forward = GetNormal(_joints[Landmarks.LeftShoulder].Prediction,
+                                _joints[Landmarks.LeftHip].Prediction,
+                                _joints[Landmarks.RightHip].Prediction);
+    _joints[Landmarks.Hip].SetRotation(Quaternion.LookRotation(forward));
+
+    foreach (Joint jointPoint in _joints) {
+      if (jointPoint != null && jointPoint.Child != null) {
+        Vector3 jointForward = forward;
+        if (jointPoint.Parent != null) {
+          jointForward = jointPoint.Parent.Prediction - jointPoint.Prediction;
+        }
+        jointPoint.SetRotation(
+            Quaternion.LookRotation(jointPoint.Prediction - jointPoint.Child.Prediction,
+                                    jointForward));
+      }
+    }
+
+    Vector3 gaze = _joints[Landmarks.Nose].Prediction - _joints[Landmarks.Head].Prediction;
+    Vector3 upwards = GetNormal(_joints[Landmarks.Nose].Prediction,
+                                _joints[Landmarks.RightEar].Prediction,
+                                _joints[Landmarks.LeftEar].Prediction);
+    _joints[Landmarks.Head].SetRotation(Quaternion.LookRotation(gaze, upwards));
   }
 
+  // <summary>
+  // Updates each joint on their predicted location. It is called when a new prediction is given.
+  // </summary>
   public void SetPoseLandmarks(NormalizedLandmarkList poseLandmarks) {
+    for (int i = 0; i < Landmarks.PoseCount; i++) {
+      NormalizedLandmark landmark = poseLandmarks.Landmark[i];
+      if (_joints[i] != null) {
+        // Apply a z-scale since the BlazePose model doesn't generate good z-axis estimates.
+        // Multiple a width since the landmark coordinates are normalized to [0, 1].
+        _joints[i].Prediction = new Vector3(-landmark.X * Width,
+                                            landmark.Y * Height,
+                                            landmark.Z * _zScale * Width);
+      }
+    }
+    UpdateSpecialJoints();
+  }
+
+  // <summary>
+  // Updates special joints that doesn't have direct mappings in the predicted ones.
+  // </summary>
+  private void UpdateSpecialJoints() {
+    // Hip, neck, spine and head are not directly generated by BlazePose. We use simple heuristics
+    // to find a good estimate with Unity-Chan. They need adjustment for custom models.
+    _joints[Landmarks.Hip].Prediction =
+        Vector3.Lerp(_joints[Landmarks.LeftHip].Prediction,
+                     _joints[Landmarks.RightHip].Prediction, 0.5f);
+    _joints[Landmarks.Neck].Prediction =
+        Vector3.Lerp(_joints[Landmarks.LeftShoulder].Prediction,
+                     _joints[Landmarks.RightShoulder].Prediction, 0.5f);
+    _joints[Landmarks.Spine].Prediction =
+        Vector3.Lerp(_joints[Landmarks.Hip].Prediction,
+                     _joints[Landmarks.Neck].Prediction, 0.2f);
+
+    var midEar = Vector3.Lerp(_joints[Landmarks.LeftEar].Prediction,
+                              _joints[Landmarks.RightEar].Prediction, 0.5f);
+    _joints[Landmarks.Head].Prediction =
+        Vector3.Lerp(_joints[Landmarks.Neck].Prediction,
+                     midEar, 0.85f);
   }
 
   // <summary>

Assets/Scripts/Tracker/PoseTracker.cs

@@ -64,6 +64,8 @@ public IEnumerator Start() {
     _graphRunner.StartRunAsync(sidePacket).AssertOk();
     Logger.LogInfo(_TAG, "Graph Runner started in async mode!");
 
+    _modelAnimator.Width = _sourceTexture.width;
+    _modelAnimator.Height = _sourceTexture.height;
     _textureFramePool.ResizeTexture(_sourceTexture.width, _sourceTexture.height, TextureFormat.RGBA32);
     _coroutine = StartCoroutine(ProcessImage(_sourceTexture));
   }